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)

{-# INLINEABLE 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

{-# INLINEABLE 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