large-records-0.2.1.0: src/Data/Record/Internal/Plugin/CodeGen.hs
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
-- | The core of the plugin implementation.
module Data.Record.Internal.Plugin.CodeGen (genLargeRecord) where
import Data.List (nubBy)
import Data.List.NonEmpty (NonEmpty(..))
import qualified Data.Generics as SYB
import Data.Record.Internal.GHC.Fresh
import Data.Record.Internal.GHC.Shim hiding (mkTyVar)
import Data.Record.Internal.GHC.TemplateHaskellStyle
import Data.Record.Internal.Plugin.Options
import Data.Record.Internal.Plugin.Record
import qualified Data.Record.Internal.Plugin.Names.GhcGenerics as GHC
import qualified Data.Record.Internal.Plugin.Names.Runtime as RT
{-------------------------------------------------------------------------------
Top-level
-------------------------------------------------------------------------------}
-- | Generate all large-records definitions for a record.
genLargeRecord :: MonadFresh m => Record -> m [LHsDecl GhcPs]
genLargeRecord r@Record{..} = concatM [
(:[]) <$> genDatatype r
, genVectorConversions r
, genIndexedAccessor r
, genUnsafeSetIndex r
, genStockInstances r
, mapM (genHasFieldInstance r) recordFields
, sequence [
genConstraintsClass r
, genConstraintsInstance r
, genGenericInstance r
, genGHCGeneric r
]
]
{-------------------------------------------------------------------------------
The type itself and conversion to and from vectors
NOTE: All generation exampleshask assume as example
> data T a b = MkT {
> tWord :: Word
> , tBool :: Bool
> , tChar :: Char
> , tA :: a
> , tListB :: [b]
> }
> deriving (Eq, Show)
-------------------------------------------------------------------------------}
-- | Generate the datatype that will represent the record
--
-- Currently this generates something like
--
-- > data T a b =
-- > forall f0 f1 f2 f3 f4. (
-- > f0 ~ Word
-- > , f1 ~ Bool
-- > , f2 ~ Char
-- > , f3 ~ a
-- > , f4 ~ [b]
-- > )
-- > => MkT {
-- > tInt :: f0
-- > , tBool :: f1
-- > , tChar :: f2
-- > , tA :: f3
-- > , tListB :: f4
-- > }
-- > deriving anyclass C -- where applicable
--
-- (possibly with strict fields). This representation accomplishes two things:
--
-- 1. The use of the existentials with type equalities prevents ghc from
-- generating field accessors.
-- 2. It can still be used in the normal way to construct record values and
-- to pattern match on records.
--
-- TODO: From ghc 9.2 and up, we should generate
--
-- > newtype T a b = TFromVector {vectorFromT :: SmallArray Any}
-- > deriving anyclass C -- where applicable
--
-- instead, along with a pattern synonym.
genDatatype :: MonadFresh m => Record -> m (LHsDecl GhcPs)
genDatatype Record{..} = pure $
DataD
recordTyName
recordTyVars
[ forallRecC
vars
(zipWith fieldContext vars recordFields)
recordConName
(zipWith fieldExistentialType vars recordFields)
]
[ DerivClause (Just (noLoc AnyclassStrategy)) (c :| [])
| DeriveAnyClass c <- recordDerivings
]
where
-- There is no need to generate fresh va riables here, as these type vars
-- cannot clash with anything else (no other type vars can be in scope).
vars :: [LRdrName]
vars = [
mkTyVar recordAnnLoc ("lr_f" <> show i)
| (i, _) <- zip [1 :: Int ..] recordFields
]
optionalBang :: LHsType GhcPs -> LHsType GhcPs
optionalBang = if allFieldsStrict recordOptions then bangType else id
fieldContext :: LRdrName -> Field -> LHsType GhcPs
fieldContext var fld = equalP (VarT var) (fieldType fld)
fieldExistentialType :: LRdrName -> Field -> (LRdrName, LHsType GhcPs)
fieldExistentialType var fld = (fieldName fld, optionalBang $ VarT var)
-- | Generate conversion to and from an array
--
-- Generates something like
--
-- > vectorFromT :: T a b -> SmallArray Any
-- > vectorFromT = \x ->
-- > case x of
-- > MkT f0 f1 f2 f3 f4 -> smallArrayFromList [
-- > unsafeCoerce f0
-- > , unsafeCoerce f1
-- > , unsafeCoerce f2
-- > , unsafeCoerce f3
-- > , unsafeCoerce f4
-- > ]
-- >
-- > vectorToT :: SmallArray Any -> T a b
-- > vectorToT = \x ->
-- > case smallArrayToList x of
-- > [f0, f1, f2, f3, f4] ->
-- > MkT (unsafeCoerce f0)
-- > (unsafeCoerce f1)
-- > (unsafeCoerce f2)
-- > (unsafeCoerce f3)
-- > (unsafeCoerce f4)
-- > _ -> error "Pattern match failure in vectorToT: vector with invalid number of elements."
--
-- TODO: From ghc 9.2, these could be identity functions. See 'genDatatype'
-- for details.
genVectorConversions :: forall m. MonadFresh m => Record -> m [LHsDecl GhcPs]
genVectorConversions r@Record{..} = concatM [
fromVector
, toVector
]
where
fromVector :: m [LHsDecl GhcPs]
fromVector = do
args <- mapM (freshName . fieldName) recordFields
return [
sigD name $
funT
(recordTypeT r)
(ConT RT.type_SmallArray `appT` ConT RT.type_Any)
, valD name $
lamE1 (conP recordConName (map varP args)) $
appE
(VarE RT.smallArrayFromList)
(listE [ VarE RT.unsafeCoerce `appE` VarE arg
| arg <- args
]
)
]
where
name :: LRdrName
name = nameVectorFrom r
toVector :: m [LHsDecl GhcPs]
toVector = do
x <- freshName $ mkExpVar recordAnnLoc "x"
args <- mapM (freshName . fieldName) recordFields
return $ [
sigD name $
funT
(ConT RT.type_SmallArray `appT` ConT RT.type_Any)
(recordTypeT r)
, valD name $
lamE1 (varP x) $
caseE
(VarE RT.smallArrayToList `appE` VarE x)
[ ( listP (map varP args)
, appsE
(ConE recordConName)
[ VarE RT.unsafeCoerce `appE` VarE arg
| arg <- args
]
)
, ( wildP
, VarE RT.error `appE` stringE matchErr
)
]
]
where
name :: LRdrName
name = nameVectorTo r
matchErr :: String
matchErr = concat [
"Pattern match failure in "
, nameBase (nameVectorTo r)
, ": vector with invalid number of elements."
]
{-------------------------------------------------------------------------------
Field accessors and 'HasField' instance
TODO: If we had support within GHC itself for accessing fields in records,
we might be able to integrate this a lot more closely with normal GHC,
especially when combined with the @NoFieldSelectors@ extension.
See <https://gitlab.haskell.org/ghc/ghc/-/issues/17991>
-------------------------------------------------------------------------------}
-- | Generate the indexed field accessor
--
-- Generates something like
--
-- > unsafeGetIndexT :: forall x a b. Int -> T a b -> x
-- > unsafeGetIndexT = \ n t -> noInlineUnsafeCo (V.unsafeIndex (vectorFromT t) n)
genIndexedAccessor :: MonadFresh m => Record -> m [LHsDecl GhcPs]
genIndexedAccessor r@Record{..} = do
x <- freshName $ mkTyVar recordAnnLoc "x"
n <- freshName $ mkExpVar recordAnnLoc "n"
t <- freshName $ mkExpVar recordAnnLoc "t"
return [
sigD name $
funT
(ConT RT.type_Int)
(recordTypeT r `funT` VarT x)
, valD name $
lamE (varP n :| [varP t]) $
appE
(VarE RT.noInlineUnsafeCo)
(appsE
(VarE RT.indexSmallArray)
[ VarE (nameVectorFrom r) `appE` VarE t
, VarE n
]
)
]
where
name :: LRdrName
name = nameUnsafeGetIndex r
-- | Generate index field overwrite
--
-- Generates something like
--
-- > unsafeSetIndexT :: forall x a b. Int -> T a b -> x -> T a b
-- > unsafeSetIndexT = \n t val ->
-- > vectorToT (V.unsafeUpd (vectorFromT t) [(n, noInlineUnsafeCo val)])
--
-- NOTE: Like 'genTo', this function used to be more complicated, because it
-- would need to take the strictness of the fields into account. If we change
-- our internal representation, we might need to be more careful with that
-- again. See 'genTo' for further discussion.
genUnsafeSetIndex :: MonadFresh m => Record -> m [LHsDecl GhcPs]
genUnsafeSetIndex r@Record{..} = do
x <- freshName $ mkTyVar recordAnnLoc "x"
n <- freshName $ mkExpVar recordAnnLoc "n"
t <- freshName $ mkExpVar recordAnnLoc "t"
val <- freshName $ mkExpVar recordAnnLoc "val"
return [
sigD name $
ConT RT.type_Int
`funT` (recordTypeT r `funT` (VarT x `funT` recordTypeT r))
, valD name $
lamE (varP n :| [varP t, (varP val)]) $
appE
(VarE (nameVectorTo r))
(appsE
(VarE RT.updateSmallArray)
[ VarE (nameVectorFrom r) `appE` VarE t
, listE [
tupE $
VarE n
:| [VarE RT.noInlineUnsafeCo `appE` VarE val]
]
]
)
]
where
name :: LRdrName
name = nameUnsafeSetIndex r
-- | Generate 'HasField' instance for single field
--
-- Generates something like
--
-- > instance x ~ Word => HasField "tInt" (T a b) x where
-- > hasField = \t -> (unsafeSetIndexT 0 t, unsafeGetIndexT 0 t)
genHasFieldInstance :: MonadFresh m => Record -> Field -> m (LHsDecl GhcPs)
genHasFieldInstance r@Record{..} Field{..} = do
x <- freshName $ mkTyVar recordAnnLoc "x"
t <- freshName $ mkExpVar recordAnnLoc "t"
return $
instanceD
[equalP (VarT x) fieldType]
(appsT
(ConT RT.type_HasField)
[ stringT (nameBase fieldName)
, recordTypeT r
, VarT x
]
)
[ ( RT.unq_hasField
, lamE1 (varP t) $
tupE $
appsE (VarE (nameUnsafeSetIndex r)) [intE fieldIndex, VarE t]
:| [appsE (VarE (nameUnsafeGetIndex r)) [intE fieldIndex, VarE t]]
)
]
[]
{-------------------------------------------------------------------------------
Generics
-------------------------------------------------------------------------------}
-- | Generate the class we will use to instantiate 'Constraints'
--
-- Generates something like this:
--
-- > class Constraints_T a b (c :: Type -> Constraint) where
-- > dictConstraints_T :: Proxy c -> Rep (Dict c) (T a b)
--
-- NOTE: It is critical that we don't give the class any superclass constraints
-- like
--
-- > class (c Word, c Bool, c Char, c a, c [b])
-- > => Constraints_T a b (c :: Type -> Constraint)
--
-- because then @ghc@ would use resolve @Constraints_T@ to that tuple instead,
-- and use lots of "tuple constraint extractor" functions, each of which have
-- the same size as the number of constraints (another example of a
-- @case f of { T x1 x2 x3 .. -> xn@ function, but now at the dictionary level).
genConstraintsClass :: MonadFresh m => Record -> m (LHsDecl GhcPs)
genConstraintsClass r@Record{..} = do
c <- freshName $ mkTyVar recordAnnLoc "c"
return $ classD
[]
(nameConstraints r)
(recordTyVars ++ [kindedTV c cKind])
[ ( nameDictConstraints r
, funT
(ConT RT.type_Proxy `appT` VarT c)
(appsT
(ConT RT.type_Rep)
[ ConT RT.type_Dict `appT` VarT c
, recordTypeT r
]
)
)
]
where
cKind :: LHsType GhcPs
cKind = ConT RT.type_Type `funT` ConT RT.type_Constraint
-- | Superclass constraints required by the constraints class instance
--
-- Generates something like
--
-- > (c Word, c Bool, c Char, c a, c [b])
--
-- However, we filter out constraints that are type variable free, so if we
-- pass, say, @Show@ for @c@, then we generate
--
-- > (Show a, Show [b])
--
-- instead. This avoids @ghc@ complaining about
--
-- > Redundant constraints: (Show Word, Show Bool, Show Char)
genRequiredConstraints :: Record -> LHsType GhcPs -> [LHsType GhcPs]
genRequiredConstraints Record{..} c =
nubBy sameType $ filter hasTypeVar $ map constrainField recordFields
where
constrainField :: Field -> LHsType GhcPs
constrainField Field{..} = c `appT` fieldType
sameType :: LHsType GhcPs -> LHsType GhcPs -> Bool
sameType = compareHs
hasTypeVar :: LHsType GhcPs -> Bool
hasTypeVar = not . null . allTyVars
allTyVars :: LHsType GhcPs -> [String]
allTyVars = SYB.everything (++) (SYB.mkQ [] isTypeVar)
isTypeVar :: LHsType GhcPs -> [String]
isTypeVar (VarT (TyVar name)) = [name]
isTypeVar _otherwise = []
-- | Generate the dictionary creation function ('dict')
--
-- Generates something like
--
-- > \p -> Rep (V.fromList [
-- > noInlineUnsafeCo (dictFor p (Proxy :: Proxy Word))
-- > , noInlineUnsafeCo (dictFor p (Proxy :: Proxy Bool))
-- > , noInlineUnsafeCo (dictFor p (Proxy :: Proxy Char))
-- > , noInlineUnsafeCo (dictFor p (Proxy :: Proxy a))
-- > , noInlineUnsafeCo (dictFor p (Proxy :: Proxy [b]))
-- > ])
genDict :: MonadFresh m => Record -> m (LHsExpr GhcPs)
genDict Record{..} = do
p <- freshName $ mkExpVar recordAnnLoc "p"
return $
lamE1 (varP p) $
appE
(ConE RT.con_Rep)
(appE
(VarE RT.smallArrayFromList)
(listE (map (dictForField p) recordFields))
)
where
dictForField :: LRdrName -> Field -> LHsExpr GhcPs
dictForField p Field{..} =
appE
(VarE RT.noInlineUnsafeCo)
(VarE RT.dictFor `appsE` [VarE p, proxyE fieldType])
-- | Generate (one and only) instance of the constraints class
--
-- Generates something like
--
-- > instance (..) => Constraints_T a b c where
-- > dictConstraints_T = ..
--
-- where the body of @dictConstraints_T@ is generated by 'genDict'.
genConstraintsInstance :: MonadFresh m => Record -> m (LHsDecl GhcPs)
genConstraintsInstance r@Record{..} = do
body <- genDict r
c <- freshName $ mkTyVar recordAnnLoc "c"
return $
instanceD
(genRequiredConstraints r (VarT c))
(appsT
(ConT (nameConstraints r))
([VarT (tyVarBndrName v) | v <- recordTyVars] ++ [VarT c]))
[(nameDictConstraints r, body)]
[]
-- | Generate metadata
--
-- Generates something like
--
-- > \_p -> Metadata {
-- > recordName = "T"
-- > , recordConstructor = "MkT"
-- > , recordSize = 5
-- > , recordFieldMetadata = Rep $ V.fromList [
-- > FieldMetadata (Proxy :: Proxy "tInt")) FieldLazy
-- > , FieldMetadata (Proxy :: Proxy "tBool")) FieldLazy
-- > , FieldMetadata (Proxy :: Proxy "tChar")) FieldLazy
-- > , FieldMetadata (Proxy :: Proxy "tA")) FieldLazy
-- > , FieldMetadata (Proxy :: Proxy "tListB")) FieldLazy
-- > ]
-- > }
genMetadata :: MonadFresh m => Record -> m (LHsExpr GhcPs)
genMetadata r@Record{..} = do
p <- freshName $ mkExpVar recordAnnLoc "p"
return $
lamE1 (varP p) $
recConE
RT.con_Metadata [
( RT.recordName
, stringE (nameRecord r)
)
, ( RT.recordConstructor
, stringE (nameBase recordConName)
)
, ( RT.recordSize
, intE (length recordFields)
)
, ( RT.recordFieldMetadata
, appE
(ConE RT.con_Rep)
(appE
(VarE RT.smallArrayFromList)
(listE (map auxField recordFields))
)
)
]
where
auxField :: Field -> LHsExpr GhcPs
auxField Field{..} =
appsE
(ConE RT.con_FieldMetadata)
[ proxyE (stringT (nameBase fieldName))
, ConE $ if allFieldsStrict recordOptions
then RT.con_FieldStrict
else RT.con_FieldLazy
]
-- | Generate definition for `from` in the `Generic` instance
--
-- Generates something like
--
-- > repFromVectorStrict . vectorFromT
genFrom :: MonadFresh m => Record -> m (LHsExpr GhcPs)
genFrom r@Record{..} = do
x <- freshName $ mkExpVar recordAnnLoc "x"
return $
lamE1 (varP x) $
VarE RT.repFromVector `appE` (VarE (nameVectorFrom r) `appE` VarE x)
-- | Generate definition for `to` in the `Generic` instance
--
-- > vectorToT . repToVector
--
-- NOTE: This function used to be more complicated. When the internal
-- representation of a record /is/ a vector, then we have to be very careful
-- with the strictness of the fields here. However, since we currently use a
-- " normal " record as our internal representation (albeit with strange types),
-- and the fields of that record have their own strictness annotation, we don't
-- have to worry about strictness here.
genTo :: MonadFresh m => Record -> m (LHsExpr GhcPs)
genTo r@Record{..} = do
x <- freshName $ mkExpVar recordAnnLoc "x"
return $
lamE1 (varP x) $
VarE (nameVectorTo r) `appE` (VarE RT.repToVector `appE` VarE x)
-- | Generate an instance of large-records 'Data.Record.Generic'.
--
-- In the sample instance below, @vectorFromT@ and @vectorToT@ are generated
-- per-record by 'genVectorFrom' and 'genVectorTo' respectively.
--
-- > instance Generic (T ...) where
-- > type Constraints (T ...) = Constraints_T ...
-- > type MetadataOf (T ...) = '[ '("field1", fieldType1), ... ]
-- >
-- > from = ..
-- > to = ..
-- > dict = dictConstraints_T
-- > metadata = ..
genGenericInstance :: MonadFresh m => Record -> m (LHsDecl GhcPs)
genGenericInstance r@Record{..} = do
metadata <- genMetadata r
from <- genFrom r
to <- genTo r
return $
instanceD
[]
(ConT RT.type_Generic `appT` recordTypeT r)
[ ( RT.unq_from , from )
, ( RT.unq_to , to )
, ( RT.unq_dict , VarE (nameDictConstraints r) )
, ( RT.unq_metadata , metadata )
]
[ tySynEqn RT.unq_type_Constraints [recordTypeT r] $
appsT
(ConT (nameConstraints r))
[VarT (tyVarBndrName v) | v <- recordTyVars]
, tySynEqn RT.unq_type_MetadataOf [recordTypeT r] $
listT [
tupT $ stringT (nameBase fieldName) :| [fieldType]
| Field{..} <- recordFields
]
]
where
{-------------------------------------------------------------------------------
"Stock" instances
-------------------------------------------------------------------------------}
-- | Generate stock instances
genStockInstances :: MonadFresh m => Record -> m [LHsDecl GhcPs]
genStockInstances r@Record{..} = concatM [
genStockInstance r d
| DeriveStock d <- recordDerivings
]
-- | For a record type @T@ and stock-derivable class @C@, generate
--
-- > instance $(genRequiredConstraints T C) => C T where
-- > $(method) = $(generic implementation)
--
-- NOTE: All of these instances depend on the 'Data.Record.Generics.Generics'
-- instance.
--
-- TODO: For 'Generic' we currently don't do anything. We could change this so
-- that we generate the 'GHC.Generics' instance only when the user asks for a
-- 'Generics' instance?
genStockInstance :: MonadFresh m => Record -> StockDeriving -> m [LHsDecl GhcPs]
genStockInstance r = pure . \case
Show -> [mkInstance RT.type_Show RT.unq_showsPrec RT.gshowsPrec]
Eq -> [mkInstance RT.type_Eq RT.unq_eq RT.geq ]
Ord -> [mkInstance RT.type_Ord RT.unq_compare RT.gcompare ]
Generic -> []
where
mkInstance :: LRdrName -> LRdrName -> LRdrName -> LHsDecl GhcPs
mkInstance cls mthd gen =
instanceD
(genRequiredConstraints r (ConT cls))
(ConT cls `appT` recordTypeT r)
[(mthd, VarE gen)]
[]
{-------------------------------------------------------------------------------
GHC generics
-------------------------------------------------------------------------------}
-- | Generate GHC generics instance
--
-- Generates something like
--
-- > instance GHC.Generic ExampleRecord where
-- > type Rep ExampleRecord = ThroughLRGenerics ExampleRecord
-- >
-- > from = WrapThroughLRGenerics
-- > to = unwrapThroughLRGenerics
--
-- See 'ThroughLRGenerics' for documentation.
genGHCGeneric :: MonadFresh m => Record -> m (LHsDecl GhcPs)
genGHCGeneric r = pure $
instanceD
[]
(ConT GHC.type_Generic `appT` recordTypeT r)
[ ( GHC.unq_from , ConE RT.con_WrapThroughLRGenerics )
, ( GHC.unq_to , VarE RT.unwrapThroughLRGenerics )
]
[ tySynEqn GHC.unq_type_Rep [recordTypeT r] $
ConT RT.type_ThroughLRGenerics `appT` recordTypeT r
]
{-------------------------------------------------------------------------------
Auxiliary functions for dealing with records
-------------------------------------------------------------------------------}
-- | The saturated type of the record (that is, with all type vars applied)
recordTypeT :: Record -> LHsType GhcPs
recordTypeT Record{..} =
ConT recordTyName `appsT` [VarT (tyVarBndrName f) | f <- recordTyVars]
{-------------------------------------------------------------------------------
Pick names for generated code
-------------------------------------------------------------------------------}
nameRecord :: Record -> String
nameRecord Record{..} = nameBase recordTyName
-- | Make name derived from the name of the record
mkDerived :: (SrcSpan -> String -> LRdrName) -> String -> Record -> LRdrName
mkDerived f prefix r = f (recordAnnLoc r) (prefix <> nameRecord r)
nameVectorFrom :: Record -> LRdrName
nameVectorTo :: Record -> LRdrName
nameUnsafeGetIndex :: Record -> LRdrName
nameUnsafeSetIndex :: Record -> LRdrName
nameConstraints :: Record -> LRdrName
nameDictConstraints :: Record -> LRdrName
nameVectorFrom = mkDerived mkExpVar "vectorFrom"
nameVectorTo = mkDerived mkExpVar "vectorTo"
nameUnsafeGetIndex = mkDerived mkExpVar "unsafeGetIndex"
nameUnsafeSetIndex = mkDerived mkExpVar "unsafeSetIndex"
nameConstraints = mkDerived mkTyCon "Constraints_"
nameDictConstraints = mkDerived mkExpVar "dictConstraints_"
{-------------------------------------------------------------------------------
Auxiliary
-------------------------------------------------------------------------------}
-- | Generate a Proxy expression for the given type.
--
-- @proxyE [t|ty|]@ will result in a @Proxy :: Proxy ty@.
proxyE :: LHsType GhcPs -> LHsExpr GhcPs
proxyE ty = sigE (ConE RT.con_Proxy) (ConT RT.type_Proxy `appT` ty)
concatM :: Applicative m => [m [a]] -> m [a]
concatM = fmap concat . sequenceA