module PrimitiveExtras.UnliftedArray
where

import PrimitiveExtras.Prelude


{-# INLINE at #-}
at :: PrimUnlifted element => UnliftedArray element -> Int -> forall result. result -> (element -> result) -> result
at :: forall element.
PrimUnlifted element =>
UnliftedArray element
-> Int -> forall result. result -> (element -> result) -> result
at UnliftedArray element
ua Int
index result
none element -> result
some =
  if forall e. UnliftedArray e -> Int
sizeofUnliftedArray UnliftedArray element
ua forall a. Ord a => a -> a -> Bool
<= Int
index
    then result
none
    else element -> result
some (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 :: forall a. PrimUnlifted a => Int -> IO a -> IO (UnliftedArray a)
replicateIO Int
size IO a
elementIO =
  do
    MutableUnliftedArray RealWorld a
array <- 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 forall a. Ord a => a -> a -> Bool
< Int
size
          then do
            a
element <- IO a
elementIO
            forall (m :: * -> *) a.
(PrimMonad m, PrimUnlifted a) =>
MutableUnliftedArray (PrimState m) a -> Int -> a -> m ()
writeUnliftedArray MutableUnliftedArray RealWorld a
array Int
index a
element
            Int -> IO (UnliftedArray a)
loop (forall a. Enum a => a -> a
succ Int
index)
          else forall (m :: * -> *) a.
PrimMonad m =>
MutableUnliftedArray (PrimState m) a -> m (UnliftedArray a)
unsafeFreezeUnliftedArray MutableUnliftedArray RealWorld a
array
      in Int -> IO (UnliftedArray a)
loop Int
0

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

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