git.iko.soy/streamly
on 79df417 help sign in
streamly/core/src/Streamly/Internal/Data/Serialize/TH.hs
513 lines 18 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
10--
11module Streamly.Internal.Data.Serialize.TH
12 ( deriveSerialize
13 , Config(..)
14 , defaultConfig
15 , deriveSerializeWith
16 ) where
18--------------------------------------------------------------------------------
19-- Imports
20--------------------------------------------------------------------------------
22import Data.List (foldl')
23import Data.Word (Word16, Word32, Word64, Word8)
25import Language.Haskell.TH
26import Language.Haskell.TH.Syntax
27import Streamly.Internal.Data.Serialize.Type
29import Streamly.Internal.Data.Unbox.TH
30 ( DataCon(..)
31 , DataType(..)
32 , appsT
33 , reifyDataType
34 )
36import qualified Streamly.Internal.Data.Serialize.TH.RecHeader as RecHeader
38import Streamly.Internal.Data.Serialize.TH.Bottom
39import Streamly.Internal.Data.Serialize.TH.Common
41--------------------------------------------------------------------------------
42-- Domain specific helpers
43--------------------------------------------------------------------------------
45exprGetSize :: Q Exp -> (Int, Type) -> Q Exp
46exprGetSize acc (i, _) = [|size $(acc) $(varE (mkFieldName i))|]
48getTagSize :: Int -> Int
49getTagSize numConstructors
50 | numConstructors == 1 = 0
51 | fromIntegral (maxBound :: Word8) >= numConstructors = 1
52 | fromIntegral (maxBound :: Word16) >= numConstructors = 2
53 | fromIntegral (maxBound :: Word32) >= numConstructors = 4
54 | fromIntegral (maxBound :: Word64) >= numConstructors = 8
55 | otherwise = error "Too many constructors"
57getTagType :: Int -> Name
58getTagType numConstructors
59 | numConstructors == 1 = error "No tag for 1 constructor"
60 | fromIntegral (maxBound :: Word8) >= numConstructors = ''Word8
61 | fromIntegral (maxBound :: Word16) >= numConstructors = ''Word16
62 | fromIntegral (maxBound :: Word32) >= numConstructors = ''Word32
63 | fromIntegral (maxBound :: Word64) >= numConstructors = ''Word64
64 | otherwise = error "Too many constructors"
66--------------------------------------------------------------------------------
67-- Size
68--------------------------------------------------------------------------------
70getNameBaseLen :: Name -> Word8
71getNameBaseLen cname =
72 let x = length (nameBase cname)
73 in if x > 63
74 then error "Max Constructor Len: 63 characters"
75 else fromIntegral x
77conEncLen :: Name -> Word8
78conEncLen cname = getNameBaseLen cname + 1
80--------------------------------------------------------------------------------
81-- Size
82--------------------------------------------------------------------------------
84mkSizeOfExpr :: Bool -> Bool -> TypeOfType -> Q Exp
85mkSizeOfExpr True False tyOfTy =
86 case tyOfTy of
87 UnitType cname ->
88 lamE
89 [varP _acc, wildP]
90 [|$(varE _acc) + $(litIntegral (conEncLen cname))|]
91 TheType con ->
92 lamE
93 [varP _acc, varP _x]
94 (caseE (varE _x) [matchCons (varE _acc) con])
95 MultiType constructors -> sizeOfHeadDt constructors
97 where
99 sizeOfFields acc fields =
100 foldl' exprGetSize acc $ zip [0..] fields
102 matchCons acc (SimpleDataCon cname fields) =
103 let a = litIntegral (conEncLen cname)
104 b = sizeOfFields acc (map snd fields)
105 expr = [|$(a) + $(b)|]
106 in matchConstructor cname (length fields) expr
108 sizeOfHeadDt cons =
109 let acc = [|$(varE _acc)|]
110 in lamE
111 [varP _acc, varP _x]
112 (caseE (varE _x) (fmap (matchCons acc) cons))
114mkSizeOfExpr False False tyOfTy =
115 case tyOfTy of
116 UnitType _ -> lamE [varP _acc, wildP] [|$(varE _acc) + 1|]
117 TheType con ->
118 lamE
119 [varP _acc, varP _x]
120 (caseE (varE _x) [matchCons (varE _acc) con])
121 MultiType constructors -> sizeOfHeadDt constructors
123 where
125 tagSizeExp numConstructors =
126 litE (IntegerL (fromIntegral (getTagSize numConstructors)))
128 -- XXX fields of the same type can be folded together, will reduce the code
129 -- size when there are many fields of the same type.
130 -- XXX const size fields can be calculated statically.
131 -- XXX This can result in large compilation times due to nesting when there
132 -- are many constructors. We can create a list and sum the list at run time
133 -- to avoid that depending on the number of constructors. Or using a let
134 -- statement for each case may help?
135 -- appE (varE 'sum) (listE (acc : map (exprGetSize (litE (IntegerL 0))) (zip [0..] fields)))
136 sizeOfFields acc fields =
137 foldl' exprGetSize acc $ zip [0..] fields
139 matchCons acc (SimpleDataCon cname fields) =
140 let expr = sizeOfFields acc (map snd fields)
141 in matchConstructor cname (length fields) expr
143 -- XXX We fix VarSize for simplicity. Should be changed later.
144 sizeOfHeadDt cons =
145 let numCons = length cons
146 acc = [|$(varE _acc) + $(tagSizeExp numCons)|]
147 in lamE
148 [varP _acc, varP _x]
149 (caseE (varE _x) (fmap (matchCons acc) cons))
151mkSizeOfExpr False True (TheType con) = RecHeader.mkSizeOfExpr con
153mkSizeOfExpr _ _ _ = errorUnimplemented
155mkSizeDec :: Config -> Type -> [DataCon] -> Q [Dec]
156mkSizeDec (Config {..}) headTy cons = do
157 -- INLINE on sizeOf actually worsens some benchmarks, and improves none
158 sizeOfMethod <-
159 mkSizeOfExpr
160 constructorTagAsString
161 recordSyntaxWithHeader
162 (typeOfType headTy cons)
163 pure
164 [ PragmaD (InlineP 'size inlineSize FunLike AllPhases)
165 , FunD 'size [Clause [] (NormalB sizeOfMethod) []]
166 ]
168--------------------------------------------------------------------------------
169-- Peek
170--------------------------------------------------------------------------------
172mkDeserializeExpr :: Bool -> Bool -> Type -> TypeOfType -> Q Exp
173mkDeserializeExpr True False headTy tyOfTy =
174 case tyOfTy of
175 UnitType cname -> deserializeConsExpr [SimpleDataCon cname []]
176 TheType con -> deserializeConsExpr [con]
177 MultiType cons -> deserializeConsExpr cons
179 where
181 deserializeConsExpr cons = do
182 conLen <- newName "conLen"
183 off1 <- newName "off1"
184 [|do ($(varP off1), $(varP conLen) :: Word8) <-
185 deserialize
186 $(varE _initialOffset)
187 $(varE _arr)
188 $(varE _endOffset)
189 $(multiIfE (map (guardCon conLen off1) cons ++ [catchAll]))|]
191 catchAll =
192 normalGE
193 [|True|]
194 [|error
195 ("Found invalid tag while peeking (" ++
196 $(lift (pprint headTy)) ++ ")")|]
198 guardCon conLen off con@(SimpleDataCon cname _) = do
199 let lenCname = getNameBaseLen cname
200 tag = map c2w (nameBase cname)
201 normalGE
202 [|($(litIntegral lenCname) == $(varE conLen))
203 && $(xorCmp tag off _arr)|]
204 [|let $(varP (makeI 0)) = $(varE off) + $(litIntegral lenCname)
205 in $(mkDeserializeExprOne 'deserialize con)|]
207mkDeserializeExpr False False headTy tyOfTy =
208 case tyOfTy of
209 -- Unit constructor
210 UnitType cname ->
211 [|pure ($(varE _initialOffset) + 1, $(conE cname))|]
212 -- Product type
213 TheType con ->
214 letE
215 [valD (varP (mkName "i0")) (normalB (varE _initialOffset)) []]
216 (mkDeserializeExprOne 'deserialize con)
217 -- Sum type
218 MultiType cons -> do
219 let lenCons = length cons
220 tagType = getTagType lenCons
221 doE
222 [ bindS
223 (tupP [varP (mkName "i0"), varP _tag])
224 [|deserialize $(varE _initialOffset) $(varE _arr) $(varE _endOffset)|]
225 , noBindS
226 (caseE
227 (sigE (varE _tag) (conT tagType))
228 (map peekMatch (zip [0 ..] cons) ++ [peekErr]))
229 ]
230 where
231 peekMatch (i, con) =
232 match
233 (litP (IntegerL i))
234 (normalB (mkDeserializeExprOne 'deserialize con)) []
235 peekErr =
236 match
237 wildP
238 (normalB
239 -- XXX Print the tag
240 [|error
241 ("Found invalid tag while peeking (" ++
242 $(lift (pprint headTy)) ++ ")")|])
243 []
245mkDeserializeExpr False True _ (TheType con@(SimpleDataCon _ fields)) = do
246 deserializeWithKeys <- newName "deserializeWithKeys"
247 updateFunc <- newName "updateFunc"
248 updateFuncDec <- RecHeader.conUpdateFuncDec updateFunc fields
249 deserializeWithKeysDec <-
250 RecHeader.mkDeserializeKeysDec deserializeWithKeys updateFunc con
251 letE
252 (pure <$> (deserializeWithKeysDec ++ updateFuncDec))
253 (RecHeader.mkDeserializeExpr
254 _initialOffset
255 _endOffset
256 deserializeWithKeys
257 con)
259mkDeserializeExpr _ _ _ _ = errorUnimplemented
261mkDeserializeDec :: Config -> Type -> [DataCon] -> Q [Dec]
262mkDeserializeDec (Config {..}) headTy cons = do
263 peekMethod <-
264 mkDeserializeExpr
265 constructorTagAsString
266 recordSyntaxWithHeader
267 headTy
268 (typeOfType headTy cons)
269 pure
270 [ PragmaD (InlineP 'deserialize inlineDeserialize FunLike AllPhases)
271 , FunD
272 'deserialize
273 [ Clause
274 (if isUnitType cons && not constructorTagAsString
275 then [VarP _initialOffset, WildP, WildP]
276 else [VarP _initialOffset, VarP _arr, VarP _endOffset])
277 (NormalB peekMethod)
278 []
279 ]
280 ]
282--------------------------------------------------------------------------------
283-- Poke
284--------------------------------------------------------------------------------
286mkSerializeExprTag :: Name -> Int -> Q Exp
287mkSerializeExprTag tagType tagVal =
288 [|serialize
289 $(varE _initialOffset)
290 $(varE _arr)
291 $((sigE (litE (IntegerL (fromIntegral tagVal))) (conT tagType)))|]
293mkSerializeExpr :: Bool -> Bool -> TypeOfType -> Q Exp
294mkSerializeExpr True False tyOfTy =
295 case tyOfTy of
296 -- Unit type
297 UnitType cname ->
298 caseE
299 (varE _val)
300 [serializeDataCon (SimpleDataCon cname [])]
301 -- Product type
302 (TheType con) ->
303 caseE
304 (varE _val)
305 [serializeDataCon con]
306 -- Sum type
307 (MultiType cons) ->
308 caseE
309 (varE _val)
310 (map serializeDataCon cons)
312 where
314 serializeDataCon (SimpleDataCon cname fields) = do
315 let tagLen8 = getNameBaseLen cname
316 conEnc = tagLen8 : map c2w (nameBase cname)
317 matchConstructor
318 cname
319 (length fields)
320 (doE [ bindS
321 (varP (mkName "i0"))
322 (serializeW8List _initialOffset _arr conEnc)
323 , noBindS (mkSerializeExprFields 'serialize fields)
324 ])
326mkSerializeExpr False False tyOfTy =
327 case tyOfTy of
328 -- Unit type
329 UnitType _ -> [|pure ($(varE _initialOffset) + 1)|]
330 -- Product type
331 (TheType (SimpleDataCon cname fields)) ->
332 letE
333 [valD (varP (mkName "i0")) (normalB (varE _initialOffset)) []]
334 (caseE
335 (varE _val)
336 [ matchConstructor
337 cname
338 (length fields)
339 (mkSerializeExprFields 'serialize fields)
340 ])
341 -- Sum type
342 (MultiType cons) -> do
343 let lenCons = length cons
344 tagType = getTagType lenCons
345 caseE
346 (varE _val)
347 (map (\(tagVal, (SimpleDataCon cname fields)) ->
348 matchConstructor
349 cname
350 (length fields)
351 (doE [ bindS
352 (varP (mkName "i0"))
353 (mkSerializeExprTag tagType tagVal)
354 , noBindS
355 (mkSerializeExprFields
356 'serialize
357 fields)
358 ]))
359 (zip [0 ..] cons))
361mkSerializeExpr False True (TheType con) =
362 RecHeader.mkSerializeExpr _initialOffset con
364mkSerializeExpr _ _ _ = errorUnimplemented
366mkSerializeDec :: Config -> Type -> [DataCon] -> Q [Dec]
367mkSerializeDec (Config {..}) headTy cons = do
368 pokeMethod <-
369 mkSerializeExpr
370 constructorTagAsString
371 recordSyntaxWithHeader
372 (typeOfType headTy cons)
373 pure
374 [ PragmaD (InlineP 'serialize inlineSerialize FunLike AllPhases)
375 , FunD
376 'serialize
377 [ Clause
378 (if isUnitType cons && not constructorTagAsString
379 then [VarP _initialOffset, WildP, WildP]
380 else [VarP _initialOffset, VarP _arr, VarP _val])
381 (NormalB pokeMethod)
382 []
383 ]
384 ]
386--------------------------------------------------------------------------------
387-- Main
388--------------------------------------------------------------------------------
390-- | A general function to derive Serialize instances where you can control
391-- which Constructors of the datatype to consider and what the Context for the
392-- 'Serialize' instance would be.
393--
394-- Consider the datatype:
395-- @
396-- data CustomDataType a b
397-- = CDTConstructor1
398-- | CDTConstructor2 Bool
399-- | CDTConstructor3 Bool b
400-- deriving (Show, Eq)
401-- @
402--
403-- Usage:
404-- @
405-- $(deriveSerializeInternal
406-- defaultConfig
407-- [AppT (ConT ''Serialize) (VarT (mkName "b"))]
408-- (AppT
409-- (AppT (ConT ''CustomDataType) (VarT (mkName "a")))
410-- (VarT (mkName "b")))
411-- [ DataCon 'CDTConstructor1 [] [] []
412-- , DataCon 'CDTConstructor2 [] [] [(Nothing, (ConT ''Bool))]
413-- , DataCon
414-- 'CDTConstructor3
415-- []
416-- []
417-- [(Nothing, (ConT ''Bool)), (Nothing, (VarT (mkName "b")))]
418-- ])
419-- @
420deriveSerializeInternal ::
421 Config -> Type -> [DataCon] -> ([Dec] -> Q [Dec]) -> Q [Dec]
422deriveSerializeInternal conf headTy cons next = do
423 sizeDec <- mkSizeDec conf headTy cons
424 peekDec <- mkDeserializeDec conf headTy cons
425 pokeDec <- mkSerializeDec conf headTy cons
426 let methods = concat [sizeDec, peekDec, pokeDec]
427 next methods
429-- | Similar to 'deriveSerialize' but take a 'Config' to control how the
430-- instance is generated.
431--
432-- Usage:
433-- @
434-- $(deriveSerializeWith
435-- config
436-- [d|instance Serialize a => Serialize (Maybe a)|])
437-- @
438deriveSerializeWith :: Config -> Q [Dec] -> Q [Dec]
439deriveSerializeWith conf mDecs = do
440 dec <- mDecs
441 case dec of
442 [inst@(InstanceD _ _ headTyWC [])] -> do
443 let headTy = unwrap dec headTyWC
444 dt <- reifyDataType (getMainTypeName dec headTy)
445 let cons = dtCons dt
446 deriveSerializeInternal conf headTy cons (next inst)
447 _ -> errorUnsupported (errorMessage dec)
448 where
450 next (InstanceD mo preds headTyWC []) methods =
451 pure [InstanceD mo preds headTyWC methods]
452 next dec _ = errorUnsupported (errorMessage dec)
454 errorMessage dec =
455 error $ unlines
456 [ "Error: deriveSerializeWith:"
457 , ""
458 , ">> " ++ pprint dec
459 , ""
460 , "The above is not a valid instance declaration."
461 , "Any haskell instance declaration without a body is valid."
462 , ""
463 , "Examples:"
464 , "instance Serialize (Proxy a)"
465 , "instance Serialize a => Serialize (Identity a)"
466 , "instance Serialize (TableT Identity)"
467 ]
469 unwrap _ (AppT (ConT _) r) = r
470 unwrap dec _ = errorMessage dec
472 getMainTypeName dec = go
474 where
476 go (ConT nm) = nm
477 go (AppT l _) = go l
478 go _ = errorMessage dec
481-- | Template haskell helper to create instances of 'Serialize' automatically.
482--
483-- Consider the datatype:
484-- @
485-- data CustomDataType a b c = ...
486-- @
487--
488-- Usage: @$(deriveSerialize ''CustomDataType)@
489--
490-- Note: All type variables automatcally get an "Serialize" constraint.
491-- The derived code will look like the following,
492-- @
493-- instance (Serialize a, Serialize b, Serialize c) => Serialize (CustomDataType a b c) where
494-- ...
495-- @
496--
497-- To control which type variables don't get the Serialize constraint, use
498-- 'deriveSerializeWith'.
499deriveSerialize :: Name -> Q [Dec]
500deriveSerialize name = do
501 dt <- reifyDataType name
502 let preds = map (unboxPred . VarT) (dtTvs dt)
503 headTy = appsT (ConT name) (map VarT (dtTvs dt))
504 cons = dtCons dt
505 deriveSerializeInternal defaultConfig headTy cons (next preds headTy)
508 where
510 next preds headTy methods =
511 pure [InstanceD Nothing preds (AppT (ConT ''Serialize) headTy) methods]
513 unboxPred ty = AppT (ConT ''Serialize) ty