{-# LANGUAGE ScopedTypeVariables #-}

-- SPDX-License-Identifier: MPL-2.0

{- | Pinned CPU storage with shared mutable views. Views retain the same
ForeignPtr, so garbage collection frees storage exactly once after the last
view/callback dies. Freeze/thaw copy; slices and buffer access are zero-copy.
Values have their Storable ABI alignment. Callers synchronize concurrent writes.
-}
module HMath.Tensor.Mutable (
    MutableTensor,
    newMutable,
    thawTensor,
    freezeTensor,
    mutableShape,
    readMutable,
    writeMutable,
    modifyMutable,
    sliceMutable,
    transposeMutable,
    withMutableBuffer,
    mutableAlignment,
) where

import Data.Proxy (Proxy)
import Foreign.ForeignPtr (ForeignPtr, mallocForeignPtrArray, withForeignPtr)
import Foreign.Marshal.Array (advancePtr, pokeArray)
import Foreign.Ptr (Ptr)
import Foreign.Storable (Storable, alignment, peekElemOff, pokeElemOff, sizeOf)
import HMath.Error
import HMath.Tensor

{- | Pinned CPU tensor with shared mutable views and garbage-collected ForeignPtr ownership.
Freeze/thaw copy; callers synchronize concurrent writes.
-}
data MutableTensor a = MutableTensor [Int] [Int] Int Int (ForeignPtr a)

{- | Return current mutable-view dimensions in axis order. Slices/transposes can change the
view shape while sharing storage.
-}
mutableShape :: MutableTensor a -> [Int]
mutableShape :: forall a. MutableTensor a -> [Int]
mutableShape (MutableTensor [Int]
dims [Int]
_ Int
_ Int
_ ForeignPtr a
_) = [Int]
dims

{- | Given a type Proxy, return the Storable ABI alignment in bytes. This reports the
element type requirement, not a user-selectable alignment.
-}
mutableAlignment :: forall a. (Storable a) => Proxy a -> Int
mutableAlignment :: forall a. Storable a => Proxy a -> Int
mutableAlignment Proxy a
_ = a -> Int
forall a. Storable a => a -> Int
alignment (a
forall a. HasCallStack => a
undefined :: a)

{- | Given checked dimensions then fill value, allocate pinned row-major storage filled with
that value. Shape/byte-size overflow returns Left MathError; allocation errors remain
IO exceptions. Empty tensors are valid.
-}
newMutable :: forall a. (Storable a) => [Int] -> a -> IO (Either MathError (MutableTensor a))
newMutable :: forall a.
Storable a =>
[Int] -> a -> IO (Either MathError (MutableTensor a))
newMutable [Int]
dims a
value = case [Int] -> Either MathError Int
shapeSize [Int]
dims of
    Left MathError
e -> Either MathError (MutableTensor a)
-> IO (Either MathError (MutableTensor a))
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (MathError -> Either MathError (MutableTensor a)
forall a b. a -> Either a b
Left MathError
e)
    Right Int
size
        | Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
size Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Int -> Integer
forall a. Integral a => a -> Integer
toInteger (a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined :: a)) 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) ->
            Either MathError (MutableTensor a)
-> IO (Either MathError (MutableTensor a))
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (MathError -> Either MathError (MutableTensor a)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"newMutable: byte allocation overflow"))
    Right Int
size -> do
        storage <- Int -> IO (ForeignPtr a)
forall a. Storable a => Int -> IO (ForeignPtr a)
mallocForeignPtrArray (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
1 Int
size)
        withForeignPtr storage $ \Ptr a
pointer -> Ptr a -> [a] -> IO ()
forall a. Storable a => Ptr a -> [a] -> IO ()
pokeArray Ptr a
pointer (Int -> a -> [a]
forall a. Int -> a -> [a]
replicate Int
size a
value)
        let strides = if Int
size Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 then 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]
dims) Int
0 else Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
drop Int
1 ((Int -> Int -> Int) -> Int -> [Int] -> [Int]
forall a b. (a -> b -> b) -> b -> [a] -> [b]
scanr Int -> Int -> Int
forall a. Num a => a -> a -> a
(*) Int
1 [Int]
dims)
        pure (Right (MutableTensor dims strides 0 size storage))

{- | Copy logical immutable tensor values to new contiguous pinned mutable storage.
Byte-size overflow returns Left InvalidArgument; allocation failures remain IO
exceptions. The source stays unchanged.
-}
thawTensor :: forall a. (Storable a) => Tensor a -> IO (Either MathError (MutableTensor a))
thawTensor :: forall a.
Storable a =>
Tensor a -> IO (Either MathError (MutableTensor a))
thawTensor Tensor a
t
    | Int -> Integer
