{-# LANGUAGE CPP #-}
{-# LANGUAGE TemplateHaskell #-}

module Data.Struct.TH (makeStruct) where

import           Control.Monad (when, zipWithM)
import           Control.Monad.Primitive (PrimMonad, PrimState)
import           Data.Either (partitionEithers)
import           Data.Primitive
import           Data.Struct
import           Data.Struct.Internal (Dict(Dict), initializeUnboxedField, st)
import           Data.List (groupBy, nub)
import           Language.Haskell.TH
import           Language.Haskell.TH.Datatype.TyVarBndr
import           Language.Haskell.TH.Syntax (VarStrictType)

#ifdef HLINT
{-# ANN module "HLint: ignore Use ." #-}
#endif

data StructRep = StructRep
  { StructRep -> Name
srState       :: Name
  , StructRep -> Name
srName        :: Name
  , StructRep -> [TyVarBndrUnit]
srTyVars      :: [TyVarBndrUnit]
#if MIN_VERSION_template_haskell(2,12,0)
  , StructRep -> [DerivClause]
srDerived     :: [DerivClause]
#else
  , srDerived     :: Cxt
#endif
  , StructRep -> Cxt
srCxt         :: Cxt
  , StructRep -> Name
srConstructor :: Name
  , StructRep -> [Member]
srMembers     :: [Member]
  } deriving Int -> StructRep -> ShowS
[StructRep] -> ShowS
StructRep -> String
(Int -> StructRep -> ShowS)
-> (StructRep -> String)
-> ([StructRep] -> ShowS)
-> Show StructRep
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [StructRep] -> ShowS
$cshowList :: [StructRep] -> ShowS
show :: StructRep -> String
$cshow :: StructRep -> String
showsPrec :: Int -> StructRep -> ShowS
$cshowsPrec :: Int -> StructRep -> ShowS
Show

data Member = Member
  { Member -> Representation
_memberRep :: Representation
  , Member -> Name
memberName :: Name
  , Member -> Type
_memberType :: Type
  }
  deriving Int -> Member -> ShowS
[Member] -> ShowS
Member -> String
(Int -> Member -> ShowS)
-> (Member -> String) -> ([Member] -> ShowS) -> Show Member
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Member] -> ShowS
$cshowList :: [Member] -> ShowS
show :: Member -> String
$cshow :: Member -> String
showsPrec :: Int -> Member -> ShowS
$cshowsPrec :: Int -> Member -> ShowS
Show

data Representation = BoxedField | UnboxedField | Slot
  deriving Int -> Representation -> ShowS
[Representation] -> ShowS
Representation -> String
(Int -> Representation -> ShowS)
-> (Representation -> String)
-> ([Representation] -> ShowS)
-> Show Representation
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Representation] -> ShowS
$cshowList :: [Representation] -> ShowS
show :: Representation -> String
$cshow :: Representation -> String
showsPrec :: Int -> Representation -> ShowS
$cshowsPrec :: Int -> Representation -> ShowS
Show

-- | Generate allocators, slots, fields, unboxed fields, Eq instances,
-- and Struct instances for the given "data types".
--
-- Inputs are expected to be "data types" parameterized by a state
-- type. Strict fields are considered to be slots, Non-strict fields
-- are considered to be boxed types, Unpacked fields are considered
-- to be unboxed primitives.
--
-- The data type should use record syntax and have a single constructor.
-- The field names will be used to generate slot, field, and unboxedField
-- values of the same name.
--
-- An allocator for the struct is generated by prefixing "alloc" to the
-- data type name.
makeStruct :: DecsQ -> DecsQ
makeStruct :: DecsQ -> DecsQ
makeStruct DecsQ
dsq =
  do [Dec]
ds   <- DecsQ
dsq
     ([Dec]
passthrough, [StructRep]
reps) <- [Either Dec StructRep] -> ([Dec], [StructRep])
forall a b. [Either a b] -> ([a], [b])
partitionEithers ([Either Dec StructRep] -> ([Dec], [StructRep]))
-> Q [Either Dec StructRep] -> Q ([Dec], [StructRep])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Dec -> Q (Either Dec StructRep))
-> [Dec] -> Q [Either Dec StructRep]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse Dec -> Q (Either Dec StructRep)
computeRep [Dec]
ds
     [[Dec]]
