{-# LANGUAGE DeriveFoldable #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}

-- SPDX-License-Identifier: MPL-2.0

{- | Checked containers and copying conversion boundaries. Boxed containers can
retain lazy elements; strictBoxed forces each element to WHNF, not deep NF.
Unboxed storage specializes to Double. ST regions cannot escape runSTBuffer;
IO storage has one ForeignPtr owner shared by aliases and is garbage collected.
IO operations synchronize on the buffer lock. Pointer callbacks borrow storage
only for their dynamic extent and must not retain the pointer or reenter the
same buffer's locking operations. Freeze and thaw always copy.
-}
module HMath.API.Buffer (
    Boxed,
    lazyBoxed,
    strictBoxed,
    boxedValues,
    boxedLength,
    mapStrict,
    UnboxedDouble,
    unboxedDouble,
    unboxedValues,
    STBuffer,
    runSTBuffer,
    thawSTBuffer,
    freezeSTBuffer,
    readSTBuffer,
    writeSTBuffer,
    IOBuffer,
    newIOBuffer,
    thawIOBuffer,
    freezeIOBuffer,
    ioBufferLength,
    readIOBuffer,
    writeIOBuffer,
    withIOBuffer,
    unboxedToBoxed,
    boxedToUnboxed,
) where

import Control.Concurrent.MVar
import Control.Monad.ST (ST, runST)
import Data.Array (Array, elems, listArray)
import Data.Array.ST (STArray, getElems, newListArray, readArray, writeArray)
import qualified Data.Array.Unboxed as U
import Foreign.ForeignPtr (ForeignPtr, mallocForeignPtrArray, withForeignPtr)
import Foreign.Marshal.Array (peekArray, pokeArray)
import Foreign.Ptr (Ptr)
import Foreign.Storable (Storable, peekElemOff, pokeElemOff, sizeOf)
import HMath.Error

{- | Immutable boxed array preserving element laziness unless built with strictBoxed;
conversions preserve logical order.
-}
newtype Boxed a = Boxed (Array Int a) deriving (Boxed a -> Boxed a -> Bool
(Boxed a -> Boxed a -> Bool)
-> (Boxed a -> Boxed a -> Bool) -> Eq (Boxed a)
forall a. Eq a => Boxed a -> Boxed a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall a. Eq a => Boxed a -> Boxed a -> Bool
== :: Boxed a -> Boxed a -> Bool
$c/= :: forall a. Eq a => Boxed a -> Boxed a -> Bool
/= :: Boxed a -> Boxed a -> Bool
Eq, Int -> Boxed a -> ShowS
[Boxed a] -> ShowS
Boxed a -> String
(Int -> Boxed a -> ShowS)
-> (Boxed a -> String) -> ([Boxed a] -> ShowS) -> Show (Boxed a)
forall a. Show a => Int -> Boxed a -> ShowS
forall a. Show a => [Boxed a] -> ShowS
forall a. Show a => Boxed a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall a. Show a => Int -> Boxed a -> ShowS
showsPrec :: Int -> Boxed a -> ShowS
$cshow :: forall a. Show a => Boxed a -> String
show :: Boxed a -> String
$cshowList :: forall a. Show a => [Boxed a] -> ShowS
showList :: [Boxed a] -> ShowS
Show, (forall a b. (a -> b) -> Boxed a -> Boxed b)
-> (forall a b. a -> Boxed b -> Boxed a) -> Functor Boxed
forall a b. a -> Boxed b -> Boxed a
forall a b. (a -> b) -> Boxed a -> Boxed b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall a b. (a -> b) -> Boxed a -> Boxed b
fmap :: forall a b. (a -> b) -> Boxed a -> Boxed b
$c<$ :: forall a b. a -> Boxed b -> Boxed a
<$ :: forall a b. a -> Boxed b -> Boxed a
Functor, (forall m. Monoid m => Boxed m -> m)
-> (forall m a. Monoid m => (a -> m) -> Boxed a -> m)
-> (forall m a. Monoid m => (a -> m) -> Boxed a -> m)
-> (forall a b. (a -> b -> b) -> b -> Boxed a -> b)
-> (forall a b. (a -> b -> b) -> b -> Boxed a -> b)
-> (forall b a. (b -> a -> b) -> b -> Boxed a -> b)
-> (forall b a. (b -> a -> b) -> b -> Boxed a -> b)
-> (forall a. (a -> a -> a) -> Boxed a -> a)
-> (forall a. (a -> a -> a) -> Boxed a -> a)
-> (forall a. Boxed a -> [a])
-> (forall a. Boxed a -> Bool)
-> (forall a. Boxed a -> Int)
-> (forall a. Eq a => a -> Boxed a -> Bool)
-> (forall a. Ord a => Boxed a -> a)
-> (forall a. Ord a => Boxed a -> a)
-> (forall a. Num a => Boxed a -> a)
-> (forall a. Num a => Boxed a -> a)
-> Foldable Boxed
forall a. Eq a => a -> Boxed a -> Bool
forall a. Num a => Boxed a -> a
forall a. Ord a => Boxed a -> a
forall m. Monoid m => Boxed m -> m
forall a. Boxed a -> Bool
forall a. Boxed a -> Int
forall a. Boxed a -> [a]
forall a. (a -> a -> a) -> Boxed a -> a
forall m a. Monoid m => (a -> m) -> Boxed a -> m
forall b a. (b -> a -> b) -> b -> Boxed a -> b
forall a b. (a -> b -> b) -> b -> Boxed a -> b
forall (t :: * -> *).
(forall m. Monoid m => t m -> m)
-> (forall m a. Monoid m => (a -> m) -> t a -> m)
-> (forall m a. Monoid m => (a -> m) -> t a -> m)
-> (forall a b. (a -> b -> b) -> b -> t a -> b)
-> (forall a b. (a -> b -> b) -> b -> t a -> b)
-> (forall b a. (b -> a -> b) -> b -> t a -> b)
-> (forall b a. (b -> a -> b) -> b -> t a -> b)
-> (forall a. (a -> a -> a) -> t a -> a)
-> (forall a. (a -> a -> a) -> t a -> a)
-> (forall a. t a -> [a])
-> (forall a. t a -> Bool)
-> (forall a. t a -> Int)
-> (forall a. Eq a => a -> t a -> Bool)
-> (forall a. Ord a => t a -> a)
-> (forall a. Ord a => t a -> a)
-> (forall a. Num a => t a -> a)
-> (forall a. Num a => t a -> a)
-> Foldable t
$cfold :: forall m. Monoid m => Boxed m -> m
fold :: forall m. Monoid m => Boxed m -> m
$cfoldMap :: forall m a. Monoid m => (a -> m) -> Boxed a -> m
foldMap :: forall m a. Monoid m => (a -> m) -> Boxed a -> m
$cfoldMap' :: forall m a. Monoid m => (a -> m) -> Boxed a -> m
foldMap' :: forall m a. Monoid m => (a -> m) -> Boxed a -> m
$cfoldr :: forall a b. (a -> b -> b) -> b -> Boxed a -> b
foldr :: forall a b. (a -> b -> b) -> b -> Boxed a -> b
$cfoldr' :: forall a b. (a -> b -> b) -> b -> Boxed a -> b
foldr' :: forall a b. (a -> b -> b) -> b -> Boxed a -> b
$cfoldl :: forall b a. (b -> a -> b) -> b -> Boxed a -> b
foldl :: forall b a. (b -> a -> b) -> b -> Boxed a -> b
$cfoldl' :: forall b a. (b -> a -> b) -> b -> Boxed a -> b
foldl' :: forall b a. (b -> a -> b) -> b -> Boxed a -> b
$cfoldr1 :: forall a. (a -> a -> a) -> Boxed a -> a
foldr1 :: forall a. (a -> a -> a) -> Boxed a -> a
$cfoldl1 :: forall a. (a -> a -> a) -> Boxed a -> a
foldl1 :: forall a. (a -> a -> a) -> Boxed a -> a
$ctoList :: forall a. Boxed a -> [a]
toList :: forall a. Boxed a -> [a]
$cnull :: forall a. Boxed a -> Bool
null :: forall a. Boxed a -> Bool
$clength :: forall a. Boxed a -> Int
length :: forall a. Boxed a -> Int
$celem :: forall a. Eq a => a -> Boxed a -> Bool
elem :: forall a. Eq a => a -> Boxed a -> Bool
$cmaximum :: forall a. Ord a => Boxed a -> a
maximum :: forall a. Ord a => Boxed a -> a
$cminimum :: forall a. Ord a => Boxed a -> a
minimum :: forall a. Ord a => Boxed a -> a
$csum :: forall a. Num a => Boxed a -> a
sum :: forall a. Num a => Boxed a -> a
$cproduct :: forall a. Num a => Boxed a -> a
product :: forall a. Num a => Boxed a -> a
Foldable, Functor Boxed
Foldable Boxed
(Functor Boxed, Foldable Boxed) =>
(forall (f :: * -> *) a b.
 Applicative f =>
 (a -> f b) -> Boxed a -> f (Boxed b))
-> (forall (f :: * -> *) a.
    Applicative f =>
    Boxed (f a) -> f (Boxed a))
-> (forall (m :: * -> *) a b.
    Monad m =>
    (a -> m b) -> Boxed a -> m (Boxed b))
-> (forall (m :: * -> *) a. Monad m => Boxed (m a) -> m (Boxed a))
-> Traversable Boxed
forall (t :: * -> *).
(Functor t, Foldable t) =>
(forall (f :: * -> *) a b.
 Applicative f =>
 (a -> f b) -> t a -> f (t b))
-> (forall (f :: * -> *) a. Applicative f => t (f a) -> f (t a))
-> (forall (m :: * -> *) a b.
    Monad m =>
    (a -> m b) -> t a -> m (t b))
-> (forall (m :: * -> *) a. Monad m => t (m a) -> m (t a))
-> Traversable t
forall (m :: * -> *) a. Monad m => Boxed (m a) -> m (Boxed a)
forall (f :: * -> *) a. Applicative f => Boxed (f a) -> f (Boxed a)
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Boxed a -> m (Boxed b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Boxed a -> f (Boxed b)
$ctraverse :: forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Boxed a -> f (Boxed b)
traverse :: forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Boxed a -> f (Boxed b)
$csequenceA :: forall (f :: * -> *) a. Applicative f => Boxed (f a) -> f (Boxed a)
sequenceA :: forall (f :: * -> *) a. Applicative f => Boxed (f a) -> f (Boxed a)
$cmapM :: forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Boxed a -> m (Boxed b)
mapM :: forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Boxed a -> m (Boxed b)
$csequence :: forall (m :: * -> *) a. Monad m => Boxed (m a) -> m (Boxed a)
sequence :: forall (m :: * -> *) a. Monad m => Boxed (m a) -> m (Boxed a)
Traversable)

-- | Finite lists only. The list spine is consumed; elements remain lazy.
lazyBoxed :: [a] -> Boxed a
lazyBoxed :: forall a. [a] -> Boxed a
lazyBoxed [a]
xs = Array Int a -> Boxed a
forall a. Array Int a -> Boxed a
Boxed ((Int, Int) -> [a] -> Array Int a
forall i e. Ix i => (i, i) -> [e] -> Array i e
listArray (Int
0, [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) [a]
xs)

{- | Construct a boxed container from a finite list, forcing every element to weak head
normal form (not deep normal form). Preserve order; element evaluation exceptions
propagate.
-}
strictBoxed :: [a] -> Boxed a
strictBoxed :: forall a. [a] -> Boxed a
strictBoxed [a]
xs = (a -> () -> ()) -> () -> [a] -> ()
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr a -> () -> ()
forall a b. a -> b -> b
seq () [a]
xs () -> Boxed a -> Boxed a
forall a b. a -> b -> b
`seq` [a] -> Boxed a
forall a. [a] -> Boxed a
lazyBoxed [a]
xs

{- | Return boxed elements in zero-based array order as a list. Elements retain the
laziness/strictness chosen at construction.
-}
boxedValues :: Boxed a -> [a]
boxedValues :: forall a. Boxed a -> [a]
boxedValues (Boxed Array Int a
a) = Array Int a -> [a]
forall i e. Array i e -> [e]
elems Array Int a
a

-- | Return the number of elements in a boxed container; empty storage has length zero.
boxedLength :: Boxed a -> Int
boxedLength :: forall a. Boxed a -> Int
boxedLength = [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([a] -> Int) -> (Boxed a -> [a]) -> Boxed a -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Boxed a -> [a]
forall a. Boxed a -> [a]
boxedValues

{- | Given element transformation then boxed input, return a boxed result with every mapped
element forced to WHNF. Preserve length/order; evaluation exceptions propagate.
-}
mapStrict :: (a -> b) -> Boxed a -> Boxed b
mapStrict :: forall a b. (a -> b) -> Boxed a -> Boxed b
mapStrict a -> b
f = [b] -> Boxed b
forall a. [a] -> Boxed a
strictBoxed ([b] -> Boxed b) -> (Boxed a -> [b]) -> Boxed a -> Boxed b
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> b) -> [a] -> [b]
forall a b. (a -> b) -> [a] -> [b]
map a -> b
f ([a] -> [b]) -> (Boxed a -> [a]) -> Boxed a -> [b]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Boxed a -> [a]
forall a. Boxed a -> [a]
boxedValues

-- | Immutable unboxed Double storage. Construct with unboxedDouble; boxed conversions copy.
newtype UnboxedDouble = UnboxedDouble (U.UArray Int Double) deriving (UnboxedDouble -> UnboxedDouble -> Bool
(UnboxedDouble -> UnboxedDouble -> Bool)
-> (UnboxedDouble -> UnboxedDouble -> Bool) -> Eq UnboxedDouble
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: UnboxedDouble -> UnboxedDouble -> Bool
== :: UnboxedDouble -> UnboxedDouble -> Bool
$c/= :: UnboxedDouble -> UnboxedDouble -> Bool
/= :: UnboxedDouble -> UnboxedDouble -> Bool
Eq, Int -> UnboxedDouble -> ShowS
[UnboxedDouble] -> ShowS
UnboxedDouble -> String
(Int -> UnboxedDouble -> ShowS)
-> (UnboxedDouble -> String)
-> ([UnboxedDouble] -> ShowS)
-> Show UnboxedDouble
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> UnboxedDouble -> ShowS
showsPrec :: Int -> UnboxedDouble -> ShowS
$cshow :: UnboxedDouble -> String
show :: UnboxedDouble -> String
$cshowList :: [UnboxedDouble] -> ShowS
showList :: [UnboxedDouble] -> ShowS
Show)

{- | Copy a finite Double list into an immutable unboxed array in list order. Values are
evaluated by unboxing; no finiteness validation is added.
-}
unboxedDouble :: [Double] -> UnboxedDouble
unboxedDouble :: [Double] -> UnboxedDouble
unboxedDouble [Double]
xs = UArray Int Double -> UnboxedDouble
UnboxedDouble ((Int, Int) -> [Double] -> UArray Int Double
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
(i, i) -> [e] -> a i e
U.listArray (Int
0, [Double] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Double]
xs Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) [Double]
xs)

-- | Copy/expand unboxed Doubles into a list in array order, preserving length.
unboxedValues :: UnboxedDouble -> [Double]
unboxedValues :: UnboxedDouble -> [Double]
unboxedValues (UnboxedDouble UArray Int Double
a) = UArray Int Double -> [Double]
forall (a :: * -> * -> *) e i. (IArray a e, Ix i) => a i e -> [e]
U.elems UArray Int Double
a

{- | Convert unboxed Double storage to a new strict boxed container, preserving order and
values.
-}
unboxedToBoxed :: UnboxedDouble -> Boxed Double
unboxedToBoxed :: UnboxedDouble -> Boxed Double
unboxedToBoxed = [Double] -> Boxed Double
forall a. [a] -> Boxed a
strictBoxed ([Double] -> Boxed Double)
-> (UnboxedDouble -> [Double]) -> UnboxedDouble -> Boxed Double
forall b c a. (b -> c) -> (a -> b) -> a -> c
. UnboxedDouble -> [Double]
unboxedValues

{- | Convert boxed Double values to a new unboxed array, forcing values during copying.
Preserve order and length.
-}
boxedToUnboxed :: Boxed Double -> UnboxedDouble
boxedToUnboxed :: Boxed Double -> UnboxedDouble
boxedToUnboxed = [Double] -> UnboxedDouble
unboxedDouble ([Double] -> UnboxedDouble)
-> (Boxed Double -> [Double]) -> Boxed Double -> UnboxedDouble
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Boxed Double -> [Double]
forall a. Boxed a -> [a]
boxedValues

{- | Mutable boxed storage confined by the ST state parameter. Freeze/thaw copy; use
runSTBuffer to prevent region escape.
-}
data STBuffer s a = STBuffer !Int !(STArray s Int a)

-- | The rank-2 state parameter prevents a mutable region from escaping.
runSTBuffer :: Boxed a -> (forall s. STBuffer s a -> ST s b) -> b
runSTBuffer :: forall a b. Boxed a -> (forall s. STBuffer s a -> ST s b) -> b
runSTBuffer Boxed a
boxed forall s. STBuffer s a -> ST s b
action = (forall s. ST s b) -> b
forall a. (forall s. ST s a) -> a
runST (Boxed a -> ST s (STBuffer s a)
forall a s. Boxed a -> ST s (STBuffer s a)
thawSTBuffer Boxed a
boxed ST s (STBuffer s a) -> (STBuffer s a -> ST s b) -> ST s b
forall a b. ST s a -> (a -> ST s b) -> ST s b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= STBuffer s a -> ST s b
forall s. STBuffer s a -> ST s b
action)

{- | Copy immutable boxed values into new mutable storage in the current ST region. The
input remains unchanged; storage cannot escape the region.
-}
thawSTBuffer :: Boxed a -> ST s (STBuffer s a)
thawSTBuffer :: forall a s. Boxed a -> ST s (STBuffer s a)
thawSTBuffer Boxed a
boxed = let xs :: [a]
xs = Boxed a -> [a]
forall a. Boxed a -> [a]
boxedValues Boxed a
boxed; n :: Int
n = [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs in Int -> STArray s Int a -> STBuffer s a
forall s a. Int -> STArray s Int a -> STBuffer s a
STBuffer Int
n (STArray s Int a -> STBuffer s a)
-> ST s (STArray s Int a) -> ST s (STBuffer s a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Int, Int) -> [a] -> ST s (STArray s Int a)
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
(i, i) -> [e] -> m (a i e)
newListArray (Int
0, Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) [a]
xs

{- | Copy current mutable ST values into an immutable boxed container; later writes do not
modify the frozen result.
-}
freezeSTBuffer :: STBuffer s a -> ST s (Boxed a)
freezeSTBuffer :: forall s a. STBuffer s a -> ST s (Boxed a)
freezeSTBuffer (STBuffer Int
_ STArray s Int a
array) = [a] -> Boxed a
forall a. [a] -> Boxed a
lazyBoxed ([a] -> Boxed a) -> ST s [a] -> ST s (Boxed a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> STArray s Int a -> ST s [a]
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> m [e]
getElems STArray s Int a
array

indexCheck :: Int -> Int -> Either MathError ()
indexCheck :: Int -> Int -> Either MathError ()
indexCheck Int
n Int
i
    | 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
n = MathError -> Either MathError ()
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"buffer index out of bounds")
    | Bool
otherwise = () -> Either MathError ()
forall a b. b -> Either a b
Right ()

{- | Given ST buffer then zero-based index, read the element. Negative/out-of-range indices
return Left InvalidArgument within ST, rather than indexing the array.
-}
readSTBuffer :: STBuffer s a -> Int -> ST s (Either MathError a)
readSTBuffer :: forall s a. STBuffer s a -> Int -> ST s (Either MathError a)
readSTBuffer (STBuffer Int
n STArray s Int a
array) Int
i = case Int -> Int -> Either MathError ()
indexCheck Int
n Int
i of Left MathError
e -> Either MathError a -> ST s (Either MathError a)
forall a. a -> ST s 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 () -> a -> Either MathError a
forall a b. b -> Either a b
Right (a -> Either MathError a) -> ST s a -> ST s (Either MathError a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> STArray s Int a -> Int -> ST s a
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> m e
readArray STArray s Int a
array Int
i

{- | Given ST buffer, zero-based index, then value, update one element and return Right ().
Invalid index returns Left InvalidArgument without writing.
-}
writeSTBuffer :: STBuffer s a -> Int -> a -> ST s (Either MathError ())
writeSTBuffer :: forall s a. STBuffer s a -> Int -> a -> ST s (Either MathError ())
writeSTBuffer (STBuffer Int
n STArray s Int a
array) Int
i a
x = case Int -> Int -> Either MathError ()
indexCheck Int
n Int
i of Left MathError
e -> Either MathError () -> ST s (Either MathError ())
forall a. a -> ST s a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (MathError -> Either MathError ()
forall a b. a -> Either a b
Left MathError
e); Right () -> () -> Either MathError ()
forall a b. b -> Either a b
Right (() -> Either MathError ())
-> ST s () -> ST s (Either MathError ())
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> STArray s Int a -> Int -> a -> ST s ()
forall (a :: * -> * -> *) e (m :: * -> *) i.
(MArray a e m, Ix i) =>
a i e -> i -> e -> m ()
writeArray STArray s Int a
array Int
i a
x

{- | Pinned garbage-collected Storable buffer with synchronized IO access. Pointer callbacks
borrow storage only for their dynamic lifetime.
-}
data IOBuffer a = IOBuffer !Int !(ForeignPtr a) !(MVar ())

{- | Return the immutable element count of a pinned IO buffer. This query neither borrows
its pointer nor takes ownership.
-}
ioBufferLength :: IOBuffer a -> Int
ioBufferLength :: forall a. IOBuffer a -> Int
ioBufferLength (IOBuffer Int
n ForeignPtr a
_ MVar ()
_) = Int
n

{- | Allocation overflow is checked before allocation. Allocation/OS exceptions
remain IO exceptions; MathError describes input and numeric contract failures.
-}
newIOBuffer :: forall a. (Storable a) => Int -> a -> IO (Either MathError (IOBuffer a))
newIOBuffer :: forall a.
Storable a =>
Int -> a -> IO (Either MathError (IOBuffer a))
newIOBuffer Int
n a
x
    | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 Bool -> Bool -> Bool
|| Int -> Integer
forall a. Integral a => a -> Integer
toInteger (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
1 Int
n) 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 (IOBuffer a) -> IO (Either MathError (IOBuffer a))
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (MathError -> Either MathError (IOBuffer a)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"IO buffer allocation size overflow"))
    | Bool
otherwise = do
        fp <- 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
n)
        withForeignPtr fp (\Ptr a
p -> Ptr a -> [a] -> IO ()
forall a. Storable a => Ptr a -> [a] -> IO ()
pokeArray Ptr a
p (Int -> a -> [a]
forall a. Int -> a -> [a]
replicate Int
n a
x))
        lock <- newMVar ()
        pure (Right (IOBuffer n fp lock))

{- | Copy boxed Storable values into new pinned IO storage. Return Left InvalidArgument for
byte-size overflow; allocation/OS/evaluation failures remain IO exceptions.
-}
thawIOBuffer :: forall a. (Storable a) => Boxed a -> IO (Either MathError (IOBuffer a))
thawIOBuffer :: forall a.
Storable a =>
Boxed a -> IO (Either MathError (IOBuffer a))
thawIOBuffer Boxed a
boxed
    | Int -> Integer
forall a. Integral a => a -> Integer
toInteger (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
1 Int
n) 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 (IOBuffer a) -> IO (Either MathError (IOBuffer a))
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (MathError -> Either MathError (IOBuffer a)
forall a b. a -> Either a b
Left (String -> MathError
InvalidArgument String
"IO buffer allocation size overflow"))
    | Bool
otherwise = do
        fp <- 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
n)
        withForeignPtr fp (\Ptr a
p -> Ptr a -> [a] -> IO ()
forall a. Storable a => Ptr a -> [a] -> IO ()
pokeArray Ptr a
p [a]
xs)
        lock <- newMVar ()
        pure (Right (IOBuffer n fp lock))
  where
    xs :: [a]
xs = Boxed a -> [a]
forall a. Boxed a -> [a]
boxedValues Boxed a
boxed; n :: Int
n = [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs

{- | Given IO buffer then callback, hold its lock and lend a Ptr plus element count for the
callback duration. Do not retain the pointer or reenter locking operations on this same
buffer. Callback exceptions propagate while releasing the lock.
-}
withIOBuffer :: (Storable a) => IOBuffer a -> (Ptr a -> Int -> IO b) -> IO b
withIOBuffer :: forall a b.
Storable a =>
IOBuffer a -> (Ptr a -> Int -> IO b) -> IO b
withIOBuffer (IOBuffer Int
n ForeignPtr a
fp MVar ()
lock) Ptr a -> Int -> IO b
action = MVar () -> (() -> IO b) -> IO b
forall a b. MVar a -> (a -> IO b) -> IO b
withMVar MVar ()
lock (\() -> ForeignPtr a -> (Ptr a -> IO b) -> IO b
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr a
fp (\Ptr a
p -> Ptr a -> Int -> IO b
action Ptr a
p Int
n))

{- | Lock and copy pinned IO values to a strict immutable boxed container. The result is
independent of later buffer writes; allocation/IO exceptions propagate.
-}
freezeIOBuffer :: (Storable a) => IOBuffer a -> IO (Boxed a)
freezeIOBuffer :: forall a. Storable a => IOBuffer a -> IO (Boxed a)
freezeIOBuffer IOBuffer a
buffer = IOBuffer a -> (Ptr a -> Int -> IO (Boxed a)) -> IO (Boxed a)
forall a b.
Storable a =>
IOBuffer a -> (Ptr a -> Int -> IO b) -> IO b
withIOBuffer IOBuffer a
buffer (\Ptr a
p Int
n -> [a] -> Boxed a
forall a. [a] -> Boxed a
strictBoxed ([a] -> Boxed a) -> IO [a] -> IO (Boxed a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Ptr a -> IO [a]
forall a. Storable a => Int -> Ptr a -> IO [a]
peekArray Int
n Ptr a
p)

{- | Given IO buffer then zero-based index, read under the buffer lock. Invalid index
returns Left InvalidArgument; IO exceptions propagate.
-}
readIOBuffer :: (Storable a) => IOBuffer a -> Int -> IO (Either MathError a)
readIOBuffer :: forall a.
Storable a =>
IOBuffer a -> Int -> IO (Either MathError a)
readIOBuffer IOBuffer a
buffer Int
i = case Int -> Int -> Either MathError ()
indexCheck (IOBuffer a -> Int
forall a. IOBuffer a -> Int
ioBufferLength IOBuffer a
buffer) Int
i 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 () -> IOBuffer a
-> (Ptr a -> Int -> IO (Either MathError a))
-> IO (Either MathError a)
forall a b.
Storable a =>
IOBuffer a -> (Ptr a -> Int -> IO b) -> IO b
withIOBuffer IOBuffer a
buffer (\Ptr a
p Int
_ -> 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
<$> Ptr a -> Int -> IO a
forall a. Storable a => Ptr a -> Int -> IO a
peekElemOff Ptr a
p Int
i)

{- | Given IO buffer, zero-based index, then value, write under its lock and return Right
(). Invalid index returns Left InvalidArgument without modification.
-}
writeIOBuffer :: (Storable a) => IOBuffer a -> Int -> a -> IO (Either MathError ())
writeIOBuffer :: forall a.
Storable a =>
IOBuffer a -> Int -> a -> IO (Either MathError ())
writeIOBuffer IOBuffer a
buffer Int
i a
x = case Int -> Int -> Either MathError ()
indexCheck (IOBuffer a -> Int
forall a. IOBuffer a -> Int
ioBufferLength IOBuffer a
buffer) Int
i 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 () -> IOBuffer a
-> (Ptr a -> Int -> IO (Either MathError ()))
-> IO (Either MathError ())
forall a b.
Storable a =>
IOBuffer a -> (Ptr a -> Int -> IO b) -> IO b
withIOBuffer IOBuffer a
buffer (\Ptr a
p Int
_ -> () -> 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
<$> Ptr a -> Int -> a -> IO ()
forall a. Storable a => Ptr a -> Int -> a -> IO ()
pokeElemOff Ptr a
p Int
i a
x)