forall a. Integral a => a -> Integer
toInteger (Tensor a -> Int
forall a. Tensor a -> Int
numel Tensor a
t) Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Int -> Integer
forall a. Integral a => a -> Integer
toInteger (a -> Int
forall a. Storable a => a -> Int
sizeOf (a
forall a. HasCallStack => a
undefined :: a)) 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) =
        Either MathError (MutableTensor a)
-> IO (Either MathError (MutableTensor a))
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (MathError -> Either MathError (MutableTensor a)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"thawTensor: byte allocation overflow"))
    | Bool
otherwise = do
        let xs :: [a]
xs = Tensor a -> [a]
forall a. Tensor a -> [a]
tensorValues Tensor a
t; size :: Int
size = Tensor a -> Int
forall a. Tensor a -> Int
numel Tensor a
t; dims :: [Int]
dims = Tensor a -> [Int]
forall a. Tensor a -> [Int]
tensorShape Tensor a
t
        storage <- Int -> IO (ForeignPtr a)
forall a. Storable a => Int -> IO (ForeignPtr a)
mallocForeignPtrArray (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
1 Int
size)
        withForeignPtr storage $ \Ptr a
pointer -> Ptr a -> [a] -> IO ()
forall a. Storable a => Ptr a -> [a] -> IO ()
pokeArray Ptr a
pointer [a]
xs
        pure
            ( Right
                ( MutableTensor
                    dims
                    (if size == 0 then replicate (length dims) 0 else drop 1 (scanr (*) 1 dims))
                    0
                    size
                    storage
                )
            )

indices :: [Int] -> [[Int]]
indices :: [Int] -> [[Int]]
indices [Int]
dims | 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]
dims = []
indices [] = [[]]
indices (Int
d : [Int]
ds) = [Int
i Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: [Int]
is | Int
i <- [Int
0 .. Int
d Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1], [Int]
is <- [Int] -> [[Int]]
indices [Int]
ds]
address :: MutableTensor a -> [Int] -> Either MathError Int
address :: forall a. MutableTensor a -> [Int] -> Either MathError Int
address (MutableTensor [Int]
dims [Int]
strides Int
offset Int
_ ForeignPtr a
_) [Int]
ix
    | [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
ix Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
dims Bool -> Bool -> Bool
|| [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
or ((Int -> Int -> Bool) -> [Int] -> [Int] -> [Bool]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\Int
i Int
d -> Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
d) [Int]
ix [Int]
dims) =
        MathError -> Either MathError Int
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"mutable tensor: index out of bounds")
    | Bool
