{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE ExtendedLiterals #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE UnboxedTuples #-}

module Data.Packed.TH.Case (caseFName, genCase) where

import Data.Packed.Reader hiding (return)
import Data.Packed.TH.Flag
import Data.Packed.TH.Pattern (patternFName)
import Data.Packed.TH.Utils (getBranchesTyList, getNameAndBangTypesFromCon, resolveAppliedType, sanitizeConName)
import Language.Haskell.TH

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

-- TODO Use generated patterns

-- | Generates a function to allow pattern matching a packed data type using the data constructors
--
--  __Example:__
--
-- For the 'Tree' data type, it generates the following function:
--
-- @
-- caseTree ::
--     ('Data.Packed.PackedReader' '[a] r b) ->
--     ('Data.Packed.PackedReader' '[Tree a, Tree a] r b) ->
--     'Data.Packed.PackedReader' '[Tree a] r b
-- caseTree leafCase nodeCase = 'Data.Packed.Reader.mkPackedReader' $ \case
--      PackedLeaf l -> 'runReaderStep' leafCase l
--      PackedNode l -> 'runReaderStep' nodeCase l
-- @
genCase ::
    [PackingFlag] ->
    -- | The name of the type to generate the function for
    Name ->
    Q [Dec]
genCase :: [PackingFlag] -> Name -> Q [Dec]
genCase [PackingFlag]
flags Name
tyName = do
    (TyConI (DataD _ _ _ _ cs _)) <- Name -> Q Info
reify Name
tyName
    -- For each data constructor, we build names for the pattern for the case functions
    -- Example: leafCase, nodeCase, etc.
    let casePatterns = Con -> (Name, Name)
buildCaseCallbackNameAndPattern (Con -> (Name, Name)) -> [Con] -> [(Name, Name)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Con]
cs
    body <- [|PackedReader $ $(buildCaseExpression casePatterns)|]
    signature <- genCaseSignature flags tyName
    return
        [ PragmaD $ InlineP (caseFName tyName) Inline FunLike AllPhases
        , signature
        , FunD
            (caseFName tyName)
            [Clause (VarP . snd <$> casePatterns) (NormalB body) []]
        ]
  where
    -- for dataconstructor Leaf, will be 'leafCase'
    buildCaseCallbackNameAndPattern :: Con -> (Name, Name)
buildCaseCallbackNameAndPattern = (\Name
name -> (Name -> Name
patternFName Name
name, Name -> Name
conNameToCaseFunctionName Name
name)) (Name -> (Name, Name)) -> (Con -> Name) -> Con -> (Name, 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
    conNameToCaseFunctionName :: Name -> Name
conNameToCaseFunctionName Name
conName = String -> Name
mkName (String -> Name) -> String -> Name
forall a b. (a -> b) -> a -> b
$ Char
'c' Char -> String -> String
forall a. a -> [a] -> [a]
: (Name -> String
sanitizeConName Name
conName) String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
"Case"

    -- Build the case .. of ... expression using the list of available xxxCase, the flag and bytestring
    buildCaseExpression :: [(Name, Name)] -> Q Exp
    buildCaseExpression :: [(Name, Name)] -> Q Exp
buildCaseExpression [(Name, Name)]
casePatterns =
        -- For each xxxCase, we build a branch for the case expression
        let matches :: [Q Match]
matches =
                ( \(Name
patternName, Name
caseCallbackName) -> do
                    let r :: Name
r = String -> Name
mkName String
"r"
                    body <- [|runReaderStep $(Name -> Q Exp
forall (m :: * -> *). Quote m => Name -> m Exp
varE Name
caseCallbackName) $(Name -> Q Exp
forall (m :: * -> *). Quote m => Name -> m Exp
varE Name
r)|]
                    let pat = Name -> Cxt -> [Pat] -> Pat
ConP Name
patternName [] [Name -> Pat
VarP Name
r]
                    return $ Match pat (NormalB body) []
                )
                    ((Name, Name) -> Q Match) -> [(Name, Name)] -> [Q Match]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(Name, Name)]
casePatterns
         in [Q Match] -> Q Exp
forall (m :: * -> *). Quote m => [m Match] -> m Exp
lamCaseE [Q Match]
matches

-- For a type 'Tree', generates the following signature
-- caseTree ::
--     ('Data.Packed.PackedReader' '[a] r b) ->
--     ('Data.Packed.PackedReader' '[Tree a, Tree a] r b) ->
--     'Data.Packed.PackedReader' '[Tree a] r b
genCaseSignature :: [PackingFlag] -> Name -> Q Dec
genCaseSignature :: [PackingFlag] -> Name -> Q Dec
genCaseSignature [PackingFlag]
flags Name
tyName = do
    (sourceType, _) <- Name -> Q (Kind, [Name])
resolveAppliedType Name
tyName
    bVar <- newName "b"
    rVar <- newName "r"
    branchesTypes <- getBranchesTyList tyName flags
    let
        bType = Name -> Q Kind
forall (m :: * -> *). Quote m => Name -> m Kind
varT Name
bVar
        rType = Name -> Q Kind
forall (m :: * -> *). Quote m => Name -> m Kind
varT Name
rVar
        lambdaTypes = (\Cxt
branchTypes -> Cxt -> Q Kind -> Q Kind -> Q Kind
buildLambdaType Cxt
branchTypes Q Kind
bType Q Kind
rType) (Cxt -> Q Kind) -> [Cxt] -> [Q Kind]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Cxt]
branchesTypes
        outType = [t|PackedReader '[$(Kind -> Q Kind
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return Kind
sourceType)] $Q Kind
rType $Q Kind
bType|]
    signature <- foldr (\Q Kind
lambda Q Kind
out -> [t|$Q Kind
lambda -> $Q Kind
out|]) outType lambdaTypes
    return $ SigD (caseFName tyName) signature
  where
    -- From a constructor (say Leaf a), build type PackedReader '[a] r b
    buildLambdaType :: [Type] -> Q Type -> Q Type -> Q Type
    buildLambdaType :: Cxt -> Q Kind -> Q Kind -> Q Kind
buildLambdaType Cxt
branchType Q Kind
returnType Q Kind
restType = do
        let branchTypeList :: Q Kind
branchTypeList = (Kind -> Q Kind -> Q Kind) -> Q Kind -> Cxt -> Q Kind
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\Kind
a Q Kind
rest -> [t|$(Kind -> Q Kind
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return Kind
a) ': $Q Kind
rest|]) [t|'[]|] Cxt
branchType
        [t|PackedReader $Q Kind
branchTypeList $Q Kind
restType $Q Kind
returnType|]