{-# 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 qualified Data.List.NonEmpty as NE
import           Data.List.NonEmpty (NonEmpty(..))
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 -> [TyVarBndrVis]
srTyVars      :: [TyVarBndrVis]
#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
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
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
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) <- forall a b. [Either a b] -> ([a], [b])
partitionEithers forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> 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 <- 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
     forall (m :: * -> *) a. Monad m => a -> m a
return ([Dec]
passthrough forall a. [a] -> [a] -> [a]
++ 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" 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" 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 [TyVarBndrVis]
vs Maybe Type
_ [Con]
cs [DerivClause]
ds) =
  do Name
state <- [TyVarBndrVis] -> Q Name
validateStateType [TyVarBndrVis]
vs
     (Name
conname, [VarStrictType]
confields) <- [Con] -> Q (Name, [VarStrictType])
validateContructor [Con]
cs
     [Member]
members <- 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

     forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a b. b -> Either a b
Right StructRep
       { srState :: Name
srState = Name
state
       , srName :: Name
srName  = Name
n
       , srTyVars :: [TyVarBndrVis]
srTyVars = [TyVarBndrVis]
vs
       , srConstructor :: Name
srConstructor = Name
conname
       , srMembers :: [Member]
srMembers = [Member]
members
       , srDerived :: [DerivClause]
srDerived = [DerivClause]
ds
       , srCxt :: Cxt
srCxt = Cxt
c
       }
computeRep Dec
d = forall (m :: * -> *) a. Monad m => a -> m a
return (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] = forall (m :: * -> *) a. Monad m => a -> m a
return (Name
name,[VarStrictType]
fields)
validateContructor [Con
_] = forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"Expected a record constructor"
validateContructor [Con]
xs = forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String
"Expected 1 constructor, got " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show (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 :: [TyVarBndrVis] -> Q Name
validateStateType :: [TyVarBndrVis] -> Q Name
validateStateType [TyVarBndrVis]
xs =
  do forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (forall (t :: * -> *) a. Foldable t => t a -> Bool
null [TyVarBndrVis]
xs) (forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
"state type expected but no type variables found")
     forall r flag.
(Name -> r) -> (Name -> Type -> r) -> TyVarBndr_ flag -> r
elimTV forall (m :: * -> *) a. Monad m => a -> m a
return Name -> Type -> Q Name
validateKindedTV (forall a. [a] -> a
last [TyVarBndrVis]
xs)
  where
    validateKindedTV :: Name -> Kind -> Q Name
    validateKindedTV :: Name -> Type -> Q Name
validateKindedTV Name
n Type
k
      | Type
k forall a. Eq a => a -> a -> Bool
== Type
starK = forall (m :: * -> *) a. Monad m => a -> m a
return Name
n
      | Bool
otherwise  = 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 forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Name -> Type -> Bool
occurs Name
s Type
fieldtype)
       (forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String
"state type may not occur in field `" forall a. [a] -> [a] -> [a]
++ Name -> String
nameBase Name
fieldname forall a. [a] -> [a] -> [a]
++ String
"`"))
     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
     forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Name -> Type -> Bool
occurs Name
s Type
f)
       (forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String
"state type may only occur in final position in slot `" forall a. [a] -> [a] -> [a]
++ Name -> String
nameBase Name
fieldname forall a. [a] -> [a] -> [a]
++ String
"`"))
     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 forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Name -> Type -> Bool
occurs Name
s Type
fieldtype)
       (forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String
"state type may not occur in unpacked field `" forall a. [a] -> [a] -> [a]
++ Name -> String
nameBase Name
fieldname forall a. [a] -> [a] -> [a]
++ String
"`"))
     forall (m :: * -> *) a. Monad m => a -> m a
return (Representation -> Name -> Type -> Member
Member Representation
UnboxedField Name
fieldname Type
fieldtype)
validateMember Name
_ VarStrictType
_ = 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 forall a. Eq a => a -> a -> Bool
== Name
y = forall (m :: * -> *) a. Monad m => a -> m a
return Type
f
unapplyType Type
t Name
n =
  forall (m :: * -> *) a. MonadFail m => String -> m a
