From 9572af05ad3342596a715919b757e96a4d435674 Mon Sep 17 00:00:00 2001
From: Nikita Volkov <nikita.y.volkov@mail.ru>
Date: Sat, 27 Dec 2014 07:18:04 +0300
Subject: [PATCH] Clean up whitespace during the quasi-quote parsing

---
 hspec/Main.hs            | 20 +++++++++++++-
 library/Hasql.hs         |  5 ++--
 library/Hasql/QParser.hs | 56 ++++++++++++++++++++++++----------------
 3 files changed, 56 insertions(+), 25 deletions(-)

diff --git a/hspec/Main.hs b/hspec/Main.hs
index 64fefc2..c6cb7f3 100644
--- a/hspec/Main.hs
+++ b/hspec/Main.hs
@@ -22,5 +22,23 @@ main =
     context "Quasi quoter" $ do
       it "generates a proper statement" $ do
         (flip shouldBe)
-          (HB.Stmt " SELECT (? + ?) " (V.fromList [HB.encodeValue 'a', HB.encodeValue 'b']) True)
+          (HB.Stmt "SELECT (? + ?)" (V.fromList [HB.encodeValue 'a', HB.encodeValue 'b']) True)
           ([H.stmt| SELECT (? + ?) |] 'a' 'b' :: HB.Stmt X)
+      it "does not drop quotes" $ do
+        let 
+          HB.Stmt t _ _ =
+            [H.stmt| SELECT "a", 'b' |]
+        (flip shouldBe)
+          "SELECT \"a\", 'b'"
+          t
+      it "cleans whitespace" $ do
+        let 
+          HB.Stmt t _ _ =
+            [H.stmt| CREATE TABLE data (
+                       field1    DECIMAL NOT NULL,
+                       field2    BIGINT  NOT NULL,
+                       PRIMARY KEY (field1)
+                     ) |]
+        (flip shouldBe)
+          "CREATE TABLE data ( field1 DECIMAL NOT NULL, field2 BIGINT NOT NULL, PRIMARY KEY (field1) )"
+          t
diff --git a/library/Hasql.hs b/library/Hasql.hs
index 386c9c0..8e84596 100644
--- a/library/Hasql.hs
+++ b/library/Hasql.hs
@@ -60,6 +60,7 @@ import qualified Hasql.RowParser as RowParser
 import qualified Hasql.QParser as QParser
 import qualified ListT
 import qualified Data.Pool as Pool
+import qualified Data.Text as Text
 import qualified Data.Vector as Vector
 import qualified Data.Vector.Mutable as MVector
 import qualified Language.Haskell.TH as TH
@@ -337,8 +338,8 @@ stmt =
   where
     parseExp s =
       do
-        n <- either (fail . showString "Parsing failure: ") return (QParser.parse (fromString s))
-        return $ statementF s (fromIntegral n)
+        (t, n) <- either (fail . showString "Parsing failure: ") return (QParser.parse (fromString s))
+        return $ statementF (Text.unpack t) (fromIntegral n)
     statementF s n =
       TH.LamE
         (map TH.VarP argNames)
diff --git a/library/Hasql/QParser.hs b/library/Hasql/QParser.hs
index a0cd3e4..3dd3fc8 100644
--- a/library/Hasql/QParser.hs
+++ b/library/Hasql/QParser.hs
@@ -2,37 +2,49 @@ module Hasql.QParser where
 
 import Hasql.Prelude hiding (takeWhile)
 import Data.Attoparsec.Text hiding (Result)
-import qualified Data.Text as Text
+import qualified Data.Text as T
+import qualified Data.Text.Lazy as TL
+import qualified Data.Text.Lazy.Builder as TLB
 
 
 -- |
--- The amount of placeholders.
-type Result =
-  (Word)
-
-parse :: Text -> Either String Result
+-- Produces a whitespace-cleaned text and a count of placeholders in it.
+parse :: Text -> Either String (Text, Int)
 parse = 
-  parseOnly $ flip execStateT 0 $
-    statement *>
-    (lift endOfInput <|> 
-     (void (lift (char ';')) <* fail "A semicolon detected. Only single statements are allowed"))
-  where
-    statement =
-      skipMany1 $ 
-        void (lift stringLit) <|> 
-        void (lift (char '?') <* modify succ) <|>
-        void (lift (notChar ';'))
+  parseOnly $ singleTemplate
+
+singleTemplate :: Parser (Text, Int)
+singleTemplate =
+  template <* 
+  ((endOfInput) <|>
+   (() <$ skipSpace <* char ';' <* fail "A semicolon detected. Only single statements are allowed"))
+
+template :: Parser (Text, Int)
+template =
+  flip runStateT 0 $ do
+    lift $ skipSpace
+    fmap (TL.toStrict . TLB.toLazyText . mconcat) $ 
+      many $ 
+        (mempty <$ lift trailingWhitespace) <|>
+        (TLB.singleton ' ' <$ lift (takeWhile1 isSpace)) <|>
+        (TLB.fromText <$> lift stringLit) <|>
+        (TLB.singleton <$> lift (char '?') <* modify succ) <|>
+        (TLB.singleton <$> lift (notChar ';'))
+
+trailingWhitespace :: Parser ()
+trailingWhitespace =
+  () <$ takeWhile1 isSpace <* endOfInput
 
 stringLit :: Parser Text
 stringLit =
   do
     quote <- 
       char '"' <|> char '\''
-    text <- 
+    content <- 
       fmap mconcat $ many $ 
-        string "\\\\" <|> 
-        string (fromString ['\\', quote]) <|> 
-        (Text.singleton <$> notChar quote)
+        TLB.fromText <$> string "\\\\" <|> 
+        TLB.fromText <$> string (fromString ['\\', quote]) <|> 
+        TLB.singleton <$> notChar quote
     char quote
-    return text
-
+    return $ TL.toStrict . TLB.toLazyText $
+      TLB.singleton quote <> content <> TLB.singleton quote