{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE ScopedTypeVariables #-}

-- SPDX-License-Identifier: MPL-2.0

{- | Fixed-width, typed array interchange. Strides and offsets count bytes;
negative/zero strides, unaligned samples, scalars and empty axes are valid.
The immutable backing bytes may contain padding. Construction validates the
entire reachable interval with Integer arithmetic before indexing. This is
an interchange representation, not the boxed numerical kernel storage.
-}
module HMath.Interchange.Array (
    Array,
    DType (..),
    Endian (..),
    Order (..),
    Limits (..),
    defaultLimits,
    InterchangeError (..),
    Element (..),
    array,
    contiguousArray,
    validateShape,
    arrayDType,
    arrayEndian,
    arrayShape,
    arrayStrides,
    arrayOffset,
    arrayBytes,
    arrayMetadata,
    withMetadata,
    elementBytes,
    elementCount,
    denseStrides,
    fromValues,
    toValues,
    fromTensor,
    toTensor,
    fromMatrix,
    toMatrix,
    repack,
    rawEncode,
    rawDecode,
    byteOffsets,
    guardI,
    checkLimits,
    boundedInt,
) where

import Control.Monad (unless)
import Data.Binary.Get
import Data.Binary.Put
import qualified Data.ByteString as B
import qualified Data.ByteString.Lazy as L
import Data.Complex (Complex ((:+)))
import Data.Int
import Data.List (nub)
import Data.Proxy
import Data.Word
import qualified HMath.Matrix as M
import qualified HMath.Tensor as T

-- | Numeric wire types; no platform-dependent Int, object or record dtype.
data DType = Bool8 | I8 | U8 | I16 | U16 | I32 | U32 | I64 | U64 | F32 | F64 | C64 | C128
    deriving (DType -> DType -> Bool
(DType -> DType -> Bool) -> (DType -> DType -> Bool) -> Eq DType
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: DType -> DType -> Bool
== :: DType -> DType -> Bool
$c/= :: DType -> DType -> Bool
/= :: DType -> DType -> Bool
Eq, Eq DType
Eq DType =>
(DType -> DType -> Ordering)
-> (DType -> DType -> Bool)
-> (DType -> DType -> Bool)
-> (DType -> DType -> Bool)
-> (DType -> DType -> Bool)
-> (DType -> DType -> DType)
-> (DType -> DType -> DType)
-> Ord DType
DType -> DType -> Bool
DType -> DType -> Ordering
DType -> DType -> DType
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: DType -> DType -> Ordering
compare :: DType -> DType -> Ordering
$c< :: DType -> DType -> Bool
< :: DType -> DType -> Bool
$c<= :: DType -> DType -> Bool
<= :: DType -> DType -> Bool
$c> :: DType -> DType -> Bool
> :: DType -> DType -> Bool
$c>= :: DType -> DType -> Bool
>= :: DType -> DType -> Bool
$cmax :: DType -> DType -> DType
max :: DType -> DType -> DType
$cmin :: DType -> DType -> DType
min :: DType -> DType -> DType
Ord, Int -> DType -> ShowS
[DType] -> ShowS
DType -> String
(Int -> DType -> ShowS)
-> (DType -> String) -> ([DType] -> ShowS) -> Show DType
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> DType -> ShowS
showsPrec :: Int -> DType -> ShowS
$cshow :: DType -> String
show :: DType -> String
$cshowList :: [DType] -> ShowS
showList :: [DType] -> ShowS
Show, Int -> DType
DType -> Int
DType -> [DType]
DType -> DType
DType -> DType -> [DType]
DType -> DType -> DType -> [DType]
(DType -> DType)
-> (DType -> DType)
-> (Int -> DType)
-> (DType -> Int)
-> (DType -> [DType])
-> (DType -> DType -> [DType])
-> (DType -> DType -> [DType])
-> (DType -> DType -> DType -> [DType])
-> Enum DType
forall a.
(a -> a)
-> (a -> a)
-> (Int -> a)
-> (a -> Int)
-> (a -> [a])
-> (a -> a -> [a])
-> (a -> a -> [a])
-> (a -> a -> a -> [a])
-> Enum a
$csucc :: DType -> DType
succ :: DType -> DType
$cpred :: DType -> DType
pred :: DType -> DType
$ctoEnum :: Int -> DType
toEnum :: Int -> DType
$cfromEnum :: DType -> Int
fromEnum :: DType -> Int
$cenumFrom :: DType -> [DType]
enumFrom :: DType -> [DType]
$cenumFromThen :: DType -> DType -> [DType]
enumFromThen :: DType -> DType -> [DType]
$cenumFromTo :: DType -> DType -> [DType]
enumFromTo :: DType -> DType -> [DType]
$cenumFromThenTo :: DType -> DType -> DType -> [DType]
enumFromThenTo :: DType -> DType -> DType -> [DType]
Enum, DType
DType -> DType -> Bounded DType
forall a. a -> a -> Bounded a
$cminBound :: DType
minBound :: DType
$cmaxBound :: DType
maxBound :: DType
Bounded)

-- | Explicit byte order. Single-byte types ignore it.
data Endian = LittleEndian | BigEndian deriving (Endian -> Endian -> Bool
(Endian -> Endian -> Bool)
-> (Endian -> Endian -> Bool) -> Eq Endian
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Endian -> Endian -> Bool
== :: Endian -> Endian -> Bool
$c/= :: Endian -> Endian -> Bool
/= :: Endian -> Endian -> Bool
Eq, Int -> Endian -> ShowS
[Endian] -> ShowS
Endian -> String
(Int -> Endian -> ShowS)
-> (Endian -> String) -> ([Endian] -> ShowS) -> Show Endian
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Endian -> ShowS
showsPrec :: Int -> Endian -> ShowS
$cshow :: Endian -> String
show :: Endian -> String
$cshowList :: [Endian] -> ShowS
showList :: [Endian] -> ShowS
Show)

-- | Physical order of a packed output; logical values are always row major.
data Order = COrder | FOrder deriving (Order -> Order -> Bool
(Order -> Order -> Bool) -> (Order -> Order -> Bool) -> Eq Order
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Order -> Order -> Bool
== :: Order -> Order -> Bool
$c/= :: Order -> Order -> Bool
/= :: Order -> Order -> Bool
Eq, Int -> Order -> ShowS
[Order] -> ShowS
Order -> String
(Int -> Order -> ShowS)
-> (Order -> String) -> ([Order] -> ShowS) -> Show Order
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Order -> ShowS
showsPrec :: Int -> Order -> ShowS
$cshow :: Order -> String
show :: Order -> String
$cshowList :: [Order] -> ShowS
showList :: [Order] -> ShowS
Show)

{- | Caller-configurable finite resource budgets, checked before allocations.
Byte budgets include metadata/file overhead as well as raw payload. The
archive budget applies to the SUM of uncompressed members, preventing bombs.
-}
data Limits = Limits
    { Limits -> Int
maxBytes :: Int
    -- ^ Maximum aggregate encoded/input/backing bytes, including relevant format overhead.
    , Limits -> Int
maxElements :: Int
    -- ^ Maximum logical scalar element count.
    , Limits -> Int
maxRank :: Int
    -- ^ Maximum array dimension count (implementation cap 255).
    , Limits -> Int
maxMetadata :: Int
    -- ^ Maximum descriptive metadata/header byte budget.
    , Limits -> Int
maxArchiveEntries :: Int
    -- ^ Maximum archive member count; aggregate uncompressed bytes are also bounded.
    }
    deriving (Limits -> Limits -> Bool
(Limits -> Limits -> Bool)
-> (Limits -> Limits -> Bool) -> Eq Limits
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Limits -> Limits -> Bool
== :: Limits -> Limits -> Bool
$c/= :: Limits -> Limits -> Bool
/= :: Limits -> Limits -> Bool
Eq, Int -> Limits -> ShowS
[Limits] -> ShowS
Limits -> String
(Int -> Limits -> ShowS)
-> (Limits -> String) -> ([Limits] -> ShowS) -> Show Limits
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Limits -> ShowS
showsPrec :: Int -> Limits -> ShowS
$cshow :: Limits -> String
show :: Limits -> String
$cshowList :: [Limits] -> ShowS
showList :: [Limits] -> ShowS
Show)

-- | 128 MiB of bytes, 8 million samples, rank 16, 64 KiB metadata, 256 members.
defaultLimits :: Limits
defaultLimits :: Limits
defaultLimits = Int -> Int -> Int -> Int -> Int -> Limits
Limits (Int
128 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
1024 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
1024) (Int
8 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
1024 Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
1024) Int
16 Int
65536 Int
256

-- | Decoders never throw on malformed bytes; IO failures have their own case.
data InterchangeError
    = InvalidData String
    | Unsupported String
    | LimitExceeded String
    | InterchangeIOError String
    deriving (InterchangeError -> InterchangeError -> Bool
(InterchangeError -> InterchangeError -> Bool)
-> (InterchangeError -> InterchangeError -> Bool)
-> Eq InterchangeError
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: InterchangeError -> InterchangeError -> Bool
== :: InterchangeError -> InterchangeError -> Bool
$c/= :: InterchangeError -> InterchangeError -> Bool
/= :: InterchangeError -> InterchangeError -> Bool
Eq, Int -> InterchangeError -> ShowS
[InterchangeError] -> ShowS
InterchangeError -> String
(Int -> InterchangeError -> ShowS)
-> (InterchangeError -> String)
-> ([InterchangeError] -> ShowS)
-> Show InterchangeError
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> InterchangeError -> ShowS
showsPrec :: Int -> InterchangeError -> ShowS
$cshow :: InterchangeError -> String
show :: InterchangeError -> String
$cshowList :: [InterchangeError] -> ShowS
showList :: [InterchangeError] -> ShowS
Show)

-- | Opaque validated bytes and descriptive UTF-8 metadata (key/value strings).
data Array = Array DType Endian [Int] [Int] Int B.ByteString [(String, String)] deriving (Array -> Array -> Bool
(Array -> Array -> Bool) -> (Array -> Array -> Bool) -> Eq Array
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Array -> Array -> Bool
== :: Array -> Array -> Bool
$c/= :: Array -> Array -> Bool
/= :: Array -> Array -> Bool
Eq, Int -> Array -> ShowS
[Array] -> ShowS
Array -> String
(Int -> Array -> ShowS)
-> (Array -> String) -> ([Array] -> ShowS) -> Show Array
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Array -> ShowS
showsPrec :: Int -> Array -> ShowS
$cshow :: Array -> String
show :: Array -> String
$cshowList :: [Array] -> ShowS
showList :: [Array] -> ShowS
Show)

{- | Return the fixed-width wire dtype of an Array; this is independent of Haskell host
Int/Double defaults.
-}
arrayDType :: Array -> DType
arrayDType :: Array -> DType
arrayDType (Array DType
x Endian
_ [Int]
_ [Int]
_ Int
_ ByteString
_ [(String, String)]
_) = DType
x

-- | Return explicit payload byte order; single-byte dtypes are endian-independent.
arrayEndian :: Array -> Endian
arrayEndian :: Array -> Endian
arrayEndian (Array DType
_ Endian
x [Int]
_ [Int]
_ Int
_ ByteString
_ [(String, String)]
_) = Endian
x

-- | Return logical dimensions in axis order; [] is scalar and zero axes are valid.
arrayShape :: Array -> [Int]
arrayShape :: Array -> [Int]
arrayShape (Array DType
_ Endian
_ [Int]
x [Int]
_ Int
_ ByteString
_ [(String, String)]
_) = [Int]
x

{- | Return signed byte strides for each axis. Unlike HMath.Tensor strides, these are bytes,
not elements.
-}
arrayStrides :: Array -> [Int]
arrayStrides :: Array -> [Int]
arrayStrides (Array DType
_ Endian
_ [Int]
_ [Int]
x Int
_ ByteString
_ [(String, String)]
_) = [Int]
x

-- | Return byte offset of logical origin within arrayBytes backing storage.
arrayOffset :: Array -> Int
arrayOffset :: Array -> Int
arrayOffset (Array DType
_ Endian
_ [Int]
_ [Int]
_ Int
x ByteString
_ [(String, String)]
_) = Int
x

{- | Return immutable backing bytes, which may include padding or data not reached by the
logical strided view. Use toValues or repack for logical contents.
-}
arrayBytes :: Array -> B.ByteString
arrayBytes :: Array -> ByteString
arrayBytes (Array DType
_ Endian
_ [Int]
_ [Int]
_ Int
_ ByteString
x [(String, String)]
_) = ByteString
x

{- | Return descriptive UTF-8-compatible key/value metadata strings in stored order;
metadata is not executable.
-}
arrayMetadata :: Array -> [(String, String)]
arrayMetadata :: Array -> [(String, String)]
arrayMetadata (Array DType
_ Endian
_ [Int]
_ [Int]
_ Int
_ ByteString
_ [(String, String)]
x) = [(String, String)]
x

{- | Given condition then InterchangeError, return Right () when True and Left that error
otherwise. This helper performs no IO.
-}
guardI :: Bool -> InterchangeError -> Either InterchangeError ()
guardI :: Bool -> InterchangeError -> Either InterchangeError ()
guardI Bool
p InterchangeError
e = Bool -> Either InterchangeError () -> Either InterchangeError ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless Bool
p (InterchangeError -> Either InterchangeError ()
forall a b. a -> Either a b
Left InterchangeError
e)

{- | Validate nonnegative finite resource caps, rank\<=255 and Int-safe
byte/element/metadata/archive budgets. Invalid settings return InvalidData.
-}
checkLimits :: Limits -> Either InterchangeError ()
checkLimits :: Limits -> Either InterchangeError ()
checkLimits Limits
lim =
    Bool -> InterchangeError -> Either InterchangeError ()
guardI
        ( Limits -> Int
maxBytes Limits
lim Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0
            Bool -> Bool -> Bool
&& Limits -> Int
maxBytes Limits
lim Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
forall a. Bounded a => a
maxBound
            Bool -> Bool -> Bool
&& Limits -> Int
maxElements Limits
lim Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0
            Bool -> Bool -> Bool
&& Limits -> Int
maxElements Limits
lim Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
forall a. Bounded a => a
maxBound
            Bool -> Bool -> Bool
&& Limits -> Int
maxRank Limits
lim Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0
            Bool -> Bool -> Bool
&& Limits -> Int
maxRank Limits
lim Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
255
            Bool -> Bool -> Bool
&& Limits -> Int
maxMetadata Limits
lim Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0
            Bool -> Bool -> Bool
&& Limits -> Int
maxMetadata Limits
lim Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
forall a. Bounded a => a
maxBound
            Bool -> Bool -> Bool
&& Limits -> Int
maxArchiveEntries Limits
lim Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0
            Bool -> Bool -> Bool
&& Limits -> Int
maxArchiveEntries Limits
lim Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
forall a. Bounded a => a
maxBound
        )
        (String -> InterchangeError
InvalidData String
"invalid resource limits")

{- | Given diagnostic name,nonnegative upper bound,then Integer value, convert only if value
lies in [0,upper] and fits Int. Otherwise return LimitExceeded with that name.
-}
boundedInt :: String -> Integer -> Integer -> Either InterchangeError Int
boundedInt :: String -> Integer -> Integer -> Either InterchangeError Int
boundedInt String
name Integer
upper Integer
x = do
    Bool -> InterchangeError -> Either InterchangeError ()
guardI (Integer
x Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
>= Integer
0 Bool -> Bool -> Bool
&& Integer
x Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<= Integer
upper Bool -> Bool -> Bool
&& Integer
x Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<= Int -> Integer
forall a. Integral a => a -> Integer
toInteger (Int
forall a. Bounded a => a
maxBound :: Int)) (String -> InterchangeError
LimitExceeded String
name)
    Int -> Either InterchangeError Int
forall a b. b -> Either a b
Right (Integer -> Int
forall a. Num a => Integer -> a
fromInteger Integer
x)

{- | Return fixed byte width of a DType, including complex pairs (C64=8,C128=16) and
Bool8=1.
-}
elementBytes :: DType -> Int
elementBytes :: DType -> Int
elementBytes DType
dt = case DType
dt of
    DType
Bool8 -> Int
1
    DType
I8 -> Int
1
    DType
U8 -> Int
1
    DType
I16 -> Int
2
    DType
U16 -> Int
2
    DType
I32 -> Int
4
    DType
U32 -> Int
4
    DType
I64 -> Int
8
    DType
U64 -> Int
8
    DType
F32 -> Int
4
    DType
F64 -> Int
8
    DType
C64 -> Int
8
    DType
C128 -> Int
16

{- | Given limits then dimensions, validate rank, nonnegative axes and element product,
returning Int element count. Invalid data or exceeded limits return InterchangeError.
-}
validateShape :: Limits -> [Int] -> Either InterchangeError Int
validateShape :: Limits -> [Int] -> Either InterchangeError Int
validateShape Limits
lim [Int]
sh = do
    Limits -> Either InterchangeError ()
checkLimits Limits
lim
    Bool -> InterchangeError -> Either InterchangeError ()
guardI ([Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
take (Limits -> Int
maxRank Limits
lim Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) [Int]
sh) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Limits -> Int
maxRank Limits
lim) (String -> InterchangeError
LimitExceeded String
"array rank")
    Bool -> InterchangeError -> Either InterchangeError ()
guardI ((Int -> Bool) -> [Int] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0) [Int]
sh) (String -> InterchangeError
InvalidData String
"negative dimension")
    String -> Integer -> Integer -> Either InterchangeError Int
