SSZ serialize FP

SSZ.hs

module SSZ (
    SSZItem(..),
    zeroVal,
    serialize,
) where

import Data.Word (Word8,Word16,Word32,Word64)
import Data.Serialize (Serialize,encode,decode)
import qualified Data.ByteString as B
import Data.ByteString (ByteString)

-- Constants.
bytesPerLengthOffset :: Int
bytesPerLengthOffset = 4

-- Type class for SSZ items.
data SSZItem a =
    Uint64 Word64 | Uint32 Word32 | Uint16 Word16 | Uint8 Word8 | Bool Bool |
    List Int [SSZItem a] | Vector Int [SSZItem a] | Container [SSZItem a] |
    Bitlist Int [Bool] | Bitvector Int [Bool]
    deriving (Show, Eq)

-- Simple helper to map over a list zipped with each item's index.
mapWithIdx :: (Num a1, Enum a1) => (a1 -> b -> a2) -> [b] -> [a2]
mapWithIdx _ [] = []
mapWithIdx f xs = zipWith f [0..] xs

-- The cereal package's encoder is big-endian, so we create our 
-- own little-endian encoder using a simple composition. 
littleEncoder :: (Serialize a) => a -> B.ByteString
littleEncoder = B.reverse . encode

-- Serialization.
serialize :: SSZItem a -> B.ByteString
serialize (Bool a) = littleEncoder a
serialize (Uint64 a) = littleEncoder a
serialize (Uint32 a) = littleEncoder a
serialize (Uint16 a) = littleEncoder a
serialize (Uint8 a) = littleEncoder a
serialize (Vector n xs) = serializeFoldable xs
serialize (List n xs) = serializeFoldable xs
serialize (Container xs) = serializeFoldable xs
serialize _ = encode False

-- Implements the SSZ serialization algorithm for a Haskell list of SSZ items.
serializeFoldable :: [SSZItem a] -> B.ByteString
serializeFoldable xs = let
    -- When serializing fixed items, we use the Maybe monad
    fixedSerialize = map (\x ->
        if isFixed x then Just $ serialize x else Nothing)
    -- We serialize the fixed and variables parts.
    fixedParts = fixedSerialize xs
    variableParts = map serialize (filter isVariable xs)
    -- We determine the fixed and variable lengths.
    fixedLengths = map determineFixedLength fixedParts
    varLengths = map B.length variableParts
    -- Interleave the necessary offsets.
    varOffsets = let l = length xs in
        map (serialize . Uint32 . addFixedVar fixedLengths varLengths)[0..l]
    interleavedFixed = mapWithIdx (\i x -> construct x varOffsets i) fixedParts
    in
        B.intercalate B.empty (interleavedFixed ++ variableParts)

-- Adds a fixed length + variable length sizes for an ith index.
addFixedVar :: [Int] -> [Int] -> Int -> Word32
addFixedVar fixedLengths varLengths i = fromIntegral . sum $ fixedLengths ++ take (i-1) varLengths

construct x xs i = case x of
    Nothing -> xs !! i
    Just a -> a

determineVarOffset :: [ByteString] -> Int -> ByteString
determineVarOffset varOffsets i = varOffsets !! i

determineFixedLength :: Maybe ByteString -> Int
determineFixedLength Nothing = bytesPerLengthOffset
determineFixedLength (Just a) = B.length a

-- Gets the length of an SSZ item in bytes.
itemLen :: SSZItem a -> Maybe Int
itemLen (Vector n _) = Just n
itemLen (List _ xs) = Just $ length xs
itemLen _ = Nothing

-- Gets the default, zero value of an SSZ item type.
zeroVal :: (Num b) => SSZItem a -> SSZItem b
zeroVal (Bool _) = Bool False
zeroVal (List n _) = List n []
zeroVal (Vector n xs) = Vector n (replicate n (zeroVal $ head xs))
zeroVal (Bitlist n _) = Bitlist n []
zeroVal (Bitvector n xs) = Bitvector n (replicate n False)
zeroVal (Container xs) = Container (replicate (length xs) (zeroVal $ head xs))
zeroVal (Uint64 _) = Uint64 0
zeroVal (Uint32 _) = Uint32 0
zeroVal (Uint16 _) = Uint16 0
zeroVal (Uint8 _) = Uint8 0

-- Size checks for SSZ items.
isVariable :: SSZItem a -> Bool
isVariable (List _ _) = True
isVariable (Vector _ _) = False
isVariable (Bitlist _ _) = True
isVariable (Bitvector _ _) = False
isVariable (Container xs) = all isVariable xs
isVariable _ = False

isFixed :: SSZItem a -> Bool
isFixed = not . isVariable

-- An SSZ item is zeroed if it equals to the zero value of its type.
-- Because the SSZItem type derives Eq, this is a trivial check.
isZero :: (Eq a, Num a) => SSZItem a -> Bool
isZero item = item == zeroVal item