{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE MagicHash #-}

module Data.Primitive.Contiguous.Shim
  ( errorThunk
  , resizeArray
  , resizeSmallArray
  , replicateSmallMutableArray
  , resizeUnliftedArray
  , replicateMutablePrimArray
  , clonePrimArrayShim
  , cloneMutablePrimArrayShim
  , freezePrimArrayShim
  ) where

import Control.Monad (when)
import Control.Monad.ST.Run (runPrimArrayST)
import Data.Primitive hiding (fromList, fromListN)
import Data.Primitive.Unlifted.Array
import Prelude hiding (all, any, elem, filter, foldMap, foldl, foldr, map, mapM, mapM_, maximum, minimum, null, read, replicate, reverse, scanl, sequence, sequence_, traverse, zip, zipWith, (<$))

import Control.Monad.Primitive (PrimMonad (..), PrimState)
import Data.Primitive.Unlifted.Class (PrimUnlifted)

errorThunk :: a
errorThunk :: forall a. a
errorThunk = [Char] -> a
forall a. HasCallStack => [Char] -> a
error [Char]
"Contiguous typeclass: unitialized element"
{-# NOINLINE errorThunk #-}

resizeArray :: (PrimMonad m) => MutableArray (PrimState m) a -> Int -> m (MutableArray (PrimState m) a)
resizeArray :: forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a
-> Int -> m (MutableArray (PrimState m) a)
resizeArray !MutableArray (PrimState m) a
src !Int
sz = do
  MutableArray (PrimState m) a
dst <- Int -> a -> m (MutableArray (PrimState m) a)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (MutableArray (PrimState m) a)
newArray Int
sz a
forall a. a
errorThunk
  MutableArray (PrimState m) a
-> Int -> MutableArray (PrimState m) a -> Int -> Int -> m ()
forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a
-> Int -> MutableArray (PrimState m) a -> Int -> Int -> m ()
copyMutableArray MutableArray (PrimState m) a
dst Int
0 MutableArray (PrimState m) a
src Int
0 (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
sz (MutableArray (PrimState m) a -> Int
forall s a. MutableArray s a -> Int
sizeofMutableArray MutableArray (PrimState m) a
src))
  MutableArray (PrimState m) a -> m (MutableArray (PrimState m) a)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure MutableArray (PrimState m) a
dst
{-# INLINE resizeArray #-}

resizeSmallArray :: (PrimMonad m) => SmallMutableArray (PrimState m) a -> Int -> m (SmallMutableArray (PrimState m) a)
resizeSmallArray :: forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a
-> Int -> m (SmallMutableArray (PrimState m) a)
resizeSmallArray !SmallMutableArray (PrimState m) a
src !Int
sz = do
  SmallMutableArray (PrimState m) a
dst <- Int -> a -> m (SmallMutableArray (PrimState m) a)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
sz a
forall a. a
errorThunk
  SmallMutableArray (PrimState m) a
-> Int -> SmallMutableArray (PrimState m) a -> Int -> Int -> m ()
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a
-> Int -> SmallMutableArray (PrimState m) a -> Int -> Int -> m ()
copySmallMutableArray SmallMutableArray (PrimState m) a
dst Int
0 SmallMutableArray (PrimState m) a
src Int
0 (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
sz (SmallMutableArray (PrimState m) a -> Int
forall s a. SmallMutableArray s a -> Int
sizeofSmallMutableArray SmallMutableArray (PrimState m) a
src))
  SmallMutableArray (PrimState m) a
-> m (SmallMutableArray (PrimState m) a)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SmallMutableArray (PrimState m) a
dst
{-# INLINE resizeSmallArray #-}

replicateSmallMutableArray ::
  (PrimMonad m) =>
  Int ->
  a ->
  m (SmallMutableArray (PrimState m) a)
replicateSmallMutableArray :: forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
replicateSmallMutableArray Int
len a
a = do
  SmallMutableArray (PrimState m) a
marr <- Int -> a -> m (SmallMutableArray (PrimState m) a)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
len a
forall a. a
errorThunk
  let go :: Int -> m ()
go !Int
ix = Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
ix Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
len) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ do
        SmallMutableArray (PrimState m) a -> Int -> a -> m ()
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
writeSmallArray SmallMutableArray (PrimState m) a
marr Int
ix a
a
        Int -> m ()
go (Int
ix Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
  Int -> m ()
go Int
0
  SmallMutableArray (PrimState m) a
-> m (SmallMutableArray (PrimState m) a)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SmallMutableArray (PrimState m) a
marr
{-# INLINE replicateSmallMutableArray #-}

resizeUnliftedArray :: (PrimMonad m, PrimUnlifted a) => MutableUnliftedArray (PrimState m) a -> Int -> m (MutableUnliftedArray (PrimState m) a)
resizeUnliftedArray :: forall (m :: * -> *) a.
(PrimMonad m, PrimUnlifted a) =>
MutableUnliftedArray (PrimState m) a
-> Int -> m (MutableUnliftedArray (PrimState m) a)
resizeUnliftedArray !MutableUnliftedArray (PrimState m) a
src !Int
sz = do
  MutableUnliftedArray (PrimState m) a
dst <- Int -> m (MutableUnliftedArray (PrimState m) a)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> m (MutableUnliftedArray (PrimState m) a)
unsafeNewUnliftedArray Int
sz
  MutableUnliftedArray (PrimState m) a
-> Int
-> MutableUnliftedArray (PrimState m) a
-> Int
-> Int
-> m ()
forall (m :: * -> *) a.
PrimMonad m =>
MutableUnliftedArray (PrimState m) a
-> Int
-> MutableUnliftedArray (PrimState m) a
-> Int
-> Int
-> m ()
copyMutableUnliftedArray MutableUnliftedArray (PrimState m) a
dst Int
0 MutableUnliftedArray (PrimState m) a
src Int
0 (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
sz (MutableUnliftedArray (PrimState m) a -> Int
forall s e. MutableUnliftedArray s e -> Int
sizeofMutableUnliftedArray MutableUnliftedArray (PrimState m) a
src))
  MutableUnliftedArray (PrimState m) a
-> m (MutableUnliftedArray (PrimState m) a)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure MutableUnliftedArray (PrimState m) a
dst
{-# INLINE resizeUnliftedArray #-}

replicateMutablePrimArray ::
  (PrimMonad m, Prim a) =>
  -- | length
  Int ->
  -- | element
  a ->
  m (MutablePrimArray (PrimState m) a)
replicateMutablePrimArray :: forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Int -> a -> m (MutablePrimArray (PrimState m) a)
replicateMutablePrimArray Int
len a
a = do
  MutablePrimArray (PrimState m) a
marr <- Int -> m (MutablePrimArray (PrimState m) a)
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newPrimArray Int
len
  MutablePrimArray (PrimState m) a -> Int -> Int -> a -> m ()
forall a (m :: * -> *).
(Prim a, PrimMonad m) =>
MutablePrimArray (PrimState m) a -> Int -> Int -> a -> m ()
setPrimArray MutablePrimArray (PrimState m) a
marr Int
0 Int
len a
a
  MutablePrimArray (PrimState m) a
-> m (MutablePrimArray (PrimState m) a)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure MutablePrimArray (PrimState m) a
marr
{-# INLINE replicateMutablePrimArray #-}

clonePrimArrayShim :: (Prim a) => PrimArray a -> Int -> Int -> PrimArray a
clonePrimArrayShim :: forall a. Prim a => PrimArray a -> Int -> Int -> PrimArray a
clonePrimArrayShim !PrimArray a
arr !Int
off !Int
len = (forall s. ST s (PrimArray a)) -> PrimArray a
forall a. (forall s. ST s (PrimArray a)) -> PrimArray a
runPrimArrayST ((forall s. ST s (PrimArray a)) -> PrimArray a)
-> (forall s. ST s (PrimArray a)) -> PrimArray a
forall a b. (a -> b) -> a -> b
$ do
  MutablePrimArray s a
marr <- Int -> ST s (MutablePrimArray (PrimState (ST s)) a)
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newPrimArray Int
len
  MutablePrimArray (PrimState (ST s)) a
-> Int -> PrimArray a -> Int -> Int -> ST s ()
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
MutablePrimArray (PrimState m) a
-> Int -> PrimArray a -> Int -> Int -> m ()
copyPrimArray MutablePrimArray s a
MutablePrimArray (PrimState (ST s)) a
marr Int
0 PrimArray a
arr Int
off Int
len
  MutablePrimArray (PrimState (ST s)) a -> ST s (PrimArray a)
forall (m :: * -> *) a.
PrimMonad m =>
MutablePrimArray (PrimState m) a -> m (PrimArray a)
unsafeFreezePrimArray MutablePrimArray s a
MutablePrimArray (PrimState (ST s)) a
marr
{-# INLINE clonePrimArrayShim #-}

cloneMutablePrimArrayShim :: (PrimMonad m, Prim a) => MutablePrimArray (PrimState m) a -> Int -> Int -> m (MutablePrimArray (PrimState m) a)
cloneMutablePrimArrayShim :: forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
MutablePrimArray (PrimState m) a
-> Int -> Int -> m (MutablePrimArray (PrimState m) a)
cloneMutablePrimArrayShim !MutablePrimArray (PrimState m) a
arr !Int
off !Int
len = do
  MutablePrimArray (PrimState m) a
marr <- Int -> m (MutablePrimArray (PrimState m) a)
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newPrimArray Int
len
  MutablePrimArray (PrimState m) a
-> Int -> MutablePrimArray (PrimState m) a -> Int -> Int -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
MutablePrimArray (PrimState m) a
-> Int -> MutablePrimArray (PrimState m) a -> Int -> Int -> m ()
copyMutablePrimArray MutablePrimArray (PrimState m) a
marr Int
0 MutablePrimArray (PrimState m) a
arr Int
off Int
len
  MutablePrimArray (PrimState m) a
-> m (MutablePrimArray (PrimState m) a)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure MutablePrimArray (PrimState m) a
marr
{-# INLINE cloneMutablePrimArrayShim #-}

freezePrimArrayShim :: (PrimMonad m, Prim a) => MutablePrimArray (PrimState m) a -> Int -> Int -> m (PrimArray a)
freezePrimArrayShim :: forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
MutablePrimArray (PrimState m) a -> Int -> Int -> m (PrimArray a)
freezePrimArrayShim !MutablePrimArray (PrimState m) a
src !Int
off !Int
len = do
  MutablePrimArray (PrimState m) a
dst <- Int -> m (MutablePrimArray (PrimState m) a)
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newPrimArray Int
len
  MutablePrimArray (PrimState m) a
-> Int -> MutablePrimArray (PrimState m) a -> Int -> Int -> m ()
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
MutablePrimArray (PrimState m) a
-> Int -> MutablePrimArray (PrimState m) a -> Int -> Int -> m ()
copyMutablePrimArray MutablePrimArray (PrimState m) a
dst Int
0 MutablePrimArray (PrimState m) a
src Int
off Int
len
  MutablePrimArray (PrimState m) a -> m (PrimArray a)
forall (m :: * -> *) a.
PrimMonad m =>
MutablePrimArray (PrimState m) a -> m (PrimArray a)
unsafeFreezePrimArray MutablePrimArray (PrimState m) a
dst
{-# INLINE freezePrimArrayShim #-}