boundedInt String
"array element count" (Int -> Integer
forall a. Integral a => a -> Integer
toInteger (Limits -> Int
maxElements Limits
lim)) ([Integer] -> Integer
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product ((Int -> Integer) -> [Int] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map Int -> Integer
forall a. Integral a => a -> Integer
toInteger [Int]
sh))

-- | Return logical shape product; scalar arrays have one element and empty axes have zero.
elementCount :: Array -> Int
elementCount :: Array -> Int
elementCount = [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product ([Int] -> Int) -> (Array -> [Int]) -> Array -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Array -> [Int]
arrayShape

{- | Dense strides use Integer arithmetic before conversion. Call only on
validated shape/byte sizes; zero axes use zero strides to avoid overflow.
-}
denseStrides :: Order -> Int -> [Int] -> [Int]
denseStrides :: Order -> Int -> [Int] -> [Int]
denseStrides Order
_ Int
_ [Int]
sh | Int
0 Int -> [Int] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Int]
sh = Int -> Int -> [Int]
forall a. Int -> a -> [a]
replicate ([Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
sh) Int
0
denseStrides Order
COrder Int
w [Int]
sh = (Integer -> Int) -> [Integer] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map Integer -> Int
forall a. Num a => Integer -> a
fromInteger (Int -> [Integer] -> [Integer]
forall a. Int -> [a] -> [a]
drop Int
1 ((Integer -> Integer -> Integer)
-> Integer -> [Integer] -> [Integer]
forall a b. (a -> b -> b) -> b -> [a] -> [b]
scanr Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
(*) (Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
w) ((Int -> Integer) -> [Int] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map Int -> Integer
forall a. Integral a => a -> Integer
toInteger [Int]
sh)))
denseStrides Order
FOrder Int
w [Int]
sh = (Integer -> Int) -> [Integer] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map Integer -> Int
forall a. Num a => Integer -> a
fromInteger (Int -> [Integer] -> [Integer]
forall a. Int -> [a] -> [a]
take ([Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
sh) ((Integer -> Integer -> Integer)
-> Integer -> [Integer] -> [Integer]
forall b a. (b -> a -> b) -> b -> [a] -> [b]
scanl Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
(*) (Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
w) ((Int -> Integer) -> [Int] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map Int -> Integer
forall a. Integral a => a -> Integer
toInteger [Int]
sh)))

{- | Construct a bounded view. Bool8 bytes reachable through the view must be
0 or 1. Floating payloads preserve IEEE NaNs, signed zero and infinities.
-}
array ::
    Limits -> DType -> Endian -> [Int] -> [Int] -> Int -> B.ByteString -> Either InterchangeError Array
array :: Limits
-> DType
-> Endian
-> [Int]
-> [Int]
-> Int
-> ByteString
-> Either InterchangeError Array
array Limits
lim DType
dt Endian
en [Int]
sh [Int]
st Int
off ByteString
bs = do
    n <- Limits -> [Int] -> Either InterchangeError Int
validateShape Limits
lim [Int]
sh
    guardI (length (take (length sh + 1) st) == length sh) (InvalidData "stride rank mismatch")
    guardI (B.length bs <= maxBytes lim) (LimitExceeded "array backing bytes")
    guardI
        (toInteger n * toInteger (elementBytes dt) <= toInteger (maxBytes lim))
        (LimitExceeded "logical array bytes")
    let extent = (Int -> Int -> Integer) -> [Int] -> [Int] -> [Integer]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\Int
d Int
s -> Int -> Integer
forall a. Integral a => a -> Integer
toInteger (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
0 (Int
d Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)) Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
s) [Int]
sh [Int]
st
        low = Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
off Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ [Integer] -> Integer
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((Integer -> Integer) -> [Integer] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map (Integer -> Integer -> Integer
forall a. Ord a => a -> a -> a
min Integer
0) [Integer]
extent)
        high = Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
off Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ [Integer] -> Integer
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((Integer -> Integer) -> [Integer] -> [Integer]
forall a b. (a -> b) -> [a] -> [b]
map (Integer -> Integer -> Integer
forall a. Ord a => a -> a -> a
max Integer
0) [Integer]
extent) Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Int -> Integer
forall a. Integral a => a -> Integer
toInteger (DType -> Int
elementBytes DType
dt)
    guardI
        (off >= 0 && if n == 0 then off <= B.length bs else low >= 0 && high <= toInteger (B.length bs))
        (InvalidData "strided view exceeds backing bytes")
    let a = DType
-> Endian
-> [Int]
-> [Int]
-> Int
-> ByteString
-> [(String, String)]
-> Array
Array DType
dt Endian
en [Int]
sh [Int]
st Int
off ByteString
bs []
    guardI
        (dt /= Bool8 || all (\Int
i -> HasCallStack => ByteString -> Int -> Word8
ByteString -> Int -> Word8
B.index ByteString
bs Int
i Word8 -> Word8 -> Bool
forall a. Ord a => a -> a -> Bool
<= Word8
1) (byteOffsets a))
        (InvalidData "noncanonical boolean")
    Right a

{- | Given limits,dtype,endian,physical order,shape,then exact packed bytes, construct a
canonical byte-strided array. Require exact byte count; invalid shape/data/limits fail.
-}
contiguousArray ::
    Limits -> DType -> Endian -> Order -> [Int] -> B.ByteString -> Either InterchangeError Array
contiguousArray :: Limits
-> DType
-> Endian
-> Order
-> [Int]
-> ByteString
-> Either InterchangeError Array
contiguousArray Limits
lim DType
dt Endian
en Order
order [Int]
sh ByteString
bs = do
    n <- Limits -> [Int] -> Either InterchangeError Int
validateShape Limits
lim [Int]
sh
    guardI
        (toInteger n * toInteger (elementBytes dt) == toInteger (B.length bs))
        (InvalidData "raw byte count does not match shape/dtype")
    array lim dt en sh (denseStrides order (elementBytes dt) sh) 0 bs

{- | Metadata is UTF-8 in binary files; keys must be unique and nonempty. The
four-byte-per-codepoint bound is intentionally conservative before encoding.
-}
withMetadata :: Limits -> [(String, String)] -> Array -> Either InterchangeError Array
withMetadata :: Limits
-> [(String, String)] -> Array -> Either InterchangeError Array
withMetadata Limits
lim [(String, String)]
meta (Array DType
dt Endian
en [Int]
sh [Int]
st Int
off ByteString
bs [(String, String)]
_) = do
    Limits -> Either InterchangeError ()
checkLimits Limits
lim
    let charBudget :: Int
charBudget = Limits -> Int
maxMetadata Limits
lim Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
4
        bounded :: [a] -> Int
bounded [a]
s = [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
take (Int
charBudget Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) [a]
s)
        count :: Integer
count = [Integer] -> Integer
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum [Int -> Integer
forall a. Integral a => a -> Integer
toInteger (String -> Int
forall a. [a] -> Int
bounded String
k) Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Int -> Integer
forall a. Integral a => a -> Integer
toInteger (String -> Int
forall a. [a] -> Int
bounded String
v) | (String
k, String
v) <- Int -> [(String, String)] -> [(String, String)]
forall a. Int -> [a] -> [a]
take (Int
charBudget Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) [(String, String)]
meta]
    Bool -> InterchangeError -> Either InterchangeError ()
guardI
        ([(String, String)] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (Int -> [(String, String)] -> [(String, String)]
forall a. Int -> [a] -> [a]
take (Int
charBudget Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) [(String, String)]
meta) Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
charBudget Bool -> Bool -> Bool
&& Integer
count Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<= Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
charBudget)
        (String -> InterchangeError
LimitExceeded String
"metadata")
    Bool -> InterchangeError -> Either InterchangeError ()
guardI
        (((String, String) -> Bool) -> [(String, String)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (Bool -> Bool
not (Bool -> Bool)
-> ((String, String) -> Bool) -> (String, String) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null (String -> Bool)
-> ((String, String) -> String) -> (String, String) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (String, String) -> String
forall a b. (a, b) -> a
fst) [(String, String)]
meta Bool -> Bool -> Bool
&& [String] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([String] -> [String]
forall a. Eq a => [a] -> [a]
nub (((String, String) -> String) -> [(String, String)] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map (String, String) -> String
forall a b. (a, b) -> a
fst [(String, String)]
meta)) Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== [(String, String)] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(String, String)]
meta)
        (String -> InterchangeError
InvalidData String
"empty/duplicate metadata key")
    Bool -> InterchangeError -> Either InterchangeError ()
guardI
        ((String -> Bool) -> [String] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all ((Char -> Bool) -> String -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (\Char
c -> Char
c Char -> Char -> Bool
forall a. Ord a => a -> a -> Bool
< Char
'\xD800' Bool -> Bool -> Bool
|| Char
c Char -> Char -> Bool
forall a. Ord a => a -> a -> Bool
> Char
'\xDFFF')) (((String, String) -> [String]) -> [(String, String)] -> [String]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (\(String
k, String
v) -> [String
k, String
v]) [(String, String)]
meta))
        (String -> InterchangeError
InvalidData String
"surrogate in metadata")
    Array -> Either InterchangeError Array
forall a b. b -> Either a b
Right (DType
-> Endian
-> [Int]
-> [Int]
-> Int
-> ByteString
-> [(String, String)]
-> Array
Array DType
dt Endian
en [Int]
sh [Int]
st Int
off ByteString
bs [(String, String)]
meta)

{- | Return backing byte offsets for each logical element in row-major coordinate order,
honoring signed strides and offset.
-}
byteOffsets :: Array -> [Int]
-- An inner zero axis must short-circuit before traversing outer dimensions.
-- Those outer dimensions can be huge even when the logical size is zero.
byteOffsets :: Array -> [Int]
byteOffsets Array
a | Array -> Int
elementCount Array
a Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = []
byteOffsets Array
a = Int -> [(Int, Int)] -> [Int]
forall {t} {a} {a}.
(Integral t, Integral a, Integral a) =>
t -> [(a, a)] -> [t]
offsets (Array -> Int
arrayOffset Array
a) ([Int] -> [Int] -> [(Int, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip (Array -> [Int]
arrayShape Array
a) (Array -> [Int]
arrayStrides Array
a))
  where
    offsets :: t -> [(a, a)] -> [t]
offsets t
p [] = [t
p]
    offsets t
p ((a
d, a
s) : [(a, a)]
rest) =
        [t
q | a
i <- [a
0 .. a
d a -> a -> a
forall a. Num a => a -> a -> a
- a
1], t
q <- t -> [(a, a)] -> [t]
offsets (Integer -> t
forall a. Num a => Integer -> a
fromInteger (t -> Integer
forall a. Integral a => a -> Integer
toInteger t
p Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ a -> Integer
forall a. Integral a => a -> Integer
toInteger a
i Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* a -> Integer
forall a. Integral a => a -> Integer
toInteger a
s)) [(a, a)]
rest]

{- | Fixed-width element encoding. Instances deliberately exclude machine Int.
Custom instances must write exactly elementBytes bytes and use a bounded
decoder. Built-in instances preserve floating-point bit representations.
-}
class Element a where
    -- | Given a Haskell Element type proxy, return its exact fixed-width wire DType.
    elementType :: proxy a -> DType

    -- | Given byte order then one typed Element, emit its binary representation with Data.Binary.Put.
    putElement :: Endian -> a -> Put

    -- | Given byte order, decode one typed Element with Data.Binary.Get; surrounding array decoders enforce byte/storage budgets.
    getElement :: Endian -> Get a

put16 :: Endian -> Word16 -> Put
put16 :: Endian -> Word16 -> Put
put16 Endian
LittleEndian = Word16 -> Put
putWord16le
put16 Endian
BigEndian = Word16 -> Put
putWord16be
put32 :: Endian -> Word32 -> Put
put32 :: Endian -> Word32 -> Put
put32 Endian
LittleEndian = Word32 -> Put
putWord32le
put32 Endian
BigEndian = Word32 -> Put
putWord32be
put64 :: Endian -> Word64 -> Put
put64 :: Endian -> Word64 -> Put
put64 Endian
LittleEndian = Word64 -> Put
putWord64le
put64 Endian
BigEndian = Word64 -> Put
putWord64be
get16 :: Endian -> Get Word16
get16 :: Endian -> Get Word16
get16 Endian
LittleEndian = Get Word16
getWord16le
get16 Endian
BigEndian = Get Word16
getWord16be
get32 :: Endian -> Get Word32
get32 :: Endian -> Get Word32
get32 Endian
LittleEndian = Get Word32
getWord32le
get32 Endian
BigEndian = Get Word32
getWord32be
get64 :: Endian -> Get Word64
get64 :: Endian -> Get Word64
get64 Endian
LittleEndian = Get Word64
getWord64le
get64 Endian
BigEndian = Get Word64
getWord64be
instance Element Bool where
    elementType :: forall (proxy :: * -> *). proxy Bool -> DType
elementType proxy Bool
_ = DType
Bool8
    putElement :: Endian -> Bool -> Put
putElement Endian
_ Bool
x = Word8 -> Put
putWord8 (if Bool
x then Word8
1 else Word8
0)
    getElement :: Endian -> Get Bool
getElement Endian
_ = do
        x <- Get Word8
getWord8
        if x <= 1 then pure (x == 1) else fail "invalid boolean"
instance Element Word8 where
    elementType :: forall (proxy :: * -> *). proxy Word8 -> DType
elementType proxy Word8
_ = DType
U8
    putElement :: Endian -> Word8 -> Put
putElement Endian
_ = Word8 -> Put
putWord8
    getElement :: Endian -> Get Word8
getElement Endian
_ = Get Word8
getWord8
instance Element Int8 where
    elementType :: forall (proxy :: * -> *). proxy Int8 -> DType
elementType proxy Int8
_ = DType
I8
    putElement :: Endian -> Int8 -> Put
putElement Endian
_ = Int8 -> Put
putInt8
    getElement :: Endian -> Get Int8
getElement Endian
_ = Get Int8
getInt8
instance Element Word16 where
    elementType :: forall (proxy :: * -> *). proxy Word16 -> DType
elementType proxy Word16
_ = DType
U16
    putElement :: Endian -> Word16 -> Put
putElement = Endian -> Word16 -> Put
put16
    getElement :: Endian -> Get Word16
getElement = Endian -> Get Word16
get16
instance Element Int16 where
    elementType :: forall (proxy :: * -> *). proxy Int16 -> DType
elementType proxy Int16
_ = DType
I16
    putElement :: Endian -> Int16 -> Put
putElement Endian
e = Endian -> Word16 -> Put
put16 Endian
e (Word16 -> Put) -> (Int16 -> Word16) -> Int16 -> Put
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int16 -> Word16
forall a b. (Integral a, Num b) => a -> b
fromIntegral
    getElement :: Endian -> Get Int16
getElement Endian
e = Word16 -> Int16
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word16 -> Int16) -> Get Word16 -> Get Int16
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Endian -> Get Word16
get16 Endian
e
instance Element Word32 where
    elementType :: forall (proxy :: * -> *). proxy Word32 -> DType
elementType proxy Word32
_ = DType
U32
    putElement :: Endian -> Word32 -> Put
putElement = Endian -> Word32 -> Put
put32
    getElement :: Endian -> Get Word32
getElement = Endian -> Get Word32
get32
instance Element Int32 where
    elementType :: forall (proxy :: * -> *). proxy Int32 -> DType
elementType proxy Int32
_ = DType
I32
    putElement :: Endian -> Int32 -> Put
putElement Endian
e = Endian -> Word32 -> Put
put32 Endian
e (Word32 -> Put) -> (Int32 -> Word32) -> Int32 -> Put
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int32 -> Word32
forall a b. (Integral a, Num b) => a -> b
fromIntegral
    getElement :: Endian -> Get Int32
getElement Endian
e = Word32 -> Int32
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word32 -> Int32) -> Get Word32 -> Get Int32
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Endian -> Get Word32
get32 Endian
e
instance Element Word64 where
    elementType :: forall (proxy :: * -> *). proxy Word64 -> DType
elementType proxy Word64
_ = DType
U64
    putElement :: Endian -> Word64 -> Put
putElement = Endian -> Word64 -> Put
put64
    getElement :: Endian -> Get Word64
getElement = Endian -> Get Word64
get64
instance Element Int64 where
    elementType :: forall (proxy :: * -> *). proxy Int64 -> DType
elementType proxy Int64
_ = DType
I64
    putElement :: Endian -> Int64 -> Put
putElement Endian
e = Endian -> Word64 -> Put
put64 Endian
e (Word64 -> Put) -> (Int64 -> Word64) -> Int64 -> Put
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int64 -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral
    getElement :: Endian -> Get Int64
getElement Endian
e = Word64 -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word64 -> Int64) -> Get Word64 -> Get Int64
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Endian -> Get Word64
get64 Endian
e
instance Element Float where
    elementType :: forall (proxy :: * -> *). proxy Float -> DType