ds's <- (StructRep -> DecsQ) -> [StructRep] -> Q [[Dec]]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse ([Dec] -> StructRep -> DecsQ
generateCode [Dec]
passthrough) [StructRep]
reps
     [Dec] -> DecsQ
forall (m :: * -> *) a. Monad m => a -> m a
return ([Dec]
passthrough [Dec] -> [Dec] -> [Dec]
forall a. [a] -> [a] -> [a]
++ [[Dec]] -> [Dec]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[Dec]]
ds's)

mkAllocName :: StructRep -> Name
mkAllocName :: StructRep -> Name
mkAllocName StructRep
rep = String -> Name
mkName (String
"alloc" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Name -> String
nameBase (StructRep -> Name
srName StructRep
rep))

mkInitName :: StructRep -> Name
mkInitName :: StructRep -> Name
mkInitName StructRep
rep = String -> Name
mkName (String
"new" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Name -> String
nameBase (StructRep -> Name
srName StructRep
rep))

------------------------------------------------------------------------
-- Input validation
------------------------------------------------------------------------

computeRep :: Dec -> Q (Either Dec StructRep)
computeRep :: Dec -> Q (Either Dec StructRep)
computeRep (DataD Cxt
c Name
n [TyVarBndrUnit]
vs Maybe Type
_ [Con]
cs [DerivClause]
ds) =
  do Name
state <- [TyVarBndrUnit] -> Q Name
validateStateType [TyVarBndrUnit]
vs
     (Name
conname, [VarStrictType]
confields) <- [Con] -> Q (Name, [VarStrictType])
validateContructor [Con]
cs
     [Member]
members <- (VarStrictType -> Q Member) -> [VarStrictType] -> Q [Member]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (Name -> VarStrictType -> Q Member
validateMember Name
state) [VarStrictType]
confields

     Either Dec StructRep -> Q (Either Dec StructRep)
forall (m :: * -> *) a. Monad m => a -> m a
return (Either Dec StructRep -> Q (Either Dec StructRep))
-> Either Dec StructRep -> Q (Either Dec StructRep)
forall a b. (a -> b) -> a -> b
$ StructRep -> Either Dec StructRep
forall a b. b -> Either a b
Right StructRep :: Name
-> Name
-> [TyVarBndrUnit]
-> [DerivClause]
-> Cxt
-> Name
-> [Member]
-> StructRep
StructRep
       { srState :: Name
srState = Name
state
       , srName :: Name
srName  = Name
n
       , srTyVars :: [TyVarBndrUnit]
srTyVars = [TyVarBndrUnit]
vs
       , srConstructor :: Name
srConstructor = Name
conname
       , srMembers :: [Member]
srMembers = [Member]
members
       , srDerived :: [DerivClause]
srDerived = [DerivClause]
ds
       , srCxt :: Cxt
srCxt = Cxt
c
       }
computeRep Dec
d = Either Dec StructRep -> Q (Either Dec StructRep)
forall (m :: * -> *) a. Monad m => a -> m a
return (Dec -> Either Dec StructRep
forall a b. a -> Either a b
Left Dec
d)

-- | Check that only a single data constructor was provided and
-- that it was a record constructor.
validateContructor :: [Con] -> Q (Name,[VarStrictType])
validateContructor :: [Con] -> Q (Name, [VarStrictType])
validateContructor [RecC Name
name [VarStrictType]
fields] = (Name, [VarStrictType]) -> Q (Name, [VarStrictType])
forall (m :: * -> *) a. Monad m => a -> m a
return (Name
name,[VarStrictType]
fields)
validateContructor [Con
_] = String -> Q (Name, [VarStrictType])
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"Expected a record constructor"
validateContructor [Con]
xs = String -> Q (Name, [VarStrictType])
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String
"Expected 1 constructor, got " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show ([Con] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Con]
xs))

-- A struct type's final type variable should be suitable for
-- use as the ('PrimState' m) argument.
validateStateType :: [TyVarBndrUnit] -> Q Name
validateStateType :: [TyVarBndrUnit] -> Q Name
validateStateType [TyVarBndrUnit]
xs =
  do Bool -> Q () -> Q ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when ([TyVarBndrUnit] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [TyVarBndrUnit]
xs) (String -> Q ()
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"state type expected but no type variables found")
     (Name -> Q Name)
