@@ -19,13 +19,18 @@ module Streamly.Internal.Data.Serialize.TH
1919-- Imports
2020--------------------------------------------------------------------------------
2121
22+import Data.Char (ord)
2223import Data.List (foldl')
2324import Data.Word (Word16, Word32, Word64, Word8)
25+import Data.Bits (Bits, (.|.), shiftL, zeroBits, xor)
26+import Streamly.Internal.System.IO (unsafeInlineIO)
2427
2528import Language.Haskell.TH
2629import Language.Haskell.TH.Syntax
2730import Streamly.Internal.Data.Serialize
2831
32+import qualified Streamly.Internal.Data.Unbox as Unbox
33+
2934import Streamly.Internal.Data.Unbox.TH
3035 ( DataCon(..)
3136 , DataType(..)
@@ -67,6 +72,9 @@ mkFieldName i = mkName ("field" ++ show i)
6772makeI :: Int -> Name
6873makeI i = mkName $ "i" ++ show i
6974
75+makeN :: Int -> Name
76+makeN i = mkName $ "n" ++ show i
77+
7078makeA :: Int -> Name
7179makeA i = mkName $ "a" ++ show i
7280
@@ -136,8 +144,185 @@ isUnitType _ = False
136144-- Size
137145--------------------------------------------------------------------------------
138146
139−mkSizeOfExpr :: TypeOfType -> Q Exp
140−mkSizeOfExpr tyOfTy =
147+{-
148+int_w8 :: Int -> Word8
149+int_w8 = fromIntegral
150+-}
151+
152+w8_w16 :: Word8 -> Word16
153+w8_w16 = fromIntegral
154+
155+w8_w32 :: Word8 -> Word32
156+w8_w32 = fromIntegral
157+
158+w8_w64 :: Word8 -> Word64
159+w8_w64 = fromIntegral
160+
161+{-
162+w32_int :: Word32 -> Int
163+w32_int = fromIntegral
164+-}
165+
166+c2w :: Char -> Word8
167+c2w = fromIntegral . ord
168+
169+shiftAdd :: Bits a => (b -> a) -> [b] -> a
170+shiftAdd conv xs =
171+ foldl' (.|.) zeroBits $
172+ map (\(j, x) -> shiftL x (j * 8)) $ zip [0 ..] $ map conv xs
173+
174+-- Note: This only works in little endian machines
175+xorCmp :: [Word8] -> Name -> Name -> Q Exp
176+xorCmp tag off arr =
177+ case tagLen of
178+ x | x < 2 -> [|$(go8 0) == zeroBits|]
179+ x | x < 4 -> [|$(go16 0) == zeroBits|]
180+ x | x < 8 -> [|$(go32 0) == zeroBits|]
181+ _ -> [|$(go64 0) == zeroBits|]
182+ where
183+ tagLen = length tag
184+ go8 i | i >= tagLen = [|zeroBits|]
185+ go8 i = do
186+ let wIntegral = litIntegral i
187+ [|xor (unsafeInlineIO
188+ (Unbox.peekByteIndex
189+ ($(varE off) + $(litIntegral i))
190+ $(varE arr)))
191+ ($(wIntegral) :: Word8) .|.
192+ $(go8 (i + 1))|]
193+ go16 i
194+ | i >= tagLen = [|zeroBits|]
195+ go16 i
196+ | tagLen - i < 2 = go16 (tagLen - 2)
197+ go16 i = do
198+ let wIntegral =
199+ litIntegral
200+ (shiftAdd w8_w16 [tag !! i, tag !! (i + 1)] :: Word16)
201+ [|xor (unsafeInlineIO
202+ (Unbox.peekByteIndex
203+ ($(varE off) + $(litIntegral i))
204+ $(varE arr)))
205+ ($(wIntegral) :: Word16) .|.
206+ $(go16 (i + 2))|]
207+ go32 i
208+ | i >= tagLen = [|zeroBits|]
209+ go32 i
210+ | tagLen - i < 4 = go32 (tagLen - 4)
211+ go32 i = do
212+ let wIntegral =
213+ litIntegral
214+ (shiftAdd
215+ w8_w32
216+ [ tag !! i
217+ , tag !! (i + 1)
218+ , tag !! (i + 2)
219+ , tag !! (i + 3)
220+ ] :: Word32)
221+ [|xor (unsafeInlineIO
222+ (Unbox.peekByteIndex
223+ ($(varE off) + $(litIntegral i))
224+ $(varE arr)))
225+ ($(wIntegral) :: Word32) .|.
226+ $(go32 (i + 4))|]
227+ go64 i
228+ | i >= tagLen = [|zeroBits|]
229+ go64 i
230+ | tagLen - i < 8 = go64 (tagLen - 8)
231+ go64 i = do
232+ let wIntegral =
233+ litIntegral
234+ (shiftAdd
235+ w8_w64
236+ [ tag !! i
237+ , tag !! (i + 1)
238+ , tag !! (i + 2)
239+ , tag !! (i + 3)
240+ , tag !! (i + 4)
241+ , tag !! (i + 5)
242+ , tag !! (i + 6)
243+ , tag !! (i + 7)
244+ ])
245+ [|xor (unsafeInlineIO
246+ (Unbox.peekByteIndex
247+ ($(varE off) + $(litIntegral i))
248+ $(varE arr)))
249+ ($(wIntegral) :: Word64) .|.
250+ $(go64 (i + 8))|]
251+
252+--------------------------------------------------------------------------------
253+-- Primitive serialization
254+--------------------------------------------------------------------------------
255+
256+-- XXX Batch serialize? See xorCmp
257+serializeW8List :: Name -> Name -> [Word8] -> Q Exp
258+serializeW8List off arr w8List = do
259+ [|let $(varP (makeN 0)) = $(varE off)
260+ in $(doE (fmap makeBind [0 .. (lenW8List - 1)] ++
261+ [noBindS ([|pure $(varE (makeN lenW8List))|])]))|]
262+
263+ where
264+
265+ lenW8List = length w8List
266+ makeBind i =
267+ bindS
268+ (varP (makeN (i + 1)))
269+ [|$(varE 'serialize)
270+ $(varE (makeN i))
271+ $(varE arr)
272+ ($(litIntegral (w8List !! i)) :: Word8)|]
273+
274+--------------------------------------------------------------------------------
275+-- Size
276+--------------------------------------------------------------------------------
277+
278+litIntegral :: Integral a => a -> Q Exp
279+litIntegral = litE . IntegerL . fromIntegral
280+
281+getNameBaseLen :: Name -> Word8
282+getNameBaseLen cname =
283+ let x = length (nameBase cname)
284+ in if x > 63
285+ then error "Max Constructor Len: 63 characters"
286+ else fromIntegral x
287+
288+conEncLen :: Name -> Word8
289+conEncLen cname = getNameBaseLen cname + 1
290+
291+--------------------------------------------------------------------------------
292+-- Size
293+--------------------------------------------------------------------------------
294+
295+mkSizeOfExpr :: Bool -> TypeOfType -> Q Exp
296+mkSizeOfExpr True tyOfTy =
297+ case tyOfTy of
298+ UnitType cname ->
299+ lamE
300+ [varP _acc, wildP]
301+ [|$(varE _acc) + $(litIntegral (conEncLen cname))|]
302+ TheType con ->
303+ lamE
304+ [varP _acc, varP _x]
305+ (caseE (varE _x) [matchCons (varE _acc) con])
306+ MultiType constructors -> sizeOfHeadDt constructors
307+
308+ where
309+
310+ sizeOfFields acc fields =
311+ foldl' exprGetSize acc $ zip [0..] fields
312+
313+ matchCons acc (SimpleDataCon cname fields) =
314+ let a = litIntegral (conEncLen cname)
315+ b = sizeOfFields acc (map snd fields)
316+ expr = [|$(a) + $(b)|]
317+ in matchConstructor cname (length fields) expr
318+
319+ sizeOfHeadDt cons =
320+ let acc = [|$(varE _acc)|]
321+ in lamE
322+ [varP _acc, varP _x]
323+ (caseE (varE _x) (fmap (matchCons acc) cons))
324+
325+mkSizeOfExpr False tyOfTy =
141326 case tyOfTy of
142327 UnitType _ -> lamE [varP _acc, wildP] [|$(varE _acc) + 1|]
143328 TheType con ->
@@ -207,8 +392,42 @@ mkDeserializeExprOne (SimpleDataCon cname fields) =
207392 [|deserialize $(varE (makeI i)) $(varE _arr) $(varE _endOffset)|]
208393
209394
210−mkDeserializeExpr :: Type -> TypeOfType -> Q Exp
211−mkDeserializeExpr headTy tyOfTy =
395+mkDeserializeExpr :: Bool -> Type -> TypeOfType -> Q Exp
396+mkDeserializeExpr True headTy tyOfTy =
397+ case tyOfTy of
398+ UnitType cname -> deserializeConsExpr [SimpleDataCon cname []]
399+ TheType con -> deserializeConsExpr [con]
400+ MultiType cons -> deserializeConsExpr cons
401+
402+ where
403+
404+ deserializeConsExpr cons = do
405+ conLen <- newName "conLen"
406+ off1 <- newName "off1"
407+ [|do ($(varP off1), $(varP conLen) :: Word8) <-
408+ deserialize
409+ $(varE _initialOffset)
410+ $(varE _arr)
411+ $(varE _endOffset)
412+ $(multiIfE (map (guardCon conLen off1) cons ++ [catchAll]))|]
413+
414+ catchAll =
415+ normalGE
416+ [|True|]
417+ [|error
418+ ("Found invalid tag while peeking (" ++
419+ $(lift (pprint headTy)) ++ ")")|]
420+
421+ guardCon conLen off con@(SimpleDataCon cname _) = do
422+ let lenCname = getNameBaseLen cname
423+ tag = map c2w (nameBase cname)
424+ normalGE
425+ [|($(litIntegral lenCname) == $(varE conLen))
426+ && $(xorCmp tag off _arr)|]
427+ [|let $(varP (makeI 0)) = $(varE off) + $(litIntegral lenCname)
428+ in $(mkDeserializeExprOne con)|]
429+
430+mkDeserializeExpr False headTy tyOfTy =
212431 case tyOfTy of
213432 -- Unit constructor
214433 UnitType cname ->
@@ -271,8 +490,40 @@ mkSerializeExprFields fields =
271490 (varP (makeI (i + 1)))
272491 [|serialize $(varE (makeI i)) $(varE _arr) $(varE (mkFieldName i))|]
273492
274−mkSerializeExpr :: TypeOfType -> Q Exp
275−mkSerializeExpr tyOfTy =
493+mkSerializeExpr :: Bool -> TypeOfType -> Q Exp
494+mkSerializeExpr True tyOfTy =
495+ case tyOfTy of
496+ -- Unit type
497+ UnitType cname ->
498+ caseE
499+ (varE _val)
500+ [serializeDataCon (SimpleDataCon cname [])]
501+ -- Product type
502+ (TheType con) ->
503+ caseE
504+ (varE _val)
505+ [serializeDataCon con]
506+ -- Sum type
507+ (MultiType cons) ->
508+ caseE
509+ (varE _val)
510+ (map serializeDataCon cons)
511+
512+ where
513+
514+ serializeDataCon (SimpleDataCon cname fields) = do
515+ let tagLen8 = getNameBaseLen cname
516+ conEnc = tagLen8 : map c2w (nameBase cname)
517+ matchConstructor
518+ cname
519+ (length fields)
520+ (doE [ bindS
521+ (varP (mkName "i0"))
522+ (serializeW8List _initialOffset _arr conEnc)
523+ , noBindS (mkSerializeExprFields fields)
524+ ])
525+
526+mkSerializeExpr False tyOfTy =
276527 case tyOfTy of
277528 -- Unit type
278529 UnitType _ -> [|pure ($(varE _initialOffset) + 1)|]
@@ -341,9 +592,11 @@ mkSerializeExpr tyOfTy =
341592deriveSerializeInternal ::
342593 Config -> Cxt -> Type -> [DataCon] -> Q [Dec]
343594deriveSerializeInternal (Config{..}) preds headTy cons = do
344− sizeOfMethod <- mkSizeOfExpr (typeOfType headTy cons)
345− peekMethod <- mkDeserializeExpr headTy (typeOfType headTy cons)
346− pokeMethod <- mkSerializeExpr (typeOfType headTy cons)
595+ sizeOfMethod <- mkSizeOfExpr constructorTagAsString (typeOfType headTy cons)
596+ peekMethod <-
597+ mkDeserializeExpr constructorTagAsString headTy (typeOfType headTy cons)
598+ pokeMethod <-
599+ mkSerializeExpr constructorTagAsString (typeOfType headTy cons)
347600 let methods =
348601 -- INLINE on sizeOf actually worsens some benchmarks, and improves
349602 -- none
@@ -353,7 +606,7 @@ deriveSerializeInternal (Config{..}) preds headTy cons = do
353606 , FunD
354607 'deserialize
355608 [ Clause
356− (if isUnitType cons
609+ (if isUnitType cons && not constructorTagAsString
357610 then [VarP _initialOffset, WildP, WildP]
358611 else [VarP _initialOffset, VarP _arr, VarP _endOffset])
359612 (NormalB peekMethod)
@@ -363,7 +616,7 @@ deriveSerializeInternal (Config{..}) preds headTy cons = do
363616 , FunD
364617 'serialize
365618 [ Clause
366− (if isUnitType cons
619+ (if isUnitType cons && not constructorTagAsString
367620 then [VarP _initialOffset, WildP, WildP]
368621 else [VarP _initialOffset, VarP _arr, VarP _val])
369622 (NormalB pokeMethod)