-- SPDX-License-Identifier: MPL-2.0

{- | HMath array wire format version 1. All structural integers are little
endian, independent of the explicit payload endian. Header fields: 8-byte
magic @HMARRAY\\0@, u16 version, u8 dtype/endian/rank/flags, u64 offset and
backing length, u32 metadata length, rank*(u64 shape,i64 byte stride), UTF-8
length-prefixed metadata, backing bytes, u32 CRC-32/ISO-HDLC of everything
preceding the checksum. Unknown versions, flags and trailing bytes fail.
CRC detects accidental corruption, not malicious modification.
-}
module HMath.Interchange.Binary (encodeBinary, decodeBinary, crc32, runDecoder) where

import Control.Monad (replicateM, unless, when)
import Data.Array.Unboxed (UArray, listArray, (!))
import Data.Binary.Get
import Data.Binary.Put
import Data.Bits
import qualified Data.ByteString as B
import qualified Data.ByteString.Char8 as C
import qualified Data.ByteString.Lazy as L
import qualified Data.Text as T
import qualified Data.Text.Encoding as E
import Data.Word
import HMath.Interchange.Array

-- | Standard CRC-32, compatible with ZIP and zlib, with a fixed 256-entry table.
crc32 :: B.ByteString -> Word32
crc32 :: ByteString -> Word32
crc32 = Word32 -> Word32
forall a. Bits a => a -> a
complement (Word32 -> Word32)
-> (ByteString -> Word32) -> ByteString -> Word32
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Word32 -> Word8 -> Word32) -> Word32 -> ByteString -> Word32
forall a. (a -> Word8 -> a) -> a -> ByteString -> a
B.foldl' Word32 -> Word8 -> Word32
forall {a}. Integral a => Word32 -> a -> Word32
step Word32
0xffffffff
  where
    step :: Word32 -> a -> Word32
step Word32
c a
x = (Word32
c Word32 -> Int -> Word32
forall a. Bits a => a -> Int -> a
`shiftR` Int
8) Word32 -> Word32 -> Word32
forall a. Bits a => a -> a -> a
`xor` (UArray Int Word32
crcTable UArray Int Word32 -> Int -> Word32
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> i -> e
! Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral ((Word32
c Word32 -> Word32 -> Word32
forall a. Bits a => a -> a -> a
`xor` a -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral a
x) Word32 -> Word32 -> Word32
forall a. Bits a => a -> a -> a
.&. Word32
255))

crcTable :: UArray Int Word32
crcTable :: UArray Int Word32
crcTable = (Int, Int) -> [Word32] -> UArray Int Word32
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
(i, i) -> [e] -> a i e
listArray (Int
0, Int
255) [(Word32 -> Word32) -> Word32 -> [Word32]
forall a. (a -> a) -> a -> [a]
iterate Word32 -> Word32
forall {a}. (Bits a, Num a) => a -> a
bitStep (Int -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n) [Word32] -> Int -> Word32
forall a. HasCallStack => [a] -> Int -> a
!! Int
8 | Int
n <- [Int
0 :: Int .. Int
255]]
  where
    bitStep :: a -> a
bitStep a
x = if a
x a -> a -> a
forall a. Bits a => a -> a -> a
.&. a
1 a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
1 then (a
x a -> Int -> a
forall a. Bits a => a -> Int -> a
`shiftR` Int
1) a -> a -> a
forall a. Bits a => a -> a -> a
`xor` a
0xedb88320 else a
x a -> Int -> a
forall a. Bits a => a -> Int -> a
`shiftR` Int
1

-- | Strict, exact decoder: trailing bytes are an error rather than ignored.
runDecoder :: Get a -> B.ByteString -> Either InterchangeError a
runDecoder :: forall a. Get a -> ByteString -> Either InterchangeError a
runDecoder Get a
parser ByteString
bytes = case Get a
-> ByteString
-> Either
     (ByteString, ByteOffset, String) (ByteString, ByteOffset, a)
forall a.
Get a
-> ByteString
-> Either
     (ByteString, ByteOffset, String) (ByteString, ByteOffset, a)
runGetOrFail Get a
parser (ByteString -> ByteString
L.fromStrict ByteString
bytes) of
    Left (ByteString
_, ByteOffset
offset, String
err) -> InterchangeError -> Either InterchangeError a
forall a b. a -> Either a b
Left (String -> InterchangeError
InvalidData (String
"byte " String -> String -> String
forall a. [a] -> [a] -> [a]
++ ByteOffset -> String
forall a. Show a => a -> String
show ByteOffset
offset String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
": " String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
err))
    Right (ByteString
rest, ByteOffset
_, a
value) | ByteString -> Bool
L.null ByteString
rest -> a -> Either InterchangeError a
forall a b. b -> Either a b
Right a
value
    Either (ByteString, ByteOffset, String) (ByteString, ByteOffset, a)
_ -> InterchangeError -> Either InterchangeError a
forall a b. a -> Either a b
Left (String -> InterchangeError
InvalidData String
"trailing bytes")

