git.iko.soy/streamly
on 79df417 help sign in
streamly/core/src/Streamly/Internal/Data/Serialize/TH/RecHeader.hs
411 lines 15 kB history raw
1{-# LANGUAGE TemplateHaskell #-}
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--
11module Streamly.Internal.Data.Serialize.TH.RecHeader
12 ( mkSerializeExpr
13 , mkDeserializeExpr
14 , mkSizeOfExpr
15 , conUpdateFuncDec
16 , mkDeserializeKeysDec
17 ) where
19--------------------------------------------------------------------------------
20-- Imports
21--------------------------------------------------------------------------------
23import Control.Monad (void)
24import Data.List (foldl')
25import Data.Word (Word32, Word8)
26import Data.Maybe (fromJust)
27import Language.Haskell.TH
28import Streamly.Internal.Data.Serialize.Type (Serialize(..))
29import Data.Foldable (foldlM)
30import Streamly.Internal.Data.Unbox (MutableByteArray)
31import Data.Proxy (Proxy(..))
33import qualified Streamly.Internal.Data.Unbox as Unbox
35import Streamly.Internal.Data.Serialize.TH.Bottom
36import Streamly.Internal.Data.Serialize.TH.Common
38--------------------------------------------------------------------------------
39-- Notes
40--------------------------------------------------------------------------------
42-- Compatibility Algorithm
43-- =======================
44--
45-- The algorithm is written without any low level implementation details. See
46-- the code for any low level implementation details.
47--
48-- Serialization:
49-- --------------
50--
51-- To serialize the data,
52--
53-- * Get the list of keys for the record as @keyList@.
54-- * Serialize the @keyList@.
55-- * Serialize the @fields@ one-by-one after serializing the @keyList@.
56--
57-- Deserialization:
58-- ----------------
59--
60-- To deserialize the data to type @T@,
61--
62-- __Checking for type match__:
63--
64-- * Get the list of keys for type @T@ as @targetKeyList@.
65-- * Get the list of keys encoded as @encodedKeyList@.
66-- * If @targetKeyList == encodedKeyList@ see the __Type Match__ section else
67-- see the __No Type Match__ section.
68--
69-- __Type Match__:
70--
71-- * Decode the fields one-by-one and construct the type @T@ in the end.
72--
73-- __No Type Match__:
74--
75-- * Decode the list of keys encoded into @encodedKeyList@.
76-- * Get the list of keys for type @T@ as @targetKeyList@.
77-- * Loop through @encodedKeyList@ and start deserializing the encoded data.
78-- * If the key is present in @encodedKeyList@ and not in @targetKeyList@
79-- then skip parsing the corresponding value.
80-- * If the key is present in @targetKeyList@ and not in @encodedKeyList@
81-- then set the value for that key as @Nothing@.
82-- * If the key is present in both @encodedKeyList@ and in @targetKeyList@
83-- parse the value.
84-- * Construct @T@ after parsing all the data.
86-- Developer Notes
87-- ===============
88--
89-- * Record update syntax is not robust across language extensions and common
90-- record plugins (like record-dot-processor, large-records, etc.).
92--------------------------------------------------------------------------------
93-- Compact lists
94--------------------------------------------------------------------------------
96-- Like haskell list but the maximum length of the list is 255
97newtype CompactList a =
98 CompactList
99 { unCompactList :: [a]
100 }
102-- We use 'Word8' to encode the length, hence the maximim number of elements in
103-- the list is 255.
104instance forall a. Serialize a => Serialize (CompactList a) where
106 -- {-# INLINE size #-}
107 size acc (CompactList xs) =
108 foldl' size (acc + (Unbox.sizeOf (Proxy :: Proxy Word8))) xs
110 -- Inlining this causes large compilation times for tests
111 {-# INLINABLE deserialize #-}
112 deserialize off arr sz = do
113 (off1, len8) <- deserialize off arr sz :: IO (Int, Word8)
114 let len = w8_int len8
115 peekList f o i | i >= 3 = do
116 -- Unfold the loop three times
117 (o1, x1) <- deserialize o arr sz
118 (o2, x2) <- deserialize o1 arr sz
119 (o3, x3) <- deserialize o2 arr sz
120 peekList (f . (\xs -> x1:x2:x3:xs)) o3 (i - 3)
121 peekList f o 0 = pure (o, f [])
122 peekList f o i = do
123 (o1, x) <- deserialize o arr sz
124 peekList (f . (x:)) o1 (i - 1)
125 (nextOff, lst) <- peekList id off1 len
126 pure (nextOff, CompactList lst)
128 -- Inlining this causes large compilation times for tests
129 {-# INLINABLE serialize #-}
130 serialize off arr (CompactList val) = do
131 void $ serialize off arr (int_w8 (length val) :: Word8)
132 let off1 = off + Unbox.sizeOf (Proxy :: Proxy Word8)
133 let pokeList o [] = pure o
134 pokeList o (x:xs) = do
135 o1 <- serialize o arr x
136 pokeList o1 xs
137 pokeList off1 val
139--------------------------------------------------------------------------------
140-- Helpers
141--------------------------------------------------------------------------------
143fieldToNameBase :: Field -> String
144fieldToNameBase = nameBase . fromJust . fst
146isMaybeType :: Type -> Bool
147isMaybeType (AppT (ConT m) _) = m == ''Maybe
148isMaybeType _ = False
150--------------------------------------------------------------------------------
151-- Size
152--------------------------------------------------------------------------------
154-- We add 4 here because we use 'serializeWithSize' for serializing.
155exprGetSize :: Q Exp -> (Int, Type) -> Q Exp
156exprGetSize acc (i, _) =
157 [|size $(acc) $(varE (mkFieldName i)) + 4|]
159sizeOfHeader :: SimpleDataCon -> Int
160sizeOfHeader (SimpleDataCon _ fields) =
161 sizeForFinalOff + sizeForHeaderLength + sizeForNumFields
162 + sum (map ((+ sizeForFieldLen) . length . fieldToNameBase) fields)
164 where
166 sizeForFinalOff = 4
167 sizeForHeaderLength = 4 -- Max header length is (255 * (255 + 1) + 1) and
168 -- hence 2 bytes is enough to store it. But we still
169 -- use 4 bytes as using 2 bytes introduces
170 -- regression.
171 sizeForNumFields = 1 -- At max 255 fields in the record constructor
172 sizeForFieldLen = 1 -- At max 255 letters in the key
174mkSizeOfExpr :: SimpleDataCon -> Q Exp
175mkSizeOfExpr con = do
176 n_acc <- newName "acc"
177 n_x <- newName "x"
178 (lamE
179 [varP n_acc, varP n_x]
180 [|$(litIntegral hlen) +
181 $(caseE (varE n_x) [matchCons (varE n_acc) con])|])
183 where
185 hlen = sizeOfHeader con
186 sizeOfFields acc fields = foldl' exprGetSize acc $ zip [0 ..] fields
187 matchCons acc (SimpleDataCon cname fields) =
188 let expr = sizeOfFields acc (map snd fields)
189 in matchConstructor cname (length fields) expr
191--------------------------------------------------------------------------------
192-- Header
193--------------------------------------------------------------------------------
195headerValue :: SimpleDataCon -> [Word8]
196headerValue (SimpleDataCon _ fields) =
197 int_w8 numFields : concat (fmap lengthPrependedFieldEncoding fields)
199 where
201 -- Error out if the number of fields or the length of key is >= 256. We use
202 -- Word8 for encoding the info and hence the max value is 255.
203 numFields =
204 let lenFields = length fields
205 in if lenFields <= 255
206 then lenFields
207 else errorUnsupported
208 "Number of fields in the record should be <= 255."
209 lengthPrependedFieldEncoding field =
210 let fEnc =
211 let fEnc_ = map c2w (fieldToNameBase field)
212 lenFEnc = length fEnc_
213 in if lenFEnc <= 255
214 then fEnc_
215 else
216 errorUnsupported
217 "Length of any key should be <= 255."
218 in (int_w8 (length fEnc)) : fEnc
220--------------------------------------------------------------------------------
221-- Peek
222--------------------------------------------------------------------------------
224-- Encoding the size is required if we want to skip the field without knowing
225-- its type. We encode the size as 'Word32' hence there is a 4 bytes increase
226-- in size.
227{-# INLINE serializeWithSize #-}
228serializeWithSize :: Serialize a => Int -> MutableByteArray -> a -> IO Int
229serializeWithSize off arr val = do
230 off1 <- serialize (off + 4) arr val
231 Unbox.pokeByteIndex off arr (int_w32 (off1 - off - 4) :: Word32)
232 pure off1
234mkSerializeExpr :: Name -> SimpleDataCon -> Q Exp
235mkSerializeExpr initialOffset (con@(SimpleDataCon cname fields)) = do
236 afterHLen <- newName "afterHLen"
237 -- Encoding the header length is required.
238 -- We first compare the header length encoded and the current header
239 -- length. Only if the header lengths match, we compare the headers.
240 [|do $(varP afterHLen) <-
241 serialize
242 ($(varE initialOffset) + 4)
243 $(varE _arr)
244 ($(litIntegral hlen) :: Word32)
245 $(varP (makeI 0)) <- $(serializeW8List afterHLen _arr hval)
246 let $(openConstructor cname (length fields)) = $(varE _val)
247 finalOff <- $(mkSerializeExprFields 'serializeWithSize fields)
248 Unbox.pokeByteIndex
249 $(varE initialOffset)
250 $(varE _arr)
251 ((fromIntegral :: Int -> Word32)
252 (finalOff - $(varE initialOffset)))
253 pure finalOff|]
255 where
257 hval = headerValue con
258 hlen = length hval
260--------------------------------------------------------------------------------
261-- Poke
262--------------------------------------------------------------------------------
264{-# INLINE deserializeWithSize #-}
265deserializeWithSize ::
266 Serialize a => Int -> MutableByteArray -> Int -> IO (Int, a)
267deserializeWithSize off arr endOff = deserialize (off + 4) arr endOff
269conUpdateFuncDec :: Name -> [Field] -> Q [Dec]
270conUpdateFuncDec funcName fields = do
271 prevAcc <- newName "prevAcc"
272 curOff <- newName "curOff"
273 endOff <- newName "endOff"
274 arr <- newName "arr"
275 key <- newName "key"
276 method <-
277 (caseE
278 (varE key)
279 (concat
280 [ map (matchField arr endOff (prevAcc, curOff)) fnames
281 , [ match
282 wildP
283 (normalB
284 [|do (valOff, valLen :: Word32) <-
285 deserialize
286 $(varE curOff)
287 $(varE arr)
288 $(varE endOff)
289 pure
290 ( $(varE prevAcc)
291 , valOff + w32_int valLen)|])
292 []
293 ]
294 ]))
295 pure
296 [ PragmaD (InlineP funcName NoInline FunLike AllPhases)
297 , FunD
298 funcName
299 [ Clause
300 [ VarP arr
301 , VarP endOff
302 , TupP [VarP prevAcc, VarP curOff]
303 , VarP key
304 ]
305 (NormalB method)
306 []
307 ]
308 ]
310 where
312 fnames = fmap (fromJust . fst) fields
313 matchField :: Name -> Name -> (Name, Name) -> Name -> Q Match
314 matchField arr endOff (acc, currOff) fname = do
315 let fnameLit = StringL (nameBase fname)
316 match
317 (litP fnameLit)
318 (normalB
319 [|do (valOff, valLen :: Word32) <-
320 deserialize
321 $(varE currOff)
322 $(varE arr)
323 $(varE endOff)
324 pure
325 ( ($(litE fnameLit), $(varE currOff)) : $(varE acc)
326 , valOff + w32_int valLen)|])
327 []
329mkDeserializeKeysDec :: Name -> Name -> SimpleDataCon -> Q [Dec]
330mkDeserializeKeysDec funcName updateFunc (SimpleDataCon cname fields) = do
331 hOff <- newName "hOff"
332 finalOff <- newName "finalOff"
333 arr <- newName "arr"
334 endOff <- newName "endOff"
335 kvEncoded <- newName "kvEncoded"
336 finalRec <- newName "finalRec"
337 let deserializeFieldExpr (Just name, ty) = do
338 let nameLit = litE (StringL (nameBase name))
339 [|case lookup $(nameLit) $(varE kvEncoded) of
340 Nothing -> $(emptyTy name ty)
341 Just off -> do
342 val <- deserializeWithSize off $(varE arr) $(varE endOff)
343 pure $ snd val|]
344 deserializeFieldExpr _ =
345 errorUnsupported "The datatype should use record syntax."
346 method <-
347 [|do (dataOff, hlist :: CompactList (CompactList Word8)) <-
348 deserialize $(varE hOff) $(varE arr) $(varE endOff)
349 let keys = wListToString . unCompactList <$> unCompactList hlist
350 ($(varP kvEncoded), _) <-
351 foldlM
352 ($(varE updateFunc) $(varE arr) $(varE endOff))
353 ([], dataOff)
354 keys
355 $(varP finalRec) <-
356 $(foldl
357 (\acc i ->
358 [|$(acc) <*>
359 $(deserializeFieldExpr i)|])
360 [|pure $(conE cname)|]
361 fields)
362 pure ($(varE finalOff), $(varE finalRec))|]
363 pure
364 [ PragmaD (InlineP funcName NoInline FunLike AllPhases)
365 , FunD
366 funcName
367 [ Clause
368 [ VarP hOff
369 , VarP finalOff
370 , VarP arr
371 , VarP endOff
372 ]
373 (NormalB method)
374 []
375 ]
376 ]
378 where
380 emptyTy k ty =
381 if isMaybeType ty
382 then [|pure Nothing|]
383 else [|error $(litE (StringL (nameBase k ++ " is not found.")))|]
386mkDeserializeExpr :: Name -> Name -> Name -> SimpleDataCon -> Q Exp
387mkDeserializeExpr initialOff endOff deserializeWithKeys con = do
388 hOff <- newName "hOff"
389 let sizeForFinalOff = 4 -- Word32
390 sizeForHeaderLength = 4 -- Word32
391 sizePreData = sizeForFinalOff + sizeForHeaderLength + hlen
392 [|do (hlenOff, encLen :: Word32) <-
393 deserialize $(varE initialOff) $(varE _arr) $(varE endOff)
394 ($(varP hOff), hlen1 :: Word32) <-
395 deserialize hlenOff $(varE _arr) $(varE endOff)
396 if (hlen1 == $(litIntegral hlen)) && $(xorCmp hval hOff _arr)
397 then do
398 let $(varP (makeI 0)) =
399 $(varE initialOff) +
400 $(litIntegral sizePreData)
401 $(mkDeserializeExprOne 'deserializeWithSize con)
402 else $(varE deserializeWithKeys)
403 $(varE hOff)
404 ($(varE initialOff) + w32_int encLen)
405 $(varE _arr)
406 $(varE endOff)|]
408 where
410 hval = headerValue con
411 hlen = length hval