-- |
-- Module      : Streamly.Internal.Data.SmallArray
-- Copyright   : (c) 2019 Composewell Technologies
--
-- License     : BSD-3-Clause
-- Maintainer  : streamly@composewell.com
-- Stability   : experimental
-- Portability : GHC

{-# OPTIONS_GHC -fno-warn-orphans #-}
{-# LANGUAGE UnboxedTuples #-}

#include "inline.hs"

module Streamly.Internal.Data.SmallArray
  (
    -- XXX should it be just Array instead? We should be able to replace one
    -- array type with another easily.
    SmallArray(..)

  , foldl'
  , foldr

  , length

  , writeN

  , toStreamD
  , toStreamDRev

  , toStream
  , toStreamRev
  , read

  , fromListN
  , fromStreamDN
  , fromStreamN

  , streamFold
  , fold
  )
where

import Prelude hiding (foldr, length, read)
import Control.DeepSeq (NFData(..))
import Control.Monad (when)
import Control.Monad.IO.Class (MonadIO, liftIO)
import GHC.IO (unsafePerformIO)
import Data.Functor.Identity (runIdentity)

import Streamly.Internal.Data.SmallArray.Type

import Streamly.Internal.Data.Tuple.Strict (Tuple'(..))
import Streamly.Internal.Data.Unfold.Type (Unfold(..))
import Streamly.Internal.Data.Fold.Type (Fold(..))
import Streamly.Internal.Data.Stream.Serial (SerialT(..))

import qualified Streamly.Internal.Data.Stream.StreamD as D
import qualified Streamly.Internal.Data.Fold.Type as FL

{-# NOINLINE bottomElement #-}
bottomElement :: a
bottomElement :: a
bottomElement = a
forall a. HasCallStack => a
undefined

{-# INLINE length #-}
length :: SmallArray a -> Int
length :: SmallArray a -> Int
length = SmallArray a -> Int
forall a. SmallArray a -> Int
sizeofSmallArray

{-# INLINE_NORMAL toStreamD #-}
toStreamD :: Monad m => SmallArray a -> D.Stream m a
toStreamD :: SmallArray a -> Stream m a
toStreamD SmallArray a
arr = (State Stream m a -> Int -> m (Step Int a)) -> Int -> Stream m a
forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
D.Stream State Stream m a -> Int -> m (Step Int a)
forall (m :: * -> *) p. Monad m => p -> Int -> m (Step Int a)
step Int
0
  where
    {-# INLINE_LATE step #-}
    step :: p -> Int -> m (Step Int a)
step p
_ Int
i
        | Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== SmallArray a -> Int
forall a. SmallArray a -> Int
length SmallArray a
arr = Step Int a -> m (Step Int a)
forall (m :: * -> *) a. Monad m => a -> m a
return Step Int a
forall s a. Step s a
D.Stop
        | Bool
otherwise =
            Step Int a -> m (Step Int a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step Int a -> m (Step Int a)) -> Step Int a -> m (Step Int a)
forall a b. (a -> b) -> a -> b
$
            case SmallArray a -> Int -> (# a #)
forall a. SmallArray a -> Int -> (# a #)
indexSmallArray## SmallArray a
arr Int
i of
                (# a
x #) -> a -> Int -> Step Int a
forall s a. a -> s -> Step s a
D.Yield a
x (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)

{-# INLINE_NORMAL toStreamDRev #-}
toStreamDRev :: Monad m => SmallArray a -> D.Stream m a
toStreamDRev :: SmallArray a -> Stream m a
toStreamDRev SmallArray a
arr = (State Stream m a -> Int -> m (Step Int a)) -> Int -> Stream m a
forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
D.Stream State Stream m a -> Int -> m (Step Int a)
forall (m :: * -> *) p. Monad m => p -> Int -> m (Step Int a)
step (SmallArray a -> Int
forall a. SmallArray a -> Int
length SmallArray a
arr Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
  where
    {-# INLINE_LATE step #-}
    step :: p -> Int -> m (Step Int a)
step p
_ Int
i
        | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 = Step Int a -> m (Step Int a)
forall (m :: * -> *) a. Monad m => a -> m a
return Step Int a
forall s a. Step s a
D.Stop
        | Bool
otherwise =
            Step Int a -> m (Step Int a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step Int a -> m (Step Int a)) -> Step Int a -> m (Step Int a)
forall a b. (a -> b) -> a -> b
$
            case SmallArray a -> Int -> (# a #)
forall a. SmallArray a -> Int -> (# a #)
indexSmallArray## SmallArray a
arr Int
i of
                (# a
x #) -> a -> Int -> Step Int a
forall s a. a -> s -> Step s a
D.Yield a
x (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)

{-# INLINE_NORMAL foldl' #-}
foldl' :: (b -> a -> b) -> b -> SmallArray a -> b
foldl' :: (b -> a -> b) -> b -> SmallArray a -> b
foldl' b -> a -> b
f b
z SmallArray a
arr = Identity b -> b
forall a. Identity a -> a
runIdentity (Identity b -> b) -> Identity b -> b
forall a b. (a -> b) -> a -> b
$ (b -> a -> b) -> b -> Stream Identity a -> Identity b
forall (m :: * -> *) b a.
Monad m =>
(b -> a -> b) -> b -> Stream m a -> m b
D.foldl' b -> a -> b
f b
z (Stream Identity a -> Identity b)
-> Stream Identity a -> Identity b
forall a b. (a -> b) -> a -> b
$ SmallArray a -> Stream Identity a
forall (m :: * -> *) a. Monad m => SmallArray a -> Stream m a
toStreamD SmallArray a
arr

{-# INLINE_NORMAL foldr #-}
foldr :: (a -> b -> b) -> b -> SmallArray a -> b
foldr :: (a -> b -> b) -> b -> SmallArray a -> b
foldr a -> b -> b
f b
z SmallArray a
arr = Identity b -> b
forall a. Identity a -> a
runIdentity (Identity b -> b) -> Identity b -> b
forall a b. (a -> b) -> a -> b
$ (a -> b -> b) -> b -> Stream Identity a -> Identity b
forall (m :: * -> *) a b.
Monad m =>
(a -> b -> b) -> b -> Stream m a -> m b
D.foldr a -> b -> b
f b
z (Stream Identity a -> Identity b)
-> Stream Identity a -> Identity b
forall a b. (a -> b) -> a -> b
$ SmallArray a -> Stream Identity a
forall (m :: * -> *) a. Monad m => SmallArray a -> Stream m a
toStreamD SmallArray a
arr

-- | @writeN n@ folds a maximum of @n@ elements from the input stream to an
-- 'SmallArray'.
--
-- Since we are folding to a 'SmallArray' @n@ should be <= 128, for larger number
-- of elements use an 'Array' from either "Streamly.Data.Array" or "Streamly.Data.Array.Foreign".
{-# INLINE_NORMAL writeN #-}
writeN :: MonadIO m => Int -> Fold m a (SmallArray a)
writeN :: Int -> Fold m a (SmallArray a)
writeN Int
len = (Tuple' (SmallMutableArray RealWorld a) Int
 -> a
 -> m (Step
         (Tuple' (SmallMutableArray RealWorld a) Int) (SmallArray a)))
-> m (Step
        (Tuple' (SmallMutableArray RealWorld a) Int) (SmallArray a))
-> (Tuple' (SmallMutableArray RealWorld a) Int -> m (SmallArray a))
-> Fold m a (SmallArray a)
forall (m :: * -> *) a b s.
(s -> a -> m (Step s b))
-> m (Step s b) -> (s -> m b) -> Fold m a b
Fold Tuple' (SmallMutableArray RealWorld a) Int
-> a
-> m (Step
        (Tuple' (SmallMutableArray RealWorld a) Int) (SmallArray a))
forall (m :: * -> *) a.
MonadIO m =>
Tuple' (SmallMutableArray RealWorld a) Int
-> a
-> m (Step
        (Tuple' (SmallMutableArray RealWorld a) Int) (SmallArray a))
step m (Step
     (Tuple' (SmallMutableArray RealWorld a) Int) (SmallArray a))
forall a.
m (Step
     (Tuple' (SmallMutableArray RealWorld a) Int) (SmallArray a))
initial Tuple' (SmallMutableArray RealWorld a) Int -> m (SmallArray a)
forall (m :: * -> *) a.
MonadIO m =>
Tuple' (SmallMutableArray RealWorld a) Int -> m (SmallArray a)
extract

    where

    {-# INLINE next #-}
    next :: SmallMutableArray RealWorld a
-> Int
-> m (Step
        (Tuple' (SmallMutableArray RealWorld a) Int) (SmallArray a))
next SmallMutableArray RealWorld a
marr Int
i = do
        let i1 :: Int
i1 = Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1
            st :: Tuple' (SmallMutableArray RealWorld a) Int
st = SmallMutableArray RealWorld a
-> Int -> Tuple' (SmallMutableArray RealWorld a) Int
forall a b. a -> b -> Tuple' a b
Tuple' SmallMutableArray RealWorld a
marr Int
i1
        if Int
len Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
i1
        then Step (Tuple' (SmallMutableArray RealWorld a) Int) (SmallArray a)
-> m (Step
        (Tuple' (SmallMutableArray RealWorld a) Int) (SmallArray a))
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (Tuple' (SmallMutableArray RealWorld a) Int) (SmallArray a)
 -> m (Step
         (Tuple' (SmallMutableArray RealWorld a) Int) (SmallArray a)))
-> Step (Tuple' (SmallMutableArray RealWorld a) Int) (SmallArray a)
-> m (Step
        (Tuple' (SmallMutableArray RealWorld a) Int) (SmallArray a))
forall a b. (a -> b) -> a -> b
$ Tuple' (SmallMutableArray RealWorld a) Int
-> Step (Tuple' (SmallMutableArray RealWorld a) Int) (SmallArray a)
forall s b. s -> Step s b
FL.Partial Tuple' (SmallMutableArray RealWorld a) Int
st
        else (SmallArray a
 -> Step
      (Tuple' (SmallMutableArray RealWorld a) Int) (SmallArray a))
-> m (SmallArray a)
-> m (Step
        (Tuple' (SmallMutableArray RealWorld a) Int) (SmallArray a))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap SmallArray a
-> Step (Tuple' (SmallMutableArray RealWorld a) Int) (SmallArray a)
forall s b. b -> Step s b
FL.Done (m (SmallArray a)
 -> m (Step
         (Tuple' (SmallMutableArray RealWorld a) Int) (SmallArray a)))
-> m (SmallArray a)
-> m (Step
        (Tuple' (SmallMutableArray RealWorld a) Int) (SmallArray a))
forall a b. (a -> b) -> a -> b
$ Tuple' (SmallMutableArray RealWorld a) Int -> m (SmallArray a)
forall (m :: * -> *) a.
MonadIO m =>
Tuple' (SmallMutableArray RealWorld a) Int -> m (SmallArray a)
extract Tuple' (SmallMutableArray RealWorld a) Int
st

    initial :: m (Step
     (Tuple' (SmallMutableArray RealWorld a) Int) (SmallArray a))
initial = do
        SmallMutableArray RealWorld a
marr <- IO (SmallMutableArray RealWorld a)
-> m (SmallMutableArray RealWorld a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (SmallMutableArray RealWorld a)
 -> m (SmallMutableArray RealWorld a))
-> IO (SmallMutableArray RealWorld a)
-> m (SmallMutableArray RealWorld a)
forall a b. (a -> b) -> a -> b
$ Int -> a -> IO (SmallMutableArray (PrimState IO) a)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray Int
len a
forall a. a
bottomElement
        SmallMutableArray RealWorld a
-> Int
-> m (Step
        (Tuple' (SmallMutableArray RealWorld a) Int) (SmallArray a))
forall (m :: * -> *) a.
MonadIO m =>
SmallMutableArray RealWorld a
-> Int
-> m (Step
        (Tuple' (SmallMutableArray RealWorld a) Int) (SmallArray a))
next SmallMutableArray RealWorld a
marr (-Int
1)

    step :: Tuple' (SmallMutableArray RealWorld a) Int
-> a
-> m (Step
        (Tuple' (SmallMutableArray RealWorld a) Int) (SmallArray a))
step (Tuple' SmallMutableArray RealWorld a
marr Int
i) a
x = do
        IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ SmallMutableArray (PrimState IO) a -> Int -> a -> IO ()
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
writeSmallArray SmallMutableArray RealWorld a
SmallMutableArray (PrimState IO) a
marr Int
i a
x
        SmallMutableArray RealWorld a
-> Int
-> m (Step
        (Tuple' (SmallMutableArray RealWorld a) Int) (SmallArray a))
forall (m :: * -> *) a.
MonadIO m =>
SmallMutableArray RealWorld a
-> Int
-> m (Step
        (Tuple' (SmallMutableArray RealWorld a) Int) (SmallArray a))
next SmallMutableArray RealWorld a
marr Int
i

    extract :: Tuple' (SmallMutableArray RealWorld a) Int -> m (SmallArray a)
extract (Tuple' SmallMutableArray RealWorld a
marr Int
l) = IO (SmallArray a) -> m (SmallArray a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (SmallArray a) -> m (SmallArray a))
-> IO (SmallArray a) -> m (SmallArray a)
forall a b. (a -> b) -> a -> b
$ SmallMutableArray (PrimState IO) a
-> Int -> Int -> IO (SmallArray a)
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> Int -> m (SmallArray a)
freezeSmallArray SmallMutableArray RealWorld a
SmallMutableArray (PrimState IO) a
marr Int
0 Int
l


{-# INLINE_NORMAL fromStreamDN #-}
fromStreamDN :: MonadIO m => Int -> D.Stream m a -> m (SmallArray a)
fromStreamDN :: Int -> Stream m a -> m (SmallArray a)
fromStreamDN Int
limit Stream m a
str = do
    SmallMutableArray RealWorld a
marr <- IO (SmallMutableArray RealWorld a)
-> m (SmallMutableArray RealWorld a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (SmallMutableArray RealWorld a)
 -> m (SmallMutableArray RealWorld a))
-> IO (SmallMutableArray RealWorld a)
-> m (SmallMutableArray RealWorld a)
forall a b. (a -> b) -> a -> b
$ Int -> a -> IO (SmallMutableArray (PrimState IO) a)
forall (m :: * -> *) a.
PrimMonad m =>
Int -> a -> m (SmallMutableArray (PrimState m) a)
newSmallArray (Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
limit Int
0) a
forall a. a
bottomElement
    Int
i <-
        (Int -> a -> m Int) -> m Int -> Stream m a -> m Int
forall (m :: * -> *) b a.
Monad m =>
(b -> a -> m b) -> m b -> Stream m a -> m b
D.foldlM'
            (\Int
i a
x -> Int
i Int -> m Int -> m Int
`seq` IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (SmallMutableArray (PrimState IO) a -> Int -> a -> IO ()
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> a -> m ()
writeSmallArray SmallMutableArray RealWorld a
SmallMutableArray (PrimState IO) a
marr Int
i a
x) m () -> m Int -> m Int
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Int -> m Int
forall (m :: * -> *) a. Monad m => a -> m a
return (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1))
            (Int -> m Int
forall (m :: * -> *) a. Monad m => a -> m a
return Int
0) (Stream m a -> m Int) -> Stream m a -> m Int
forall a b. (a -> b) -> a -> b
$
        Int -> Stream m a -> Stream m a
forall (m :: * -> *) a.
Applicative m =>
Int -> Stream m a -> Stream m a
D.take Int
limit Stream m a
str
    IO (SmallArray a) -> m (SmallArray a)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (SmallArray a) -> m (SmallArray a))
-> IO (SmallArray a) -> m (SmallArray a)
forall a b. (a -> b) -> a -> b
$ SmallMutableArray (PrimState IO) a
-> Int -> Int -> IO (SmallArray a)
forall (m :: * -> *) a.
PrimMonad m =>
SmallMutableArray (PrimState m) a -> Int -> Int -> m (SmallArray a)
freezeSmallArray SmallMutableArray RealWorld a
SmallMutableArray (PrimState IO) a
marr Int
0 Int
i

-- | Create a 'SmallArray' from the first @n@ elements of a list. The
-- array may hold less than @n@ elements if the length of the list <=
-- @n@.
--
-- It is recommended to use a value of @n@ <= 128. For larger sized
-- arrays, use an 'Array' from "Streamly.Data.Array" or
-- "Streamly.Data.Array.Foreign"
{-# INLINABLE fromListN #-}
fromListN :: Int -> [a] -> SmallArray a
fromListN :: Int -> [a] -> SmallArray a
fromListN Int
n [a]
xs = IO (SmallArray a) -> SmallArray a
forall a. IO a -> a
unsafePerformIO (IO (SmallArray a) -> SmallArray a)
-> IO (SmallArray a) -> SmallArray a
forall a b. (a -> b) -> a -> b
$ Int -> Stream IO a -> IO (SmallArray a)
forall (m :: * -> *) a.
MonadIO m =>
Int -> Stream m a -> m (SmallArray a)
fromStreamDN Int
n (Stream IO a -> IO (SmallArray a))
-> Stream IO a -> IO (SmallArray a)
forall a b. (a -> b) -> a -> b
$ [a] -> Stream IO a
forall (m :: * -> *) a. Applicative m => [a] -> Stream m a
D.fromList [a]
xs

instance NFData a => NFData (SmallArray a) where
    {-# INLINE rnf #-}
    rnf :: SmallArray a -> ()
rnf = (() -> a -> ()) -> () -> SmallArray a -> ()
forall b a. (b -> a -> b) -> b -> SmallArray a -> b
foldl' (\()
_ a
x -> a -> ()
forall a. NFData a => a -> ()
rnf a
x) ()

-- | Create a 'SmallArray' from the first @n@ elements of a stream. The
-- array is allocated to size @n@, if the stream terminates before @n@
-- elements then the array may hold less than @n@ elements.
--
-- For optimal performance use this with @n@ <= 128.
{-# INLINE fromStreamN #-}
fromStreamN :: MonadIO m => Int -> SerialT m a -> m (SmallArray a)
fromStreamN :: Int -> SerialT m a -> m (SmallArray a)
fromStreamN Int
n (SerialT Stream m a
m) = do
    Bool -> m () -> m ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0) (m () -> m ()) -> m () -> m ()
forall a b. (a -> b) -> a -> b
$ [Char] -> m ()
forall a. HasCallStack => [Char] -> a
error [Char]
"fromStreamN: negative write count specified"
    Int -> Stream m a -> m (SmallArray a)
forall (m :: * -> *) a.
MonadIO m =>
Int -> Stream m a -> m (SmallArray a)
fromStreamDN Int
n (Stream m a -> m (SmallArray a)) -> Stream m a -> m (SmallArray a)
forall a b. (a -> b) -> a -> b
$ Stream m a -> Stream m a
forall (m :: * -> *) a. Applicative m => Stream m a -> Stream m a
D.fromStreamK Stream m a
m

{-# INLINE_EARLY toStream #-}
toStream :: Monad m => SmallArray a -> SerialT m a
toStream :: SmallArray a -> SerialT m a
toStream = Stream m a -> SerialT m a
forall (m :: * -> *) a. Stream m a -> SerialT m a
SerialT (Stream m a -> SerialT m a)
-> (SmallArray a -> Stream m a) -> SmallArray a -> SerialT m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Stream m a -> Stream m a
forall (m :: * -> *) a. Monad m => Stream m a -> Stream m a
D.toStreamK (Stream m a -> Stream m a)
-> (SmallArray a -> Stream m a) -> SmallArray a -> Stream m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SmallArray a -> Stream m a
forall (m :: * -> *) a. Monad m => SmallArray a -> Stream m a
toStreamD

{-# INLINE_EARLY toStreamRev #-}
toStreamRev :: Monad m => SmallArray a -> SerialT m a
toStreamRev :: SmallArray a -> SerialT m a
toStreamRev = Stream m a -> SerialT m a
forall (m :: * -> *) a. Stream m a -> SerialT m a
SerialT (Stream m a -> SerialT m a)
-> (SmallArray a -> Stream m a) -> SmallArray a -> SerialT m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Stream m a -> Stream m a
forall (m :: * -> *) a. Monad m => Stream m a -> Stream m a
D.toStreamK (Stream m a -> Stream m a)
-> (SmallArray a -> Stream m a) -> SmallArray a -> Stream m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SmallArray a -> Stream m a
forall (m :: * -> *) a. Monad m => SmallArray a -> Stream m a
toStreamDRev

{-# INLINE fold #-}
fold :: Monad m => Fold m a b -> SmallArray a -> m b
fold :: Fold m a b -> SmallArray a -> m b
fold Fold m a b
f SmallArray a
arr = Fold m a b -> Stream m a -> m b
forall (m :: * -> *) a b.
Monad m =>
Fold m a b -> Stream m a -> m b
D.fold Fold m a b
f (SmallArray a -> Stream m a
forall (m :: * -> *) a. Monad m => SmallArray a -> Stream m a
toStreamD SmallArray a
arr)

{-# INLINE streamFold #-}
streamFold :: Monad m => (SerialT m a -> m b) -> SmallArray a -> m b
streamFold :: (SerialT m a -> m b) -> SmallArray a -> m b
streamFold SerialT m a -> m b
f SmallArray a
arr = SerialT m a -> m b
f (SmallArray a -> SerialT m a
forall (m :: * -> *) a. Monad m => SmallArray a -> SerialT m a
toStream SmallArray a
arr)

{-# INLINE_NORMAL read #-}
read :: Monad m => Unfold m (SmallArray a) a
read :: Unfold m (SmallArray a) a
read = ((SmallArray a, Int) -> m (Step (SmallArray a, Int) a))
-> (SmallArray a -> m (SmallArray a, Int))
-> Unfold m (SmallArray a) a
forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
Unfold (SmallArray a, Int) -> m (Step (SmallArray a, Int) a)
forall (m :: * -> *) a.
Monad m =>
(SmallArray a, Int) -> m (Step (SmallArray a, Int) a)
step SmallArray a -> m (SmallArray a, Int)
forall (m :: * -> *) b a. (Monad m, Num b) => a -> m (a, b)
inject
  where
    inject :: a -> m (a, b)
inject a
arr = (a, b) -> m (a, b)
forall (m :: * -> *) a. Monad m => a -> m a
return (a
arr, b
0)
    step :: (SmallArray a, Int) -> m (Step (SmallArray a, Int) a)
step (SmallArray a
arr, Int
i)
        | Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== SmallArray a -> Int
forall a. SmallArray a -> Int
length SmallArray a
arr = Step (SmallArray a, Int) a -> m (Step (SmallArray a, Int) a)
forall (m :: * -> *) a. Monad m => a -> m a
return Step (SmallArray a, Int) a
forall s a. Step s a
D.Stop
        | Bool
otherwise =
            Step (SmallArray a, Int) a -> m (Step (SmallArray a, Int) a)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (SmallArray a, Int) a -> m (Step (SmallArray a, Int) a))
-> Step (SmallArray a, Int) a -> m (Step (SmallArray a, Int) a)
forall a b. (a -> b) -> a -> b
$
            case SmallArray a -> Int -> (# a #)
forall a. SmallArray a -> Int -> (# a #)
indexSmallArray## SmallArray a
arr Int
i of
                (# a
x #) -> a -> (SmallArray a, Int) -> Step (SmallArray a, Int) a
forall s a. a -> s -> Step s a
D.Yield a
x (SmallArray a
arr, Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)