@@ -0,0 +1,328 @@
1+{-# LANGUAGE TemplateHaskell #-}
4+-- Module : Streamly.Internal.Data.Serialize.TH.RecHeader
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.RecHeader
16+ , mkDeserializeKeysExpr
19+--------------------------------------------------------------------------------
21+--------------------------------------------------------------------------------
23+import Control.Monad (void)
24+import Data.List (foldl')
25+import Data.Word (Word16, Word32, Word8)
26+import Data.Maybe (fromJust)
27+import Language.Haskell.TH
28+import Streamly.Internal.Data.Serialize.Type (Serialize(..))
29+import Data.Foldable (foldlM)
30+import Streamly.Internal.Data.Unbox (MutableByteArray)
31+import Data.Proxy (Proxy(..))
33+import qualified Streamly.Internal.Data.Unbox as Unbox
35+import Streamly.Internal.Data.Serialize.TH.Bottom
36+import Streamly.Internal.Data.Serialize.TH.Common
38+--------------------------------------------------------------------------------
40+--------------------------------------------------------------------------------
42+-- Encoding the header length is required to check if we want to compare the
43+-- header. We should not compare the headers if the lengths of the headers are
44+-- not equal as this might lead to a runtime error.
46+--------------------------------------------------------------------------------
48+--------------------------------------------------------------------------------
50+-- Like haskell list but the maximum length of the list is 255
51+newtype CompactList a =
53+ { unCompactList :: [a]
56+-- We use 'Word8' to encode the length, hence the maximim number of elements in
58+instance forall a. Serialize a => Serialize (CompactList a) where
60+ -- {-# INLINE size #-}
61+ size acc (CompactList xs) =
62+ foldl' size (acc + (Unbox.sizeOf (Proxy :: Proxy Word8))) xs
64+ -- Inlining this causes large compilation times for tests
65+ {-# INLINABLE deserialize #-}
66+ deserialize off arr sz = do
67+ (off1, len8) <- deserialize off arr sz :: IO (Int, Word8)
68+ let len = w8_int len8
69+ peekList f o i | i >= 3 = do
70+ -- Unfold the loop three times
71+ (o1, x1) <- deserialize o arr sz
72+ (o2, x2) <- deserialize o1 arr sz
73+ (o3, x3) <- deserialize o2 arr sz
74+ peekList (f . (\xs -> x1:x2:x3:xs)) o3 (i - 3)
75+ peekList f o 0 = pure (o, f [])
77+ (o1, x) <- deserialize o arr sz
78+ peekList (f . (x:)) o1 (i - 1)
79+ (nextOff, lst) <- peekList id off1 len
80+ pure (nextOff, CompactList lst)
82+ -- Inlining this causes large compilation times for tests
83+ {-# INLINABLE serialize #-}
84+ serialize off arr (CompactList val) = do
85+ void $ serialize off arr (int_w8 (length val) :: Word8)
86+ let off1 = off + Unbox.sizeOf (Proxy :: Proxy Word8)
87+ let pokeList o [] = pure o
88+ pokeList o (x:xs) = do
89+ o1 <- serialize o arr x
93+--------------------------------------------------------------------------------
95+--------------------------------------------------------------------------------
97+fieldToNameBase :: Field -> String
98+fieldToNameBase = nameBase . fromJust . fst
100+isMaybeType :: Type -> Bool
101+isMaybeType (AppT (ConT m) _) = m == ''Maybe
102+isMaybeType _ = False
104+--------------------------------------------------------------------------------
106+--------------------------------------------------------------------------------
108+-- We add 4 here because we use 'serializeWithSize' for serializing.
109+exprGetSize :: Q Exp -> (Int, Type) -> Q Exp
110+exprGetSize acc (i, _) =
111+ [|size $(acc) $(varE (mkFieldName i)) + 4|]
113+sizeOfHeader :: SimpleDataCon -> Int
114+sizeOfHeader (SimpleDataCon _ fields) =
115+ sizeForFinalOff + sizeForHeaderLength + sizeForNumFields
116+ + sum (map ((+ sizeForFieldLen) . length . fieldToNameBase) fields)
121+ sizeForHeaderLength = 2 -- Max header length is (255 * (255 + 1) + 1) and
122+ -- hence 2 bytes is enough to store it.
123+ sizeForNumFields = 1 -- At max 255 fields in the record constructor
124+ sizeForFieldLen = 1 -- At max 255 letters in the key
126+mkSizeOfExpr :: SimpleDataCon -> Q Exp
127+mkSizeOfExpr con = do
128+ n_acc <- newName "acc"
131+ [varP n_acc, varP n_x]
132+ [|$(litIntegral hlen) +
133+ $(caseE (varE n_x) [matchCons (varE n_acc) con])|])
137+ hlen = sizeOfHeader con
138+ sizeOfFields acc fields = 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+--------------------------------------------------------------------------------
145+--------------------------------------------------------------------------------
147+headerValue :: SimpleDataCon -> [Word8]
148+headerValue (SimpleDataCon _ fields) =
149+ int_w8 numFields : concat (fmap lengthPrependedFieldEncoding fields)
153+ -- Error out if the number of fields or the length of key is >= 256. We use
154+ -- Word8 for encoding the info and hence the max value is 255.
156+ let lenFields = length fields
157+ in if lenFields <= 255
159+ else errorUnsupported
160+ lengthPrependedFieldEncoding field =
162+ let fEnc_ = map c2w (fieldToNameBase field)
163+ lenFEnc = length fEnc_
164+ in if lenFEnc <= 255
166+ else errorUnsupported
167+ in (int_w8 (length fEnc)) : fEnc
169+--------------------------------------------------------------------------------
171+--------------------------------------------------------------------------------
173+-- Encoding the size is required if we want to skip the field.
174+-- We encode the size as 'Word32' hence there is a 4 bytes increase in size.
175+{-# INLINE serializeWithSize #-}
176+serializeWithSize :: Serialize a => Int -> MutableByteArray -> a -> IO Int
177+serializeWithSize off arr val = do
178+ off1 <- serialize (off + 4) arr val
179+ Unbox.pokeByteIndex off arr (int_w32 (off1 - off - 4) :: Word32)
182+mkSerializeExpr :: Name -> SimpleDataCon -> Q Exp
183+mkSerializeExpr initialOffset (con@(SimpleDataCon cname fields)) = do
184+ afterHLen <- newName "afterHLen"
185+ [|do $(varP afterHLen) <-
187+ ($(varE initialOffset) + 4)
189+ ($(litIntegral hlen) :: Word16)
190+ $(varP (makeI 0)) <- $(serializeW8List afterHLen _arr hval)
191+ let $(openConstructor cname (length fields)) = $(varE _val)
192+ finalOff <- $(mkSerializeExprFields 'serializeWithSize fields)
194+ $(varE initialOffset)
196+ ((fromIntegral :: Int -> Word32)
197+ (finalOff - $(varE initialOffset)))
202+ hval = headerValue con
205+--------------------------------------------------------------------------------
207+--------------------------------------------------------------------------------
209+{-# INLINE deserializeWithSize #-}
210+deserializeWithSize ::
211+ Serialize a => Int -> MutableByteArray -> Int -> IO (Int, a)
212+deserializeWithSize off arr endOff = deserialize (off + 4) arr endOff
214+conUpdateFuncExpr :: Name -> Type -> [Field] -> Q [Dec]
215+conUpdateFuncExpr funcName headTy fields = do
216+ prevRec <- newName "prevRec"
217+ curOff <- newName "curOff"
218+ endOff <- newName "endOff"
219+ arr <- newName "arr"
220+ key <- newName "key"
225+ [ map (matchField arr endOff (prevRec, curOff)) fnames
229+ [|do (valOff, valLen :: Word32) <-
236+ , valOff + w32_int valLen)|])
241+ [ PragmaD (InlineP funcName NoInline FunLike AllPhases)
247+ , TupP [VarP prevRec, VarP curOff]
257+ fnames = fmap (fromJust . fst) fields
258+ matchField :: Name -> Name -> (Name, Name) -> Name -> Q Match
259+ matchField arr endOff (prevRec, currOff) fname = do
260+ val <- newName "val"
261+ newOff <- newName "newOff"
263+ (litP (StringL (nameBase fname)))
265+ [|do ($(varP newOff), $(varP val)) <-
271+ ( $(recUpdE (varE prevRec) [pure (fname, VarE val)])
273+ , $(varE newOff))|])
276+mkDeserializeKeysExpr ::
277+ Name -> Name -> Name -> Name -> Name -> SimpleDataCon -> Q Exp
278+mkDeserializeKeysExpr
279+ hOff finalOff arr endOff updateFunc (SimpleDataCon cname fields) = do
281+ [|do (dataOff, hlist :: CompactList (CompactList Word8)) <-
282+ deserialize $(varE hOff) $(varE arr) $(varE endOff)
283+ let keys = wListToString . unCompactList <$> unCompactList hlist
286+ ($(varE updateFunc) $(varE arr) $(varE endOff))
287+ ($(emptyConExpr), dataOff)
289+ pure ($(varE finalOff), finalRec)|]
296+ else [|error $(litE (StringL (nameBase k ++ " is not found.")))|]
297+ fieldToEmptyRecC (Just name, ty) = (name, ) <$> emptyTy name ty
298+ fieldToEmptyRecC _ = errorUnsupported
299+ emptyConExpr = recConE cname (fmap fieldToEmptyRecC fields)
301+mkDeserializeExpr :: Name -> Name -> Name -> SimpleDataCon -> Q Exp
302+mkDeserializeExpr initialOff endOff deserializeWithKeys con = do
303+ hOff <- newName "hOff"
304+ let sizeForFinalOff = 4 -- Word32
305+ sizeForHeaderLength = 2 -- Word16
306+ sizePreData = sizeForFinalOff + sizeForHeaderLength + hlen
307+ [|do (hlenOff, encLen :: Word32) <-
308+ deserialize $(varE initialOff) $(varE _arr) $(varE endOff)
309+ ($(varP hOff), hlen1 :: Word16) <-
310+ deserialize hlenOff $(varE _arr) $(varE endOff)
311+ if hlen1 == $(litIntegral hlen)
312+ then if $(xorCmp hval hOff _arr)
314+ let $(varP (makeI 0)) =
315+ $(varE initialOff) +
316+ $(litIntegral sizePreData)
317+ $(mkDeserializeExprOne 'deserializeWithSize con)
318+ else $(varE deserializeWithKeys)
320+ ($(varE initialOff) + w32_int encLen)
321+ else $(varE deserializeWithKeys)
323+ ($(varE initialOff) + w32_int encLen)|]
327+ hval = headerValue con