module Data.Array.Comfort.Storable.Mutable (
   Array,
   MutArray.STArray,
   MutArray.IOArray,
   shape,

   MutArray.new,
   read, MutArray.readMaybe,
   write,
   update,
   toList,
   fromList,
   vectorFromList,

   MutArray.thaw,
   MutArray.freeze,
   ) where

import qualified Data.Array.Comfort.Storable.Mutable.Unchecked as MutArray
import qualified Data.Array.Comfort.Shape as Shape
import Data.Array.Comfort.Storable.Mutable.Unchecked (Array)
import Data.Maybe (fromMaybe)

import Foreign.Marshal.Array (pokeArray)
import Foreign.Storable (Storable)

import Control.Monad.Primitive (PrimMonad)

import Prelude hiding (read)


shape :: Array m sh a -> sh
shape :: Array m sh a -> sh
shape = Array m sh a -> sh
forall (m :: * -> *) sh a. Array m sh a -> sh
MutArray.shape


read ::
   (PrimMonad m, Shape.Indexed sh, Storable a) =>
   Array m sh a -> Shape.Index sh -> m a
read :: Array m sh a -> Index sh -> m a
read Array m sh a
arr Index sh
ix =
   m a -> Maybe (m a) -> m a
forall a. a -> Maybe a -> a
fromMaybe
      ([Char] -> m a
forall a. HasCallStack => [Char] -> a
error [Char]
"Array.Comfort.Storable.Mutable.read: index out of bounds")
      (Array m sh a -> Index sh -> Maybe (m a)
forall (m :: * -> *) sh a.
(PrimMonad m, Indexed sh, Storable a) =>
Array m sh a -> Index sh -> Maybe (m a)
MutArray.readMaybe Array m sh a
arr Index sh
ix)

write ::
   (PrimMonad m, Shape.Indexed sh, Storable a) =>
   Array m sh a -> Shape.Index sh -> a -> m ()
write :: Array m sh a -> Index sh -> a -> m ()
write Array m sh a
arr Index sh
ix a
a =
   if sh -> Index sh -> Bool
forall sh. Indexed sh => sh -> Index sh -> Bool
Shape.inBounds (Array m sh a -> sh
forall (m :: * -> *) sh a. Array m sh a -> sh
shape Array m sh a
arr) Index sh
ix
      then Array m sh a -> Index sh -> a -> m ()
forall (m :: * -> *) sh a.
(PrimMonad m, Indexed sh, Storable a) =>
Array m sh a -> Index sh -> a -> m ()
MutArray.write Array m sh a
arr Index sh
ix a
a
      else [Char] -> m ()
forall a. HasCallStack => [Char] -> a
error [Char]
"Array.Comfort.Storable.Mutable.write: index out of bounds"

update ::
   (PrimMonad m, Shape.Indexed sh, Storable a) =>
   Array m sh a -> Shape.Index sh -> (a -> a) -> m ()
update :: Array m sh a -> Index sh -> (a -> a) -> m ()
update Array m sh a
arr Index sh
ix a -> a
f =
   if sh -> Index sh -> Bool
forall sh. Indexed sh => sh -> Index sh -> Bool
Shape.inBounds (Array m sh a -> sh
forall (m :: * -> *) sh a. Array m sh a -> sh
shape Array m sh a
arr) Index sh
ix
      then Array m sh a -> Index sh -> (a -> a) -> m ()
forall (m :: * -> *) sh a.
(PrimMonad m, Indexed sh, Storable a) =>
Array m sh a -> Index sh -> (a -> a) -> m ()
MutArray.update Array m sh a
arr Index sh
ix a -> a
f
      else [Char] -> m ()
forall a. HasCallStack => [Char] -> a
error [Char]
"Array.Comfort.Storable.Mutable.update: index out of bounds"

toList :: (PrimMonad m, Shape.C sh, Storable a) => Array m sh a -> m [a]
toList :: Array m sh a -> m [a]
toList = Array m sh a -> m [a]
forall (m :: * -> *) sh a.
(PrimMonad m, C sh, Storable a) =>
Array m sh a -> m [a]
MutArray.toList

fromList ::
   (PrimMonad m, Shape.C sh, Storable a) => sh -> [a] -> m (Array m sh a)
fromList :: sh -> [a] -> m (Array m sh a)
fromList sh
sh [a]
xs =
   sh -> (Int -> Ptr a -> IO ()) -> m (Array m sh a)
forall (m :: * -> *) sh a.
(PrimMonad m, C sh, Storable a) =>
sh -> (Int -> Ptr a -> IO ()) -> m (Array m sh a)
MutArray.unsafeCreateWithSize sh
sh ((Int -> Ptr a -> IO ()) -> m (Array m sh a))
-> (Int -> Ptr a -> IO ()) -> m (Array m sh a)
forall a b. (a -> b) -> a -> b
$ \Int
size Ptr a
ptr ->
      Ptr a -> [a] -> IO ()
forall a. Storable a => Ptr a -> [a] -> IO ()
pokeArray Ptr a
ptr ([a] -> IO ()) -> [a] -> IO ()
forall a b. (a -> b) -> a -> b
$ Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
take Int
size ([a] -> [a]) -> [a] -> [a]
forall a b. (a -> b) -> a -> b
$
      [a]
xs [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++
      a -> [a]
forall a. a -> [a]
repeat ([Char] -> a
forall a. HasCallStack => [Char] -> a
error [Char]
"Array.Comfort.Storable.fromList: list too short for shape")

vectorFromList ::
   (PrimMonad m, Storable a) => [a] -> m (Array m (Shape.ZeroBased Int) a)
vectorFromList :: [a] -> m (Array m (ZeroBased Int) a)
vectorFromList = [a] -> m (Array m (ZeroBased Int) a)
forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
[a] -> m (Array m (ZeroBased Int) a)
MutArray.vectorFromList