elementType proxy Float
_ = DType
F32
    putElement :: Endian -> Float -> Put
putElement Endian
LittleEndian = Float -> Put
putFloatle
    putElement Endian
BigEndian = Float -> Put
putFloatbe
    getElement :: Endian -> Get Float
getElement Endian
LittleEndian = Get Float
getFloatle
    getElement Endian
BigEndian = Get Float
getFloatbe
instance Element Double where
    elementType :: forall (proxy :: * -> *). proxy Double -> DType
elementType proxy Double
_ = DType
F64
    putElement :: Endian -> Double -> Put
putElement Endian
LittleEndian = Double -> Put
putDoublele
    putElement Endian
BigEndian = Double -> Put
putDoublebe
    getElement :: Endian -> Get Double
getElement Endian
LittleEndian = Get Double
getDoublele
    getElement Endian
BigEndian = Get Double
getDoublebe
instance Element (Complex Float) where
    elementType :: forall (proxy :: * -> *). proxy (Complex Float) -> DType
elementType proxy (Complex Float)
_ = DType
C64
    putElement :: Endian -> Complex Float -> Put
putElement Endian
en (Float
a :+ Float
b) = Endian -> Float -> Put
forall a. Element a => Endian -> a -> Put
putElement Endian
en Float
a Put -> Put -> Put
forall a b. PutM a -> PutM b -> PutM b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Endian -> Float -> Put
forall a. Element a => Endian -> a -> Put
putElement Endian
en Float
b
    getElement :: Endian -> Get (Complex Float)
