{-# LANGUAGE MagicHash #-}
{-# LANGUAGE UnboxedTuples #-}

module Data.Packed.TH.Pattern (genPatterns, patternFName, tagReaderFName) where

import Control.Monad (forM)
import Data.Packed.Reader (PackedFragment (..))
import Data.Packed.TH.Flag (PackingFlag (..))
import Data.Packed.TH.Utils
import GHC.Exts
import GHC.Ptr
import Language.Haskell.TH

-- | Generates patterns for packed values of the given type
--
--  __Example:__
--
-- For the 'Tree' data type, it generates the following patterns:
--
-- @
-- PackedLeaf :: (PackedFragment (a ': r)) -> PackedFragment (Tree a ': r)
--
-- PackedNode :: (PackedFragment (Tree A ': Tree a ': r)) -> PackedFragment (Tree a ': r)
-- @
--
-- More specifically, it generates the following code:
--
-- @
--
-- {-# INLINE readTreeTag #-}
-- readTreeTag :: PackedFragment (Tree a ': r) -> (# Word8#, PackedFragment n #)
-- readTreeTag (PF ptr@(Ptr addr) i) = case readWord8OffAddr# addr 0# realWorld# of
--     (# _, t #) -> case t of
--         0 -> (# t, PF (ptr `plusPtr` 1) (i - 1) #)
--         1 -> (# t, PF (ptr `plusPtr` 1) (i - 1) #)
--         _ -> error $ "Bag tag: Got " ++ show (W8# t)
--
-- {-# INLINE PackedLeaf #-}
-- pattern PackedLeaf :: PackedFragment (a ': r) -> 'PackedFragment' (Tree a ': r)
-- pattern PackedLeaf pf <- (readTreeTag -> (# 0, pf #))
--
-- {-# INLINE PackedNode #-}
-- pattern PackedNode :: 'PackedFragment' (Tree a ': Tree a ': r) -> 'PackedFragment' (Tree a ': r)
-- pattern PackedNode pf <- (readTreeTag -> (# 1, pf #))
--
-- {-# COMPLETE PackedLeaf, PackedNode #-}
-- @
genPatterns :: [PackingFlag] -> Name -> Q [Dec]
genPatterns :: [PackingFlag] -> Name -> Q [Dec]
genPatterns [PackingFlag]
flags Name
tyName = do
    (TyConI (DataD _ _ _ _ cs _)) <- Name -> Q Info
reify Name
tyName
    patterns <- concat <$> forM (zip cs [0 ..]) (uncurry (genPattern flags tyName))
    tagReader <- genTagReader flags tyName
    let completePragma =
            let
                patternNames :: [Name]
patternNames = Name -> Name
patternFName (Name -> Name) -> (Con -> Name) -> Con -> Name
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, [BangType]) -> Name
forall a b. (a, b) -> a
fst ((Name, [BangType]) -> Name)
-> (Con -> (Name, [BangType])) -> Con -> Name
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Con -> (Name, [BangType])
getNameAndBangTypesFromCon (Con -> Name) -> [Con] -> [Name]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Con]
cs
             in
                [Name] -> Maybe Name -> Pragma
CompleteP [Name]
patternNames Maybe Name
forall a. Maybe a
Nothing
    return $ tagReader ++ patterns ++ [PragmaD completePragma]

genPattern :: [PackingFlag] -> Name -> Con -> Integer -> Q [Dec]
genPattern :: [PackingFlag] -> Name -> Con -> Integer -> Q [Dec]
genPattern [PackingFlag]
flags Name
tyName Con
con Integer
conIdx = do
    let patternName :: Name
patternName = Name -> Name
patternFName (Name -> Name) -> Name -> Name
forall a b. (a -> b) -> a -> b
$ (Name, [BangType]) -> Name
forall a b. (a, b) -> a
fst ((Name, [BangType]) -> Name) -> (Name, [BangType]) -> Name
forall a b. (a -> b) -> a -> b
$ Con -> (Name, [BangType])
getNameAndBangTypesFromCon Con
con
    pfName <- String -> Q Name
forall (m :: * -> *). Quote m => String -> m Name
newName String
"pf"
    (ty, _) <- resolveAppliedType tyName
    rVar <- newName "r"
    conArgTypes <- getBranchTyList con flags
    patternType <- do
        patternFieldType <- foldr (\Kind
curr Q Kind
rest -> [t|$(Kind -> Q Kind
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return Kind
curr) ': $Q Kind
rest|]) (varT rVar) conArgTypes
        fromType <- [t|$(return ty) ': $(varT rVar)|]
        [t|PackedFragment ($(return patternFieldType)) -> PackedFragment ($(return fromType))|]
    let patSynSig =
            Name -> Kind -> Dec
PatSynSigD
                Name
patternName
                Kind
patternType
    patSyn <- do
        let viewPat = [p|(# $(Lit -> Q Pat
forall (m :: * -> *). Quote m => Lit -> m Pat
litP (Lit -> Q Pat) -> Lit -> Q Pat
forall a b. (a -> b) -> a -> b
$ Integer -> Lit
WordPrimL Integer
conIdx), $(Name -> Q Pat
forall (m :: * -> *). Quote m => Name -> m Pat
varP Name
pfName) #)|]
        patSynD
            patternName
            (prefixPatSyn [pfName])
            unidir
            (viewP (varE $ tagReaderFName tyName) viewPat)

    return
        [ PragmaD (InlineP patternName Inline FunLike AllPhases)
        , patSynSig
        , patSyn
        ]

patternFName :: Name -> Name
patternFName :: Name -> Name
patternFName Name
name = String -> Name
mkName (String -> Name) -> String -> Name
forall a b. (a -> b) -> a -> b
$ String
"Packed" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Name -> String
sanitizeConName Name
name

genTagReader :: [PackingFlag] -> Name -> Q [Dec]
genTagReader :: [PackingFlag] -> Name -> Q [Dec]
genTagReader [PackingFlag]
flags Name
tyName = do
    rType <- String -> Q Name
forall (m :: * -> *). Quote m => String -> m Name
newName String
"r"
    (ty, _) <- resolveAppliedType tyName
    conTypes <- getBranchesTyList tyName flags
    let fName = Name -> Name
tagReaderFName Name
tyName
    prototype <-
        let
            inType = [t|PackedFragment ($(Kind -> Q Kind
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return Kind
ty) ': $(Name -> Q Kind
forall (m :: * -> *). Quote m => Name -> m Kind
varT Name
rType))|]
            outType = [t|(# Word#, PackedFragment $(Name -> Q Kind
forall (m :: * -> *). Quote m => Name -> m Kind
varT (Name -> Q Kind) -> Name -> Q Kind
forall a b. (a -> b) -> a -> b
$ String -> Name
mkName String
"n") #)|]
         in
            [t|$inType -> $outType|]
    (pat, body) <- do
        let addrName = String -> Name
mkName String
"addr"
            ptrName = String -> Name
mkName String
"ptr"
            lengthName = String -> Name
mkName String
"l"
            tagVarName = String -> Name
mkName String
"t"
        readWordExp <- [|readWord8OffAddr# $(varE addrName) 0# realWorld#|]
        caseBranches <- forM (zip [0 ..] conTypes) $ \(Integer
conIdx, Cxt
_) -> do
            tName <- String -> Q Name
forall (m :: * -> *). Quote m => String -> m Name
newName String
"t"
            let pat = Name -> Pat -> Pat
AsP Name
tName (Lit -> Pat
LitP (Lit -> Pat) -> Lit -> Pat
forall a b. (a -> b) -> a -> b
$ Integer -> Lit
WordPrimL Integer
conIdx)
            branchExp <- [|(# $(varE tName), PF ($(varE ptrName) `plusPtr` 1) ($(varE lengthName) - 1) #)|]
            return $ Match pat (NormalB branchExp) []
        errorBranch <- do
            tName <- newName "t"
            match (asP tName wildP) (normalB [|error $ "Bad tag, got " ++ show (W# $(varE tName))|]) []
        let caseExpr = Q Exp -> [Q Match] -> Q Exp
forall (m :: * -> *). Quote m => m Exp -> [m Match] -> m Exp
caseE [|word8ToWord# $(Name -> Q Exp
forall (m :: * -> *). Quote m => Name -> m Exp
varE Name
tagVarName)|] (Match -> Q Match
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return (Match -> Q Match) -> [Match] -> [Q Match]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ([Match]
caseBranches [Match] -> [Match] -> [Match]
forall a. [a] -> [a] -> [a]
++ [Match
errorBranch]))
        body <- [|case $(return readWordExp) of (# _, $(varP tagVarName) #) -> $caseExpr|]
        pattern_ <- do
            ptrPattern <- asP ptrName [p|Ptr $(varP addrName)|]
            [p|(PF $(return ptrPattern) $(varP lengthName))|]
        return (pattern_, body)

    return
        [ PragmaD (InlineP fName Inline FunLike AllPhases)
        , SigD fName prototype
        , FunD fName [Clause [pat] (NormalB body) []]
        ]

tagReaderFName :: Name -> Name
tagReaderFName :: Name -> Name
tagReaderFName Name
tyName = String -> Name
mkName (String -> Name) -> String -> Name
forall a b. (a -> b) -> a -> b
$ String
"read" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Name -> String
sanitizeConName Name
tyName String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"Tag"