@@ -19,13 +19,18 @@ module Streamly.Internal.Data.Serialize.TH
20--------------------------------------------------------------------------------
20--------------------------------------------------------------------------------
22+import Data.Char (ord)
22import Data.List (foldl')
23import Data.List (foldl')
23import Data.Word (Word16, Word32, Word64, Word8)
24import Data.Word (Word16, Word32, Word64, Word8)
25+import Data.Bits (Bits, (.|.), shiftL, zeroBits, xor)
26+import Streamly.Internal.System.IO (unsafeInlineIO)
25import Language.Haskell.TH
28import Language.Haskell.TH
26import Language.Haskell.TH.Syntax
29import Language.Haskell.TH.Syntax
27import Streamly.Internal.Data.Serialize
30import Streamly.Internal.Data.Serialize
32+import qualified Streamly.Internal.Data.Unbox as Unbox
29import Streamly.Internal.Data.Unbox.TH
34import Streamly.Internal.Data.Unbox.TH
30 ( DataCon(..)
35 ( DataCon(..)
31 , DataType(..)
36 , DataType(..)
@@ -67,6 +72,9 @@ mkFieldName i = mkName ("field" ++ show i)
67makeI :: Int -> Name
72makeI :: Int -> Name
68makeI i = mkName $ "i" ++ show i
73makeI i = mkName $ "i" ++ show i
76+makeN i = mkName $ "n" ++ show i
70makeA :: Int -> Name
78makeA :: Int -> Name
71makeA i = mkName $ "a" ++ show i
79makeA i = mkName $ "a" ++ show i
@@ -136,8 +144,185 @@ isUnitType _ = False
137--------------------------------------------------------------------------------
145--------------------------------------------------------------------------------
148+int_w8 :: Int -> Word8
149+int_w8 = fromIntegral
152+w8_w16 :: Word8 -> Word16
153+w8_w16 = fromIntegral
155+w8_w32 :: Word8 -> Word32
156+w8_w32 = fromIntegral
158+w8_w64 :: Word8 -> Word64
159+w8_w64 = fromIntegral
162+w32_int :: Word32 -> Int
163+w32_int = fromIntegral
167+c2w = fromIntegral . ord
169+shiftAdd :: Bits a => (b -> a) -> [b] -> a
171+ foldl' (.|.) zeroBits $
172+ map (\(j, x) -> shiftL x (j * 8)) $ zip [0 ..] $ map conv xs
174+-- Note: This only works in little endian machines
175+xorCmp :: [Word8] -> Name -> Name -> Q Exp
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|]
184+ go8 i | i >= tagLen = [|zeroBits|]
186+ let wIntegral = litIntegral i
187+ [|xor (unsafeInlineIO
188+ (Unbox.peekByteIndex
189+ ($(varE off) + $(litIntegral i))
191+ ($(wIntegral) :: Word8) .|.
194+ | i >= tagLen = [|zeroBits|]
196+ | tagLen - i < 2 = go16 (tagLen - 2)
200+ (shiftAdd w8_w16 [tag !! i, tag !! (i + 1)] :: Word16)
201+ [|xor (unsafeInlineIO
202+ (Unbox.peekByteIndex
203+ ($(varE off) + $(litIntegral i))
205+ ($(wIntegral) :: Word16) .|.
208+ | i >= tagLen = [|zeroBits|]
210+ | tagLen - i < 4 = go32 (tagLen - 4)
221+ [|xor (unsafeInlineIO
222+ (Unbox.peekByteIndex
223+ ($(varE off) + $(litIntegral i))
225+ ($(wIntegral) :: Word32) .|.
228+ | i >= tagLen = [|zeroBits|]
230+ | tagLen - i < 8 = go64 (tagLen - 8)
245+ [|xor (unsafeInlineIO
246+ (Unbox.peekByteIndex
247+ ($(varE off) + $(litIntegral i))
249+ ($(wIntegral) :: Word64) .|.
252+--------------------------------------------------------------------------------
253+-- Primitive serialization
254+--------------------------------------------------------------------------------
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))|])]))|]
265+ lenW8List = length w8List
268+ (varP (makeN (i + 1)))
269+ [|$(varE 'serialize)
272+ ($(litIntegral (w8List !! i)) :: Word8)|]
274+--------------------------------------------------------------------------------
276+--------------------------------------------------------------------------------
278+litIntegral :: Integral a => a -> Q Exp
279+litIntegral = litE . IntegerL . fromIntegral
281+getNameBaseLen :: Name -> Word8
282+getNameBaseLen cname =
283+ let x = length (nameBase cname)
285+ then error "Max Constructor Len: 63 characters"
288+conEncLen :: Name -> Word8
289+conEncLen cname = getNameBaseLen cname + 1
291+--------------------------------------------------------------------------------
293+--------------------------------------------------------------------------------
139−mkSizeOfExpr :: TypeOfType -> Q Exp
295+mkSizeOfExpr :: Bool -> TypeOfType -> Q Exp
140−mkSizeOfExpr tyOfTy =
296+mkSizeOfExpr True tyOfTy =
301+ [|$(varE _acc) + $(litIntegral (conEncLen cname))|]
304+ [varP _acc, varP _x]
305+ (caseE (varE _x) [matchCons (varE _acc) con])
306+ MultiType constructors -> sizeOfHeadDt constructors
310+ sizeOfFields acc fields =
311+ foldl' exprGetSize acc $ zip [0..] fields
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
320+ let acc = [|$(varE _acc)|]
322+ [varP _acc, varP _x]
323+ (caseE (varE _x) (fmap (matchCons acc) cons))
325+mkSizeOfExpr False tyOfTy =
141 case tyOfTy of
326 case tyOfTy of
142 UnitType _ -> lamE [varP _acc, wildP] [|$(varE _acc) + 1|]
327 UnitType _ -> lamE [varP _acc, wildP] [|$(varE _acc) + 1|]
143 TheType con ->
328 TheType con ->
@@ -207,8 +392,42 @@ mkDeserializeExprOne (SimpleDataCon cname fields) =
207 [|deserialize $(varE (makeI i)) $(varE _arr) $(varE _endOffset)|]
392 [|deserialize $(varE (makeI i)) $(varE _arr) $(varE _endOffset)|]
210−mkDeserializeExpr :: Type -> TypeOfType -> Q Exp
395+mkDeserializeExpr :: Bool -> Type -> TypeOfType -> Q Exp
211−mkDeserializeExpr headTy tyOfTy =
396+mkDeserializeExpr True headTy tyOfTy =
398+ UnitType cname -> deserializeConsExpr [SimpleDataCon cname []]
399+ TheType con -> deserializeConsExpr [con]
400+ MultiType cons -> deserializeConsExpr cons
404+ deserializeConsExpr cons = do
405+ conLen <- newName "conLen"
406+ off1 <- newName "off1"
407+ [|do ($(varP off1), $(varP conLen) :: Word8) <-
409+ $(varE _initialOffset)
412+ $(multiIfE (map (guardCon conLen off1) cons ++ [catchAll]))|]
418+ ("Found invalid tag while peeking (" ++
419+ $(lift (pprint headTy)) ++ ")")|]
421+ guardCon conLen off con@(SimpleDataCon cname _) = do
422+ let lenCname = getNameBaseLen cname
423+ tag = map c2w (nameBase cname)
425+ [|($(litIntegral lenCname) == $(varE conLen))
426+ && $(xorCmp tag off _arr)|]
427+ [|let $(varP (makeI 0)) = $(varE off) + $(litIntegral lenCname)
428+ in $(mkDeserializeExprOne con)|]
430+mkDeserializeExpr False headTy tyOfTy =
212 case tyOfTy of
431 case tyOfTy of
213 -- Unit constructor
432 -- Unit constructor
214 UnitType cname ->
433 UnitType cname ->
@@ -271,8 +490,40 @@ mkSerializeExprFields fields =
271 (varP (makeI (i + 1)))
490 (varP (makeI (i + 1)))
272 [|serialize $(varE (makeI i)) $(varE _arr) $(varE (mkFieldName i))|]
491 [|serialize $(varE (makeI i)) $(varE _arr) $(varE (mkFieldName i))|]
274−mkSerializeExpr :: TypeOfType -> Q Exp
493+mkSerializeExpr :: Bool -> TypeOfType -> Q Exp
275−mkSerializeExpr tyOfTy =
494+mkSerializeExpr True tyOfTy =
500+ [serializeDataCon (SimpleDataCon cname [])]
505+ [serializeDataCon con]
510+ (map serializeDataCon cons)
514+ serializeDataCon (SimpleDataCon cname fields) = do
515+ let tagLen8 = getNameBaseLen cname
516+ conEnc = tagLen8 : map c2w (nameBase cname)
521+ (varP (mkName "i0"))
522+ (serializeW8List _initialOffset _arr conEnc)
523+ , noBindS (mkSerializeExprFields fields)
526+mkSerializeExpr False tyOfTy =
276 case tyOfTy of
527 case tyOfTy of
277 -- Unit type
528 -- Unit type
278 UnitType _ -> [|pure ($(varE _initialOffset) + 1)|]
529 UnitType _ -> [|pure ($(varE _initialOffset) + 1)|]
@@ -341,9 +592,11 @@ mkSerializeExpr tyOfTy =
341deriveSerializeInternal ::
592deriveSerializeInternal ::
342 Config -> Cxt -> Type -> [DataCon] -> Q [Dec]
593 Config -> Cxt -> Type -> [DataCon] -> Q [Dec]
343deriveSerializeInternal (Config{..}) preds headTy cons = do
594deriveSerializeInternal (Config{..}) preds headTy cons = do
344− sizeOfMethod <- mkSizeOfExpr (typeOfType headTy cons)
595+ sizeOfMethod <- mkSizeOfExpr constructorTagAsString (typeOfType headTy cons)
345− peekMethod <- mkDeserializeExpr headTy (typeOfType headTy cons)
597+ mkDeserializeExpr constructorTagAsString headTy (typeOfType headTy cons)
346− pokeMethod <- mkSerializeExpr (typeOfType headTy cons)
599+ mkSerializeExpr constructorTagAsString (typeOfType headTy cons)
347 let methods =
600 let methods =
348 -- INLINE on sizeOf actually worsens some benchmarks, and improves
601 -- INLINE on sizeOf actually worsens some benchmarks, and improves
@@ -353,7 +606,7 @@ deriveSerializeInternal (Config{..}) preds headTy cons = do
354 'deserialize
607 'deserialize
356− (if isUnitType cons
609+ (if isUnitType cons && not constructorTagAsString
357 then [VarP _initialOffset, WildP, WildP]
610 then [VarP _initialOffset, WildP, WildP]
358 else [VarP _initialOffset, VarP _arr, VarP _endOffset])
611 else [VarP _initialOffset, VarP _arr, VarP _endOffset])
359 (NormalB peekMethod)
612 (NormalB peekMethod)
@@ -363,7 +616,7 @@ deriveSerializeInternal (Config{..}) preds headTy cons = do
364 'serialize
617 'serialize
366− (if isUnitType cons
619+ (if isUnitType cons && not constructorTagAsString
367 then [VarP _initialOffset, WildP, WildP]
620 then [VarP _initialOffset, WildP, WildP]
368 else [VarP _initialOffset, VarP _arr, VarP _val])
621 else [VarP _initialOffset, VarP _arr, VarP _val])
369 (NormalB pokeMethod)
622 (NormalB pokeMethod)