module Agda.Utils.MinimalArray.MutablePrim where

import Data.Coerce
import qualified Data.Primitive.PrimArray as A
import qualified Agda.Utils.MinimalArray.Prim as PA
import Control.Monad.Primitive
import Data.Primitive.Types

newtype Array s a = Array {forall s a. Array s a -> MutablePrimArray s a
unwrap :: A.MutablePrimArray s a}
type IOArray a = Array RealWorld a

{-# INLINE new #-}
new :: Prim a => PrimMonad m => Int -> m (Array (PrimState m) a)
new :: forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
Int -> m (Array (PrimState m) a)
new Int
n = MutablePrimArray (PrimState m) a -> Array (PrimState m) a
forall s a. MutablePrimArray s a -> Array s a
Array (MutablePrimArray (PrimState m) a -> Array (PrimState m) a)
-> m (MutablePrimArray (PrimState m) a)
-> m (Array (PrimState m) a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> m (MutablePrimArray (PrimState m) a)
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
A.newPrimArray Int
n

{-# INLINE size #-}
size :: PrimMonad m => Prim a => Array (PrimState m) a -> m Int
size :: forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Array (PrimState m) a -> m Int
size (Array MutablePrimArray (PrimState m) a
arr) = MutablePrimArray (PrimState m) a -> m Int
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
MutablePrimArray (PrimState m) a -> m Int
A.getSizeofMutablePrimArray MutablePrimArray (PrimState m) a
arr

{-# INLINE unsafeRead #-}
unsafeRead :: PrimMonad m => Prim a => Array (PrimState m) a -> Int -> m a
unsafeRead :: forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Array (PrimState m) a -> Int -> m a
unsafeRead (Array MutablePrimArray (PrimState m) a
arr) Int
i = MutablePrimArray (PrimState m) a -> Int -> m a
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> m a
A.readPrimArray MutablePrimArray (PrimState m) a
arr Int
i

{-# INLINE unsafeWrite #-}
unsafeWrite :: PrimMonad m => Prim a => Array (PrimState m) a -> Int -> a -> m ()
unsafeWrite :: forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Array (PrimState m) a -> Int -> a -> m ()
unsafeWrite (Array MutablePrimArray (PrimState m) a
arr) Int
i a
a = MutablePrimArray (PrimState m) a -> Int -> a -> m ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> a -> m ()
A.writePrimArray MutablePrimArray (PrimState m) a
arr Int
i a
a

{-# INLINE read #-}
read :: PrimMonad m => Prim a => Array (PrimState m) a -> Int -> m a
read :: forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Array (PrimState m) a -> Int -> m a
read Array (PrimState m) a
arr Int
i = do
  s <- Array (PrimState m) a -> m Int
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Array (PrimState m) a -> m Int
size Array (PrimState m) a
arr
  if 0 <= i && i < s then
    unsafeRead arr i
  else
    error "Array.read: out of bounds"

{-# INLINE write #-}
write :: PrimMonad m => Prim a => Array (PrimState m) a -> Int -> a -> m ()
write :: forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Array (PrimState m) a -> Int -> a -> m ()
write Array (PrimState m) a
arr Int
i a
a = do
  s <- Array (PrimState m) a -> m Int
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Array (PrimState m) a -> m Int
size Array (PrimState m) a
arr
  if 0 <= i && i < s then
    unsafeWrite arr i a
  else
    error "Array.write: out of bounds"

{-# INLINE freeze #-}
freeze :: PrimMonad m => Prim a => Array (PrimState m) a -> m (PA.Array a)
freeze :: forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Array (PrimState m) a -> m (Array a)
freeze (Array MutablePrimArray (PrimState m) a
arr) = do
  !s   <- Array (PrimState m) a -> m Int
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Array (PrimState m) a -> m Int
size (MutablePrimArray (PrimState m) a -> Array (PrimState m) a
forall s a. MutablePrimArray s a -> Array s a
Array MutablePrimArray (PrimState m) a
arr)
  !arr <- A.freezePrimArray arr 0 s
  pure (PA.Array arr)

{-# INLINE set #-}
set :: Prim a => PrimMonad m => Array (PrimState m) a -> a -> m ()
set :: forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
Array (PrimState m) a -> a -> m ()
set (Array MutablePrimArray (PrimState m) a
arr) a
a = do
  s <- Array (PrimState m) a -> m Int
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Array (PrimState m) a -> m Int
size (MutablePrimArray (PrimState m) a -> Array (PrimState m) a
forall s a. MutablePrimArray s a -> Array s a
Array MutablePrimArray (PrimState m) a
arr)
  A.setPrimArray arr 0 s a

{-# INLINE unsafeFreeze #-}
unsafeFreeze :: PrimMonad m => Array (PrimState m) a -> m (PA.Array a)
unsafeFreeze :: forall (m :: * -> *) a.
PrimMonad m =>
Array (PrimState m) a -> m (Array a)
unsafeFreeze (Array MutablePrimArray (PrimState m) a
arr) = PrimArray a -> Array a
forall a. PrimArray a -> Array a
PA.Array (PrimArray a -> Array a) -> m (PrimArray a) -> m (Array a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MutablePrimArray (PrimState m) a -> m (PrimArray a)
forall (m :: * -> *) a.
PrimMonad m =>
MutablePrimArray (PrimState m) a -> m (PrimArray a)
A.unsafeFreezePrimArray MutablePrimArray (PrimState m) a
arr