module Agda.Utils.MinimalArray.MutableLifted where

import qualified GHC.Arr as GHC
import qualified GHC.IOArray as GHC
import Data.Coerce
import qualified Data.Primitive.Array as A
import qualified Agda.Utils.MinimalArray.Lifted as LA
import Control.Monad.Primitive

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

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

{-# INLINE size #-}
size :: Array s a -> Int
size :: forall s a. Array s a -> Int
size (Array MutableArray s a
arr) = MutableArray s a -> Int
forall s a. MutableArray s a -> Int
A.sizeofMutableArray MutableArray s a
arr

{-# INLINE unsafeRead #-}
unsafeRead :: PrimMonad m => Array (PrimState m) a -> Int -> m a
unsafeRead :: forall (m :: * -> *) a.
PrimMonad m =>
Array (PrimState m) a -> Int -> m a
unsafeRead (Array MutableArray (PrimState m) a
arr) Int
i = MutableArray (PrimState m) a -> Int -> m a
forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> m a
A.readArray MutableArray (PrimState m) a
arr Int
i

{-# INLINE unsafeWrite #-}
unsafeWrite :: PrimMonad m => Array (PrimState m) a -> Int -> a -> m ()
unsafeWrite :: forall (m :: * -> *) a.
PrimMonad m =>
Array (PrimState m) a -> Int -> a -> m ()
unsafeWrite (Array MutableArray (PrimState m) a
arr) Int
i a
a = MutableArray (PrimState m) a -> Int -> a -> m ()
forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> a -> m ()
A.writeArray MutableArray (PrimState m) a
arr Int
i a
a

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

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

{-# INLINE freeze #-}
freeze :: PrimMonad m => Array (PrimState m) a -> m (LA.Array a)
freeze :: forall (m :: * -> *) a.
PrimMonad m =>
Array (PrimState m) a -> m (Array a)
freeze (Array MutableArray (PrimState m) a
arr) = Array a -> Array a
forall a. Array a -> Array a
LA.Array (Array a -> Array a) -> m (Array a) -> m (Array a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MutableArray (PrimState m) a -> Int -> Int -> m (Array a)
forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> Int -> Int -> m (Array a)
A.freezeArray MutableArray (PrimState m) a
arr Int
0 (Array (PrimState m) a -> Int
forall s a. Array s a -> Int
size (MutableArray (PrimState m) a -> Array (PrimState m) a
forall s a. MutableArray s a -> Array s a
Array MutableArray (PrimState m) a
arr))

{-# INLINE unsafeFreeze #-}
unsafeFreeze :: PrimMonad m => Array (PrimState m) a -> m (LA.Array a)
unsafeFreeze :: forall (m :: * -> *) a.
PrimMonad m =>
Array (PrimState m) a -> m (Array a)
unsafeFreeze (Array MutableArray (PrimState m) a
arr) = Array a -> Array a
forall a. Array a -> Array a
LA.Array (Array a -> Array a) -> m (Array a) -> m (Array a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MutableArray (PrimState m) a -> m (Array a)
forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a -> m (Array a)
A.unsafeFreezeArray MutableArray (PrimState m) a
arr

fromGHCArray :: GHC.IOArray i e -> IOArray e
fromGHCArray :: forall i e. IOArray i e -> IOArray e
fromGHCArray (GHC.IOArray (GHC.STArray i
_ i
_ Int
_ MutableArray# RealWorld e
arr)) =
  MutableArray RealWorld e -> Array RealWorld e
forall s a. MutableArray s a -> Array s a
Array (MutableArray# RealWorld e -> MutableArray RealWorld e
forall s a. MutableArray# s a -> MutableArray s a
A.MutableArray MutableArray# RealWorld e
arr)