-> (Name -> Type -> Q Name) -> TyVarBndrUnit -> Q Name
forall r flag.
(Name -> r) -> (Name -> Type -> r) -> TyVarBndrUnit -> r
elimTV Name -> Q Name
forall (m :: * -> *) a. Monad m => a -> m a
return Name -> Type -> Q Name
validateKindedTV ([TyVarBndrUnit] -> TyVarBndrUnit
forall a. [a] -> a
last [TyVarBndrUnit]
xs)
  where
    validateKindedTV :: Name -> Kind -> Q Name
    validateKindedTV :: Name -> Type -> Q Name
validateKindedTV Name
n Type
k
      | Type
k Type -> Type -> Bool
forall a. Eq a => a -> a -> Bool
== Type
starK = Name -> Q Name
forall (m :: * -> *) a. Monad m => a -> m a
return Name
n
      | Bool
otherwise  = String -> Q Name
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"state type should have kind *"

-- | Figure out which record fields are Slots and which are
-- Fields. Slots will have types ending in the state type
validateMember :: Name -> VarStrictType -> Q Member
validateMember :: Name -> VarStrictType -> Q Member
validateMember Name
s (Name
fieldname,Bang SourceUnpackedness
NoSourceUnpackedness SourceStrictness
NoSourceStrictness,Type
fieldtype) =
  do Bool -> Q () -> Q ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Name -> Type -> Bool
occurs Name
s Type
fieldtype)
       (String -> Q ()
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String
"state type may not occur in field `" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Name -> String
nameBase Name
fieldname String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"`"))
     Member -> Q Member
forall (m :: * -> *) a. Monad m => a -> m a
return (Representation -> Name -> Type -> Member
Member Representation
BoxedField Name
fieldname Type
fieldtype)
validateMember Name
s (Name
fieldname,Bang SourceUnpackedness
NoSourceUnpackedness SourceStrictness
SourceStrict,Type
fieldtype) =
  do Type
f <- Type -> Name -> Q Type
unapplyType Type
fieldtype Name
s
     Bool -> Q () -> Q ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Name -> Type -> Bool
occurs Name
s Type
f)
       (String -> Q ()
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String
"state type may only occur in final position in slot `" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Name -> String
nameBase Name
fieldname String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"`"))
     Member -> Q Member
forall (m :: * -> *) a. Monad m => a -> m a
return (Representation -> Name -> Type -> Member
Member Representation
Slot Name
fieldname Type
f)
validateMember Name
s (Name
fieldname,Bang SourceUnpackedness
SourceUnpack SourceStrictness
SourceStrict,Type
fieldtype) =
  do Bool -> Q () -> Q ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Name -> Type -> Bool
occurs Name
s Type
fieldtype)
       (String -> Q ()
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String
"state type may not occur in unpacked field `" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Name -> String
nameBase Name
fieldname String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"`"))
     Member -> Q Member
forall (m :: * -> *) a. Monad m => a -> m a
return (Representation -> Name -> Type -> Member
Member Representation
UnboxedField Name
fieldname Type
fieldtype)
validateMember Name
_ VarStrictType
_ = String -> Q Member
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"validateMember: can't unpack nonstrict fields"

unapplyType :: Type -> Name -> Q Type
unapplyType :: Type -> Name -> Q Type
unapplyType (AppT Type
f (VarT Name
x)) Name
y | Name
x Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
== Name
y = Type -> Q Type
forall (m :: * -> *) a. Monad m => a -> m a
return Type
f
unapplyType Type
t Name
n =
  String -> Q Type
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String -> Q Type) -> String -> Q Type
forall a b. (a -> b) -> a -> b
$ String
"Unable to match state type of slot: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Type -> String
forall a. Show a => a -> String
show Type
t String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" | expected: " String -> ShowS
forall a. [a] -> [a] -> [a]
++ Name -> String
nameBase Name
n

------------------------------------------------------------------------
-- Code generation
------------------------------------------------------------------------

generateCode :: [Dec] -> StructRep -> DecsQ
generateCode :: [Dec] -> StructRep -> DecsQ
generateCode [Dec]
ds StructRep
rep = [[Dec]] -> [Dec]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[Dec]] -> [Dec]) -> Q [[Dec]] -> DecsQ
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [DecsQ] -> Q [[Dec]]
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence
  [ StructRep -> DecsQ
generateDataType StructRep
rep
  , StructRep -> DecsQ
generateStructInstance StructRep
rep
  , StructRep -> DecsQ
generateMembers StructRep
rep
  , StructRep -> DecsQ
generateNew StructRep
rep
  , StructRep -> DecsQ
generateAlloc StructRep
rep
  , [Dec] -> StructRep -> DecsQ
generateRoles [Dec]
ds StructRep
rep
  ]