magic :: B.ByteString
magic :: ByteString
magic = String -> ByteString
C.pack String
"HMARRAY\0"
putMeta :: [(String, String)] -> Put
putMeta :: [(String, String)] -> Put
putMeta [(String, String)]
meta = do
    Word32 -> Put
putWord32le (Int -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([(String, String)] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(String, String)]
meta))
    ((String, String) -> Put) -> [(String, String)] -> Put
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (\(String
k, String
v) -> String -> Put
putString String
k Put -> Put -> Put
forall a b. PutM a -> PutM b -> PutM b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> String -> Put
putString String
v) [(String, String)]
meta
  where
    putString :: String -> Put
putString String
s = let b :: ByteString
b = Text -> ByteString
E.encodeUtf8 (String -> Text
T.pack String
s) in Word32 -> Put
putWord32le (Int -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral (ByteString -> Int
B.length ByteString
b)) Put -> Put -> Put
forall a b. PutM a -> PutM b -> PutM b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> ByteString -> Put
putByteString ByteString
b

{- | Given limits then Array, encode HMath wire-format v1 including
shape/strides/offset/metadata/backing storage and CRC-32. Invalid metadata/layout or
exceeded aggregate byte budget fails.
-}
encodeBinary :: Limits -> Array -> Either InterchangeError B.ByteString
encodeBinary :: Limits -> Array -> Either InterchangeError ByteString
encodeBinary Limits
lim Array
a = do
    _ <-
        Limits
-> DType
-> Endian
-> [Int]
-> [Int]
-> Int
-> ByteString
-> Either InterchangeError Array
array
            Limits
lim
            (Array -> DType
arrayDType Array
a)
            (Array -> Endian
arrayEndian Array
a)
            (Array -> [Int]
arrayShape Array
a)
            (Array -> [Int]
arrayStrides Array
a)
            (Array -> Int
arrayOffset Array
a)
            (Array -> ByteString
arrayBytes Array
a)
    _ <- withMetadata lim (arrayMetadata a) a
    let meta = ByteString -> ByteString
L.toStrict (Put -> ByteString
runPut ([(String, String)] -> Put
putMeta (Array -> [(String, String)]
arrayMetadata Array
a)))
        total =
            Integer
38
                Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
16 Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Int -> Integer
forall a. Integral a => a -> Integer
toInteger ([Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (Array -> [Int]
arrayShape Array
a))
                Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Int -> Integer
forall a. Integral a => a -> Integer
toInteger (ByteString -> Int
B.length ByteString
meta)
                Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Int -> Integer
forall a. Integral a => a -> Integer
toInteger (ByteString -> Int
B.length (Array -> ByteString
arrayBytes Array
a))
    guardI
        (B.length meta <= maxMetadata lim && total <= toInteger (maxBytes lim))
        (LimitExceeded "binary output bytes")
    let content =
            ByteString -> ByteString
L.toStrict
                ( Put -> ByteString
runPut (Put -> ByteString) -> Put -> ByteString
forall a b. (a -> b) -> a -> b
$ do
                    ByteString -> Put
putByteString ByteString
magic
                    Word16 -> Put
putWord16le Word16
1
                    Word8 -> Put
putWord8 (Int -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral (DType -> Int
forall a. Enum a => a -> Int
fromEnum (Array -> DType
arrayDType Array
a)))
                    Word8 -> Put
putWord8 (if Array -> Endian
arrayEndian Array
a Endian -> Endian -> Bool
forall a. Eq a => a -> a -> Bool
== Endian
LittleEndian then Word8
0 else Word8
1)
                    Word8 -> Put
putWord8 (Int -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral ([Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (Array -> [Int]
arrayShape Array
a)))
                    Word8 -> Put
putWord8 Word8
0
                    Word64 -> Put
putWord64le (Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Array -> Int
arrayOffset Array
a))
                    Word64 -> Put
putWord64le (Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (ByteString -> Int
B.length (Array -> ByteString
arrayBytes Array
a)))
                    Word32 -> Put
putWord32le (Int -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral (ByteString -> Int
B.length ByteString
meta))
                    [Put] -> Put
forall (t :: * -> *) (m :: * -> *) a.
(Foldable t, Monad m) =>
t (m a) -> m ()
sequence_
                        [ Word64 -> Put
putWord64le (Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
d) Put -> Put -> Put
forall a b. PutM a -> PutM b -> PutM b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> ByteOffset -> Put
putInt64le (Int -> ByteOffset
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
s)
                        | (Int
d, Int
s) <- [Int] -> [Int] -> [(Int, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip (Array -> [Int]
arrayShape Array
a) (Array -> [Int]
arrayStrides Array
a)
                        ]
                    ByteString -> Put
putByteString ByteString
meta
                    ByteString -> Put
putByteString (Array -> ByteString
arrayBytes Array
a)
                )
    Right (content <> L.toStrict (runPut (putWord32le (crc32 content))))

{- | Given limits then bytes, validate HMath v1 magic/header/flags/CRC and exact framing,
returning an Array. Truncation, corruption, unsupported metadata or exceeded budgets
return InterchangeError.
-}
decodeBinary :: Limits -> B.ByteString -> Either InterchangeError Array
decodeBinary :: Limits -> ByteString -> Either InterchangeError Array
decodeBinary Limits
lim ByteString
input = do
    Limits -> Either InterchangeError ()
checkLimits Limits
lim
    Bool -> InterchangeError -> Either InterchangeError ()
guardI (ByteString -> Int
B.length ByteString
input Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Limits -> Int
maxBytes Limits
lim) (String -> InterchangeError
LimitExceeded String
"binary input bytes")
    Bool -> InterchangeError -> Either InterchangeError ()
guardI (ByteString -> Int
B.length ByteString
input Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
42) (String -> InterchangeError
InvalidData String
"truncated binary header")
    let (ByteString
content, ByteString
checksum) = Int -> ByteString -> (ByteString, ByteString)
B.splitAt (ByteString -> Int
B.length ByteString
input Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
4) ByteString
input
    stored <- Get Word32 -> ByteString -> Either InterchangeError Word32
forall a. Get a -> ByteString -> Either InterchangeError a
runDecoder Get Word32
getWord32le ByteString
checksum
    guardI (stored == crc32 content) (InvalidData "CRC-32 mismatch")
    (dt, en, sh, st, off, payload, meta) <- runDecoder parser content
    a <- array lim dt en sh st off payload
    withMetadata lim meta a
  where
    failUnless :: Bool -> String -> f ()
failUnless Bool
b String
s = Bool -> f () -> f ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless Bool
b (String -> f ()
forall a. String -> f a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
s)
    size :: String -> Int -> Word64 -> Get Int
    size :: String -> Int -> Word64 -> Get Int
size String
name Int
cap Word64
n =
        if Word64
n Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
<= Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
cap Bool -> Bool -> Bool
&& Word64
n Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
<= Int -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
forall a. Bounded a => a
maxBound :: Int)
            then Int -> Get Int
forall a. a -> Get a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Word64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
n)
            else String -> Get Int
forall a. String -> Get a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String
name String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
" limit")
    parser :: Get
  (DType, Endian, [Int], [Int], Int, ByteString, [(String, String)])