otherwise = Int -> Either MathError Int
forall a b. b -> Either a b
Right (Int
offset Int -> Int -> Int
forall a. Num a => a -> a -> a
+ [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((Int -> Int -> Int) -> [Int] -> [Int] -> [Int]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith Int -> Int -> Int
forall a. Num a => a -> a -> a
(*) [Int]
ix [Int]
strides))

{- | Given a mutable view then full zero-based coordinate list, read one element through the
view strides. Wrong rank or out-of-range coordinates return Left InvalidArgument.
Callers synchronize concurrent writes.
-}
readMutable :: (Storable a) => MutableTensor a -> [Int] -> IO (Either MathError a)
readMutable :: forall a.
Storable a =>
MutableTensor a -> [Int] -> IO (Either MathError a)
readMutable t :: MutableTensor a
t@(MutableTensor [Int]
_ [Int]
_ Int
_ Int
_ ForeignPtr a
storage) [Int]
ix = case MutableTensor a -> [Int] -> Either MathError Int
forall a. MutableTensor a -> [Int] -> Either MathError Int
address MutableTensor a
t [Int]
ix of
    Left MathError
e -> Either MathError a -> IO (Either MathError a)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (MathError -> Either MathError a
forall a b. a -> Either a b
Left MathError
e)
    Right Int
location -> a -> Either MathError a
forall a b. b -> Either a b
Right (a -> Either MathError a) -> IO a -> IO (Either MathError a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ForeignPtr a -> (Ptr a -> IO a) -> IO a
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr a
storage (\Ptr a
pointer -> Ptr a -> Int -> IO a
forall a. Storable a => Ptr a -> Int -> IO a
peekElemOff Ptr a
pointer Int
location)

{- | Given a mutable view, full coordinate list, then value, update that storage element.
Invalid coordinates return Left InvalidArgument without writing; synchronize concurrent
access externally.
-}
writeMutable :: (Storable a) => MutableTensor a -> [Int] -> a -> IO (Either MathError ())
writeMutable :: forall a.
Storable a =>
MutableTensor a -> [Int] -> a -> IO (Either MathError ())
writeMutable t :: MutableTensor a
t@(MutableTensor [Int]
_ [Int]
_ Int
_ Int
_ ForeignPtr a
storage) [Int]
ix a
value = case MutableTensor a -> [Int] -> Either MathError Int
forall a. MutableTensor a -> [Int] -> Either MathError Int
address MutableTensor a
t [Int]
ix of
    Left MathError
e -> Either MathError () -> IO (Either MathError ())
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (MathError -> Either MathError ()
forall a b. a -> Either a b
Left MathError
e)
    Right Int
location -> () -> Either MathError ()
forall a b. b -> Either a b
Right (() -> Either MathError ()) -> IO () -> IO (Either MathError ())
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ForeignPtr a -> (Ptr a -> IO ()) -> IO ()
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr a
storage (\Ptr a
pointer -> Ptr a -> Int -> a -> IO ()
forall a. Storable a => Ptr a -> Int -> a -> IO ()
pokeElemOff Ptr a
pointer Int
location a
value)

{- | Given view, coordinates, then transformation, read and replace an element. Invalid
coordinates fail. This read-modify-write is not atomic across threads; callers provide
synchronization.
-}
modifyMutable :: (Storable a) => MutableTensor a -> [Int] -> (a -> a) -> IO (Either MathError ())
modifyMutable :: forall a.
Storable a =>
MutableTensor a -> [Int] -> (a -> a) -> IO (Either MathError ())
modifyMutable MutableTensor a
t [Int]
ix a -> a
action = do
    result <- MutableTensor a -> [Int] -> IO (Either MathError a)
forall a.
Storable a =>
MutableTensor a -> [Int] -> IO (Either MathError a)
readMutable MutableTensor a
t [Int]
ix
    case result of { Left MathError
e -> Either MathError () -> IO (Either MathError ())
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (MathError -> Either MathError ()
forall a b. a -> Either a b
Left MathError
e); Right a
value -> MutableTensor a -> [Int] -> a -> IO (Either MathError ())
forall a.
Storable a =>
MutableTensor a -> [Int] -> a -> IO (Either MathError ())
writeMutable MutableTensor a
t [Int]
ix (a -> a
action a
value) }

{- | Copy current logical view values into a new immutable tensor in row-major order. Return
checked tensor construction errors; later mutable writes do not affect the result.
-}
freezeTensor :: (Storable a) => MutableTensor a -> IO (Either MathError (Tensor a))
freezeTensor :: forall a.
Storable a =>
MutableTensor a -> IO (Either MathError (Tensor a))
freezeTensor MutableTensor a
t = do
    values <- ([Int] -> IO (Either MathError a))
-> [[Int]] -> IO [Either MathError a]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> [a] -> f [b]
traverse (MutableTensor a -> [Int] -> IO (Either MathError a)
forall a.
Storable a =>
MutableTensor a -> [Int] -> IO (Either MathError a)
readMutable MutableTensor a
t) ([Int] -> [[Int]]
indices (MutableTensor a -> [Int]
forall a. MutableTensor a -> [Int]
mutableShape MutableTensor a
t))
    pure (sequence values >>= tensor (mutableShape t))

{- | Given axis, start, count, nonzero step, then view, return a zero-copy mutable slice
retaining shared ForeignPtr ownership. Negative steps are supported if all reached
indices stay in bounds. Invalid extent/axis or stride overflow fails.
-}
sliceMutable :: Int -> Int -> Int -> Int -> MutableTensor a -> Either MathError (MutableTensor a)
sliceMutable :: forall a.
Int
-> Int
-> Int
-> Int
-> MutableTensor a
-> Either MathError (MutableTensor a)
sliceMutable Int
axis Int
start Int
count Int
step (MutableTensor [Int]
dims [Int]
strides Int
offset Int
size ForeignPtr a
storage)
    | Int
axis Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int
axis Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
dims = MathError -> Either MathError (MutableTensor a)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"sliceMutable: invalid axis")
    | Int
count Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0
        Bool -> Bool -> Bool
|| Int
step Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
        Bool -> Bool -> Bool
|| Int
start Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0
        Bool -> Bool -> Bool
|| (Int
count Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 Bool -> Bool -> Bool
&& (Int
start Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= [Int]
dims [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
axis Bool -> Bool -> Bool
|| Integer
end Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
0 Bool -> Bool -> Bool
|| Integer
end Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
>= Int -> Integer
forall a. Integral a => a -> Integer
toInteger ([Int]
dims [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
axis)))
        Bool -> Bool -> Bool
|| (Int
count Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 Bool -> Bool -> Bool
&& Int
start Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> [Int]
dims [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
axis) =
        MathError -> Either MathError (MutableTensor a)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"sliceMutable: invalid extent")
    | Integer
stride 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
minBound :: Int) Bool -> Bool -> Bool
|| Integer
stride 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) =
        MathError -> Either MathError (MutableTensor a)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"sliceMutable: stride overflow")
    | Bool