-- Generates: newtype TyCon a b c s = DataCon (Object s)
generateDataType :: StructRep -> DecsQ
generateDataType :: StructRep -> DecsQ
generateDataType StructRep
rep = [Q Dec] -> DecsQ
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence
  [ CxtQ
-> Name
-> [TyVarBndrUnit]
-> Maybe Type
-> ConQ
-> [DerivClauseQ]
-> Q Dec
newtypeD (Cxt -> CxtQ
forall (m :: * -> *) a. Monad m => a -> m a
return (StructRep -> Cxt
srCxt StructRep
rep)) (StructRep -> Name
srName StructRep
rep) (StructRep -> [TyVarBndrUnit]
srTyVars StructRep
rep)
      Maybe Type
forall a. Maybe a
Nothing
      (Name -> [BangTypeQ] -> ConQ
normalC
         (StructRep -> Name
srConstructor StructRep
rep)
         [ BangQ -> Q Type -> BangTypeQ
bangType
             (SourceUnpackednessQ -> SourceStrictnessQ -> BangQ
bang SourceUnpackednessQ
noSourceUnpackedness SourceStrictnessQ
noSourceStrictness)
             [t| Object $(varT (srState rep)) |]
         ])
#if MIN_VERSION_template_haskell(2,12,0)
      ((DerivClause -> DerivClauseQ) -> [DerivClause] -> [DerivClauseQ]
forall a b. (a -> b) -> [a] -> [b]
map DerivClause -> DerivClauseQ
forall (m :: * -> *) a. Monad m => a -> m a
return (StructRep -> [DerivClause]
srDerived StructRep
rep))
#else
      (return (srDerived rep))
#endif
  ]

generateRoles :: [Dec] -> StructRep -> DecsQ
generateRoles :: [Dec] -> StructRep -> DecsQ
generateRoles [Dec]
ds StructRep
rep
  | Bool
hasRoleAnnotation = [Dec] -> DecsQ
forall (m :: * -> *) a. Monad m => a -> m a
return []
  | Bool
otherwise = [Q Dec] -> DecsQ
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence [ Name -> [Role] -> Q Dec
roleAnnotD (StructRep -> Name
srName StructRep
rep) (StructRep -> [Role]
computeRoles StructRep
rep) ]

  where
  hasRoleAnnotation :: Bool
hasRoleAnnotation = (Dec -> Bool) -> [Dec] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any Dec -> Bool
isTargetRoleAnnot [Dec]
ds

  isTargetRoleAnnot :: Dec -> Bool
isTargetRoleAnnot (RoleAnnotD Name
n [Role]
_) = Name
n Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
== StructRep -> Name
srName StructRep
rep
  isTargetRoleAnnot Dec
_ = Bool
False

-- Currently all roles are set to nominal. A more general solution
-- should be able to infer some representional/phantom roles. To do
-- this for arbitrary types we'll need a way to query the roles of
-- existing type constructors to infer the correct roles.
computeRoles :: StructRep -> [Role]
computeRoles :: StructRep -> [Role]
computeRoles = (TyVarBndrUnit -> Role) -> [TyVarBndrUnit] -> [Role]
forall a b. (a -> b) -> [a] -> [b]
map (Role -> TyVarBndrUnit -> Role
forall a b. a -> b -> a
const Role
NominalR) ([TyVarBndrUnit] -> [Role])
-> (StructRep -> [TyVarBndrUnit]) -> StructRep -> [Role]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. StructRep -> [TyVarBndrUnit]
srTyVars

-- | Type of the object not applied to a state type. This
-- should have kind * -> *
repType1 :: StructRep -> TypeQ
repType1 :: StructRep -> Q Type
repType1 StructRep
rep = Name -> [TyVarBndrUnit] -> Q Type
repTypeHelper (StructRep -> Name
srName StructRep
rep) ([TyVarBndrUnit] -> [TyVarBndrUnit]
forall a. [a] -> [a]
init (StructRep -> [TyVarBndrUnit]
srTyVars StructRep
rep))

-- | Type of the object as originally declared, fully applied.
repType :: StructRep -> TypeQ
repType :: StructRep -> Q Type
repType StructRep
rep = Name -> [TyVarBndrUnit] -> Q Type
repTypeHelper (StructRep -> Name
srName StructRep
rep) (StructRep -> [TyVarBndrUnit]
srTyVars StructRep
rep)

