module Data.Packed.TH.Skip (genSkip, skipFName) where

import Data.Packed.FieldSize (skipWithFieldSize)
import Data.Packed.Reader (PackedReader)
import qualified Data.Packed.Reader as R
import Data.Packed.Skippable (Skippable (skip))
import Data.Packed.TH.Case (caseFName)
import Data.Packed.TH.Flag (PackingFlag (..))
import Data.Packed.TH.Utils
import Language.Haskell.TH

-- For a data type 'Tree', will generate the function name 'skipTree'
skipFName :: Name -> Name
skipFName :: Name -> Name
skipFName Name
tyName = String -> Name
mkName (String -> Name) -> String -> Name
forall a b. (a -> b) -> a -> b
$ String
"skip" String -> String -> String
forall a. [a] -> [a] -> [a]
++ Name -> String
sanitizeConName Name
tyName

-- | Generates an function to skip a value of the given type in a 'Data.Packed.Packed'
--
--  __Example:__
--
-- For the 'Tree' data type, it generates the following function:
--
-- @
-- skipTree :: ('Data.Packed.Skippable' a) => 'Data.Packed.PackedReader' '[Tree a] r ()
-- skipTree = caseTree
--      'Data.Packed.Skip.skip'
--      ('skipTree' >> 'skipTree')
-- @
genSkip :: [PackingFlag] -> Name -> Q [Dec]
genSkip :: [PackingFlag] -> Name -> Q [Dec]
genSkip [PackingFlag]
flags Name
tyName = do
    let fName :: Name
fName = Name -> Name
skipFName Name
tyName
    lambdas <- [PackingFlag] -> Name -> Q [Exp]
genSkipLambdas [PackingFlag]
flags Name
tyName
    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 <- genSkipSignature tyName
    return [signature, fun]

-- Generates all the lambda functions we will need, to skip using caseTree
genSkipLambdas :: [PackingFlag] -> Name -> Q [Exp]
genSkipLambdas :: [PackingFlag] -> Name -> Q [Exp]
genSkipLambdas [PackingFlag]
flags Name
tyName = do
    (TyConI (DataD _ _ _ _ cs _)) <- Name -> Q Info
reify Name
tyName
    mapM
        ( \Con
con ->
            let (Name
_, [BangType]
bt) = Con -> (Name, [BangType])
getNameAndBangTypesFromCon Con
con
             in [PackingFlag] -> Cxt -> Q Exp
genSkipLambda [PackingFlag]
flags (BangType -> Kind
forall a b. (a, b) -> b
snd (BangType -> Kind) -> [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 skip method
genSkipLambda :: [PackingFlag] -> [Type] -> Q Exp
genSkipLambda :: [PackingFlag] -> Cxt -> Q Exp
genSkipLambda [PackingFlag]
flags Cxt
conParameterTypes =
    (Bool -> Q Exp -> Q Exp) -> Q Exp -> [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
        ( \Bool
hasSize Q Exp
ret ->
            let
                skipFSAndSkipField :: Q Exp
skipFSAndSkipField = [|skipWithFieldSize R.>> $Q Exp
ret|]
                skipField :: Q Exp
skipField = [|skip R.>> $Q Exp
ret|]
             in
                if Bool
hasSize then Q Exp
skipFSAndSkipField else Q Exp
skipField
        )
        [|R.return ()|]
        ([Bool] -> Q Exp) -> [Bool] -> Q Exp
forall a b. (a -> b) -> a -> b
$ (\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 -> Bool) -> [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)]
  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

-- Generates the following function signature for a data type 'Tree'
-- skipTree :: ('Data.Packed.Skippable' a) => 'Data.Packed.PackedReader' '[Tree a] r ()
genSkipSignature :: Name -> Q Dec
genSkipSignature :: Name -> Q Dec
genSkipSignature Name
tyName = do
    (sourceType, typeParameterNames) <- Name -> Q (Kind, [Name])
resolveAppliedType Name
tyName
    let fName = Name -> Name
skipFName Name
tyName
        -- Type variables for Needs
        r = 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
"r"
        -- Define Skippable constraints on each of the type parameters
        constraints = (Name -> Q Kind) -> [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|Skippable $(Name -> Q Kind
forall (m :: * -> *). Quote m => Name -> m Kind
varT Name
tyVarName)|]) [Name]
typeParameterNames
        signature = [t|PackedReader '[$(Kind -> Q Kind
forall a. a -> Q a
forall (m :: * -> *) a. Monad m => a -> m a
return Kind
sourceType)] $Q Kind
r ()|]
    sigD fName (forallT [] constraints signature)