{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE QuasiQuotes #-}
{-# LANGUAGE TemplateHaskell #-}

-- | Quasiquotation for byte builders.
module Data.Bytes.Builder.Template
  ( bldr
  ) where

import Control.Monad (when)
import Data.Bytes.Builder.Class (toBuilder)
import GHC.Ptr (Ptr(Ptr))
import Language.Haskell.Meta.Parse (parseExp)
import Language.Haskell.TH (Q,Exp)
import Language.Haskell.TH.Lib (integerL,stringPrimL,litE)
import Language.Haskell.TH.Quote (QuasiQuoter(..))

import qualified Data.Bytes.Builder as Builder
import qualified Data.ByteString.Short as SBS
import qualified Data.Text.Short as TS
import qualified Language.Haskell.TH as TH

-- | A quasiquoter for builders. Haskell expressions are interpolated
-- with backticks, and the @ToBuilder@ class is used to convert them
-- to builders. Several common escape sequences for whitespace and
-- control characters are recongized. Consider the following expression,
-- where the binding @partition@ has type @Word32@:
--
-- > [templ|[WARN] Partition `partition` has invalid data.\n|]
--
-- This expression has type @Builder@ and expands to:
--
-- > Builder.cstringLen (Ptr "[WARN] Partition "#, 17) <>
-- > Builder.toBuilder partition <>
-- > Builder.cstringLen (Ptr " has invalid data.\n"#, 19)
--
-- The @ToBuilder@ instance for @Word32@ uses decimal encoding, so this
-- would result in the following if @partition@ was 42 (with a newline
-- character at the end):
--
-- > [WARN] Partition 42 has invalid data.
--
-- In the future, a more sophisticated @bbldr@ variant will be added
-- that will support expressions where the maximum length of the entire
-- builder can be computed at compile time. 
bldr :: QuasiQuoter
bldr :: QuasiQuoter
bldr = QuasiQuoter :: (String -> Q Exp)
-> (String -> Q Pat)
-> (String -> Q Type)
-> (String -> Q [Dec])
-> QuasiQuoter
QuasiQuoter
  { quoteExp :: String -> Q Exp
quoteExp = String -> Q Exp
templExp
  , quotePat :: String -> Q Pat
quotePat = String -> String -> Q Pat
forall (m :: * -> *) p a. MonadFail m => String -> p -> m a
notHandled String
"patterns"
  , quoteType :: String -> Q Type
quoteType = String -> String -> Q Type
forall (m :: * -> *) p a. MonadFail m => String -> p -> m a
notHandled String
"types"
  , quoteDec :: String -> Q [Dec]
quoteDec = String -> String -> Q [Dec]
forall (m :: * -> *) p a. MonadFail m => String -> p -> m a
notHandled String
"declarations"
  }
  where
  notHandled :: String -> p -> m a
notHandled String
things p
_ = String -> m a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> m a) -> String -> m a
forall a b. (a -> b) -> a -> b
$
    String
things String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"are not handled by the byte template quasiquoter"

templExp :: String -> Q Exp
templExp :: String -> Q Exp
templExp String
inp = do
  Q ()
checkOverloadedStrings
  Template
rawParts <- case String -> Either String Template
parse String
inp of
    Left String
err -> String -> Q Template
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
err
    Right [] -> String -> Q Template
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"empty template"
    Right Template
v -> Template -> Q Template
forall (f :: * -> *) a. Applicative f => a -> f a
pure Template
v
  let expParts :: [Q Exp]
expParts = TemplPart -> Q Exp
compile (TemplPart -> Q Exp) -> Template -> [Q Exp]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Template
rawParts
  (Q Exp -> Q Exp -> Q Exp) -> [Q Exp] -> Q Exp
forall (t :: * -> *) a. Foldable t => (a -> a -> a) -> t a -> a
foldl1 (\Q Exp
e1 Q Exp
e2 -> [| $e1 <> $e2 |]) [Q Exp]
expParts

checkOverloadedStrings :: Q ()
checkOverloadedStrings :: Q ()
checkOverloadedStrings = do
  Bool
olEnabled <- Extension -> Q Bool
TH.isExtEnabled Extension
TH.OverloadedStrings
  Bool -> Q () -> Q ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Bool -> Bool
not Bool
olEnabled) (Q () -> Q ()) -> Q () -> Q ()
forall a b. (a -> b) -> a -> b
$
    String -> Q ()
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"Byte templates require the OverloadedStrings extension enabled."

type Template = [TemplPart]
data TemplPart
  = Literal String
  | Splice String

compile :: TemplPart -> Q Exp
compile :: TemplPart -> Q Exp
compile (Literal String
lit) =
  let bytes :: [Word8]
bytes = ShortByteString -> [Word8]
SBS.unpack (ShortByteString -> [Word8])
-> (String -> ShortByteString) -> String -> [Word8]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShortText -> ShortByteString
TS.toShortByteString (ShortText -> ShortByteString)
-> (String -> ShortText) -> String -> ShortByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> ShortText
TS.pack (String -> [Word8]) -> String -> [Word8]
forall a b. (a -> b) -> a -> b
$ String
lit
      strExp :: Q Exp
strExp = Lit -> Q Exp
litE (Lit -> Q Exp) -> ([Word8] -> Lit) -> [Word8] -> Q Exp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Word8] -> Lit
stringPrimL ([Word8] -> Q Exp) -> [Word8] -> Q Exp
forall a b. (a -> b) -> a -> b
$ [Word8]
bytes
      strLen :: Q Exp
strLen = Lit -> Q Exp
litE (Lit -> Q Exp) -> (Int -> Lit) -> Int -> Q Exp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Integer -> Lit
integerL (Integer -> Lit) -> (Int -> Integer) -> Int -> Lit
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> Q Exp) -> Int -> Q Exp
forall a b. (a -> b) -> a -> b
$ [Word8] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Word8]
bytes
   in [|Builder.cstringLen (Ptr $(strExp), $(strLen))|]