otherwise =
        MutableTensor a -> Either MathError (MutableTensor a)
forall a b. b -> Either a b
Right
            ( [Int] -> [Int] -> Int -> Int -> ForeignPtr a -> MutableTensor a
forall a.
[Int] -> [Int] -> Int -> Int -> ForeignPtr a -> MutableTensor a
MutableTensor
                (Int -> Int -> [Int] -> [Int]
forall {a}. Int -> a -> [a] -> [a]
replace Int
axis Int
count [Int]
dims)
                (Int -> Int -> [Int] -> [Int]
forall {a}. Int -> a -> [a] -> [a]
replace Int
axis (Integer -> Int
forall a. Num a => Integer -> a
fromInteger Integer
stride) [Int]
strides)
                (if Int
count Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 then Int
offset else Integer -> Int
forall a. Num a => Integer -> a
fromInteger Integer
newOffset)
                Int
size
                ForeignPtr a
storage
            )
  where
    end :: Integer
end = Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
start Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Int -> Integer
forall a. Integral a => a -> Integer
toInteger (Int
count 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
step
    stride :: Integer
stride = Int -> Integer
forall a. Integral a => a -> Integer
toInteger ([Int]
strides [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
axis) Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
step
    newOffset :: Integer
newOffset = Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
offset Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Int -> Integer
forall a. Integral a => a -> Integer
toInteger Int
start Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
* Int -> Integer
forall a. Integral a => a -> Integer
toInteger ([Int]
strides [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
axis)
    replace :: Int -> a -> [a] -> [a]
replace Int
i a
x [a]
xs = Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
take Int
i [a]
xs [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
: Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
drop (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) [a]
xs

{- | Reverse the order of all shape axes and strides, returning a zero-copy mutable view.
For rank two this is a matrix transpose; storage ownership remains shared.
-}
transposeMutable :: MutableTensor a -> MutableTensor a
transposeMutable :: forall a. MutableTensor a -> MutableTensor a
transposeMutable (MutableTensor [Int]
dims [Int]
strides Int
offset Int
size ForeignPtr a
storage) = [Int] -> [Int] -> Int -> Int -> ForeignPtr a -> MutableTensor a
forall a.
[Int] -> [Int] -> Int -> Int -> ForeignPtr a -> MutableTensor a
MutableTensor ([Int] -> [Int]
forall a. [a] -> [a]
reverse [Int]
dims) ([Int] -> [Int]
forall a. [a] -> [a]
reverse [Int]
strides) Int
offset Int
size ForeignPtr a
storage

{- | The pointer is valid only during the callback. Shape/strides describe views;
callbacks must honor their extents and may not retain the borrowed pointer.
-}
withMutableBuffer :: (Storable a) => MutableTensor a -> (Ptr a -> [Int] -> [Int] -> IO b) -> IO b
withMutableBuffer :: forall a b.
Storable a =>
MutableTensor a -> (Ptr a -> [Int] -> [Int] -> IO b) -> IO b
withMutableBuffer (MutableTensor [Int]
dims [Int]
strides Int
offset Int
_ ForeignPtr a
storage) Ptr a -> [Int] -> [Int] -> IO b
action = ForeignPtr a -> (Ptr a -> IO b) -> IO b
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr a
storage ((Ptr a -> IO b) -> IO b) -> (Ptr a -> IO b) -> IO b
forall a b. (a -> b) -> a -> b
$ \Ptr a
pointer -> Ptr a -> [Int] -> [Int] -> IO b
action (Ptr a
pointer Ptr a -> Int -> Ptr a
forall a. Storable a => Ptr a -> Int -> Ptr a
`advancePtr` Int
offset) [Int]
dims [Int]
strides