parser = do
        m <- Int -> Get ByteString
getByteString Int
8
        failUnless (m == magic) "bad magic"
        version <- getWord16le
        failUnless (version == 1) "unsupported format version"
        typ <- getWord8
        failUnless (fromIntegral typ <= fromEnum (maxBound :: DType)) "unsupported dtype"
        endian <- getWord8
        failUnless (endian <= 1) "invalid endian"
        rank <- getWord8
        failUnless (fromIntegral rank <= maxRank lim) "rank limit"
        flags <- getWord8
        failUnless (flags == 0) "unsupported flags"
        off <- getWord64le >>= size "offset" (maxBytes lim)
        len <- getWord64le >>= size "payload" (maxBytes lim)
        mlen <- getWord32le >>= size "metadata" (maxMetadata lim) . fromIntegral
        ds <-
            replicateM
                (fromIntegral rank)
                ( do
                    d <- getWord64le >>= size "dimension" (maxBound :: Int)
                    s <- getInt64le
                    failUnless
                        (toInteger s >= toInteger (minBound :: Int) && toInteger s <= toInteger (maxBound :: Int))
                        "stride integer overflow"
                    pure (d, fromIntegral s)
                )
        let sh = ((Int, Int) -> Int) -> [(Int, Int)] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Int, Int) -> Int
forall a b. (a, b) -> a
fst [(Int, Int)]
ds
            st = ((Int, Int) -> Int) -> [(Int, Int)] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Int, Int) -> Int
forall a b. (a, b) -> b
snd [(Int, Int)]
ds
        either (fail . show) (const (pure ())) (validateShape lim sh)
        mb <- getByteString mlen
        meta <- either (fail . show) pure (runDecoder metaParser mb)
        payload <- getByteString len
        pure
            ( toEnum (fromIntegral typ)
            , if endian == 0 then LittleEndian else BigEndian
            , sh
            , st
            , off
            , payload
            , meta
            )
    metaParser :: Get [(String, String)]
metaParser = do
        count <- Get Word32
getWord32le
        when (toInteger count > toInteger (maxMetadata lim `div` 8)) (fail "metadata entry count")
        replicateM (fromIntegral count) ((,) <$> getString <*> getString)
    getString :: Get String
getString = do
        len <- Get Word32
getWord32le Get Word32 -> (Word32 -> Get Int) -> Get Int
forall a b. Get a -> (a -> Get b) -> Get b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= String -> Int -> Word64 -> Get Int
size String
"metadata string" (Limits -> Int
maxMetadata Limits
lim) (Word64 -> Get Int) -> (Word32 -> Word64) -> Word32 -> Get Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Word32 -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral
        bytes <- getByteString len
        either (const (fail "invalid UTF-8")) (pure . T.unpack) (E.decodeUtf8' bytes)