{-# LANGUAGE TemplateHaskell #-}
module Streamly.Internal.Data.Serialize.TH.RecHeader
( mkRecSerializeExpr
, mkRecDeserializeExpr
, mkRecSizeOfExpr
, conUpdateFuncDec
, mkDeserializeKeysDec
) where
import Control.Monad (void)
import Data.List (foldl')
import Data.Word (Word32, Word8)
import Data.Maybe (fromJust)
import Language.Haskell.TH
import Streamly.Internal.Data.Serialize.Type (Serialize(..))
import Data.Foldable (foldlM)
import Streamly.Internal.Data.MutByteArray.Type (MutByteArray)
import Data.Proxy (Proxy(..))
import qualified Streamly.Internal.Data.Unbox as Unbox
import Streamly.Internal.Data.Serialize.TH.Bottom
import Streamly.Internal.Data.Serialize.TH.Common
newtype CompactList a =
CompactList
{ forall a. CompactList a -> [a]
unCompactList :: [a]
}
instance forall a. Serialize a => Serialize (CompactList a) where
addSizeTo :: Int -> CompactList a -> Int
addSizeTo Int
acc (CompactList [a]
xs) =
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' forall a. Serialize a => Int -> a -> Int
addSizeTo (Int
acc forall a. Num a => a -> a -> a
+ (forall a. Unbox a => Proxy a -> Int
Unbox.sizeOf (forall {k} (t :: k). Proxy t
Proxy :: Proxy Word8))) [a]
xs
{-# INLINABLE deserializeAt #-}
deserializeAt :: Int -> MutByteArray -> Int -> IO (Int, CompactList a)
deserializeAt Int
off MutByteArray
arr Int
sz = do
(Int
off1, Word8
len8) <- forall a. Serialize a => Int -> MutByteArray -> Int -> IO (Int, a)
deserializeAt Int
off MutByteArray
arr Int
sz :: IO (Int, Word8)
let len :: Int
len = Word8 -> Int
w8_int Word8
len8
peekList :: ([a] -> b) -> Int -> t -> IO (Int, b)
peekList [a] -> b
f Int
o t
i | t
i forall a. Ord a => a -> a -> Bool
>= t
3 = do
(Int
o1, a
x1) <- forall a. Serialize a => Int -> MutByteArray -> Int -> IO (Int, a)
deserializeAt Int
o MutByteArray
arr Int
sz
(Int
o2, a
x2) <- forall a. Serialize a => Int -> MutByteArray -> Int -> IO (Int, a)
deserializeAt Int
o1 MutByteArray
arr Int
sz
(Int
o3, a
x3) <- forall a. Serialize a => Int -> MutByteArray -> Int -> IO (Int, a)
deserializeAt Int
o2 MutByteArray
arr Int
sz
([a] -> b) -> Int -> t -> IO (Int, b)
peekList ([a] -> b
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. (\[a]
xs -> a
x1forall a. a -> [a] -> [a]
:a
x2forall a. a -> [a] -> [a]
:a
x3forall a. a -> [a] -> [a]
:[a]
xs)) Int
o3 (t
i forall a. Num a => a -> a -> a
- t
3)
peekList [a] -> b
f Int
o t
0 = forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
o, [a] -> b
f [])
peekList [a] -> b
f Int
o t
i = do
(Int
o1, a
x) <- forall a. Serialize a => Int -> MutByteArray -> Int -> IO (Int, a)
deserializeAt Int
o MutByteArray
arr Int
sz
([a] -> b) -> Int -> t -> IO (Int, b)
peekList ([a] -> b
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a
xforall a. a -> [a] -> [a]
:)) Int
o1 (t
i forall a. Num a => a -> a -> a
- t
1)
(Int
nextOff, [a]
lst) <- forall {t} {a} {b}.
(Ord t, Num t, Serialize a) =>
([a] -> b) -> Int -> t -> IO (Int, b)
peekList forall a. a -> a
id Int
off1 Int
len
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Int
nextOff, forall a. [a] -> CompactList a
CompactList [a]
lst)
{-# INLINABLE serializeAt #-}
serializeAt :: Int -> MutByteArray -> CompactList a -> IO Int
serializeAt Int
off MutByteArray
arr (CompactList [a]
val) = do
forall (f :: * -> *) a. Functor f => f a -> f ()
void forall a b. (a -> b) -> a -> b
$ forall a. Serialize a => Int -> MutByteArray -> a -> IO Int
serializeAt Int
off MutByteArray
arr (Int -> Word8
int_w8 (forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
val) :: Word8)
let off1 :: Int
off1 = Int
off forall a. Num a => a -> a -> a
+ forall a. Unbox a => Proxy a -> Int
Unbox.sizeOf (forall {k} (t :: k). Proxy t
Proxy :: Proxy Word8)
let pokeList :: Int -> [a] -> IO Int
pokeList Int
o [] = forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
o
pokeList Int
o (a
x:[a]
xs) = do
Int
o1 <- forall a. Serialize a => Int -> MutByteArray -> a -> IO Int
serializeAt Int
o MutByteArray
arr a
x
Int -> [a] -> IO Int
pokeList Int
o1 [a]
xs
forall {a}. Serialize a => Int -> [a] -> IO Int
pokeList Int
off1 [a]
val
fieldToNameBase :: Field -> String
fieldToNameBase :: Field -> String
fieldToNameBase = Name -> String
nameBase forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. HasCallStack => Maybe a -> a
fromJust forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst
isMaybeType :: Type -> Bool
isMaybeType :: Type -> Bool
isMaybeType (AppT (ConT Name
m) Type
_) = Name
m forall a. Eq a => a -> a -> Bool
== ''Maybe
isMaybeType Type
_ = Bool
False
exprGetSize :: Q Exp -> (Int, Type) -> Q Exp
exprGetSize :: Q Exp -> (Int, Type) -> Q Exp
exprGetSize Q Exp
acc (Int
i, Type
_) =
[|addSizeTo $(acc) $(varE (mkFieldName i)) + 4|]
sizeOfHeader :: SimpleDataCon -> Int
(SimpleDataCon Name
_ [Field]
fields) =
Int
sizeForFinalOff forall a. Num a => a -> a -> a
+ Int
sizeForHeaderLength forall a. Num a => a -> a -> a
+ Int
sizeForNumFields
forall a. Num a => a -> a -> a
+ forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum (forall a b. (a -> b) -> [a] -> [b]
map ((forall a. Num a => a -> a -> a
+ Int
sizeForFieldLen) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (t :: * -> *) a. Foldable t => t a -> Int
length forall b c a. (b -> c) -> (a -> b) -> a -> c
. Field -> String
fieldToNameBase) [Field]
fields)
where
sizeForFinalOff :: Int
sizeForFinalOff = Int
4
sizeForHeaderLength :: Int
sizeForHeaderLength = Int
4
sizeForNumFields :: Int
sizeForNumFields = Int
1
sizeForFieldLen :: Int
sizeForFieldLen = Int
1
mkRecSizeOfExpr :: SimpleDataCon -> Q Exp
mkRecSizeOfExpr :: SimpleDataCon -> Q Exp
mkRecSizeOfExpr SimpleDataCon
con = do
Name
n_acc <- forall (m :: * -> *). Quote m => String -> m Name
newName String
"acc"
Name
n_x <- forall (m :: * -> *). Quote m => String -> m Name
newName String
"x"
(forall (m :: * -> *). Quote m => [m Pat] -> m Exp -> m Exp
lamE
[forall (m :: * -> *). Quote m => Name -> m Pat
varP Name
n_acc, forall (m :: * -> *). Quote m => Name -> m Pat
varP Name
n_x]
[|$(litIntegral hlen) +
$(caseE (varE n_x) [matchCons (varE n_acc) con])|])
where
hlen :: Int
hlen = SimpleDataCon -> Int
sizeOfHeader SimpleDataCon
con
sizeOfFields :: Q Exp -> [Type] -> Q Exp
sizeOfFields Q Exp
acc [Type]
fields = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Q Exp -> (Int, Type) -> Q Exp
exprGetSize Q Exp
acc forall a b. (a -> b) -> a -> b
$ forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 ..] [Type]
fields
matchCons :: Q Exp -> SimpleDataCon -> Q Match
matchCons Q Exp
acc (SimpleDataCon Name
cname [Field]
fields) =
let expr :: Q Exp
expr = Q Exp -> [Type] -> Q Exp
sizeOfFields Q Exp
acc (forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> b
snd [Field]
fields)
in Name -> Int -> Q Exp -> Q Match
matchConstructor Name
cname (forall (t :: * -> *) a. Foldable t => t a -> Int
length [Field]
fields) Q Exp
expr
headerValue :: SimpleDataCon -> [Word8]
(SimpleDataCon Name
_ [Field]
fields) =
Int -> Word8
int_w8 Int
numFields forall a. a -> [a] -> [a]
: forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Field -> [Word8]
lengthPrependedFieldEncoding [Field]
fields)
where
numFields :: Int
numFields =
let lenFields :: Int
lenFields = forall (t :: * -> *) a. Foldable t => t a -> Int
length [Field]
fields
in if Int
lenFields forall a. Ord a => a -> a -> Bool
<= Int
255
then Int
lenFields
else forall a. String -> a
errorUnsupported
String
"Number of fields in the record should be <= 255."
lengthPrependedFieldEncoding :: Field -> [Word8]
lengthPrependedFieldEncoding Field
field =
let fEnc :: [Word8]
fEnc =
let fEnc_ :: [Word8]
fEnc_ = forall a b. (a -> b) -> [a] -> [b]
map Char -> Word8
c2w (Field -> String
fieldToNameBase Field
field)
lenFEnc :: Int
lenFEnc = forall (t :: * -> *) a. Foldable t => t a -> Int
length [Word8]
fEnc_
in if Int
lenFEnc forall a. Ord a => a -> a -> Bool
<= Int
255
then [Word8]
fEnc_
else
forall a. String -> a
errorUnsupported
String
"Length of any key should be <= 255."
in (Int -> Word8
int_w8 (forall (t :: * -> *) a. Foldable t => t a -> Int
length [Word8]
fEnc)) forall a. a -> [a] -> [a]
: [Word8]
fEnc
{-# INLINE serializeWithSize #-}
serializeWithSize :: Serialize a => Int -> MutByteArray -> a -> IO Int
serializeWithSize :: forall a. Serialize a => Int -> MutByteArray -> a -> IO Int
serializeWithSize Int
off MutByteArray
arr a
val = do
Int
off1 <- forall a. Serialize a => Int -> MutByteArray -> a -> IO Int
serializeAt (Int
off forall a. Num a => a -> a -> a
+ Int
4) MutByteArray
arr a
val
forall a. Unbox a => Int -> MutByteArray -> a -> IO ()
Unbox.pokeAt Int
off MutByteArray
arr (Int -> Word32
int_w32 (Int
off1 forall a. Num a => a -> a -> a
- Int
off forall a. Num a => a -> a -> a
- Int
4) :: Word32)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
off1
mkRecSerializeExpr :: Name -> SimpleDataCon -> Q Exp
mkRecSerializeExpr :: Name -> SimpleDataCon -> Q Exp
mkRecSerializeExpr Name
initialOffset (con :: SimpleDataCon
con@(SimpleDataCon Name
cname [Field]
fields)) = do
Name
afterHLen <- forall (m :: * -> *). Quote m => String -> m Name
newName String
"afterHLen"
[|do $(varP afterHLen) <-
serializeAt
($(varE initialOffset) + 4)
$(varE _arr)
($(litIntegral hlen) :: Word32)
$(varP (makeI 0)) <- $(serializeW8List afterHLen _arr hval)
let $(openConstructor cname (length fields)) = $(varE _val)
finalOff <- $(mkSerializeExprFields 'serializeWithSize fields)
Unbox.pokeAt
$(varE initialOffset)
$(varE _arr)
((fromIntegral :: Int -> Word32)
(finalOff - $(varE initialOffset)))
pure finalOff|]
where
hval :: [Word8]
hval = SimpleDataCon -> [Word8]
headerValue SimpleDataCon
con
hlen :: Int
hlen = forall (t :: * -> *) a. Foldable t => t a -> Int
length [Word8]
hval
{-# INLINE deserializeWithSize #-}
deserializeWithSize ::
Serialize a => Int -> MutByteArray -> Int -> IO (Int, a)
deserializeWithSize :: forall a. Serialize a => Int -> MutByteArray -> Int -> IO (Int, a)
deserializeWithSize Int
off MutByteArray
arr Int
endOff = forall a. Serialize a => Int -> MutByteArray -> Int -> IO (Int, a)
deserializeAt (Int
off forall a. Num a => a -> a -> a
+ Int
4) MutByteArray
arr Int
endOff
conUpdateFuncDec :: Name -> [Field] -> Q [Dec]
conUpdateFuncDec :: Name -> [Field] -> Q [Dec]
conUpdateFuncDec Name
funcName [Field]
fields = do
Name
prevAcc <- forall (m :: * -> *). Quote m => String -> m Name
newName String
"prevAcc"
Name
curOff <- forall (m :: * -> *). Quote m => String -> m Name
newName String
"curOff"
Name
endOff <- forall (m :: * -> *). Quote m => String -> m Name
newName String
"endOff"
Name
arr <- forall (m :: * -> *). Quote m => String -> m Name
newName String
"arr"
Name
key <- forall (m :: * -> *). Quote m => String -> m Name
newName String
"key"
Exp
method <-
(forall (m :: * -> *). Quote m => m Exp -> [m Match] -> m Exp
caseE
(forall (m :: * -> *). Quote m => Name -> m Exp
varE Name
key)
(forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat
[ forall a b. (a -> b) -> [a] -> [b]
map (Name -> Name -> (Name, Name) -> Name -> Q Match
matchField Name
arr Name
endOff (Name
prevAcc, Name
curOff)) [Name]
fnames
, [ forall (m :: * -> *).
Quote m =>
m Pat -> m Body -> [m Dec] -> m Match
match
forall (m :: * -> *). Quote m => m Pat
wildP
(forall (m :: * -> *). Quote m => m Exp -> m Body
normalB
[|do (valOff, valLen :: Word32) <-
deserializeAt
$(varE curOff)
$(varE arr)
$(varE endOff)
pure
( $(varE prevAcc)
, valOff + w32_int valLen)|])
[]
]
]))
forall (f :: * -> *) a. Applicative f => a -> f a
pure
[ Pragma -> Dec
PragmaD (Name -> Inline -> RuleMatch -> Phases -> Pragma
InlineP Name
funcName Inline
NoInline RuleMatch
FunLike Phases
AllPhases)
, Name -> [Clause] -> Dec
FunD
Name
funcName
[ [Pat] -> Body -> [Dec] -> Clause
Clause
[ Name -> Pat
VarP Name
arr
, Name -> Pat
VarP Name
endOff
, [Pat] -> Pat
TupP [Name -> Pat
VarP Name
prevAcc, Name -> Pat
VarP Name
curOff]
, Name -> Pat
VarP Name
key
]
(Exp -> Body
NormalB Exp
method)
[]
]
]
where
fnames :: [Name]
fnames = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall a. HasCallStack => Maybe a -> a
fromJust forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst) [Field]
fields
matchField :: Name -> Name -> (Name, Name) -> Name -> Q Match
matchField :: Name -> Name -> (Name, Name) -> Name -> Q Match
matchField Name
arr Name
endOff (Name
acc, Name
currOff) Name
fname = do
let fnameLit :: Lit
fnameLit = String -> Lit
StringL (Name -> String
nameBase Name
fname)
forall (m :: * -> *).
Quote m =>
m Pat -> m Body -> [m Dec] -> m Match
match
(forall (m :: * -> *). Quote m => Lit -> m Pat
litP Lit
fnameLit)
(forall (m :: * -> *). Quote m => m Exp -> m Body
normalB
[|do (valOff, valLen :: Word32) <-
deserializeAt
$(varE currOff)
$(varE arr)
$(varE endOff)
pure
( ($(litE fnameLit), $(varE currOff)) : $(varE acc)
, valOff + w32_int valLen)|])
[]
mkDeserializeKeysDec :: Name -> Name -> SimpleDataCon -> Q [Dec]
mkDeserializeKeysDec :: Name -> Name -> SimpleDataCon -> Q [Dec]
mkDeserializeKeysDec Name
funcName Name
updateFunc (SimpleDataCon Name
cname [Field]
fields) = do
Name
hOff <- forall (m :: * -> *). Quote m => String -> m Name
newName String
"hOff"
Name
finalOff <- forall (m :: * -> *). Quote m => String -> m Name
newName String
"finalOff"
Name
arr <- forall (m :: * -> *). Quote m => String -> m Name
newName String
"arr"
Name
endOff <- forall (m :: * -> *). Quote m => String -> m Name
newName String
"endOff"
Name
kvEncoded <- forall (m :: * -> *). Quote m => String -> m Name
newName String
"kvEncoded"
Name
finalRec <- forall (m :: * -> *). Quote m => String -> m Name
newName String
"finalRec"
let deserializeFieldExpr :: Field -> m Exp
deserializeFieldExpr (Just Name
name, Type
ty) = do
let nameLit :: m Exp
nameLit = forall (m :: * -> *). Quote m => Lit -> m Exp
litE (String -> Lit
StringL (Name -> String
nameBase Name
name))
[|case lookup $(nameLit) $(varE kvEncoded) of
Nothing -> $(emptyTy name ty)
Just off -> do
val <- deserializeWithSize off $(varE arr) $(varE endOff)
pure $ snd val|]
deserializeFieldExpr Field
_ =
forall a. String -> a
errorUnsupported String
"The datatype should use record syntax."
Exp
method <-
[|do (dataOff, hlist :: CompactList (CompactList Word8)) <-
deserializeAt $(varE hOff) $(varE arr) $(varE endOff)
let keys = wListToString . unCompactList <$> unCompactList hlist
($(varP kvEncoded), _) <-
foldlM
($(varE updateFunc) $(varE arr) $(varE endOff))
([], dataOff)
keys
$(varP finalRec) <-
$(foldl
(\acc i ->
[|$(acc) <*>
$(deserializeFieldExpr i)|])
[|pure $(conE cname)|]
fields)
pure ($(varE finalOff), $(varE finalRec))|]
forall (f :: * -> *) a. Applicative f => a -> f a
pure
[ Pragma -> Dec
PragmaD (Name -> Inline -> RuleMatch -> Phases -> Pragma
InlineP Name
funcName Inline
NoInline RuleMatch
FunLike Phases
AllPhases)
, Name -> [Clause] -> Dec
FunD
Name
funcName
[ [Pat] -> Body -> [Dec] -> Clause
Clause
[ Name -> Pat
VarP Name
hOff
, Name -> Pat
VarP Name
finalOff
, Name -> Pat
VarP Name
arr
, Name -> Pat
VarP Name
endOff
]
(Exp -> Body
NormalB Exp
method)
[]
]
]
where
emptyTy :: Name -> Type -> m Exp
emptyTy Name
k Type
ty =
if Type -> Bool
isMaybeType Type
ty
then [|pure Nothing|]
else [|error $(litE (StringL (nameBase k ++ " is not found.")))|]
mkRecDeserializeExpr :: Name -> Name -> Name -> SimpleDataCon -> Q Exp
mkRecDeserializeExpr :: Name -> Name -> Name -> SimpleDataCon -> Q Exp
mkRecDeserializeExpr Name
initialOff Name
endOff Name
deserializeWithKeys SimpleDataCon
con = do
Name
hOff <- forall (m :: * -> *). Quote m => String -> m Name
newName String
"hOff"
let sizeForFinalOff :: Int
sizeForFinalOff = Int
4
sizeForHeaderLength :: Int
sizeForHeaderLength = Int
4
sizePreData :: Int
sizePreData = Int
sizeForFinalOff forall a. Num a => a -> a -> a
+ Int
sizeForHeaderLength forall a. Num a => a -> a -> a
+ Int
hlen
[|do (hlenOff, encLen :: Word32) <-
deserializeAt $(varE initialOff) $(varE _arr) $(varE endOff)
($(varP hOff), hlen1 :: Word32) <-
deserializeAt hlenOff $(varE _arr) $(varE endOff)
if (hlen1 == $(litIntegral hlen)) && $(xorCmp hval hOff _arr)
then do
let $(varP (makeI 0)) =
$(varE initialOff) +
$(litIntegral sizePreData)
$(mkDeserializeExprOne 'deserializeWithSize con)
else $(varE deserializeWithKeys)
$(varE hOff)
($(varE initialOff) + w32_int encLen)
$(varE _arr)
$(varE endOff)|]
where
hval :: [Word8]
hval = SimpleDataCon -> [Word8]
headerValue SimpleDataCon
con
hlen :: Int
hlen = forall (t :: * -> *) a. Foldable t => t a -> Int
length [Word8]
hval