fail forall a b. (a -> b) -> a -> b
$ String
"Unable to match state type of slot: " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show Type
t forall a. [a] -> [a] -> [a]
++ String
" | expected: " 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 = forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> 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 = forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence
  [ forall (m :: * -> *).
Quote m =>
m Cxt
-> Name
-> [TyVarBndrVis]
-> Maybe Type
-> m Con
-> [m DerivClause]
-> m Dec
newtypeD (forall (m :: * -> *) a. Monad m => a -> m a
return (StructRep -> Cxt
srCxt StructRep
rep)) (StructRep -> Name
srName StructRep
rep) (StructRep -> [TyVarBndrVis]
srTyVars StructRep
rep)
      forall a. Maybe a
Nothing
      (forall (m :: * -> *). Quote m => Name -> [m BangType] -> m Con
normalC
         (StructRep -> Name
srConstructor StructRep
rep)
         [ forall (m :: * -> *). Quote m => m Bang -> m Type -> m BangType
bangType
             (forall (m :: * -> *).
Quote m =>
m SourceUnpackedness -> m SourceStrictness -> m Bang
bang forall (m :: * -> *). Quote m => m SourceUnpackedness
noSourceUnpackedness forall (m :: * -> *). Quote m => m SourceStrictness
noSourceStrictness)
             [t| Object $(varT (srState rep)) |]
         ])
#if MIN_VERSION_template_haskell(2,12,0)
      (forall a b. (a -> b) -> [a] -> [b]
map 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 = forall (m :: * -> *) a. Monad m => a -> m a
return []
  | Bool
otherwise = forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence [ forall (m :: * -> *). Quote m => Name -> [Role] -> m Dec
roleAnnotD (StructRep -> Name
srName StructRep
rep) (StructRep -> [Role]
computeRoles StructRep
rep) ]

  where
  hasRoleAnnotation :: Bool
hasRoleAnnotation = 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 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 = forall a b. (a -> b) -> [a] -> [b]
map (forall a b. a -> b -> a
const Role
NominalR) forall b c a. (b -> c) -> (a -> b) -> a -> c
. StructRep -> [TyVarBndrVis]
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 -> [TyVarBndrVis] -> Q Type
repTypeHelper (StructRep -> Name
srName StructRep
rep) (forall a. [a] -> [a]
init (StructRep -> [TyVarBndrVis]
srTyVars StructRep
rep))

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

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

-- Construct a 'TypeQ' from a 'TyVarBndr'
tyVarBndrT :: TyVarBndrVis -> TypeQ
tyVarBndrT :: TyVarBndrVis -> Q Type
tyVarBndrT = forall r flag.
(Name -> r) -> (Name -> Type -> r) -> TyVarBndr_ flag -> r
elimTV forall (m :: * -> *). Quote m => Name -> m Type
varT (forall (m :: * -> *). Quote m => m Type -> Type -> m Type
sigT forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *). Quote m => Name -> m 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 <- forall (m :: * -> *). Quote m => String -> m Name
newName String
"m"
     let m :: TypeQ
         m :: Q Type
m = forall (m :: * -> *). Quote m => Name -> m Type
varT Name
mName

         n :: Int
n = forall (t :: * -> *) a. Foldable t => t a -> Int
length (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
       (forall (m :: * -> *).
Quote m =>
[TyVarBndr Specificity] -> m Cxt -> m Type -> m Type
forallT [Name -> TyVarBndr Specificity
plainTVSpecified Name
mName] (forall (m :: * -> *). Quote m => [m Type] -> m Cxt
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 <- forall (m :: * -> *). Quote m => String -> m Name
newName String
"this"
     let ms :: [NonEmpty Member]
ms = forall (f :: * -> *) a.
Foldable f =>
(a -> a -> Bool) -> f a -> [NonEmpty a]
NE.groupBy Member -> Member -> Bool
isNeighbor (StructRep -> [Member]
srMembers StructRep
rep)

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

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

     let name :: Name
name = StructRep -> Name
mkInitName StructRep
rep
         body :: ExpQ
body = forall (m :: * -> *). Quote m => [m Stmt] -> m Exp
doE
                -- allocate struct
              forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *). Quote m => m Pat -> m Exp -> m Stmt
bindS (forall (m :: * -> *). Quote m => Name -> m Pat
varP Name
this) (forall (m :: * -> *). Quote m => Name -> m Exp
varE (StructRep -> Name
mkAllocName StructRep
rep))

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

                -- return initialized struct
             forall a. [a] -> [a] -> [a]
++ [ forall (m :: * -> *). Quote m => m Exp -> m Stmt
noBindS [| return $(varE this) |] ]

     forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence
       [ forall (m :: * -> *). Quote m => Name -> m Type -> m Dec
sigD Name
name (StructRep -> Q Type
newStructType StructRep
rep)
       , forall (m :: * -> *). Quote m => Name -> [m Clause] -> m Dec
funD Name
name [ forall (m :: * -> *).
Quote m =>
[m Pat] -> m Body -> [m Dec] -> m Clause
clause (forall (m :: * -> *). Quote m => Name -> m Pat
varP forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap forall a. NonEmpty a -> [a]
NE.toList [NonEmpty (Name, Member)]
msWithArgs)
                            (forall (m :: * -> *). Quote m => m Exp -> m Body
normalB [| st $body |] ) [] ]
       ]


assignN :: ExpQ -> Int -> NonEmpty (Name,Member) -> ExpQ

assignN :: ExpQ -> Int -> NonEmpty (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 NonEmpty (Name, Member)
us =
  do let n :: Int
n = forall a. NonEmpty a -> Int
NE.length NonEmpty (Name, Member)
us
     Name
mba <- forall (m :: * -> *). Quote m => String -> m Name
newName String
"mba"
     let arg0 :: Name
arg0 = forall a b. (a, b) -> a
fst (forall a. NonEmpty a -> a
NE.head NonEmpty (Name, Member)
us)
     forall (m :: * -> *). Quote m => [m Stmt] -> m Exp
doE forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *). Quote m => m Pat -> m Exp -> m Stmt
bindS (forall (m :: * -> *). Quote m => Name -> m Pat
varP Name
mba) [| initializeUnboxedField i n (sizeOf $(varE arg0)) $this |]
         forall a. a -> [a] -> [a]
: [ forall (m :: * -> *). Quote m => m Exp -> m Stmt
noBindS [| writeByteArray $(varE mba) j $(varE arg) |]
           | (Int
j,(Name
arg,Member
_)) <- forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 :: Int ..] (forall a. NonEmpty a -> [a]
NE.toList NonEmpty (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 <- forall (m :: * -> *). Quote m => String -> m Name
newName String
"m"
     let m :: TypeQ
         m :: Q Type
m = forall (m :: * -> *). Quote m => Name -> m 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) = forall (m :: * -> *) a. Monad m => a -> m a
return Type
t
         buildType (Member Representation
UnboxedField Name
_ Type
t) = 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 = 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 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 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> 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 forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *).
Quote m =>
[TyVarBndr Specificity] -> m Cxt -> m Type -> m Type
forallT [Name -> TyVarBndr Specificity
plainTVSpecified Name
mName] (forall (m :: * -> *). Quote m => [m Type] -> m Cxt
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
  = forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$>
    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..]
      (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 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 =
  forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> 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) <- forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 :: Int ..] [Member]
us
    ]
  where
  addPrimCxt :: Type -> Q Type -> Q Type
addPrimCxt (VarT Name
t) = forall (m :: * -> *).
Quote m =>
[TyVarBndr Specificity] -> m Cxt -> m Type -> m Type
forallT [] (forall (m :: * -> *). Quote m => [m Type] -> m Cxt
cxt [Name -> Q Type
primPred Name
t])
  addPrimCxt 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 =
  forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence
    [ forall (m :: * -> *). Quote m => Name -> m Type -> m Dec
sigD Name
name (StructRep -> Q Type -> Q Type
forallRepT StructRep
rep Q Type
typ)
    , Name -> ExpQ -> Q Dec
simpleValD Name
name ExpQ
def
    , forall (m :: * -> *).
Quote m =>
Name -> Inline -> RuleMatch -> Phases -> m 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 = forall (m :: * -> *).
Quote m =>
m Pat -> m Body -> [m Dec] -> m Dec
valD (forall (m :: * -> *). Quote m => Name -> m Pat
varP Name
var) (forall (m :: * -> *). Quote m => m Exp -> m Body
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 = forall (m :: * -> *).
Quote m =>
[TyVarBndr Specificity] -> m Cxt -> m Type -> m Type
forallT (forall a. [a] -> [a]
init (forall newFlag oldFlag.
newFlag -> [TyVarBndr_ oldFlag] -> [TyVarBndr_ newFlag]
changeTVFlags Specificity
SpecifiedSpec (StructRep -> [TyVarBndrVis]
srTyVars StructRep
rep))) (forall (m :: * -> *). Quote m => [m Type] -> m Cxt
cxt [])

(-->) :: TypeQ -> TypeQ -> TypeQ
Q Type
f --> :: Q Type -> Q Type -> Q Type
--> Q Type
x = forall (m :: * -> *). Quote m => m Type
arrowT forall (m :: * -> *). Quote m => m Type -> m Type -> m Type
`appT` Q Type
f forall (m :: * -> *). Quote m => m Type -> m Type -> m 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 forall a. Eq a => a -> a -> Bool
== Name
m
occurs Name
n (ForallT [TyVarBndr Specificity]
_ 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

#if !MIN_VERSION_template_haskell(2,21,0) && !MIN_VERSION_th_abstraction(0,6,0)
type TyVarBndrVis = TyVarBndrUnit
#endif

-- 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)))")