getElement Endian
en = Float -> Float -> Complex Float
forall a. a -> a -> Complex a
(:+) (Float -> Float -> Complex Float)
-> Get Float -> Get (Float -> Complex Float)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Endian -> Get Float
forall a. Element a => Endian -> Get a
getElement Endian
en Get (Float -> Complex Float) -> Get Float -> Get (Complex Float)
forall a b. Get (a -> b) -> Get a -> Get b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Endian -> Get Float
forall a. Element a => Endian -> Get a
getElement Endian
en
instance Element (Complex Double) where
    elementType :: forall (proxy :: * -> *). proxy (Complex Double) -> DType
elementType proxy (Complex Double)
_ = DType
C128
    putElement :: Endian -> Complex Double -> Put
putElement Endian
en (Double
a :+ Double
b) = Endian -> Double -> Put
forall a. Element a => Endian -> a -> Put
putElement Endian
en Double
a Put -> Put -> Put
forall a b. PutM a -> PutM b -> PutM b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Endian -> Double -> Put
forall a. Element a => Endian -> a -> Put
putElement Endian
en Double
b
    getElement :: Endian -> Get (Complex Double)
getElement Endian
en = Double -> Double -> Complex Double
forall a. a -> a -> Complex a
(:+) (Double -> Double -> Complex Double)
-> Get Double -> Get (Double -> Complex Double)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Endian -> Get Double
forall a. Element a => Endian -> Get a
getElement Endian
en Get (Double -> Complex Double)
-> Get Double -> Get (Complex Double)
forall a b. Get (a -> b) -> Get a -> Get b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Endian -> Get Double
forall a. Element a => Endian -> Get a
getElement Endian
en

