@@ -0,0 +1,328 @@
1+{-# LANGUAGE TemplateHaskell #-}
2+
3+-- |
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
9+-- Portability : GHC
10+--
11+module Streamly.Internal.Data.Serialize.TH.RecHeader
12+ ( mkSerializeExpr
13+ , mkDeserializeExpr
14+ , mkSizeOfExpr
15+ , conUpdateFuncExpr
16+ , mkDeserializeKeysExpr
17+ ) where
18+
19+--------------------------------------------------------------------------------
20+-- Imports
21+--------------------------------------------------------------------------------
22+
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(..))
32+
33+import qualified Streamly.Internal.Data.Unbox as Unbox
34+
35+import Streamly.Internal.Data.Serialize.TH.Bottom
36+import Streamly.Internal.Data.Serialize.TH.Common
37+
38+--------------------------------------------------------------------------------
39+-- Notes
40+--------------------------------------------------------------------------------
41+
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.
45+
46+--------------------------------------------------------------------------------
47+-- Compact lists
48+--------------------------------------------------------------------------------
49+
50+-- Like haskell list but the maximum length of the list is 255
51+newtype CompactList a =
52+ CompactList
53+ { unCompactList :: [a]
54+ }
55+
56+-- We use 'Word8' to encode the length, hence the maximim number of elements in
57+-- the list is 255.
58+instance forall a. Serialize a => Serialize (CompactList a) where
59+
60+ -- {-# INLINE size #-}
61+ size acc (CompactList xs) =
62+ foldl' size (acc + (Unbox.sizeOf (Proxy :: Proxy Word8))) xs
63+
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 [])
76+ peekList f o i = do
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)
81+
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
90+ pokeList o1 xs
91+ pokeList off1 val
92+
93+--------------------------------------------------------------------------------
94+-- Helpers
95+--------------------------------------------------------------------------------
96+
97+fieldToNameBase :: Field -> String
98+fieldToNameBase = nameBase . fromJust . fst
99+
100+isMaybeType :: Type -> Bool
101+isMaybeType (AppT (ConT m) _) = m == ''Maybe
102+isMaybeType _ = False
103+
104+--------------------------------------------------------------------------------
105+-- Size
106+--------------------------------------------------------------------------------
107+
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|]
112+
113+sizeOfHeader :: SimpleDataCon -> Int
114+sizeOfHeader (SimpleDataCon _ fields) =
115+ sizeForFinalOff + sizeForHeaderLength + sizeForNumFields
116+ + sum (map ((+ sizeForFieldLen) . length . fieldToNameBase) fields)
117+
118+ where
119+
120+ sizeForFinalOff = 4
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
125+
126+mkSizeOfExpr :: SimpleDataCon -> Q Exp
127+mkSizeOfExpr con = do
128+ n_acc <- newName "acc"
129+ n_x <- newName "x"
130+ (lamE
131+ [varP n_acc, varP n_x]
132+ [|$(litIntegral hlen) +
133+ $(caseE (varE n_x) [matchCons (varE n_acc) con])|])
134+
135+ where
136+
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
142+
143+--------------------------------------------------------------------------------
144+-- Header
145+--------------------------------------------------------------------------------
146+
147+headerValue :: SimpleDataCon -> [Word8]
148+headerValue (SimpleDataCon _ fields) =
149+ int_w8 numFields : concat (fmap lengthPrependedFieldEncoding fields)
150+
151+ where
152+
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.
155+ numFields =
156+ let lenFields = length fields
157+ in if lenFields <= 255
158+ then lenFields
159+ else errorUnsupported
160+ lengthPrependedFieldEncoding field =
161+ let fEnc =
162+ let fEnc_ = map c2w (fieldToNameBase field)
163+ lenFEnc = length fEnc_
164+ in if lenFEnc <= 255
165+ then fEnc_
166+ else errorUnsupported
167+ in (int_w8 (length fEnc)) : fEnc
168+
169+--------------------------------------------------------------------------------
170+-- Peek
171+--------------------------------------------------------------------------------
172+
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)
180+ pure off1
181+
182+mkSerializeExpr :: Name -> SimpleDataCon -> Q Exp
183+mkSerializeExpr initialOffset (con@(SimpleDataCon cname fields)) = do
184+ afterHLen <- newName "afterHLen"
185+ [|do $(varP afterHLen) <-
186+ serialize
187+ ($(varE initialOffset) + 4)
188+ $(varE _arr)
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)
193+ Unbox.pokeByteIndex
194+ $(varE initialOffset)
195+ $(varE _arr)
196+ ((fromIntegral :: Int -> Word32)
197+ (finalOff - $(varE initialOffset)))
198+ pure finalOff|]
199+
200+ where
201+
202+ hval = headerValue con
203+ hlen = length hval
204+
205+--------------------------------------------------------------------------------
206+-- Poke
207+--------------------------------------------------------------------------------
208+
209+{-# INLINE deserializeWithSize #-}
210+deserializeWithSize ::
211+ Serialize a => Int -> MutableByteArray -> Int -> IO (Int, a)
212+deserializeWithSize off arr endOff = deserialize (off + 4) arr endOff
213+
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"
221+ method <-
222+ (caseE
223+ (varE key)
224+ (concat
225+ [ map (matchField arr endOff (prevRec, curOff)) fnames
226+ , [ match
227+ wildP
228+ (normalB
229+ [|do (valOff, valLen :: Word32) <-
230+ deserialize
231+ $(varE curOff)
232+ $(varE arr)
233+ $(varE endOff)
234+ pure
235+ ( $(varE prevRec)
236+ , valOff + w32_int valLen)|])
237+ []
238+ ]
239+ ]))
240+ pure
241+ [ PragmaD (InlineP funcName NoInline FunLike AllPhases)
242+ , FunD
243+ funcName
244+ [ Clause
245+ [ VarP arr
246+ , VarP endOff
247+ , TupP [VarP prevRec, VarP curOff]
248+ , VarP key
249+ ]
250+ (NormalB method)
251+ []
252+ ]
253+ ]
254+
255+ where
256+
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"
262+ match
263+ (litP (StringL (nameBase fname)))
264+ (normalB
265+ [|do ($(varP newOff), $(varP val)) <-
266+ deserializeWithSize
267+ $(varE currOff)
268+ $(varE arr)
269+ $(varE endOff)
270+ pure
271+ ( $(recUpdE (varE prevRec) [pure (fname, VarE val)])
272+ :: $(pure headTy)
273+ , $(varE newOff))|])
274+ []
275+
276+mkDeserializeKeysExpr ::
277+ Name -> Name -> Name -> Name -> Name -> SimpleDataCon -> Q Exp
278+mkDeserializeKeysExpr
279+ hOff finalOff arr endOff updateFunc (SimpleDataCon cname fields) = do
280+
281+ [|do (dataOff, hlist :: CompactList (CompactList Word8)) <-
282+ deserialize $(varE hOff) $(varE arr) $(varE endOff)
283+ let keys = wListToString . unCompactList <$> unCompactList hlist
284+ (finalRec, _) <-
285+ foldlM
286+ ($(varE updateFunc) $(varE arr) $(varE endOff))
287+ ($(emptyConExpr), dataOff)
288+ keys
289+ pure ($(varE finalOff), finalRec)|]
290+
291+ where
292+
293+ emptyTy k ty =
294+ if isMaybeType ty
295+ then [|Nothing|]
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)
300+
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)
313+ then do
314+ let $(varP (makeI 0)) =
315+ $(varE initialOff) +
316+ $(litIntegral sizePreData)
317+ $(mkDeserializeExprOne 'deserializeWithSize con)
318+ else $(varE deserializeWithKeys)
319+ $(varE hOff)
320+ ($(varE initialOff) + w32_int encLen)
321+ else $(varE deserializeWithKeys)
322+ $(varE hOff)
323+ ($(varE initialOff) + w32_int encLen)|]
324+
325+ where
326+
327+ hval = headerValue con
328+ hlen = length hval