repTypeHelper :: Name -> [TyVarBndrUnit] -> TypeQ
repTypeHelper :: Name -> [TyVarBndrUnit] -> Q Type
repTypeHelper Name
c [TyVarBndrUnit]
vs = (Q Type -> Q Type -> Q Type) -> Q Type -> [Q Type] -> Q Type
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl Q Type -> Q Type -> Q Type
appT (Name -> Q Type
conT Name
c) (TyVarBndrUnit -> Q Type
tyVarBndrT (TyVarBndrUnit -> Q Type) -> [TyVarBndrUnit] -> [Q Type]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [TyVarBndrUnit]
vs)

-- Construct a 'TypeQ' from a 'TyVarBndr'
tyVarBndrT :: TyVarBndrUnit -> TypeQ
tyVarBndrT :: TyVarBndrUnit -> Q Type
tyVarBndrT = (Name -> Q Type)
-> (Name -> Type -> Q Type) -> TyVarBndrUnit -> Q Type
forall r flag.
(Name -> r) -> (Name -> Type -> r) -> TyVarBndrUnit -> r
elimTV Name -> Q Type
varT (Q Type -> Type -> Q Type
sigT (Q Type -> Type -> Q Type)
-> (Name -> Q Type) -> Name -> Type -> Q Type
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Name -> Q Type
varT)

generateStructInstance :: StructRep -> DecsQ
generateStructInstance :: StructRep -> DecsQ
generateStructInstance StructRep
rep =
  [d| instance Struct $(repType1 rep) where struct = Dict
      instance Eq     $(repType  rep) where (==)   = eqStruct
    |]

-- generates: allocDataCon = alloc <n>
generateAlloc :: StructRep -> DecsQ
generateAlloc :: StructRep -> DecsQ
generateAlloc StructRep
rep =
  do Name
mName <- String -> Q Name
newName String
"m"
     let m :: TypeQ
         m :: Q Type
m = Name -> Q Type
varT Name
mName

         n :: Int
n = [[Member]] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ((Member -> Member -> Bool) -> [Member] -> [[Member]]
forall a. (a -> a -> Bool) -> [a] -> [[a]]
groupBy Member -> Member -> Bool
isNeighbor (StructRep -> [Member]
srMembers StructRep
rep))
         allocName :: Name
allocName = StructRep -> Name
mkAllocName StructRep
rep

     StructRep -> Name -> Q Type -> ExpQ -> DecsQ
simpleDefinition StructRep
rep Name
allocName
       ([TyVarBndrUnit] -> CxtQ -> Q Type -> Q Type
forallT [Name -> TyVarBndrUnit
plainTVSpecified Name
mName] ([Q Type] -> CxtQ
cxt [])
          [t| PrimMonad $m => $m ( $(repType1 rep) (PrimState $m) ) |])
       [| alloc n |]


-- generates:
-- newDataCon a .. = do this <- alloc <n>; set field1 this a; ...; return this
generateNew :: StructRep -> DecsQ
generateNew :: StructRep -> DecsQ
generateNew StructRep
rep =
  do Name
this <- String -> Q Name
newName String
"this"
     let ms :: [[Member]]
ms = (Member -> Member -> Bool) -> [Member] -> [[Member]]
forall a. (a -> a -> Bool) -> [a] -> [[a]]
groupBy Member -> Member -> Bool
isNeighbor (StructRep -> [Member]
srMembers StructRep
rep)

         addName :: Member -> Q (Name, Member)
addName Member
m = do Name
n <- String -> Q Name
newName (Name -> String
nameBase (Member -> Name
memberName Member
m))
                        (Name, Member) -> Q (Name, Member)
forall (m :: * -> *) a. Monad m => a -> m a
return (Name
n,Member
m)

     [[(Name, Member)]]
msWithArgs <- ([Member] -> Q [(Name, Member)])
-> [[Member]] -> Q [[(Name, Member)]]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse ((Member -> Q (Name, Member)) -> [Member] -> Q [(Name, Member)]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse Member -> Q (Name, Member)
addName) [[Member]]
ms

     let name :: Name
name = StructRep -> Name
mkInitName StructRep
rep
         body :: ExpQ
body = [StmtQ] -> ExpQ
doE
                -- allocate struct
              ([StmtQ] -> ExpQ) -> [StmtQ] -> ExpQ
