{-# LANGUAGE QualifiedDo #-}

module Data.Packed.TH.Read (readFName, genRead) where

import Data.Packed.Reader hiding (return)
import qualified Data.Packed.Reader as R
import Data.Packed.TH.Case (caseFName)
import Data.Packed.TH.Flag (PackingFlag (..))
import Data.Packed.TH.Utils
import Data.Packed.Unpackable
import Language.Haskell.TH

readFName :: Name -> Name
readFName :: Name -> Name
readFName 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

-- | Generates an function to read (i.e. deserialise) the given data type.
--
--  __Example:__
--
-- For the 'Tree' data type, it generates the following function:
--
-- @
-- readTree :: ('Unpackable' a) => 'Data.Packed.PackedReader' '[Tree a] r (Tree a)
-- readTree = caseTree
--     ('Data.Packed.reader' >>= \\leafContent ->
--          'Data.Packed.Reader.return' $ Leaf leafContent
--     )
--
--     ('Data.Packed.reader' >>= \\leftContent ->
--      'Data.Packed.reader' >>= \\rightContent ->
--          'Data.Packed.Reader.return' $ Node leftContent rightContent
--     )
-- @
--
-- __Note__ We use bindings ('Data.Packed.Reader.>>=') intead of a do-notation, since 'Data.Packed.Reader' is not a monad. It's an indexed monad, meaning that the user would have to enable the 'QualifiedDo' extenstion for it to compile.
genRead ::
    [PackingFlag] ->
    Name ->
    -- | The name of the type to generate the function for
    Q [Dec]
genRead :: [PackingFlag] -> Name -> Q [Dec]
genRead [PackingFlag]
flags Name
tyName = do
    let fName :: Name
fName = Name -> Name
readFName Name
tyName
    (resolvedType, typeVariables) <- Name -> Q (Type, [Name])
resolveAppliedType Name
tyName
    lambdas <- genReadLambdas flags tyName
    -- we fold the list of lambda by applring them to `caseTree packed`
    funExpr <-
        foldl
            (\Q Exp
rest Exp
arg -> [|$Q Exp
rest $(Exp -> Q Exp
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return Exp
arg)|])
            (varE $ caseFName tyName)
            lambdas
    let fun = Name -> [Clause] -> Dec
FunD Name
fName [[Pat] -> Body -> [Dec] -> Clause
Clause [] (Exp -> Body
NormalB Exp
funExpr) []]
    signature <- genReadSignature tyName resolvedType typeVariables
    return [signature, fun]

-- Generates all the lambda functions we will need, to unpack using caseTree
genReadLambdas :: [PackingFlag] -> Name -> Q [Exp]
genReadLambdas :: [PackingFlag] -> Name -> Q [Exp]
genReadLambdas [PackingFlag]
flags Name
tyName = do
    (TyConI (DataD _ _ _ _ cs _)) <- Name -> Q Info
reify Name
tyName
    mapM
        ( \Con
con ->
            let (Name
conName, [BangType]
bt) = Con -> (Name, [BangType])
getNameAndBangTypesFromCon Con
con
             in [PackingFlag] -> Name -> Cxt -> Q Exp
genReadLambda [PackingFlag]
flags Name
conName (BangType -> Type
forall a b. (a, b) -> b
snd (BangType -> Type) -> [BangType] -> Cxt
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [BangType]
bt)
        )
        cs

-- generates a single lambda to use with caseTree for our unpack method
genReadLambda :: [PackingFlag] -> Name -> [Type] -> Q Exp
genReadLambda :: [PackingFlag] -> Name -> Cxt -> Q Exp
genReadLambda [PackingFlag]
flags Name
conName Cxt
conParameterTypes = do
    let appliedConstructor :: Exp
appliedConstructor =
            (Exp -> Name -> Exp) -> Exp -> [Name] -> Exp
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl
                (\Exp
rest Name
arg -> Exp -> Exp -> Exp
AppE Exp
rest (Exp -> Exp) -> Exp -> Exp
forall a b. (a -> b) -> a -> b
$ Name -> Exp
VarE Name
arg)
                (Name -> Exp
ConE Name
conName)
                ([Name] -> Exp) -> [Name] -> Exp
forall a b. (a -> b) -> a -> b
$ (\Int
i -> String -> Name
mkName (String -> Name) -> String -> Name
forall a b. (a -> b) -> a -> b
$ String
"arg" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
i)
                    (Int -> Name) -> [Int] -> [Name]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Int
0 .. (Cxt -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Cxt
conParameterTypes Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)]
    Exp -> Q Exp