{- | Given limits,endian,shape,then typed Element list, encode to a contiguous array with
that Element dtype. Require exact shape/count match and resource bounds; failures
return InterchangeError.
-}
fromValues ::
    forall a. (Element a) => Limits -> Endian -> [Int] -> [a] -> Either InterchangeError Array
fromValues :: forall a.
Element a =>
Limits -> Endian -> [Int] -> [a] -> Either InterchangeError Array
fromValues Limits
lim Endian
en [Int]
sh [a]
xs = do
    n <- Limits -> [Int] -> Either InterchangeError Int
validateShape Limits
lim [Int]
sh
    let dt = Proxy a -> DType
forall a (proxy :: * -> *). Element a => proxy a -> DType
forall (proxy :: * -> *). proxy a -> DType
elementType (Proxy a
forall {k} (t :: k). Proxy t
Proxy :: Proxy a)
    guardI
        (toInteger n * toInteger (elementBytes dt) <= toInteger (maxBytes lim))
        (LimitExceeded "encoded values")
    guardI (length (take (n + 1) xs) == n) (InvalidData "value count does not match shape")
    contiguousArray lim dt en COrder sh (L.toStrict (runPut (mapM_ (putElement en) xs)))

{- | Decode logical row-major values as the requested Element type, requiring an exact dtype
match. No silent numeric cast occurs; mismatch returns Unsupported/InvalidData.
-}
toValues :: forall a. (Element a) => Array -> Either InterchangeError [a]
toValues :: forall a. Element a => Array -> Either InterchangeError [a]
toValues Array
a = do
    Bool -> InterchangeError -> Either InterchangeError ()
