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

module Data.Primitive.Contiguous.Shim
  ( errorThunk
  , resizeArray
  , resizeUnliftedArray
  , resizeSmallUnliftedArray
  , replicateMutablePrimArray
  ) where

import Data.Primitive
import Data.Primitive.Unlifted.Array
import Data.Primitive.Unlifted.SmallArray
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
  let !srcSz :: Int
srcSz = MutableArray (PrimState m) a -> Int
forall s a. MutableArray s a -> Int
sizeofMutableArray MutableArray (PrimState m) a
src
  case Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Int
sz Int
srcSz of
    Ordering
EQ -> 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
src
    Ordering
LT -> MutableArray (PrimState m) a
-> Int -> Int -> m (MutableArray (PrimState m) a)
forall (m :: * -> *) a.
PrimMonad m =>
MutableArray (PrimState m) a
-> Int -> Int -> m (MutableArray (PrimState m) a)
cloneMutableArray MutableArray (PrimState m) a
src Int
0 Int
sz
    Ordering
GT -> do
      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
      copyMutableArray dst 0 src 0 srcSz
      pure dst
{-# INLINE resizeArray #-}

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
  let !srcSz :: Int
srcSz = MutableUnliftedArray (PrimState m) a -> Int
forall s e. MutableUnliftedArray s e -> Int
sizeofMutableUnliftedArray MutableUnliftedArray (PrimState m) a
src
  case Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare Int
sz Int
srcSz of
    Ordering
EQ -> 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
src
    Ordering
LT -> MutableUnliftedArray (PrimState m) a
-> Int -> Int -> m (MutableUnliftedArray (PrimState m) a)
forall (m :: * -> *) a.
PrimMonad m =>
MutableUnliftedArray (PrimState m) a
-> Int -> Int -> m (MutableUnliftedArray (PrimState m) a)
cloneMutableUnliftedArray MutableUnliftedArray (PrimState m) a
src Int
0 Int
sz
    Ordering
GT -> do
      dst <- Int -> m (MutableUnliftedArray (PrimState m) a)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> m (MutableUnliftedArray (PrimState m) a)
unsafeNewUnliftedArray Int
sz
      copyMutableUnliftedArray dst 0 src 0 srcSz
      pure dst
{-# INLINE resizeUnliftedArray #-}

resizeSmallUnliftedArray :: (PrimMonad m, PrimUnlifted a) => SmallMutableUnliftedArray (PrimState m) a -> Int -> m (SmallMutableUnliftedArray (PrimState m) a)
resizeSmallUnliftedArray :: forall (m :: * -> *) a.
(PrimMonad m, PrimUnlifted a) =>
SmallMutableUnliftedArray (PrimState m) a
-> Int -> m (SmallMutableUnliftedArray (PrimState m) a)
resizeSmallUnliftedArray !SmallMutableUnliftedArray (PrimState m) a
src !Int
sz = do
  srcSz <- SmallMutableUnliftedArray (PrimState m) a -> m Int
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableUnliftedArray (PrimState m) a -> m Int
getSizeofSmallMutableUnliftedArray SmallMutableUnliftedArray (PrimState m) a
src
  case compare sz srcSz of
    Ordering
EQ -> SmallMutableUnliftedArray (PrimState m) a
-> m (SmallMutableUnliftedArray (PrimState m) a)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SmallMutableUnliftedArray (PrimState m) a
src
    Ordering
LT -> SmallMutableUnliftedArray (PrimState m) a
-> Int -> Int -> m (SmallMutableUnliftedArray (PrimState m) a)
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableUnliftedArray (PrimState m) a
-> Int -> Int -> m (SmallMutableUnliftedArray (PrimState m) a)
cloneSmallMutableUnliftedArray SmallMutableUnliftedArray (PrimState m) a
src Int
0 Int
sz
    Ordering
GT -> do
      dst <- Int -> m (SmallMutableUnliftedArray (PrimState m) a)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> m (SmallMutableUnliftedArray (PrimState m) a)
unsafeNewSmallUnliftedArray Int
sz
      copySmallMutableUnliftedArray dst 0 src 0 srcSz
      pure dst
{-# INLINE resizeSmallUnliftedArray #-}


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
  marr <- Int -> m (MutablePrimArray (PrimState m) a)
forall (m :: * -> *) a.
(PrimMonad m, Prim a) =>
Int -> m (MutablePrimArray (PrimState m) a)
newPrimArray Int
len
  setPrimArray marr 0 len a
  pure marr
{-# INLINE replicateMutablePrimArray #-}