git.iko.soy/streamly
on 5d096df help sign in
streamly/core/src/Streamly/Internal/Data/Serialize/TH.hs
530 lines 19 kB history raw
1{-# LANGUAGE TemplateHaskell #-}
3-- |
4-- Module : Streamly.Internal.Data.Serialize.TH
5-- Copyright : (c) 2023 Composewell Technologies
6-- License : BSD3-3-Clause
7-- Maintainer : streamly@composewell.com
8-- Stability : experimental
9-- Portability : GHC
11-- XXX Instead of applying the product constructor in one go can we apply it
12-- one at a time in case of too many arguments, compiler may not have to save
13-- them in local vars.
15module Streamly.Internal.Data.Serialize.TH
16 (
17 -- Deriving
18 deriveSerialize
19 , deriveSerializeWith
21 -- Utilities
22 , module Streamly.Internal.Data.Serialize.TH.Bottom
23 -- ** Common
24 , module Streamly.Internal.Data.Serialize.TH.Common
25 -- ** RecHeader
26 , module Streamly.Internal.Data.Serialize.TH.RecHeader
27 ) where
29--------------------------------------------------------------------------------
30-- Imports
31--------------------------------------------------------------------------------
33import Data.List (foldl')
34import Data.Word (Word16, Word32, Word64, Word8)
36import Language.Haskell.TH
37import Language.Haskell.TH.Syntax
38import Streamly.Internal.Data.Serialize.Type
40import Streamly.Internal.Data.Unbox.TH
41 ( DataCon(..)
42 , DataType(..)
43 , reifyDataType
44 )
46import qualified Streamly.Internal.Data.Serialize.TH.RecHeader as RecHeader
48import Streamly.Internal.Data.Serialize.TH.Bottom
49import Streamly.Internal.Data.Serialize.TH.Common
50import Streamly.Internal.Data.Serialize.TH.RecHeader
52--------------------------------------------------------------------------------
53-- Domain specific helpers
54--------------------------------------------------------------------------------
56exprGetSize :: Q Exp -> (Int, Type) -> Q Exp
57exprGetSize acc (i, _) = [|addSizeTo $(acc) $(varE (mkFieldName i))|]
59getTagSize :: Int -> Int
60getTagSize numConstructors
61 | numConstructors == 1 = 0
62 | fromIntegral (maxBound :: Word8) >= numConstructors = 1
63 | fromIntegral (maxBound :: Word16) >= numConstructors = 2
64 | fromIntegral (maxBound :: Word32) >= numConstructors = 4
65 | fromIntegral (maxBound :: Word64) >= numConstructors = 8
66 | otherwise = error "Too many constructors"
68getTagType :: Int -> Name
69getTagType numConstructors
70 | numConstructors == 1 = error "No tag for 1 constructor"
71 | fromIntegral (maxBound :: Word8) >= numConstructors = ''Word8
72 | fromIntegral (maxBound :: Word16) >= numConstructors = ''Word16
73 | fromIntegral (maxBound :: Word32) >= numConstructors = ''Word32
74 | fromIntegral (maxBound :: Word64) >= numConstructors = ''Word64
75 | otherwise = error "Too many constructors"
77--------------------------------------------------------------------------------
78-- Size
79--------------------------------------------------------------------------------
81getNameBaseLen :: Name -> Word8
82getNameBaseLen cname =
83 let x = length (nameBase cname)
84 in if x > 63
85 then error "Max Constructor Len: 63 characters"
86 else fromIntegral x
88conEncLen :: Name -> Word8
89conEncLen cname = getNameBaseLen cname + 1
91--------------------------------------------------------------------------------
92-- Size
93--------------------------------------------------------------------------------
95mkSizeOfExpr :: Bool -> Bool -> TypeOfType -> Q Exp
96mkSizeOfExpr True False tyOfTy =
97 case tyOfTy of
98 UnitType cname ->
99 lamE
100 [varP _acc, wildP]
101 [|$(varE _acc) + $(litIntegral (conEncLen cname))|]
102 TheType con ->
103 lamE
104 [varP _acc, varP _x]
105 (caseE (varE _x) [matchCons (varE _acc) con])
106 MultiType constructors -> sizeOfHeadDt constructors
108 where
110 sizeOfFields acc fields =
111 foldl' exprGetSize acc $ zip [0..] fields
113 matchCons acc (SimpleDataCon cname fields) =
114 let a = litIntegral (conEncLen cname)
115 b = sizeOfFields acc (map snd fields)
116 expr = [|$(a) + $(b)|]
117 in matchConstructor cname (length fields) expr
119 sizeOfHeadDt cons =
120 let acc = [|$(varE _acc)|]
121 in lamE
122 [varP _acc, varP _x]
123 (caseE (varE _x) (fmap (matchCons acc) cons))
125mkSizeOfExpr False False tyOfTy =
126 case tyOfTy of
127 UnitType _ -> lamE [varP _acc, wildP] [|$(varE _acc) + 1|]
128 TheType con ->
129 lamE
130 [varP _acc, varP _x]
131 (caseE (varE _x) [matchCons (varE _acc) con])
132 MultiType constructors -> sizeOfHeadDt constructors
134 where
136 tagSizeExp numConstructors =
137 litE (IntegerL (fromIntegral (getTagSize numConstructors)))
139 -- XXX fields of the same type can be folded together, will reduce the code
140 -- size when there are many fields of the same type.
141 -- XXX const size fields can be calculated statically.
142 -- XXX This can result in large compilation times due to nesting when there
143 -- are many constructors. We can create a list and sum the list at run time
144 -- to avoid that depending on the number of constructors. Or using a let
145 -- statement for each case may help?
146 -- appE (varE 'sum) (listE (acc : map (exprGetSize (litE (IntegerL 0))) (zip [0..] fields)))
147 sizeOfFields acc fields =
148 foldl' exprGetSize acc $ zip [0..] fields
150 matchCons acc (SimpleDataCon cname fields) =
151 let expr = sizeOfFields acc (map snd fields)
152 in matchConstructor cname (length fields) expr
154 -- XXX We fix VarSize for simplicity. Should be changed later.
155 sizeOfHeadDt cons =
156 let numCons = length cons
157 acc = [|$(varE _acc) + $(tagSizeExp numCons)|]
158 in lamE
159 [varP _acc, varP _x]
160 (caseE (varE _x) (fmap (matchCons acc) cons))
162mkSizeOfExpr False True (TheType con) = RecHeader.mkRecSizeOfExpr con
164mkSizeOfExpr _ _ _ = errorUnimplemented
166mkSizeDec :: SerializeConfig -> Type -> [DataCon] -> Q [Dec]
167mkSizeDec (SerializeConfig {..}) headTy cons = do
168 -- INLINE on sizeOf actually worsens some benchmarks, and improves none
169 sizeOfMethod <-
170 mkSizeOfExpr
171 cfgConstructorTagAsString
172 cfgRecordSyntaxWithHeader
173 (typeOfType headTy cons)
174 pure
175 ( maybe
176 []
177 (\x -> [PragmaD (InlineP 'addSizeTo x FunLike AllPhases)])
178 cfgInlineSize
179 ++ [FunD 'addSizeTo [Clause [] (NormalB sizeOfMethod) []]]
180 )
182--------------------------------------------------------------------------------
183-- Peek
184--------------------------------------------------------------------------------
186mkDeserializeExpr :: Bool -> Bool -> Type -> TypeOfType -> Q Exp
187mkDeserializeExpr True False headTy tyOfTy =
188 case tyOfTy of
189 UnitType cname -> deserializeConsExpr [SimpleDataCon cname []]
190 TheType con -> deserializeConsExpr [con]
191 MultiType cons -> deserializeConsExpr cons
193 where
195 deserializeConsExpr cons = do
196 conLen <- newName "conLen"
197 off1 <- newName "off1"
198 [|do ($(varP off1), $(varP conLen) :: Word8) <-
199 deserializeAt
200 $(varE _initialOffset)
201 $(varE _arr)
202 $(varE _endOffset)
203 $(multiIfE (map (guardCon conLen off1) cons ++ [catchAll]))|]
205 catchAll =
206 normalGE
207 [|True|]
208 [|error
209 ("Found invalid tag while peeking (" ++
210 $(lift (pprint headTy)) ++ ")")|]
212 guardCon conLen off con@(SimpleDataCon cname _) = do
213 let lenCname = getNameBaseLen cname
214 tag = map c2w (nameBase cname)
215 normalGE
216 [|($(litIntegral lenCname) == $(varE conLen))
217 && $(xorCmp tag off _arr)|]
218 [|let $(varP (makeI 0)) = $(varE off) + $(litIntegral lenCname)
219 in $(mkDeserializeExprOne 'deserializeAt con)|]
221mkDeserializeExpr False False headTy tyOfTy =
222 case tyOfTy of
223 -- Unit constructor
224 -- XXX Should we peek and check if the byte value is 0?
225 UnitType cname ->
226 [|pure ($(varE _initialOffset) + 1, $(conE cname))|]
227 -- Product type
228 TheType con ->
229 letE
230 [valD (varP (mkName "i0")) (normalB (varE _initialOffset)) []]
231 (mkDeserializeExprOne 'deserializeAt con)
232 -- Sum type
233 MultiType cons -> do
234 let lenCons = length cons
235 tagType = getTagType lenCons
236 doE
237 [ bindS
238 (tupP [varP (mkName "i0"), varP _tag])
239 [|deserializeAt $(varE _initialOffset) $(varE _arr) $(varE _endOffset)|]
240 , noBindS
241 (caseE
242 (sigE (varE _tag) (conT tagType))
243 (map peekMatch (zip [0 ..] cons) ++ [peekErr]))
244 ]
245 where
246 peekMatch (i, con) =
247 match
248 (litP (IntegerL i))
249 (normalB (mkDeserializeExprOne 'deserializeAt con)) []
250 peekErr =
251 match
252 wildP
253 (normalB
254 -- XXX Print the tag
255 [|error
256 ("Found invalid tag while peeking (" ++
257 $(lift (pprint headTy)) ++ ")")|])
258 []
260mkDeserializeExpr False True _ (TheType con@(SimpleDataCon _ fields)) = do
261 deserializeWithKeys <- newName "deserializeWithKeys"
262 updateFunc <- newName "updateFunc"
263 updateFuncDec <- RecHeader.conUpdateFuncDec updateFunc fields
264 deserializeWithKeysDec <-
265 RecHeader.mkDeserializeKeysDec deserializeWithKeys updateFunc con
266 letE
267 (pure <$> (deserializeWithKeysDec ++ updateFuncDec))
268 (RecHeader.mkRecDeserializeExpr
269 _initialOffset
270 _endOffset
271 deserializeWithKeys
272 con)
274mkDeserializeExpr _ _ _ _ = errorUnimplemented
276mkDeserializeDec :: SerializeConfig -> Type -> [DataCon] -> Q [Dec]
277mkDeserializeDec (SerializeConfig {..}) headTy cons = do
278 peekMethod <-
279 mkDeserializeExpr
280 cfgConstructorTagAsString
281 cfgRecordSyntaxWithHeader
282 headTy
283 (typeOfType headTy cons)
284 pure
285 ( maybe
286 []
287 (\x -> [PragmaD (InlineP 'deserializeAt x FunLike AllPhases)])
288 cfgInlineDeserialize
289 ++
290 [ FunD
291 'deserializeAt
292 [ Clause
293 (if isUnitType cons && not cfgConstructorTagAsString
294 then [VarP _initialOffset, WildP, WildP]
295 else [VarP _initialOffset, VarP _arr, VarP _endOffset])
296 (NormalB peekMethod)
297 []
298 ]
299 ]
300 )
302--------------------------------------------------------------------------------
303-- Poke
304--------------------------------------------------------------------------------
306mkSerializeExprTag :: Name -> Int -> Q Exp
307mkSerializeExprTag tagType tagVal =
308 [|serializeAt
309 $(varE _initialOffset)
310 $(varE _arr)
311 $((sigE (litE (IntegerL (fromIntegral tagVal))) (conT tagType)))|]
313mkSerializeExpr :: Bool -> Bool -> TypeOfType -> Q Exp
314mkSerializeExpr True False tyOfTy =
315 case tyOfTy of
316 -- Unit type
317 UnitType cname ->
318 caseE
319 (varE _val)
320 [serializeDataCon (SimpleDataCon cname [])]
321 -- Product type
322 (TheType con) ->
323 caseE
324 (varE _val)
325 [serializeDataCon con]
326 -- Sum type
327 (MultiType cons) ->
328 caseE
329 (varE _val)
330 (map serializeDataCon cons)
332 where
334 serializeDataCon (SimpleDataCon cname fields) = do
335 let tagLen8 = getNameBaseLen cname
336 conEnc = tagLen8 : map c2w (nameBase cname)
337 matchConstructor
338 cname
339 (length fields)
340 (doE [ bindS
341 (varP (mkName "i0"))
342 (serializeW8List _initialOffset _arr conEnc)
343 , noBindS (mkSerializeExprFields 'serializeAt fields)
344 ])
346mkSerializeExpr False False tyOfTy =
347 case tyOfTy of
348 -- Unit type
349 UnitType _ ->
350 [|serializeAt $(varE _initialOffset) $(varE _arr) (0 :: Word8)|]
351 -- Product type
352 (TheType (SimpleDataCon cname fields)) ->
353 letE
354 [valD (varP (mkName "i0")) (normalB (varE _initialOffset)) []]
355 (caseE
356 (varE _val)
357 [ matchConstructor
358 cname
359 (length fields)
360 (mkSerializeExprFields 'serializeAt fields)
361 ])
362 -- Sum type
363 (MultiType cons) -> do
364 let lenCons = length cons
365 tagType = getTagType lenCons
366 caseE
367 (varE _val)
368 (map (\(tagVal, (SimpleDataCon cname fields)) ->
369 matchConstructor
370 cname
371 (length fields)
372 (doE [ bindS
373 (varP (mkName "i0"))
374 (mkSerializeExprTag tagType tagVal)
375 , noBindS
376 (mkSerializeExprFields
377 'serializeAt
378 fields)
379 ]))
380 (zip [0 ..] cons))
382mkSerializeExpr False True (TheType con) =
383 RecHeader.mkRecSerializeExpr _initialOffset con
385mkSerializeExpr _ _ _ = errorUnimplemented
387mkSerializeDec :: SerializeConfig -> Type -> [DataCon] -> Q [Dec]
388mkSerializeDec (SerializeConfig {..}) headTy cons = do
389 pokeMethod <-
390 mkSerializeExpr
391 cfgConstructorTagAsString
392 cfgRecordSyntaxWithHeader
393 (typeOfType headTy cons)
394 pure
395 ( maybe
396 []
397 (\x -> [PragmaD (InlineP 'serializeAt x FunLike AllPhases)])
398 cfgInlineSerialize
399 ++
400 [FunD
401 'serializeAt
402 [ Clause
403 (if isUnitType cons && not cfgConstructorTagAsString
404 then [VarP _initialOffset, VarP _arr, WildP]
405 else [VarP _initialOffset, VarP _arr, VarP _val])
406 (NormalB pokeMethod)
407 []
408 ]
409 ]
410 )
412--------------------------------------------------------------------------------
413-- Main
414--------------------------------------------------------------------------------
416-- | A general function to derive Serialize instances where you can control
417-- which Constructors of the datatype to consider and what the Context for the
418-- 'Serialize' instance would be.
419--
420-- Consider the datatype:
421-- @
422-- data CustomDataType a b
423-- = CDTConstructor1
424-- | CDTConstructor2 Bool
425-- | CDTConstructor3 Bool b
426-- deriving (Show, Eq)
427-- @
428--
429-- Usage:
430-- @
431-- $(deriveSerializeInternal
432-- serializeConfig
433-- [AppT (ConT ''Serialize) (VarT (mkName "b"))]
434-- (AppT
435-- (AppT (ConT ''CustomDataType) (VarT (mkName "a")))
436-- (VarT (mkName "b")))
437-- [ DataCon 'CDTConstructor1 [] [] []
438-- , DataCon 'CDTConstructor2 [] [] [(Nothing, (ConT ''Bool))]
439-- , DataCon
440-- 'CDTConstructor3
441-- []
442-- []
443-- [(Nothing, (ConT ''Bool)), (Nothing, (VarT (mkName "b")))]
444-- ])
445-- @
446deriveSerializeInternal ::
447 SerializeConfig -> Type -> [DataCon] -> ([Dec] -> Q [Dec]) -> Q [Dec]
448deriveSerializeInternal conf headTy cons next = do
449 sizeDec <- mkSizeDec conf headTy cons
450 peekDec <- mkDeserializeDec conf headTy cons
451 pokeDec <- mkSerializeDec conf headTy cons
452 let methods = concat [sizeDec, peekDec, pokeDec]
453 next methods
455-- | @deriveSerializeWith config-modifier instance-dec@ generates a template
456-- Haskell splice consisting of a declaration of a 'Serialize' instance.
457-- @instance-dec@ is a template Haskell declaration splice consisting of a
458-- standard Haskell instance declaration without the type class methods (e.g.
459-- @[d|instance Serialize a => Serialize (Maybe a)|]@).
460--
461-- The type class methods for the given instance are generated according to the
462-- supplied @config-modifier@ parameter. See 'SerializeConfig' for default
463-- configuration settings.
464--
465-- Usage:
466--
467-- @
468-- \$(deriveSerializeWith
469-- ( inlineSerializeAt (Just NoInline)
470-- . inlineDeserializeAt (Just NoInline)
471-- )
472-- [d|instance Serialize a => Serialize (Maybe a)|])
473-- @
474deriveSerializeWith ::
475 (SerializeConfig -> SerializeConfig) -> Q [Dec] -> Q [Dec]
476deriveSerializeWith modifier mDecs = do
477 dec <- mDecs
478 case dec of
479 [InstanceD mo preds headTyWC []] -> do
480 let headTy = unwrap dec headTyWC
481 dt <- reifyDataType (getMainTypeName dec headTy)
482 let cons = dtCons dt
483 deriveSerializeInternal
484 (modifier serializeConfig) headTy cons (next mo preds headTyWC)
485 _ -> errorMessage dec
487 where
489 next mo preds headTyWC methods = pure [InstanceD mo preds headTyWC methods]
491 errorMessage dec =
492 error $ unlines
493 [ "Error: deriveSerializeWith:"
494 , ""
495 , ">> " ++ pprint dec
496 , ""
497 , "The supplied declaration is not a valid instance declaration."
498 , "Provide a valid Haskell instance declaration without a body."
499 , ""
500 , "Examples:"
501 , "instance Serialize (Proxy a)"
502 , "instance Serialize a => Serialize (Identity a)"
503 , "instance Serialize (TableT Identity)"
504 ]
506 unwrap _ (AppT (ConT _) r) = r
507 unwrap dec _ = errorMessage dec
509 getMainTypeName dec = go
511 where
513 go (ConT nm) = nm
514 go (AppT l _) = go l
515 go _ = errorMessage dec
517-- | Given an 'Serialize' instance declaration splice without the methods (e.g.
518-- @[d|instance Serialize a => Serialize (Maybe a)|]@), generate an instance
519-- declaration including all the type class method implementations.
520--
521-- >>> deriveSerialize = deriveSerializeWith id
522--
523-- Usage:
524--
525-- @
526-- \$(deriveSerialize
527-- [d|instance Serialize a => Serialize (Maybe a)|])
528-- @
529deriveSerialize :: Q [Dec] -> Q [Dec]
530deriveSerialize = deriveSerializeWith id