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 :: forall (m :: * -> *) sh a. Array m sh a -> sh
shape = 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 :: forall (m :: * -> *) sh a.
(PrimMonad m, Indexed sh, Storable a) =>
Array m sh a -> Index sh -> m a
read Array m sh a
arr Index sh
ix =
   forall a. a -> Maybe a -> a
fromMaybe
      (forall a. HasCallStack => [Char] -> a
error [Char]
"Array.Comfort.Storable.Mutable.read: index out of bounds")
      (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 :: forall (m :: * -> *) sh a.
(PrimMonad m, Indexed sh, Storable a) =>
Array m sh a -> Index sh -> a -> m ()
write Array m sh a
arr Index sh
ix a
a =
   if forall sh. Indexed sh => sh -> Index sh -> Bool
Shape.inBounds (forall (m :: * -> *) sh a. Array m sh a -> sh
shape Array m sh a
arr) Index sh
ix
      then 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 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 :: forall (m :: * -> *) sh a.
(PrimMonad m, Indexed sh, Storable a) =>
Array m sh a -> Index sh -> (a -> a) -> m ()
update Array m sh a
arr Index sh
ix a -> a
f =
   if forall sh. Indexed sh => sh -> Index sh -> Bool
Shape.inBounds (forall (m :: * -> *) sh a. Array m sh a -> sh
shape Array m sh a
arr) Index sh
ix
      then 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 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 :: forall (m :: * -> *) sh a.
(PrimMonad m, C sh, Storable a) =>
Array m sh a -> m [a]
toList = 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 :: forall (m :: * -> *) sh a.
(PrimMonad m, C sh, Storable a) =>
sh -> [a] -> m (Array m sh a)
fromList sh
sh [a]
xs =
   forall (m :: * -> *) sh a.
(PrimMonad m, C sh, Storable a) =>
sh -> (Int -> Ptr a -> IO ()) -> m (Array m sh a)
MutArray.unsafeCreateWithSize sh
sh forall a b. (a -> b) -> a -> b
$ \Int
size Ptr a
ptr ->
      forall a. Storable a => Ptr a -> [a] -> IO ()
pokeArray Ptr a
ptr forall a b. (a -> b) -> a -> b
$ forall a. Int -> [a] -> [a]
take Int
size forall a b. (a -> b) -> a -> b
$
      [a]
xs forall a. [a] -> [a] -> [a]
++
      forall a. a -> [a]
repeat (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 :: forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
[a] -> m (Array m (ZeroBased Int) a)
vectorFromList = forall (m :: * -> *) a.
(PrimMonad m, Storable a) =>
[a] -> m (Array m (ZeroBased Int) a)
MutArray.vectorFromList