Arrays.hs

------------------------------------------------------------------------------

------------------------------------------------------------------------------

-- |
-- Module:      Database.PostgreSQL.Simple.Arrays
-- Copyright:   (c) 2012 Leon P Smith
-- License:     BSD3
-- Maintainer:  Leon P Smith <leon@melding-monads.com>
-- Stability:   experimental
--
-- A Postgres array parser and pretty-printer.
module Database.PostgreSQL.Simple.Arrays where

import Control.Applicative (Alternative (..))
import Data.Attoparsec.ByteString.Char8
import Data.ByteString.Char8 (ByteString)
import qualified Data.ByteString.Char8 as B

-- | Parse one of three primitive field formats: array, quoted and plain.
arrayFormat :: Char -> Parser ArrayFormat
arrayFormat delim =
  Array
    <$> array delim
      <|> Plain
    <$> plain delim
      <|> Quoted
    <$> quoted

data ArrayFormat
  = Array [ArrayFormat]
  | Plain ByteString
  | Quoted ByteString
  deriving (Eq, Show, Ord)

array :: Char -> Parser [ArrayFormat]
array delim = char '{' *> option [] (arrays <|> strings) <* char '}'
  where
    strings = sepBy1 (Quoted <$> quoted <|> Plain <$> plain delim) (char delim)
    arrays = sepBy1 (Array <$> array delim) (char ',')

-- NB: Arrays seem to always be delimited by commas.

-- | Recognizes a quoted string.
quoted :: Parser ByteString
quoted = char '"' *> option "" contents <* char '"'
  where
    esc' = char '\\' *> (char '\\' <|> char '"')
    unQ = takeWhile1 (notInClass "\"\\")
    contents = mconcat <$> many (unQ <|> B.singleton <$> esc')

-- | Recognizes a plain string literal, not containing quotes or brackets and
--   not containing the delimiter character.
plain :: Char -> Parser ByteString
plain delim = takeWhile1 (notInClass (delim : "\"{}"))

-- Mutually recursive 'fmt' and 'delimit' separate out value formatting
-- from the subtleties of delimiting.

-- | Format an array format item, using the delimiter character if the item is
--   itself an array.
fmt :: Char -> ArrayFormat -> ByteString
fmt = fmt' False

-- | Format a list of array format items, inserting the appropriate delimiter
--   between them. When the items are arrays, they will be delimited with
--   commas; otherwise, they are delimited with the passed-in-delimiter.
delimit :: Char -> [ArrayFormat] -> ByteString
delimit _ [] = ""
delimit c [x] = fmt' True c x
delimit c (x : y : z) = (fmt' True c x `B.snoc` c') `mappend` delimit c (y : z)
  where
    c'
      | Array _ <- x = ','
      | otherwise = c

-- | Format an array format item, using the delimiter character if the item is
--   itself an array, optionally applying quoting rules. Creates copies for
--   safety when used in 'FromField' instances.
fmt' :: Bool -> Char -> ArrayFormat -> ByteString
fmt' quoting c x =
  case x of
    Array items -> '{' `B.cons` (delimit c items `B.snoc` '}')
    Plain bytes -> B.copy bytes
    Quoted q
      | quoting -> '"' `B.cons` (esc q `B.snoc` '"')
      | otherwise -> B.copy q

-- NB: The 'snoc' and 'cons' functions always copy.

-- | Escape a string according to Postgres double-quoted string format.
esc :: ByteString -> ByteString
esc = B.concatMap f
  where
    f '"' = "\\\""
    f '\\' = "\\\\"
    f c = B.singleton c

-- TODO: Implement easy performance improvements with unfoldr.