guardI
        (Array -> DType
arrayDType Array
a DType -> DType -> Bool
forall a. Eq a => a -> a -> Bool
== Proxy a -> DType
forall a (proxy :: * -> *). Element a => proxy a -> DType
forall (proxy :: * -> *). proxy a -> DType
elementType (Proxy a
forall {k} (t :: k). Proxy t
Proxy :: Proxy a))
        (String -> InterchangeError
InvalidData String
"requested element type differs from dtype")
    (Int -> Either InterchangeError a)
-> [Int] -> Either InterchangeError [a]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM Int -> Either InterchangeError a
forall {b}. Element b => Int -> Either InterchangeError b
sample (Array -> [Int]
byteOffsets Array
a)
  where
    sample :: Int -> Either InterchangeError b
sample Int
i = case Get b
-> LazyByteString
-> Either
     (LazyByteString, Int64, String) (LazyByteString, Int64, b)
forall a.
Get a
-> LazyByteString
-> Either
     (LazyByteString, Int64, String) (LazyByteString, Int64, a)
runGetOrFail
        (Endian -> Get b
forall a. Element a => Endian -> Get a
getElement (Array -> Endian
arrayEndian Array
a))
        (ByteString -> LazyByteString
L.fromStrict (Int -> ByteString -> ByteString
B.take (DType -> Int
elementBytes (Array -> DType
arrayDType Array
a)) (Int -> ByteString -> ByteString
B.drop Int
i (Array -> ByteString
arrayBytes Array
a)))) of
        Left (LazyByteString
_, Int64
_, String
err) -> InterchangeError -> Either InterchangeError b
forall a b. a -> Either a b
Left (String -> InterchangeError
InvalidData String
err)
        Right (LazyByteString
rest, Int64
_, b
x) | LazyByteString -> Bool
L.null LazyByteString
rest -> b -> Either InterchangeError b
forall a b. b -> Either a b
Right b
x
        Either (LazyByteString, Int64, String) (LazyByteString, Int64, b)
