module PrimitiveExtras.UnliftedArray
where

import PrimitiveExtras.Prelude


{-# INLINE at #-}
at :: PrimUnlifted element => UnliftedArray element -> Int -> forall result. result -> (element -> result) -> result
at :: UnliftedArray element
-> Int -> forall result. result -> (element -> result) -> result
at UnliftedArray element
ua Int
index result
none element -> result
some =
  if UnliftedArray element -> Int
forall e. UnliftedArray e -> Int
sizeofUnliftedArray UnliftedArray element
ua Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
index
    then result
none
    else element -> result
some (UnliftedArray element -> Int -> element
forall a. PrimUnlifted a => UnliftedArray a -> Int -> a
indexUnliftedArray UnliftedArray element
ua Int
index)

{-# INLINABLE replicateIO #-}
replicateIO :: PrimUnlifted a => Int -> IO a -> IO (UnliftedArray a)
replicateIO :: Int -> IO a -> IO (UnliftedArray a)
replicateIO Int
size IO a
elementIO =
  do
    MutableUnliftedArray RealWorld a
array <- Int -> IO (MutableUnliftedArray (PrimState IO) a)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> m (MutableUnliftedArray (PrimState m) a)
unsafeNewUnliftedArray Int
size
    let
      loop :: Int -> IO (UnliftedArray a)
loop Int
index =
        if Int
index Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
size
          then do
            a
element <- IO a
elementIO
            MutableUnliftedArray (PrimState IO) a -> Int -> a -> IO ()
forall (m :: * -> *) a.
(PrimMonad m, PrimUnlifted a) =>
MutableUnliftedArray (PrimState m) a -> Int -> a -> m ()
writeUnliftedArray MutableUnliftedArray RealWorld a
MutableUnliftedArray (PrimState IO) a
array Int
index a
element
            Int -> IO (UnliftedArray a)
loop (Int -> Int
forall a. Enum a => a -> a
succ Int
index)
          else MutableUnliftedArray (PrimState IO) a -> IO (UnliftedArray a)
forall (m :: * -> *) a.
PrimMonad m =>
MutableUnliftedArray (PrimState m) a -> m (UnliftedArray a)
unsafeFreezeUnliftedArray MutableUnliftedArray RealWorld a
MutableUnliftedArray (PrimState IO) a
array
      in Int -> IO (UnliftedArray a)
loop Int
0

{-# INLINABLE generate #-}
generate :: PrimUnlifted a => Int -> (Int -> IO a) -> IO (UnliftedArray a)
generate :: Int -> (Int -> IO a) -> IO (UnliftedArray a)
generate Int
size Int -> IO a
elementIO =
  do
    MutableUnliftedArray RealWorld a
array <- Int -> IO (MutableUnliftedArray (PrimState IO) a)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> m (MutableUnliftedArray (PrimState m) a)
unsafeNewUnliftedArray Int
size
    let
      loop :: Int -> IO (UnliftedArray a)
loop Int
index =
        if Int
index Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
size
          then do
            a
element <- Int -> IO a
elementIO Int
index
            MutableUnliftedArray (PrimState IO) a -> Int -> a -> IO ()
forall (m :: * -> *) a.
(PrimMonad m, PrimUnlifted a) =>
MutableUnliftedArray (PrimState m) a -> Int -> a -> m ()
writeUnliftedArray MutableUnliftedArray RealWorld a
MutableUnliftedArray (PrimState IO) a
array Int
index a
element
            Int -> IO (UnliftedArray a)
loop (Int -> Int
forall a. Enum a => a -> a
succ Int
index)
          else MutableUnliftedArray (PrimState IO) a -> IO (UnliftedArray a)
forall (m :: * -> *) a.
PrimMonad m =>
MutableUnliftedArray (PrimState m) a -> m (UnliftedArray a)
unsafeFreezeUnliftedArray MutableUnliftedArray RealWorld a
MutableUnliftedArray (PrimState IO) a
array
      in Int -> IO (UnliftedArray a)
loop Int
0

traverse_ :: (Monad m, PrimUnlifted a) => (a -> m ()) -> UnliftedArray a -> m ()
traverse_ :: (a -> m ()) -> UnliftedArray a -> m ()
traverse_ a -> m ()
action UnliftedArray a
array =
  let
    size :: Int
size = UnliftedArray a -> Int
forall e. UnliftedArray e -> Int
sizeofUnliftedArray UnliftedArray a
array
    iterate :: Int -> m ()
iterate Int
index = if Int
index Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
size
      then do
        a -> m ()
action (UnliftedArray a -> Int -> a
forall a. PrimUnlifted a => UnliftedArray a -> Int -> a
indexUnliftedArray UnliftedArray a
array Int
index)
        Int -> m ()
iterate (Int -> Int
forall a. Enum a => a -> a
succ Int
index)
      else () -> m ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()
    in Int -> m ()
iterate Int
0