compile (Splice String
str) = case String -> Either String Exp
parseExp String
str of
  Left String
err -> String -> Q Exp
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
err
  Right Exp
hs -> [|toBuilder $(pure hs)|]

parse :: String -> Either String Template
parse :: String -> Either String Template
parse = String -> Either String Template
partsLoop
  where
  partsLoop :: String -> Either String Template
partsLoop String
"" = do
    Template -> Either String Template
forall (f :: * -> *) a. Applicative f => a -> f a
pure []
  partsLoop (Char
'`':String
inp) = do
    (!String
spl, !String
rest) <- String -> Either String (String, String)
spliceLoop String
inp
    (String -> TemplPart
Splice String
splTemplPart -> Template -> Template
forall a. a -> [a] -> [a]
:) (Template -> Template)
-> Either String Template -> Either String Template
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> Either String Template
partsLoop String
rest
  partsLoop String
inp = do
    (!String
lit, !String
rest) <- String -> String -> Either String (String, String)
litLoop String
"" String
inp
    (String -> TemplPart
Literal String
litTemplPart -> Template -> Template
forall a. a -> [a] -> [a]
:) (Template -> Template)
-> Either String Template -> Either String Template
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> Either String Template
partsLoop String
rest
  litLoop :: String -> String -> Either String (String, String)
  litLoop :: String -> String -> Either String (String, String)
litLoop !String
acc rest :: String
rest@String
"" = (String, String) -> Either String (String, String)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (String -> String
forall a. [a] -> [a]
reverse String
acc, String
rest)
  litLoop !String
acc rest :: String
rest@(Char
'`':String
_) = (String, String) -> Either String (String, String)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (String -> String
forall a. [a] -> [a]
reverse String
acc, String
rest)
  litLoop !String
acc (Char
'\\':String
next) = do
    (Char
c, String
rest) <- String -> Either String (Char, String)
parseEscape String
next
    String -> String -> Either String (String, String)
litLoop (Char
cChar -> String -> String
forall a. a -> [a] -> [a]
:String
acc) String
rest
  litLoop !String
acc (Char
c:String
rest) = String -> String -> Either String (String, String)
litLoop (Char
cChar -> String -> String
forall a. a -> [a] -> [a]
:String
acc) String
rest
  spliceLoop :: String -> Either String (String, String)
  spliceLoop :: String -> Either String (String, String)
spliceLoop String
inp = case (Char -> Bool) -> String -> (String, String)
forall a. (a -> Bool) -> [a] -> ([a], [a])
break (Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
'`') String
inp of
    ([], String
_) -> String -> Either String (String, String)
forall a b. a -> Either a b
Left String
"internal error"
    (String
hs, Char
'`':String
rest) -> (String, String) -> Either String (String, String)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (String
hs, String
rest)
    (String
_, Char
_:String
_) -> String -> Either String (String, String)
forall a b. a -> Either a b
Left String
"internal error"
    (String
_, []) -> String -> Either String (String, String)
forall a b. a -> Either a b
Left String
"unterminated interpolation"
  parseEscape :: String -> Either String (Char, String)
  parseEscape :: String -> Either String (Char, String)
parseEscape String
"" = String -> Either String (Char, String)
forall a b. a -> Either a b
Left String
"incomplete escape"
  parseEscape (Char
'\\':String
rest) = (Char, String) -> Either String (Char, String)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Char
'\\', String
rest)
  parseEscape (Char
'`':String
rest) = (Char, String) -> Either String (Char, String)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Char
'`', String
rest)
  parseEscape (Char
'\'':String
rest) = (Char, String) -> Either String (Char, String)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Char
'\'', String
rest)
  parseEscape (Char
'\"':String
rest) = (Char, String) -> Either String (Char, String)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Char
'\"', String
rest)
  parseEscape (Char
'0':String
rest) = (Char, String) -> Either String (Char, String)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Char
'\0', String
rest)
  parseEscape (Char
'a':String
rest) = (Char, String) -> Either String (Char, String)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Char
'\a', String
rest)
  parseEscape (Char
'b':String
rest) = (Char, String) -> Either String (Char, String)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Char
'\b', String
rest)
  parseEscape (Char
'f':String
rest) = (Char, String) -> Either String (Char, String)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Char
'\f', String
rest)
  parseEscape (Char
'n':String
rest) = (Char, String) -> Either String (Char, String)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Char
'\n', String
rest)
  parseEscape (Char
'r':String
rest) = (Char, String) -> Either String (Char, String)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Char
'\r', String
rest)
  parseEscape (Char
't':String
rest) = (Char, String) -> Either String (Char, String)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Char
'\t', String
rest)
  parseEscape (Char
'v':String
rest) = (Char, String) -> Either String (Char, String)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Char
'\v', String
rest)
  parseEscape (Char
c:String
_) = String -> Either String (Char, String)
forall a b. a -> Either a b
Left (String -> Either String (Char, String))
-> String -> Either String (Char, String)
forall a b. (a -> b) -> a -> b
$ String
"unrecognized escape: \\" String -> String -> String
forall a. [a] -> [a] -> [a]
++ [Char
c]