@@ -0,0 +1,416 @@
1+{-# LANGUAGE TemplateHaskell #-}
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
11+module Streamly.Internal.Data.Serialize.TH
13+ , deriveSerializeWith
16+--------------------------------------------------------------------------------
18+--------------------------------------------------------------------------------
20+import Data.Word (Word16, Word32, Word64, Word8)
22+import Language.Haskell.TH
23+import Language.Haskell.TH.Syntax
24+import Streamly.Internal.Data.Serialize
26+import Streamly.Internal.Data.Unbox.TH
34+--------------------------------------------------------------------------------
36+--------------------------------------------------------------------------------
38+type Field = (Maybe Name, Type)
49+_initialOffset :: Name
50+_initialOffset = mkName "initialOffset"
55+mkFieldName :: Int -> Name
56+mkFieldName i = mkName ("field" ++ show i)
59+makeI i = mkName $ "i" ++ show i
62+makeA i = mkName $ "a" ++ show i
64+--------------------------------------------------------------------------------
65+-- Domain specific helpers
66+--------------------------------------------------------------------------------
68+matchConstructor :: Name -> Int -> Q Exp -> Q Match
69+matchConstructor cname numFields exp0 =
71+ (conP cname (map varP (map mkFieldName [0 .. (numFields - 1)])))
75+exprGetSize :: Int -> Type -> Q Exp
78+ (sigE (varE 'size) (appT (conT ''Size) (pure ty)))
80+ (conP 'VarSize [varP _f])
81+ (normalB (appE (varE _f) (varE (mkFieldName i))))
83+ , match (conP 'ConstSize [varP _sz]) (normalB (varE _sz)) []
92+getTagSize :: Int -> Int
93+getTagSize numConstructors
94+ | numConstructors == 1 = 0
95+ | fromIntegral (maxBound :: Word8) >= numConstructors = 1
96+ | fromIntegral (maxBound :: Word16) >= numConstructors = 2
97+ | fromIntegral (maxBound :: Word32) >= numConstructors = 4
98+ | fromIntegral (maxBound :: Word64) >= numConstructors = 8
99+ | otherwise = error "Too many constructors"
101+getTagType :: Int -> Name
102+getTagType numConstructors
103+ | numConstructors == 1 = error "No tag for 1 constructor"
104+ | fromIntegral (maxBound :: Word8) >= numConstructors = ''Word8
105+ | fromIntegral (maxBound :: Word16) >= numConstructors = ''Word16
106+ | fromIntegral (maxBound :: Word32) >= numConstructors = ''Word32
107+ | fromIntegral (maxBound :: Word64) >= numConstructors = ''Word64
108+ | otherwise = error "Too many constructors"
110+--------------------------------------------------------------------------------
112+--------------------------------------------------------------------------------
114+isUnitType :: [DataCon] -> Bool
115+isUnitType [DataCon _ _ _ []] = True
118+mkSizeOfExpr :: Type -> [DataCon] -> Q Exp
119+mkSizeOfExpr headTy constructors =
120+ case constructors of
123+ ("Attempting to get size with no constructors (" ++
124+ $(lift (pprint headTy)) ++ ")")|]
125+ -- One constructor with no fields is a unit type. Size of a unit type is
127+ [constructor@(DataCon _ _ _ fields)] ->
129+ [] -> appE (conE 'ConstSize) (litE (IntegerL 1))
135+ (caseE (varE _x) [matchCons constructor]))
141+ litE (IntegerL (fromIntegral (getTagSize (length constructors))))
143+ sizeOfField (i, (_, ty)) = exprGetSize i ty
145+ sizeOfFields fields =
146+ appE (varE 'sum) (listE (map sizeOfField (zip [0..] fields)))
148+ matchCons (DataCon cname _ _ fields) =
149+ matchConstructor cname (length fields) (sizeOfFields fields)
151+ -- XXX We fix VarSize for simplicity. Should be changed later.
158+ + $(caseE (varE _x) (fmap matchCons constructors))|])
160+--------------------------------------------------------------------------------
162+--------------------------------------------------------------------------------
164+mkDeserializeExprOne :: DataCon -> Q Exp
165+mkDeserializeExprOne (DataCon cname _ _ fields) =
167+ [] -> [|pure ($(varE (mkName "i0")), $(conE cname))|]
171+ [ fmap makeBind [0 .. (numFields - 1)]
176+ [ varE (makeI numFields)
180+ [0 .. (numFields - 1)]))
185+ numFields = length fields
188+ (tupP [varP (makeI (i + 1)), varP (makeA i)])
189+ [|deserialize $(varE (makeI i)) $(varE _arr)|]
192+mkDeserializeExpr :: Type -> [DataCon] -> Q Exp
193+mkDeserializeExpr headTy cons =
197+ ("Attempting to peek type with no constructors (" ++
198+ $(lift (pprint headTy)) ++ ")")|]
199+ [(DataCon cname _ _ [])] ->
200+ [|pure ($(varE _initialOffset) + 1, $(conE cname))|]
203+ [valD (varP (mkName "i0")) (normalB (varE _initialOffset)) []]
204+ (mkDeserializeExprOne con)
208+ (tupP [varP (mkName "i0"), varP _tag])
209+ [|deserialize $(varE _initialOffset) $(varE _arr)|]
212+ (sigE (varE _tag) (conT tagType))
213+ (map peekMatch (zip [0 ..] cons) ++ [peekErr]))
216+ lenCons = length cons
217+ tagType = getTagType lenCons
218+ peekMatch (i, con) =
219+ match (litP (IntegerL i)) (normalB (mkDeserializeExprOne con)) []
225+ ("Found invalid tag while peeking (" ++
226+ $(lift (pprint headTy)) ++ ")")|])
229+--------------------------------------------------------------------------------
231+--------------------------------------------------------------------------------
233+mkSerializeExprTag :: Name -> Int -> Q Exp
234+mkSerializeExprTag tagType tagVal =
236+ $(varE _initialOffset)
238+ $((sigE (litE (IntegerL (fromIntegral tagVal))) (conT tagType)))|]
240+mkSerializeExprFields :: [Field] -> Q Exp
241+mkSerializeExprFields fields =
243+ [] -> [|pure ($(varE (mkName "i0")))|]
246+ (fmap makeBind [0 .. (numFields - 1)] ++
247+ [noBindS ([|pure $(varE (makeI numFields))|])])
249+ numFields = length fields
252+ (varP (makeI (i + 1)))
253+ [|serialize $(varE (makeI i)) $(varE _arr) $(varE (mkFieldName i))|]
255+mkSerializeExpr :: Type -> [DataCon] -> Q Exp
256+mkSerializeExpr headTy cons =
260+ ("Attempting to poke type with no constructors (" ++
261+ $(lift (pprint headTy)) ++ ")")|]
262+ [(DataCon _ _ _ [])] -> [|pure ($(varE _initialOffset) + 1)|]
263+ [(DataCon cname _ _ fields)] ->
265+ [valD (varP (mkName "i0")) (normalB (varE _initialOffset)) []]
271+ (mkSerializeExprFields fields)
276+ (map (\(tagVal, (DataCon cname _ _ fields)) ->
281+ (varP (mkName "i0"))
282+ (mkSerializeExprTag tagType tagVal)
283+ , noBindS (mkSerializeExprFields fields)
287+ lenCons = length cons
288+ tagType = getTagType lenCons
290+--------------------------------------------------------------------------------
292+--------------------------------------------------------------------------------
294+-- | A general function to derive Serialize instances where you can control
295+-- which Constructors of the datatype to consider and what the Context for the
296+-- 'Serialize' instance would be.
298+-- Consider the datatype:
300+-- data CustomDataType a b
302+-- | CDTConstructor2 Bool
303+-- | CDTConstructor3 Bool b
304+-- deriving (Show, Eq)
309+-- $(deriveSerializeInternal
310+-- [AppT (ConT ''Serialize) (VarT (mkName "b"))]
312+-- (AppT (ConT ''CustomDataType) (VarT (mkName "a")))
313+-- (VarT (mkName "b")))
314+-- [ DataCon 'CDTConstructor1 [] [] []
315+-- , DataCon 'CDTConstructor2 [] [] [(Nothing, (ConT ''Bool))]
320+-- [(Nothing, (ConT ''Bool)), (Nothing, (VarT (mkName "b")))]
323+deriveSerializeInternal :: Cxt -> Type -> [DataCon] -> Q [Dec]
324+deriveSerializeInternal preds headTy cons = do
325+ sizeOfMethod <- mkSizeOfExpr headTy cons
326+ peekMethod <- mkDeserializeExpr headTy cons
327+ pokeMethod <- mkSerializeExpr headTy cons
329+ [ FunD 'size [Clause [] (NormalB sizeOfMethod) []]
334+ then [VarP _initialOffset, WildP]
335+ else [VarP _initialOffset, VarP _arr])
336+ (NormalB peekMethod)
343+ then [VarP _initialOffset, WildP, WildP]
344+ else [VarP _initialOffset, VarP _arr, VarP _val])
345+ (NormalB pokeMethod)
349+ return [plainInstanceD preds (AppT (ConT ''Serialize) headTy) methods]
351+-- | Template haskell helper to create instances of 'Serialize' automatically.
353+-- Consider the datatype:
355+-- data CustomDataType a b c = ...
358+-- Usage: @$(deriveSerialize ''CustomDataType)@
360+-- Note: All type variables automatcally get an "Serialize" constraint.
361+-- The derived code will look like the following,
363+-- instance (Serialize a, Serialize b, Serialize c) => Serialize (CustomDataType a b c) where
367+-- To control which type variables get the Serialize constraint, use
368+-- 'deriveSerializeWith'.
369+deriveSerialize :: Name -> Q [Dec]
370+deriveSerialize name = do
371+ dt <- reifyDataType name
372+ let preds = map (unboxPred . VarT) (dtTvs dt)
373+ headTy = appsT (ConT name) (map VarT (dtTvs dt))
375+ deriveSerializeInternal preds headTy cons
380+#if MIN_VERSION_template_haskell(2,10,0)
381+ AppT (ConT ''Serialize) ty
383+ ClassP ''Serialize [ty]
386+-- | Like 'deriveSerialize' but control which types variables get the 'Serialize'
389+-- Consider the datatype:
391+-- data CustomDataType a b c = ...
394+-- Usage: @$(deriveSerializeWith ["a", "c"] ''CustomDataType)@
397+-- instance (Serialize a, Serialize c) => Serialize (CustomDataType a b c) where
401+deriveSerializeWith :: [String] -> Name -> Q [Dec]
402+deriveSerializeWith vars name = do
403+ dt <- reifyDataType name
404+ let preds = map (unboxPred . VarT) (fmap mkName vars)
405+ headTy = appsT (ConT name) (map VarT (dtTvs dt))
407+ deriveSerializeInternal preds headTy cons
412+#if MIN_VERSION_template_haskell(2,10,0)
413+ AppT (ConT ''Serialize) ty
415+ ClassP ''Serialize [ty]