{-# LANGUAGE PatternGuards #-}

------------------------------------------------------------------------------
-- |
-- 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 (Applicative(..), Alternative(..), (<$>))
import           Data.ByteString.Char8 (ByteString)
import qualified Data.ByteString.Char8 as B
import           Data.Monoid
import           Data.Attoparsec.ByteString.Char8


-- | Parse one of three primitive field formats: array, quoted and plain.
arrayFormat :: Char -> Parser ArrayFormat
arrayFormat :: Char -> Parser ArrayFormat
arrayFormat Char
delim  =  [ArrayFormat] -> ArrayFormat
Array  ([ArrayFormat] -> ArrayFormat)
-> Parser ByteString [ArrayFormat] -> Parser ArrayFormat
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Char -> Parser ByteString [ArrayFormat]
array Char
delim
                  Parser ArrayFormat -> Parser ArrayFormat -> Parser ArrayFormat
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> ByteString -> ArrayFormat
Plain  (ByteString -> ArrayFormat)
-> Parser ByteString ByteString -> Parser ArrayFormat
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Char -> Parser ByteString ByteString
plain Char
delim
                  Parser ArrayFormat -> Parser ArrayFormat -> Parser ArrayFormat
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> ByteString -> ArrayFormat
Quoted (ByteString -> ArrayFormat)
-> Parser ByteString ByteString -> Parser ArrayFormat
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Parser ByteString ByteString
quoted

data ArrayFormat = Array [ArrayFormat]
                 | Plain ByteString
                 | Quoted ByteString
                   deriving (ArrayFormat -> ArrayFormat -> Bool
(ArrayFormat -> ArrayFormat -> Bool)
-> (ArrayFormat -> ArrayFormat -> Bool) -> Eq ArrayFormat
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: ArrayFormat -> ArrayFormat -> Bool
$c/= :: ArrayFormat -> ArrayFormat -> Bool
== :: ArrayFormat -> ArrayFormat -> Bool
$c== :: ArrayFormat -> ArrayFormat -> Bool
Eq, Int -> ArrayFormat -> ShowS
[ArrayFormat] -> ShowS
ArrayFormat -> String
(Int -> ArrayFormat -> ShowS)
-> (ArrayFormat -> String)
-> ([ArrayFormat] -> ShowS)
-> Show ArrayFormat
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [ArrayFormat] -> ShowS
$cshowList :: [ArrayFormat] -> ShowS
show :: ArrayFormat -> String
$cshow :: ArrayFormat -> String
showsPrec :: Int -> ArrayFormat -> ShowS
$cshowsPrec :: Int -> ArrayFormat -> ShowS
Show, Eq ArrayFormat
Eq ArrayFormat
-> (ArrayFormat -> ArrayFormat -> Ordering)
-> (ArrayFormat -> ArrayFormat -> Bool)
-> (ArrayFormat -> ArrayFormat -> Bool)
-> (ArrayFormat -> ArrayFormat -> Bool)
-> (ArrayFormat -> ArrayFormat -> Bool)
-> (ArrayFormat -> ArrayFormat -> ArrayFormat)
-> (ArrayFormat -> ArrayFormat -> ArrayFormat)
-> Ord ArrayFormat
ArrayFormat -> ArrayFormat -> Bool
ArrayFormat -> ArrayFormat -> Ordering
ArrayFormat -> ArrayFormat -> ArrayFormat
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: ArrayFormat -> ArrayFormat -> ArrayFormat
$cmin :: ArrayFormat -> ArrayFormat -> ArrayFormat
max :: ArrayFormat -> ArrayFormat -> ArrayFormat
$cmax :: ArrayFormat -> ArrayFormat -> ArrayFormat
>= :: ArrayFormat -> ArrayFormat -> Bool
$c>= :: ArrayFormat -> ArrayFormat -> Bool
> :: ArrayFormat -> ArrayFormat -> Bool
$c> :: ArrayFormat -> ArrayFormat -> Bool
<= :: ArrayFormat -> ArrayFormat -> Bool
$c<= :: ArrayFormat -> ArrayFormat -> Bool
< :: ArrayFormat -> ArrayFormat -> Bool
$c< :: ArrayFormat -> ArrayFormat -> Bool
compare :: ArrayFormat -> ArrayFormat -> Ordering
$ccompare :: ArrayFormat -> ArrayFormat -> Ordering
$cp1Ord :: Eq ArrayFormat
Ord)

array :: Char -> Parser [ArrayFormat]
array :: Char -> Parser ByteString [ArrayFormat]
array Char
delim = Char -> Parser Char
char Char
'{' Parser Char
-> Parser ByteString [ArrayFormat]
-> Parser ByteString [ArrayFormat]
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> [ArrayFormat]
-> Parser ByteString [ArrayFormat]
-> Parser ByteString [ArrayFormat]
forall (f :: * -> *) a. Alternative f => a -> f a -> f a
option [] (Parser ByteString [ArrayFormat]
arrays Parser ByteString [ArrayFormat]
-> Parser ByteString [ArrayFormat]
-> Parser ByteString [ArrayFormat]
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> Parser ByteString [ArrayFormat]
strings) Parser ByteString [ArrayFormat]
-> Parser Char -> Parser ByteString [ArrayFormat]
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f a
<* Char -> Parser Char
char Char
'}'
  where
    strings :: Parser ByteString [ArrayFormat]
strings = Parser ArrayFormat
-> Parser Char -> Parser ByteString [ArrayFormat]
forall (f :: * -> *) a s. Alternative f => f a -> f s -> f [a]
sepBy1 (ByteString -> ArrayFormat
Quoted (ByteString -> ArrayFormat)
-> Parser ByteString ByteString -> Parser ArrayFormat
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Parser ByteString ByteString
quoted Parser ArrayFormat -> Parser ArrayFormat -> Parser ArrayFormat
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> ByteString -> ArrayFormat
Plain (ByteString -> ArrayFormat)
-> Parser ByteString ByteString -> Parser ArrayFormat
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Char -> Parser ByteString ByteString
plain Char
delim) (Char -> Parser Char
char Char
delim)
    arrays :: Parser ByteString [ArrayFormat]
arrays  = Parser ArrayFormat
-> Parser Char -> Parser ByteString [ArrayFormat]
forall (f :: * -> *) a s. Alternative f => f a -> f s -> f [a]
sepBy1 ([ArrayFormat] -> ArrayFormat
Array ([ArrayFormat] -> ArrayFormat)
-> Parser ByteString [ArrayFormat] -> Parser ArrayFormat
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Char -> Parser ByteString [ArrayFormat]
array Char
delim) (Char -> Parser Char
char Char
',')
    -- NB: Arrays seem to always be delimited by commas.

-- | Recognizes a quoted string.
quoted :: Parser ByteString
quoted :: Parser ByteString ByteString
quoted  = Char -> Parser Char
char Char
'"' Parser Char
-> Parser ByteString ByteString -> Parser ByteString ByteString
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> ByteString
-> Parser ByteString ByteString -> Parser ByteString ByteString
forall (f :: * -> *) a. Alternative f => a -> f a -> f a
option ByteString
"" Parser ByteString ByteString
contents Parser ByteString ByteString
-> Parser Char -> Parser ByteString ByteString
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f a
<* Char -> Parser Char
char Char
'"'
  where
    esc' :: Parser Char
esc' = Char -> Parser Char
char Char
'\\' Parser Char -> Parser Char -> Parser Char
forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> (Char -> Parser Char
char Char
'\\' Parser Char -> Parser Char -> Parser Char
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> Char -> Parser Char
char Char
'"')
    unQ :: Parser ByteString ByteString
unQ = (Char -> Bool) -> Parser ByteString ByteString
takeWhile1 (String -> Char -> Bool
notInClass String
"\"\\")
    contents :: Parser ByteString ByteString
contents = [ByteString] -> ByteString
forall a. Monoid a => [a] -> a
mconcat ([ByteString] -> ByteString)
-> Parser ByteString [ByteString] -> Parser ByteString ByteString
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Parser ByteString ByteString -> Parser ByteString [ByteString]
forall (f :: * -> *) a. Alternative f => f a -> f [a]
many (Parser ByteString ByteString
unQ Parser ByteString ByteString
-> Parser ByteString ByteString -> Parser ByteString ByteString
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> Char -> ByteString
B.singleton (Char -> ByteString) -> Parser Char -> Parser ByteString ByteString
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Parser Char
esc')

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

-- 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 :: Char -> ArrayFormat -> ByteString
fmt = Bool -> Char -> ArrayFormat -> ByteString
fmt' Bool
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 :: Char -> [ArrayFormat] -> ByteString
delimit Char
_      [] = ByteString
""
delimit Char
c     [ArrayFormat
x] = Bool -> Char -> ArrayFormat -> ByteString
fmt' Bool
True Char
c ArrayFormat
x
delimit Char
c (ArrayFormat
x:ArrayFormat
y:[ArrayFormat]
z) = (Bool -> Char -> ArrayFormat -> ByteString
fmt' Bool
True Char
c ArrayFormat
x ByteString -> Char -> ByteString
`B.snoc` Char
c') ByteString -> ByteString -> ByteString
forall a. Monoid a => a -> a -> a
`mappend` Char -> [ArrayFormat] -> ByteString
delimit Char
c (ArrayFormat
yArrayFormat -> [ArrayFormat] -> [ArrayFormat]
forall a. a -> [a] -> [a]
:[ArrayFormat]
z)
  where
    c' :: Char
c' | Array [ArrayFormat]
_ <- ArrayFormat
x = Char
','
       | Bool
otherwise    = Char
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' :: Bool -> Char -> ArrayFormat -> ByteString
fmt' Bool
quoting Char
c ArrayFormat
x =
  case ArrayFormat
x of
    Array [ArrayFormat]
items          -> Char
'{' Char -> ByteString -> ByteString
`B.cons` (Char -> [ArrayFormat] -> ByteString
delimit Char
c [ArrayFormat]
items ByteString -> Char -> ByteString
`B.snoc` Char
'}')
    Plain ByteString
bytes          -> ByteString -> ByteString
B.copy ByteString
bytes
    Quoted ByteString
q | Bool
quoting   -> Char
'"' Char -> ByteString -> ByteString
`B.cons` (ByteString -> ByteString
esc ByteString
q ByteString -> Char -> ByteString
`B.snoc` Char
'"')
             | Bool
otherwise -> ByteString -> ByteString
B.copy ByteString
q
    -- NB: The 'snoc' and 'cons' functions always copy.

-- | Escape a string according to Postgres double-quoted string format.
esc :: ByteString -> ByteString
esc :: ByteString -> ByteString
esc = (Char -> ByteString) -> ByteString -> ByteString
B.concatMap Char -> ByteString
f
  where
    f :: Char -> ByteString
f Char
'"'  = ByteString
"\\\""
    f Char
'\\' = ByteString
"\\\\"
    f Char
c    = Char -> ByteString
B.singleton Char
c
  -- TODO: Implement easy performance improvements with unfoldr.