forall a b. (a -> b) -> a -> b
$ PatQ -> ExpQ -> StmtQ
bindS (Name -> PatQ
varP Name
this) (Name -> ExpQ
varE (StructRep -> Name
mkAllocName StructRep
rep))

                -- initialize each member
              StmtQ -> [StmtQ] -> [StmtQ]
forall a. a -> [a] -> [a]
: (ExpQ -> StmtQ
noBindS (ExpQ -> StmtQ) -> [ExpQ] -> [StmtQ]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Int -> [(Name, Member)] -> ExpQ)
-> [Int] -> [[(Name, Member)]] -> [ExpQ]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (ExpQ -> Int -> [(Name, Member)] -> ExpQ
assignN (Name -> ExpQ
varE Name
this)) [Int
0..] [[(Name, Member)]]
msWithArgs)

                -- return initialized struct
             [StmtQ] -> [StmtQ] -> [StmtQ]
forall a. [a] -> [a] -> [a]
++ [ ExpQ -> StmtQ
noBindS [| return $(varE this) |] ]

     [Q Dec] -> DecsQ
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence
       [ Name -> Q Type -> Q Dec
sigD Name
name (StructRep -> Q Type
newStructType StructRep
rep)
       , Name -> [ClauseQ] -> Q Dec
funD Name
name [ [PatQ] -> BodyQ -> [Q Dec] -> ClauseQ
clause (Name -> PatQ
varP (Name -> PatQ)
-> ((Name, Member) -> Name) -> (Name, Member) -> PatQ
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Name, Member) -> Name
forall a b. (a, b) -> a
fst ((Name, Member) -> PatQ) -> [(Name, Member)] -> [PatQ]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [[(Name, Member)]] -> [(Name, Member)]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[(Name, Member)]]
msWithArgs)
                            (ExpQ -> BodyQ
normalB [| st $body |] ) [] ]
       ]


assignN :: ExpQ -> Int -> [(Name,Member)] -> ExpQ

assignN :: ExpQ -> Int -> [(Name, Member)] -> ExpQ
assignN ExpQ
this Int
_ [(Name
arg,Member Representation
BoxedField Name
n Type
_)] =
  [| setField $(varE n) $this $(varE arg) |]

assignN ExpQ
this Int
_ [(Name
arg,Member Representation
Slot Name
n Type
_)] =
  [| set      $(varE n) $this $(varE arg)|]

assignN ExpQ
this Int
i [(Name, Member)]
us =
  do let n :: Int
n = [(Name, Member)] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(Name, Member)]
us
     Name
mba <- String -> Q Name
newName String
"mba"
     let arg0 :: Name
arg0 = (Name, Member) -> Name
forall a b. (a, b) -> a
fst ([(Name, Member)] -> (Name, Member)
forall a. [a] -> a
head [(Name, Member)]
us)
     [StmtQ] -> ExpQ
doE ([StmtQ] -> ExpQ) -> [StmtQ] -> ExpQ
forall a b. (a -> b) -> a -> b
$ PatQ -> ExpQ -> StmtQ
bindS (Name -> PatQ
varP Name
mba) [| initializeUnboxedField i n (sizeOf $(varE arg0)) $this |]
         StmtQ -> [StmtQ] -> [StmtQ]
forall a. a -> [a] -> [a]
: [ ExpQ -> StmtQ
noBindS [| writeByteArray $(varE mba) j $(varE arg) |]
           | (Int
j,(Name
arg,Member
_)) <- [Int] -> [(Name, Member)] -> [(Int, (Name, Member))]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 :: Int ..] [(Name, Member)]
us ]

-- | The type of the struct initializer is complicated enough to
-- pull it out here.
-- generates:
-- PrimMonad m => field1 -> field2 -> ... -> m (TyName a b ... (PrimState m))
newStructType :: StructRep -> TypeQ
newStructType :: StructRep -> Q Type
newStructType StructRep
rep =
  do Name
mName <- String -> Q Name
newName String
"m"
     let m :: TypeQ
         m :: Q Type
m = Name -> Q Type
varT Name
mName

         s :: Q Type
s = [t| PrimState $m |]
         obj :: Q Type
obj = StructRep -> Q Type
repType1 StructRep
rep

         buildType :: Member -> Q Type