buildBindingExpression Exp
appliedConstructor
  where
    hasSizeFlag :: Bool
hasSizeFlag = PackingFlag
InsertFieldSize PackingFlag -> [PackingFlag] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [PackingFlag]
flags
    skipLastFieldSizeFlag :: Bool
skipLastFieldSizeFlag = PackingFlag
SkipLastFieldSize PackingFlag -> [PackingFlag] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [PackingFlag]
flags
    buildBindingExpression :: Exp -> Q Exp
    buildBindingExpression :: Exp -> Q Exp
buildBindingExpression Exp
appliedConstructor =
        ((Int, Bool) -> Q Exp -> Q Exp) -> Q Exp -> [(Int, Bool)] -> Q Exp
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr
            ( \(Int
argIndex, Bool
hasSize) Q Exp
ret ->
                let
                    skipAndUnpack :: Q Exp
skipAndUnpack = [|skip R.>> $Q Exp
unpackExpr|]
                    unpackExpr :: Q Exp
unpackExpr = [|reader R.>>= \($(Name -> Q Pat
forall (m :: * -> *). Quote m => Name -> m Pat
varP (Name -> Q Pat) -> Name -> Q Pat
forall a b. (a -> b) -> a -> b
$ String -> Name
mkName (String -> Name) -> String -> Name
forall a b. (a -> b) -> a -> b
$ String
"arg" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
argIndex)) -> $Q Exp
ret|]
                 in
                    if Bool
hasSize then Q Exp
skipAndUnpack else Q Exp
unpackExpr
            )
            [|R.return ($(Q Exp -> Q Exp
forall (m :: * -> *). Quote m => m Exp -> m Exp
parensE (Exp -> Q Exp
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return Exp
appliedConstructor)))|]
            ([(Int, Bool)] -> Q Exp) -> [(Int, Bool)] -> Q Exp
forall a b. (a -> b) -> a -> b
$ (\Int
i -> (Int
i, Bool
hasSizeFlag Bool -> Bool -> Bool
&& (Bool -> Bool
not Bool
skipLastFieldSizeFlag Bool -> Bool -> Bool
|| (Bool
skipLastFieldSizeFlag Bool -> Bool -> Bool
&& Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Cxt -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Cxt
conParameterTypes Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1))))
                (Int -> (Int, Bool)) -> [Int] -> [(Int, Bool)]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Int
0 .. (Cxt -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Cxt
conParameterTypes Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)]

-- For a type 'Tree', generates the following function signature
-- readTree :: ('Unpackable' a) => 'Data.Packed.PackedReader' '[Tree a] r (Tree a)
genReadSignature :: Name -> Type -> [Name] -> Q Dec
genReadSignature :: Name -> Type -> [Name] -> Q Dec
genReadSignature Name
tyName Type
resolvedType [Name]
typeVariables = do
    restTypeName <- String -> Q Name
forall (m :: * -> *). Quote m => String -> m Name
newName String
"r"
    let readerType = [t|PackedReader '[$(Type -> Q Type
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return Type
resolvedType)] $(Name -> Q Type
forall (m :: * -> *). Quote m => Name -> m Type
varT Name
restTypeName) ($(Type -> Q Type
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return Type
resolvedType))|]
        constraints = (Name -> Q Type) -> [Name] -> Q Cxt
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (\Name
tyVarName -> [t|Unpackable $(Name -> Q Type
forall (m :: * -> *). Quote m => Name -> m Type
varT Name
tyVarName)|]) [Name]
typeVariables
        signature = Q Type
readerType
    sigD (readFName tyName) $ forallT [] constraints signature