@@ -0,0 +1,416 @@
1+{-# LANGUAGE TemplateHaskell #-}
2+
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+--
11+module Streamly.Internal.Data.Serialize.TH
12+ ( deriveSerialize
13+ , deriveSerializeWith
14+ ) where
15+
16+--------------------------------------------------------------------------------
17+-- Imports
18+--------------------------------------------------------------------------------
19+
20+import Data.Word (Word16, Word32, Word64, Word8)
21+
22+import Language.Haskell.TH
23+import Language.Haskell.TH.Syntax
24+import Streamly.Internal.Data.Serialize
25+
26+import Streamly.Internal.Data.Unbox.TH
27+ ( DataCon(..)
28+ , DataType(..)
29+ , appsT
30+ , plainInstanceD
31+ , reifyDataType
32+ )
33+
34+--------------------------------------------------------------------------------
35+-- Helpers
36+--------------------------------------------------------------------------------
37+
38+type Field = (Maybe Name, Type)
39+
40+_x :: Name
41+_x = mkName "x"
42+
43+_arr :: Name
44+_arr = mkName "arr"
45+
46+_tag :: Name
47+_tag = mkName "tag"
48+
49+_initialOffset :: Name
50+_initialOffset = mkName "initialOffset"
51+
52+_val :: Name
53+_val = mkName "val"
54+
55+mkFieldName :: Int -> Name
56+mkFieldName i = mkName ("field" ++ show i)
57+
58+makeI :: Int -> Name
59+makeI i = mkName $ "i" ++ show i
60+
61+makeA :: Int -> Name
62+makeA i = mkName $ "a" ++ show i
63+
64+--------------------------------------------------------------------------------
65+-- Domain specific helpers
66+--------------------------------------------------------------------------------
67+
68+matchConstructor :: Name -> Int -> Q Exp -> Q Match
69+matchConstructor cname numFields exp0 =
70+ match
71+ (conP cname (map varP (map mkFieldName [0 .. (numFields - 1)])))
72+ (normalB exp0)
73+ []
74+
75+exprGetSize :: Int -> Type -> Q Exp
76+exprGetSize i ty =
77+ caseE
78+ (sigE (varE 'size) (appT (conT ''Size) (pure ty)))
79+ [ match
80+ (conP 'VarSize [varP _f])
81+ (normalB (appE (varE _f) (varE (mkFieldName i))))
82+ []
83+ , match (conP 'ConstSize [varP _sz]) (normalB (varE _sz)) []
84+ ]
85+
86+ where
87+
88+ _f = mkName $ "f"
89+ _sz = mkName $ "sz"
90+
91+
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"
100+
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"
109+
110+--------------------------------------------------------------------------------
111+-- Size
112+--------------------------------------------------------------------------------
113+
114+isUnitType :: [DataCon] -> Bool
115+isUnitType [DataCon _ _ _ []] = True
116+isUnitType _ = False
117+
118+mkSizeOfExpr :: Type -> [DataCon] -> Q Exp
119+mkSizeOfExpr headTy constructors =
120+ case constructors of
121+ [] ->
122+ [|error
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
126+ -- 1.
127+ [constructor@(DataCon _ _ _ fields)] ->
128+ case fields of
129+ [] -> appE (conE 'ConstSize) (litE (IntegerL 1))
130+ _ ->
131+ appE
132+ (conE 'VarSize)
133+ (lamE
134+ [varP _x]
135+ (caseE (varE _x) [matchCons constructor]))
136+ _ -> sizeOfHeadDt
137+
138+ where
139+
140+ tagSizeExp =
141+ litE (IntegerL (fromIntegral (getTagSize (length constructors))))
142+
143+ sizeOfField (i, (_, ty)) = exprGetSize i ty
144+
145+ sizeOfFields fields =
146+ appE (varE 'sum) (listE (map sizeOfField (zip [0..] fields)))
147+
148+ matchCons (DataCon cname _ _ fields) =
149+ matchConstructor cname (length fields) (sizeOfFields fields)
150+
151+ -- XXX We fix VarSize for simplicity. Should be changed later.
152+ sizeOfHeadDt =
153+ appE
154+ (conE 'VarSize)
155+ (lamE
156+ [varP _x]
157+ [|$(tagSizeExp)
158+ + $(caseE (varE _x) (fmap matchCons constructors))|])
159+
160+--------------------------------------------------------------------------------
161+-- Peek
162+--------------------------------------------------------------------------------
163+
164+mkDeserializeExprOne :: DataCon -> Q Exp
165+mkDeserializeExprOne (DataCon cname _ _ fields) =
166+ case fields of
167+ [] -> [|pure ($(varE (mkName "i0")), $(conE cname))|]
168+ _ ->
169+ doE
170+ (concat
171+ [ fmap makeBind [0 .. (numFields - 1)]
172+ , [ noBindS
173+ (appE
174+ (varE 'pure)
175+ (tupE
176+ [ varE (makeI numFields)
177+ , appsE
178+ (conE cname :
179+ (map (varE . makeA)
180+ [0 .. (numFields - 1)]))
181+ ]))
182+ ]
183+ ])
184+ where
185+ numFields = length fields
186+ makeBind i =
187+ bindS
188+ (tupP [varP (makeI (i + 1)), varP (makeA i)])
189+ [|deserialize $(varE (makeI i)) $(varE _arr)|]
190+
191+
192+mkDeserializeExpr :: Type -> [DataCon] -> Q Exp
193+mkDeserializeExpr headTy cons =
194+ case cons of
195+ [] ->
196+ [|error
197+ ("Attempting to peek type with no constructors (" ++
198+ $(lift (pprint headTy)) ++ ")")|]
199+ [(DataCon cname _ _ [])] ->
200+ [|pure ($(varE _initialOffset) + 1, $(conE cname))|]
201+ [con] ->
202+ letE
203+ [valD (varP (mkName "i0")) (normalB (varE _initialOffset)) []]
204+ (mkDeserializeExprOne con)
205+ _ ->
206+ doE
207+ [ bindS
208+ (tupP [varP (mkName "i0"), varP _tag])
209+ [|deserialize $(varE _initialOffset) $(varE _arr)|]
210+ , noBindS
211+ (caseE
212+ (sigE (varE _tag) (conT tagType))
213+ (map peekMatch (zip [0 ..] cons) ++ [peekErr]))
214+ ]
215+ where
216+ lenCons = length cons
217+ tagType = getTagType lenCons
218+ peekMatch (i, con) =
219+ match (litP (IntegerL i)) (normalB (mkDeserializeExprOne con)) []
220+ peekErr =
221+ match
222+ wildP
223+ (normalB
224+ [|error
225+ ("Found invalid tag while peeking (" ++
226+ $(lift (pprint headTy)) ++ ")")|])
227+ []
228+
229+--------------------------------------------------------------------------------
230+-- Poke
231+--------------------------------------------------------------------------------
232+
233+mkSerializeExprTag :: Name -> Int -> Q Exp
234+mkSerializeExprTag tagType tagVal =
235+ [|serialize
236+ $(varE _initialOffset)
237+ $(varE _arr)
238+ $((sigE (litE (IntegerL (fromIntegral tagVal))) (conT tagType)))|]
239+
240+mkSerializeExprFields :: [Field] -> Q Exp
241+mkSerializeExprFields fields =
242+ case fields of
243+ [] -> [|pure ($(varE (mkName "i0")))|]
244+ _ ->
245+ doE
246+ (fmap makeBind [0 .. (numFields - 1)] ++
247+ [noBindS ([|pure $(varE (makeI numFields))|])])
248+ where
249+ numFields = length fields
250+ makeBind i =
251+ bindS
252+ (varP (makeI (i + 1)))
253+ [|serialize $(varE (makeI i)) $(varE _arr) $(varE (mkFieldName i))|]
254+
255+mkSerializeExpr :: Type -> [DataCon] -> Q Exp
256+mkSerializeExpr headTy cons =
257+ case cons of
258+ [] ->
259+ [|error
260+ ("Attempting to poke type with no constructors (" ++
261+ $(lift (pprint headTy)) ++ ")")|]
262+ [(DataCon _ _ _ [])] -> [|pure ($(varE _initialOffset) + 1)|]
263+ [(DataCon cname _ _ fields)] ->
264+ letE
265+ [valD (varP (mkName "i0")) (normalB (varE _initialOffset)) []]
266+ (caseE
267+ (varE _val)
268+ [ matchConstructor
269+ cname
270+ (length fields)
271+ (mkSerializeExprFields fields)
272+ ])
273+ _ ->
274+ caseE
275+ (varE _val)
276+ (map (\(tagVal, (DataCon cname _ _ fields)) ->
277+ matchConstructor
278+ cname
279+ (length fields)
280+ (doE [ bindS
281+ (varP (mkName "i0"))
282+ (mkSerializeExprTag tagType tagVal)
283+ , noBindS (mkSerializeExprFields fields)
284+ ]))
285+ (zip [0 ..] cons))
286+ where
287+ lenCons = length cons
288+ tagType = getTagType lenCons
289+
290+--------------------------------------------------------------------------------
291+-- Main
292+--------------------------------------------------------------------------------
293+
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.
297+--
298+-- Consider the datatype:
299+-- @
300+-- data CustomDataType a b
301+-- = CDTConstructor1
302+-- | CDTConstructor2 Bool
303+-- | CDTConstructor3 Bool b
304+-- deriving (Show, Eq)
305+-- @
306+--
307+-- Usage:
308+-- @
309+-- $(deriveSerializeInternal
310+-- [AppT (ConT ''Serialize) (VarT (mkName "b"))]
311+-- (AppT
312+-- (AppT (ConT ''CustomDataType) (VarT (mkName "a")))
313+-- (VarT (mkName "b")))
314+-- [ DataCon 'CDTConstructor1 [] [] []
315+-- , DataCon 'CDTConstructor2 [] [] [(Nothing, (ConT ''Bool))]
316+-- , DataCon
317+-- 'CDTConstructor3
318+-- []
319+-- []
320+-- [(Nothing, (ConT ''Bool)), (Nothing, (VarT (mkName "b")))]
321+-- ])
322+-- @
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
328+ let methods =
329+ [ FunD 'size [Clause [] (NormalB sizeOfMethod) []]
330+ , FunD
331+ 'deserialize
332+ [ Clause
333+ (if isUnitType cons
334+ then [VarP _initialOffset, WildP]
335+ else [VarP _initialOffset, VarP _arr])
336+ (NormalB peekMethod)
337+ []
338+ ]
339+ , FunD
340+ 'serialize
341+ [ Clause
342+ (if isUnitType cons
343+ then [VarP _initialOffset, WildP, WildP]
344+ else [VarP _initialOffset, VarP _arr, VarP _val])
345+ (NormalB pokeMethod)
346+ []
347+ ]
348+ ]
349+ return [plainInstanceD preds (AppT (ConT ''Serialize) headTy) methods]
350+
351+-- | Template haskell helper to create instances of 'Serialize' automatically.
352+--
353+-- Consider the datatype:
354+-- @
355+-- data CustomDataType a b c = ...
356+-- @
357+--
358+-- Usage: @$(deriveSerialize ''CustomDataType)@
359+--
360+-- Note: All type variables automatcally get an "Serialize" constraint.
361+-- The derived code will look like the following,
362+-- @
363+-- instance (Serialize a, Serialize b, Serialize c) => Serialize (CustomDataType a b c) where
364+-- ...
365+-- @
366+--
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))
374+ cons = dtCons dt
375+ deriveSerializeInternal preds headTy cons
376+
377+ where
378+
379+ unboxPred ty =
380+#if MIN_VERSION_template_haskell(2,10,0)
381+ AppT (ConT ''Serialize) ty
382+#else
383+ ClassP ''Serialize [ty]
384+#endif
385+
386+-- | Like 'deriveSerialize' but control which types variables get the 'Serialize'
387+-- constraint.
388+--
389+-- Consider the datatype:
390+-- @
391+-- data CustomDataType a b c = ...
392+-- @
393+--
394+-- Usage: @$(deriveSerializeWith ["a", "c"] ''CustomDataType)@
395+--
396+-- @
397+-- instance (Serialize a, Serialize c) => Serialize (CustomDataType a b c) where
398+-- ...
399+-- @
400+--
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))
406+ cons = dtCons dt
407+ deriveSerializeInternal preds headTy cons
408+
409+ where
410+
411+ unboxPred ty =
412+#if MIN_VERSION_template_haskell(2,10,0)
413+ AppT (ConT ''Serialize) ty
414+#else
415+ ClassP ''Serialize [ty]
416+#endif