buildType (Member Representation
BoxedField   Name
_ Type
t) = Type -> Q Type
forall (m :: * -> *) a. Monad m => a -> m a
return Type
t
         buildType (Member Representation
UnboxedField Name
_ Type
t) = Type -> Q Type
forall (m :: * -> *) a. Monad m => a -> m a
return Type
t
         buildType (Member Representation
Slot         Name
_ Type
f) = [t| $(return f) $s |]

         r :: Q Type
r = (Q Type -> Q Type -> Q Type) -> Q Type -> [Q Type] -> Q Type
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr Q Type -> Q Type -> Q Type
(-->)
               [t| $m ($obj $s) |]
               (Member -> Q Type
buildType (Member -> Q Type) -> [Member] -> [Q Type]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> StructRep -> [Member]
srMembers StructRep
rep)

         primPreds :: [Q Type]
primPreds = Name -> Q Type
primPred (Name -> Q Type) -> [Name] -> [Q Type]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Name] -> [Name]
forall a. Eq a => [a] -> [a]
nub [ Name
t | Member Representation
UnboxedField Name
_ (VarT Name
t) <- StructRep -> [Member]
srMembers StructRep
rep ]

     StructRep -> Q Type -> Q Type
forallRepT StructRep
rep (Q Type -> Q Type) -> Q Type -> Q Type
forall a b. (a -> b) -> a -> b
$ [TyVarBndrUnit] -> CxtQ -> Q Type -> Q Type
forallT [Name -> TyVarBndrUnit
plainTVSpecified Name
mName] ([Q Type] -> CxtQ
cxt [Q Type]
primPreds)
       [t| PrimMonad $m => $r |]

-- generates a slot, field, or unboxedField definition per member
generateMembers :: StructRep -> DecsQ
generateMembers :: StructRep -> DecsQ
generateMembers StructRep
rep
  = [[Dec]] -> [Dec]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[Dec]] -> [Dec]) -> Q [[Dec]] -> DecsQ
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>
    (Int -> [Member] -> DecsQ) -> [Int] -> [[Member]] -> Q [[Dec]]
forall (m :: * -> *) a b c.
Applicative m =>
(a -> b -> m c) -> [a] -> [b] -> m [c]
zipWithM
      (StructRep -> Int -> [Member] -> DecsQ
generateMember1 StructRep
rep)
      [Int
0..]
      ((Member -> Member -> Bool) -> [Member] -> [[Member]]
forall a. (a -> a -> Bool) -> [a] -> [[a]]
groupBy Member -> Member -> Bool
isNeighbor (StructRep -> [Member]
srMembers StructRep
rep))

isNeighbor :: Member -> Member -> Bool
isNeighbor :: Member -> Member -> Bool
isNeighbor (Member Representation
UnboxedField Name
_ Type
t) (Member Representation
UnboxedField Name
_ Type
u) = Type
t Type -> Type -> Bool
forall a. Eq a => a -> a -> Bool
== Type
u
isNeighbor Member
_ Member
_ = Bool
False

------------------------------------------------------------------------

generateMember1 :: StructRep -> Int -> [Member] -> DecsQ

-- generates: fieldname = field <n>
generateMember1 :: StructRep -> Int -> [Member] -> DecsQ
generateMember1 StructRep
rep Int
n [Member Representation
BoxedField Name
fieldname Type
fieldtype] =
  StructRep -> Name -> Q Type -> ExpQ -> DecsQ
simpleDefinition StructRep
rep Name
fieldname
    [t| Field $(repType1 rep) $(return fieldtype) |]
    [| field n |]

-- generates: slotname = slot <n>
generateMember1 StructRep
rep Int
n [Member Representation
Slot Name
slotname Type
slottype] =
  StructRep -> Name -> Q Type -> ExpQ -> DecsQ
simpleDefinition StructRep
rep Name
slotname
    [t| Slot $(repType1 rep) $(return slottype) |]
    [| slot n |]

-- It the first type patterns didn't hit then we expect a list
-- of unboxed fields due to the call to groupBy in generateMembers
-- generates: fieldname = unboxedField <n> <i>
generateMember1 StructRep
rep Int
n [Member]
us =
  [[Dec]] -> [Dec]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[Dec]] -> [Dec]) -> Q [[Dec]] -> DecsQ
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [DecsQ] -> Q [[Dec]]
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence
    [ StructRep -> Name -> Q Type -> ExpQ -> DecsQ
simpleDefinition StructRep
rep Name
fieldname
          (Type -> Q Type -> Q Type
addPrimCxt Type
fieldtype
             [t| Field $(repType1 rep) $(return fieldtype) |])
          [| unboxedField n i |]

    | (Int
i,Member Representation
UnboxedField Name
fieldname Type
fieldtype) <- [Int] -> [Member] -> [(Int, Member)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 :: Int ..] [Member]
us
    ]
  where
  addPrimCxt :: Type -> Q Type -> Q Type
