module Data.Packed.TH.Start (startFName, genStart) where

import Data.Packed.FieldSize
import Data.Packed.Needs
import Data.Packed.Packable (write)
import Data.Packed.TH.Flag (PackingFlag (..))
import Data.Packed.TH.Utils
import Data.Word (Word8)
import Language.Haskell.TH

-- | For a constructor 'Leaf', will generate the function name 'startLeaf'
startFName :: Name -> Name
startFName :: Name -> Name
startFName Name
conName = String -> Name
mkName (String -> Name) -> String -> Name
forall a b. (a -> b) -> a -> b
$ String
"start" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Name -> String
sanitizeConName Name
conName

-- | Generates a function that prepares a 'Data.Packed.Needs' to receive values from a data constructor.
--
-- __Example:__
--
-- For the 'Tree' data type, it generates the following functions
--
-- @
-- startLeaf :: NeedsBuilder (Tree a ': r) t (a ': r) t
-- startLeaf = 'mkNeedsBuilder' (\n -> runBuilder (write (0 :: Word8) ('unsafeCastNeeds' n)))
--
-- startNode :: NeedsBuilder (Tree a ': r) t (Tree a ': Tree a ': r) t
-- startNode = 'mkNeedsBuilder' (\n -> runBuilder (write (1 :: Word8) ('unsafeCastNeeds' n)))
-- @
genStart ::
    [PackingFlag] ->
    -- | The name of the data constructor to generate the function for
    Name ->
    -- | The 'Tag' (byte) to write for this constructor
    Tag ->
    -- | The list of 'Type's of the data constructor's arguments
    [Type] ->
    Q [Dec]
genStart :: [PackingFlag] -> Name -> Tag -> [Type] -> Q [Dec]
genStart [PackingFlag]
flags Name
conName Tag
tag [Type]
paramTypeList = do
    let fName :: Name
fName = Name -> Name
startFName Name
conName
        constructorParamTypes :: [Q Type]
constructorParamTypes = Type -> Q Type
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return (Type -> Q Type) -> [Type] -> [Q Type]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Type]
paramTypeList
    (DataConI _ conType _) <- Name -> Q Info
reify Name
conName
    sig <-
        let r = Name -> Q Type
forall (m :: * -> *). Quote m => Name -> m Type
varT (Name -> Q Type) -> Name -> Q Type
forall a b. (a -> b) -> a -> b
$ String -> Name
mkName String
"r"
            t = Name -> Q Type
forall (m :: * -> *). Quote m => Name -> m Type
varT (Name -> Q Type) -> Name -> Q Type
forall a b. (a -> b) -> a -> b
$ String -> Name
mkName String
"t"
            insertFieldSizes = 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
            skipLastFieldSize = 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
            -- From the list of the constructor's parameters, generate the correct type for 'Data.Packed.Needs'
            -- For Leaf a, we will obtain Needs (a ': r)
            -- For node, if size flag is enabled, we will get Needs (FieldSize ': Tree a ': FieldSize ': Tree a ': r)
            destNeedsTypeParams =
                ((Int, Q Type) -> Q Type -> Q Type)
-> Q Type -> [(Int, Q Type)] -> Q Type
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr
                    ( \(Int
i, Q Type
x) Q Type
xs ->
                        if Bool
insertFieldSizes Bool -> Bool -> Bool
&& (Bool -> Bool
not Bool
skipLastFieldSize Bool -> Bool -> Bool
|| (Bool
skipLastFieldSize Bool -> Bool -> Bool
&& Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
1))
                            then [t|FieldSize ': $Q Type
x ': $Q Type
xs|]
                            else [t|$Q Type
x ': $Q Type
xs|]
                    )
                    Q Type
r
                    ([(Int, Q Type)] -> Q Type) -> [(Int, Q Type)] -> Q Type
forall a b. (a -> b) -> a -> b
$ [Int] -> [Q Type] -> [(Int, Q Type)]
forall a b. [a] -> [b] -> [(a, b)]
zip ([Int] -> [Int]
forall a. [a] -> [a]
reverse [Int
0 .. [Q Type] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Q Type]
constructorParamTypes]) [Q Type]
constructorParamTypes
         in [t|NeedsBuilder ($(return $ getParentTypeFromConstructorType conType) ': $r) $t $destNeedsTypeParams $t|]
    expr <- [|mkNeedsBuilder (\n -> runBuilder (write (tag :: Word8)) (unsafeCastNeeds n))|]
    return
        [ SigD fName sig
        , FunD fName [Clause [] (NormalB expr) []]
        ]