_ -> InterchangeError -> Either InterchangeError b
forall a b. a -> Either a b
Left (String -> InterchangeError
InvalidData String
"element decoder size mismatch")

{- | Given limits,endian,then typed Haskell Tensor, encode logical row-major values and
preserve shape. The Element instance determines wire dtype; resource bounds can fail.
-}
fromTensor :: (Element a) => Limits -> Endian -> T.Tensor a -> Either InterchangeError Array
fromTensor :: forall a.
Element a =>
Limits -> Endian -> Tensor a -> Either InterchangeError Array
fromTensor Limits
lim Endian
en Tensor a
t =
    Limits -> Endian -> [Int] -> [a] -> Either InterchangeError Array
forall a.
Element a =>
Limits -> Endian -> [Int] -> [a] -> Either InterchangeError Array
fromValues Limits
lim Endian
en (Tensor a -> [Int]
forall a. Tensor a -> [Int]
T.tensorShape Tensor a
t) (if Tensor a -> Int
forall a. Tensor a -> Int
T.numel Tensor a
t Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 then [] else Tensor a -> [a]
forall a. Tensor a -> [a]
T.tensorValues Tensor a
t)

{- | Decode an Array of matching Element dtype to a contiguous Haskell Tensor, preserving
logical shape/order. Dtype or checked tensor-construction failure returns
InterchangeError.
-}
toTensor :: (Element a) => Array -> Either InterchangeError (T.Tensor a)
toTensor :: forall a. Element a => Array -> Either InterchangeError (Tensor a)
toTensor Array
a = Array -> Either InterchangeError [a]
forall a. Element a => Array -> Either InterchangeError [a]
toValues Array
a Either InterchangeError [a]
-> ([a] -> Either InterchangeError (Tensor a))
-> Either InterchangeError (Tensor a)
forall a b.
Either InterchangeError a
-> (a -> Either InterchangeError b) -> Either InterchangeError b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (MathError -> Either InterchangeError (Tensor a))
-> (Tensor a -> Either InterchangeError (Tensor a))
-> Either MathError (Tensor a)
-> Either InterchangeError (Tensor a)
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (InterchangeError -> Either InterchangeError (Tensor a)
forall a b. a -> Either a b
Left (InterchangeError -> Either InterchangeError (Tensor a))
-> (MathError -> InterchangeError)
-> MathError
-> Either InterchangeError (Tensor a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> InterchangeError
InvalidData (String -> InterchangeError)
-> (MathError -> String) -> MathError -> InterchangeError
forall b c a. (b -> c) -> (a -> b) -> a -> c
. MathError -> String
forall a. Show a => a -> String
show) Tensor a -> Either InterchangeError (Tensor a)
forall a b. b -> Either a b
Right (Either MathError (Tensor a) -> Either InterchangeError (Tensor a))
-> ([a] -> Either MathError (Tensor a))
-> [a]
-> Either InterchangeError (Tensor a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Int] -> [a] -> Either MathError (Tensor a)
forall a. [Int] -> [a] -> Either MathError (Tensor a)
T.tensor (Array -> [Int]
arrayShape Array
a)

{- | Given limits,endian,then typed Matrix, encode its row-major values as a rank-two Array.
Dtype comes from Element; resource limits can fail.
-}
fromMatrix :: (Element a) => Limits -> Endian -> M.Matrix a -> Either InterchangeError Array
fromMatrix :: forall a.
Element a =>
Limits -> Endian -> Matrix a -> Either InterchangeError Array
fromMatrix Limits
lim Endian
en Matrix a
m = let (Int
r, Int
c) = Matrix a -> (Int, Int)
forall a. Matrix a -> (Int, Int)
M.shape Matrix a
m in Limits -> Endian -> [Int] -> [a] -> Either InterchangeError Array
forall a.
Element a =>
Limits -> Endian -> [Int] -> [a] -> Either InterchangeError Array
fromValues Limits
lim Endian
en [Int
r, Int
c] (Matrix a -> [a]
forall a. Matrix a -> [a]
M.toList Matrix a
m)

{- | Decode a matching typed rank-two Array to Matrix. Rank/dtype or checked shape
conversion failures return InterchangeError.
-}
toMatrix :: (Element a) => Array -> Either InterchangeError (M.Matrix a)
toMatrix :: forall a. Element a => Array -> Either InterchangeError (Matrix a)
toMatrix Array
a = case Array -> [Int]
arrayShape Array
a of
    [Int
r, Int
c] -> Array -> Either InterchangeError [a]
forall a. Element a => Array -> Either InterchangeError [a]
toValues Array
a Either InterchangeError [a]
-> ([a] -> Either InterchangeError (Matrix a))
-> Either InterchangeError (Matrix a)
forall a b.
Either InterchangeError a
-> (a -> Either InterchangeError b) -> Either InterchangeError b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (MathError -> Either InterchangeError (Matrix a))
-> (Matrix a -> Either InterchangeError (Matrix a))
-> Either MathError (Matrix a)
-> Either InterchangeError (Matrix a)
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either (InterchangeError -> Either InterchangeError (Matrix a)
forall a b. a -> Either a b
Left (InterchangeError -> Either InterchangeError (Matrix a))
-> (MathError -> InterchangeError)
-> MathError
-> Either InterchangeError (Matrix a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> InterchangeError
InvalidData (String -> InterchangeError)
-> (MathError -> String) -> MathError -> InterchangeError
forall b c a. (b -> c) -> (a -> b) -> a -> c
. MathError -> String
forall a. Show a => a -> String
show) Matrix a -> Either InterchangeError (Matrix a)
forall a b. b -> Either a b
Right (Either MathError (Matrix a) -> Either InterchangeError (Matrix a))
-> ([a] -> Either MathError (Matrix a))
-> [a]
-> Either InterchangeError (Matrix a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Int -> [a] -> Either MathError (Matrix a)
forall a. Int -> Int -> [a] -> Either MathError (Matrix a)
M.fromListChecked Int
r Int
c
    [Int]
_ -> InterchangeError -> Either InterchangeError (Matrix a)
forall a b. a -> Either a b
Left (String -> InterchangeError
InvalidData String
"matrix requires rank two")

{- | Pack arbitrary views, preserving dtype and exact component bits. Complex
byte swapping reverses each real/imaginary component separately.
-}
repack :: Limits -> Endian -> Order -> Array -> Either InterchangeError Array
repack :: Limits -> Endian -> Order -> Array -> Either InterchangeError Array
repack Limits
lim Endian
en Order
order Array
a = do
    let dt :: DType
dt = Array -> DType
arrayDType Array
a
        width :: Int
width = DType -> Int
elementBytes DType
dt
        component :: Int
component = if DType
dt DType -> DType -> Bool
forall a. Eq a => a -> a -> Bool
== DType
C64 Bool -> Bool -> Bool
|| DType
dt DType -> DType -> Bool
forall a. Eq a => a -> a -> Bool
== DType
C128 then Int
width Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2 else Int
width
        swap :: ByteString -> ByteString
swap ByteString
b
            | Endian
en Endian -> Endian -> Bool
forall a. Eq a => a -> a -> Bool
== Array -> Endian
arrayEndian Array
a Bool -> Bool -> Bool
|| Int
component Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1 = ByteString
b
            | Bool
otherwise =
                [ByteString] -> ByteString
B.concat [ByteString -> ByteString
B.reverse (Int -> ByteString -> ByteString
B.take Int
component (Int -> ByteString -> ByteString
B.drop Int
k ByteString
b)) | Int
k <- [Int
0, Int
component .. Int
width Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]]
        offsets :: [Int]
offsets = case Order
order of
            Order
COrder -> Array -> [Int]
byteOffsets Array
a
            Order
FOrder ->
                Array -> [Int]
byteOffsets
                    ( DType
-> Endian
-> [Int]
-> [Int]
-> Int
-> ByteString
-> [(String, String)]
-> Array
Array
                        DType
dt
                        (Array -> Endian
arrayEndian Array
a)
                        ([Int] -> [Int]
forall a. [a] -> [a]
reverse (Array -> [Int]
arrayShape Array
a))
                        ([Int] -> [Int]
forall a. [a] -> [a]
reverse (Array -> [Int]
arrayStrides Array
a))
                        (Array -> Int
arrayOffset Array
a)
                        (Array -> ByteString
arrayBytes Array
a)
                        []
                    )
    _ <- Limits -> [Int] -> Either InterchangeError Int
validateShape Limits
lim (Array -> [Int]
arrayShape Array
a)
    guardI
        (toInteger (elementCount a) * toInteger width <= toInteger (maxBytes lim))
        (LimitExceeded "repacked bytes")
    b <-
        contiguousArray
            lim
            dt
            en
            order
            (arrayShape a)
            (B.concat [swap (B.take width (B.drop i (arrayBytes a))) | i <- offsets])
    withMetadata lim (arrayMetadata a) b

{- | Given limits,destination endian,physical order,then Array, repack its logical payload
and return raw bytes without shape/header metadata. Checked resource/layout failures
propagate.
-}
rawEncode :: Limits -> Endian -> Order -> Array -> Either InterchangeError B.ByteString
rawEncode :: Limits
-> Endian -> Order -> Array -> Either InterchangeError ByteString
rawEncode Limits
lim Endian
en Order
order = (Array -> ByteString)
-> Either InterchangeError Array
-> Either InterchangeError ByteString
forall a b.
(a -> b) -> Either InterchangeError a -> Either InterchangeError b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Array -> ByteString
arrayBytes (Either InterchangeError Array
 -> Either InterchangeError ByteString)
-> (Array -> Either InterchangeError Array)
-> Array
-> Either InterchangeError ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Limits -> Endian -> Order -> Array -> Either InterchangeError Array
repack Limits
lim Endian
en Order
order

{- | Given limits,dtype,endian,order,shape,then bytes, construct a packed Array. Caller
supplies all metadata; exact byte-count and shape/resource validation is enforced.
-}
rawDecode ::
    Limits -> DType -> Endian -> Order -> [Int] -> B.ByteString -> Either InterchangeError Array
rawDecode :: Limits
-> DType
-> Endian
-> Order
-> [Int]
-> ByteString
-> Either InterchangeError Array
rawDecode = Limits
-> DType
-> Endian
-> Order
-> [Int]
-> ByteString
-> Either InterchangeError Array
contiguousArray