addPrimCxt (VarT Name
t) = [TyVarBndrUnit] -> CxtQ -> Q Type -> Q Type
forallT [] ([Q Type] -> CxtQ
cxt [Name -> Q Type
primPred Name
t])
  addPrimCxt Type
_        = Q Type -> Q Type
forall a. a -> a
id

-- Generate code for definitions without arguments, with type variables
-- quantified over those in the struct rep, including an inline pragma
simpleDefinition :: StructRep -> Name -> TypeQ -> ExpQ -> DecsQ
simpleDefinition :: StructRep -> Name -> Q Type -> ExpQ -> DecsQ
simpleDefinition StructRep
rep Name
name Q Type
typ ExpQ
def =
  [Q Dec] -> DecsQ
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence
    [ Name -> Q Type -> Q Dec
sigD Name
name (StructRep -> Q Type -> Q Type
forallRepT StructRep
rep Q Type
typ)
    , Name -> ExpQ -> Q Dec
simpleValD Name
name ExpQ
def
    , Name -> Inline -> RuleMatch -> Phases -> Q Dec
pragInlD Name
name Inline
Inline RuleMatch
FunLike Phases
AllPhases
    ]

------------------------------------------------------------------------

-- Simple use of 'valD' bind an expression to a name
simpleValD :: Name -> ExpQ -> DecQ
simpleValD :: Name -> ExpQ -> Q Dec
simpleValD Name
var ExpQ
val = PatQ -> BodyQ -> [Q Dec] -> Q Dec
valD (Name -> PatQ
varP Name
var) (ExpQ -> BodyQ
normalB ExpQ
val) []

-- Quantifies over all of the type variables in a struct data type
-- except the state variable which is likely to be ('PrimState' s)
forallRepT :: StructRep -> TypeQ -> TypeQ
forallRepT :: StructRep -> Q Type -> Q Type
forallRepT StructRep
rep = [TyVarBndrUnit] -> CxtQ -> Q Type -> Q Type
forallT ([TyVarBndrUnit] -> [TyVarBndrUnit]
forall a. [a] -> [a]
init (Specificity -> [TyVarBndrUnit] -> [TyVarBndrUnit]
forall newFlag oldFlag.
newFlag -> [TyVarBndrUnit] -> [TyVarBndrUnit]
changeTVFlags Specificity
SpecifiedSpec (StructRep -> [TyVarBndrUnit]
srTyVars StructRep
rep))) ([Q Type] -> CxtQ
cxt [])

(-->) :: TypeQ -> TypeQ -> TypeQ
Q Type
f --> :: Q Type -> Q Type -> Q Type
--> Q Type
x = Q Type
arrowT Q Type -> Q Type -> Q Type
`appT` Q Type
f Q Type -> Q Type -> Q Type
`appT` Q Type
x

primPred :: Name -> PredQ
primPred :: Name -> Q Type
primPred Name
t = [t| Prim $(varT t) |]

occurs :: Name -> Type -> Bool
occurs :: Name -> Type -> Bool
occurs Name
n (AppT Type
f Type
x) = Name -> Type -> Bool
occurs Name
n Type
f Bool -> Bool -> Bool
|| Name -> Type -> Bool
occurs Name
n Type
x
occurs Name
n (VarT Name
m) = Name
n Name -> Name -> Bool
forall a. Eq a => a -> a -> Bool
== Name
m
occurs Name
n (ForallT [TyVarBndrUnit]
_ Cxt
_ Type
t) = Name -> Type -> Bool
occurs Name
n Type
t -- all names are fresh in quoted code, see below
occurs Name
n (SigT Type
t Type
_) = Name -> Type -> Bool
occurs Name
n Type
t
occurs Name
_ Type
_ = Bool
False

-- Prelude Language.Haskell.TH> runQ (stringE . show =<< [t| forall a. a -> (forall a. a) |])
-- LitE (StringL "ForallT [PlainTV a_0] [] (AppT (AppT ArrowT (VarT a_0)) (ForallT [PlainTV a_1] [] (VarT a_1)))")