#include "inline.hs"

-- |
-- Module      : Streamly.Internal.Data.Stream.StreamD.Type
-- Copyright   : (c) 2018 Composewell Technologies
--               (c) Roman Leshchinskiy 2008-2010
-- License     : BSD-3-Clause
-- Maintainer  : streamly@composewell.com
-- Stability   : experimental
-- Portability : GHC

-- The stream type is inspired by the vector package.  A few functions in this
-- module have been originally adapted from the vector package (c) Roman
-- Leshchinskiy. See the notes in specific functions.

module Streamly.Internal.Data.Stream.StreamD.Type
    (
    -- * The stream type
      Step (..)
    -- XXX UnStream is exported to avoid a performance issue in concatMap if we
    -- use the pattern synonym "Stream".
    , Stream (Stream, UnStream)

    -- * Primitives
    , nilM
    , consM
    , uncons

    -- * From Unfold
    , unfold

    -- * From Values
    , fromPure
    , fromEffect

    -- * From Containers
    , fromList

    -- * Conversions From/To
    , fromStreamK
    , toStreamK

    -- * Running a 'Fold'
    , fold
    , fold_
    , foldOn

    -- * Right Folds
    , foldrT
    , foldrM
    , foldrMx
    , foldr
    , foldrS

    -- * Left Folds
    , foldl'
    , foldlM'
    , foldlx'
    , foldlMx'

    -- * Special Folds
    , drain

    -- * To Containers
    , toList

    -- * Multi-stream folds
    , eqBy
    , cmpBy

    -- * Transformations
    , map
    , mapM
    , take
    , takeWhile
    , takeWhileM
    , takeEndBy
    , takeEndByM

    -- * Nesting
    , ConcatMapUState (..)
    , unfoldMany
    , concatMap
    , concatMapM
    , FoldMany (..) -- for inspection testing
    , FoldManyPost (..)
    , foldMany
    , foldManyPost
    , refoldMany
    , chunksOf
    )
where

import Control.Applicative (liftA2)
import Control.Monad.Catch (MonadThrow, throwM)
import Control.Monad.Trans.Class (lift, MonadTrans)
import Data.Functor (($>))
import Data.Functor.Identity (Identity(..))
import Fusion.Plugin.Types (Fuse(..))
import GHC.Base (build)
import GHC.Types (SPEC(..))
import Prelude hiding (map, mapM, foldr, take, concatMap, takeWhile)

import Streamly.Internal.Data.Fold.Type (Fold(..))
import Streamly.Internal.Data.Refold.Type (Refold(..))
import Streamly.Internal.Data.Stream.StreamD.Step (Step (..))
import Streamly.Internal.Data.SVar.Type (State, adaptState, defState)
import Streamly.Internal.Data.Unfold.Type (Unfold(..))

import qualified Streamly.Internal.Data.Fold.Type as FL
import qualified Streamly.Internal.Data.Stream.StreamK.Type as K

------------------------------------------------------------------------------
-- The direct style stream type
------------------------------------------------------------------------------

-- gst = global state
-- | A stream consists of a step function that generates the next step given a
-- current state, and the current state.
data Stream m a =
    forall s. UnStream (State K.Stream m a -> s -> m (Step s a)) s

unShare :: Stream m a -> Stream m a
unShare :: forall (m :: * -> *) a. Stream m a -> Stream m a
unShare (UnStream State Stream m a -> s -> m (Step s a)
step s
state) = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
UnStream forall {m :: * -> *} {a}. State Stream m a -> s -> m (Step s a)
step' s
state
    where step' :: State Stream m a -> s -> m (Step s a)
step' State Stream m a
gst = State Stream m a -> s -> m (Step s a)
step (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst)

pattern Stream :: (State K.Stream m a -> s -> m (Step s a)) -> s -> Stream m a
pattern $bStream :: forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
$mStream :: forall {r} {m :: * -> *} {a}.
Stream m a
-> (forall {s}. (State Stream m a -> s -> m (Step s a)) -> s -> r)
-> ((# #) -> r)
-> r
Stream step state <- (unShare -> UnStream step state)
    where Stream = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
UnStream

#if __GLASGOW_HASKELL__ >= 802
{-# COMPLETE Stream #-}
#endif

------------------------------------------------------------------------------
-- Primitives
------------------------------------------------------------------------------

-- | An empty 'Stream' with a side effect.
{-# INLINE_NORMAL nilM #-}
nilM :: Applicative m => m b -> Stream m a
nilM :: forall (m :: * -> *) b a. Applicative m => m b -> Stream m a
nilM m b
m = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream (\State Stream m a
_ ()
_ -> m b
m forall (f :: * -> *) a b. Functor f => f a -> b -> f b
$> forall s a. Step s a
Stop) ()

{-# INLINE_NORMAL consM #-}
consM :: Applicative m => m a -> Stream m a -> Stream m a
consM :: forall (m :: * -> *) a.
Applicative m =>
m a -> Stream m a -> Stream m a
consM m a
m (Stream State Stream m a -> s -> m (Step s a)
step s
state) = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m a -> Maybe s -> m (Step (Maybe s) a)
step1 forall a. Maybe a
Nothing

    where

    {-# INLINE_LATE step1 #-}
    step1 :: State Stream m a -> Maybe s -> m (Step (Maybe s) a)
step1 State Stream m a
_ Maybe s
Nothing = (forall s a. a -> s -> Step s a
`Yield` forall a. a -> Maybe a
Just s
state) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> m a
m
    step1 State Stream m a
gst (Just s
st) = do
          (\case
            Yield a
a s
s -> forall s a. a -> s -> Step s a
Yield a
a (forall a. a -> Maybe a
Just s
s)
            Skip  s
s   -> forall s a. s -> Step s a
Skip (forall a. a -> Maybe a
Just s
s)
            Step s a
Stop      -> forall s a. Step s a
Stop) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> State Stream m a -> s -> m (Step s a)
step State Stream m a
gst s
st

-- | Does not fuse, has the same performance as the StreamK version.
{-# INLINE_NORMAL uncons #-}
uncons :: Monad m => Stream m a -> m (Maybe (a, Stream m a))
uncons :: forall (m :: * -> *) a.
Monad m =>
Stream m a -> m (Maybe (a, Stream m a))
uncons (UnStream State Stream m a -> s -> m (Step s a)
step s
state) = s -> m (Maybe (a, Stream m a))
go s
state
  where
    go :: s -> m (Maybe (a, Stream m a))
go s
st = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step forall (t :: (* -> *) -> * -> *) (m :: * -> *) a. State t m a
defState s
st
        case Step s a
r of
            Yield a
x s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just (a
x, forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m a -> s -> m (Step s a)
step s
s)
            Skip  s
s   -> s -> m (Maybe (a, Stream m a))
go s
s
            Step s a
Stop      -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a. Maybe a
Nothing

------------------------------------------------------------------------------
-- From 'Unfold'
------------------------------------------------------------------------------

data UnfoldState s = UnfoldNothing | UnfoldJust s

-- | Convert an 'Unfold' into a 'Stream' by supplying it a seed.
--
{-# INLINE_NORMAL unfold #-}
unfold :: Applicative m => Unfold m a b -> a -> Stream m b
unfold :: forall (m :: * -> *) a b.
Applicative m =>
Unfold m a b -> a -> Stream m b
unfold (Unfold s -> m (Step s b)
ustep a -> m s
inject) a
seed = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream forall {p}. p -> UnfoldState s -> m (Step (UnfoldState s) b)
step forall s. UnfoldState s
UnfoldNothing

    where

    {-# INLINE_LATE step #-}
    step :: p -> UnfoldState s -> m (Step (UnfoldState s) b)
step p
_ UnfoldState s
UnfoldNothing = forall s a. s -> Step s a
Skip forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall s. s -> UnfoldState s
UnfoldJust forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> m s
inject a
seed
    step p
_ (UnfoldJust s
st) = do
        (\case
            Yield b
x s
s -> forall s a. a -> s -> Step s a
Yield b
x (forall s. s -> UnfoldState s
UnfoldJust s
s)
            Skip s
s    -> forall s a. s -> Step s a
Skip (forall s. s -> UnfoldState s
UnfoldJust s
s)
            Step s b
Stop      -> forall s a. Step s a
Stop) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> s -> m (Step s b)
ustep s
st

------------------------------------------------------------------------------
-- From Values
------------------------------------------------------------------------------

-- | Create a singleton 'Stream' from a pure value.
{-# INLINE_NORMAL fromPure #-}
fromPure :: Applicative m => a -> Stream m a
fromPure :: forall (m :: * -> *) a. Applicative m => a -> Stream m a
fromPure a
x = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream (\State Stream m a
_ Bool
s -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall {p}. p -> Bool -> Step Bool a
step forall a. HasCallStack => a
undefined Bool
s) Bool
True
  where
    {-# INLINE_LATE step #-}
    step :: p -> Bool -> Step Bool a
step p
_ Bool
True  = forall s a. a -> s -> Step s a
Yield a
x Bool
False
    step p
_ Bool
False = forall s a. Step s a
Stop

-- | Create a singleton 'Stream' from a monadic action.
{-# INLINE_NORMAL fromEffect #-}
fromEffect :: Applicative m => m a -> Stream m a
fromEffect :: forall (m :: * -> *) a. Applicative m => m a -> Stream m a
fromEffect m a
m = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream forall {p}. p -> Bool -> m (Step Bool a)
step Bool
True

    where

    {-# INLINE_LATE step #-}
    step :: p -> Bool -> m (Step Bool a)
step p
_ Bool
True  = (forall s a. a -> s -> Step s a
`Yield` Bool
False) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> m a
m
    step p
_ Bool
False = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall s a. Step s a
Stop

------------------------------------------------------------------------------
-- From Containers
------------------------------------------------------------------------------

-- Adapted from the vector package.
-- | Convert a list of pure values to a 'Stream'
{-# INLINE_LATE fromList #-}
fromList :: Applicative m => [a] -> Stream m a
fromList :: forall (m :: * -> *) a. Applicative m => [a] -> Stream m a
fromList = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream forall {f :: * -> *} {p} {a}.
Applicative f =>
p -> [a] -> f (Step [a] a)
step
  where
    {-# INLINE_LATE step #-}
    step :: p -> [a] -> f (Step [a] a)
step p
_ (a
x:[a]
xs) = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
x [a]
xs
    step p
_ []     = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall s a. Step s a
Stop

------------------------------------------------------------------------------
-- Conversions From/To
------------------------------------------------------------------------------

-- | Convert a CPS encoded StreamK to direct style step encoded StreamD
{-# INLINE_LATE fromStreamK #-}
fromStreamK :: Applicative m => K.Stream m a -> Stream m a
fromStreamK :: forall (m :: * -> *) a. Applicative m => Stream m a -> Stream m a
fromStreamK = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream forall {m :: * -> *} {a}.
Applicative m =>
State Stream m a -> Stream m a -> m (Step (Stream m a) a)
step
    where
    step :: State Stream m a -> Stream m a -> m (Step (Stream m a) a)
step State Stream m a
gst Stream m a
m1 =
        let stop :: m (Step s a)
stop       = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall s a. Step s a
Stop
            single :: a -> f (Step (Stream m a) a)
single a
a   = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
a forall (m :: * -> *) a. Stream m a
K.nil
            yieldk :: a -> s -> f (Step s a)
yieldk a
a s
r = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield a
a s
r
         in forall (m :: * -> *) a r.
State Stream m a
-> (a -> Stream m a -> m r)
-> (a -> m r)
-> m r
-> Stream m a
-> m r
K.foldStreamShared State Stream m a
gst forall {f :: * -> *} {a} {s}.
Applicative f =>
a -> s -> f (Step s a)
yieldk forall {f :: * -> *} {a} {m :: * -> *} {a}.
Applicative f =>
a -> f (Step (Stream m a) a)
single forall {s} {a}. m (Step s a)
stop Stream m a
m1

-- | Convert a direct style step encoded StreamD to a CPS encoded StreamK
{-# INLINE_LATE toStreamK #-}
toStreamK :: Monad m => Stream m a -> K.Stream m a
toStreamK :: forall (m :: * -> *) a. Monad m => Stream m a -> Stream m a
toStreamK (Stream State Stream m a -> s -> m (Step s a)
step s
state) = s -> Stream m a
go s
state
    where
    go :: s -> Stream m a
go s
st = forall (m :: * -> *) a.
(forall r.
 State Stream m a
 -> (a -> Stream m a -> m r) -> (a -> m r) -> m r -> m r)
-> Stream m a
K.MkStream forall a b. (a -> b) -> a -> b
$ \State Stream m a
gst a -> Stream m a -> m r
yld a -> m r
_ m r
stp ->
      let go' :: s -> m r
go' s
ss = do
           Step s a
r <- State Stream m a -> s -> m (Step s a)
step State Stream m a
gst s
ss
           case Step s a
r of
               Yield a
x s
s -> a -> Stream m a -> m r
yld a
x (s -> Stream m a
go s
s)
               Skip  s
s   -> s -> m r
go' s
s
               Step s a
Stop      -> m r
stp
      in s -> m r
go' s
st

#ifndef DISABLE_FUSION
{-# RULES "fromStreamK/toStreamK fusion"
    forall s. toStreamK (fromStreamK s) = s #-}
{-# RULES "toStreamK/fromStreamK fusion"
    forall s. fromStreamK (toStreamK s) = s #-}
#endif

------------------------------------------------------------------------------
-- Running a 'Fold'
------------------------------------------------------------------------------

{-# INLINE_NORMAL fold #-}
fold :: Monad m => Fold m a b -> Stream m a -> m b
fold :: forall (m :: * -> *) a b.
Monad m =>
Fold m a b -> Stream m a -> m b
fold Fold m a b
fld Stream m a
strm = do
    (b
b, Stream m a
_) <- forall (m :: * -> *) a b.
Monad m =>
Fold m a b -> Stream m a -> m (b, Stream m a)
fold_ Fold m a b
fld Stream m a
strm
    forall (m :: * -> *) a. Monad m => a -> m a
return b
b

{-# INLINE_NORMAL fold_ #-}
fold_ :: Monad m => Fold m a b -> Stream m a -> m (b, Stream m a)
fold_ :: forall (m :: * -> *) a b.
Monad m =>
Fold m a b -> Stream m a -> m (b, Stream m a)
fold_ (Fold s -> a -> m (Step s b)
fstep m (Step s b)
begin s -> m b
done) (Stream State Stream m a -> s -> m (Step s a)
step s
state) = do
    Step s b
res <- m (Step s b)
begin
    case Step s b
res of
        FL.Partial s
fs -> SPEC -> s -> s -> m (b, Stream m a)
go SPEC
SPEC s
fs s
state
        FL.Done b
fb -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$! (b
fb, forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m a -> s -> m (Step s a)
step s
state)

    where

    {-# INLINE go #-}
    go :: SPEC -> s -> s -> m (b, Stream m a)
go !SPEC
_ !s
fs s
st = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step forall (t :: (* -> *) -> * -> *) (m :: * -> *) a. State t m a
defState s
st
        case Step s a
r of
            Yield a
x s
s -> do
                Step s b
res <- s -> a -> m (Step s b)
fstep s
fs a
x
                case Step s b
res of
                    FL.Done b
b -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$! (b
b, forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m a -> s -> m (Step s a)
step s
s)
                    FL.Partial s
fs1 -> SPEC -> s -> s -> m (b, Stream m a)
go SPEC
SPEC s
fs1 s
s
            Skip s
s -> SPEC -> s -> s -> m (b, Stream m a)
go SPEC
SPEC s
fs s
s
            Step s a
Stop -> do
                b
b <- s -> m b
done s
fs
                forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$! (b
b, forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream (\ State Stream m a
_ ()
_ -> forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop) ())

{-# INLINE_NORMAL foldOn #-}
foldOn :: Monad m => Fold m a b -> Stream m a -> Fold m a b
foldOn :: forall (m :: * -> *) a b.
Monad m =>
Fold m a b -> Stream m a -> Fold m a b
foldOn (Fold s -> a -> m (Step s b)
fstep m (Step s b)
finitial s -> m b
fextract) (Stream State Stream m a -> s -> m (Step s a)
sstep s
state) =
    forall (m :: * -> *) a b s.
(s -> a -> m (Step s b))
-> m (Step s b) -> (s -> m b) -> Fold m a b
Fold s -> a -> m (Step s b)
fstep m (Step s b)
initial s -> m b
fextract

    where

    initial :: m (Step s b)
initial = do
        Step s b
res <- m (Step s b)
finitial
        case Step s b
res of
            FL.Partial s
fs -> SPEC -> s -> s -> m (Step s b)
go SPEC
SPEC s
fs s
state
            FL.Done b
fb -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s b. b -> Step s b
FL.Done b
fb

    {-# INLINE go #-}
    go :: SPEC -> s -> s -> m (Step s b)
go !SPEC
_ !s
fs s
st = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
sstep forall (t :: (* -> *) -> * -> *) (m :: * -> *) a. State t m a
defState s
st
        case Step s a
r of
            Yield a
x s
s -> do
                Step s b
res <- s -> a -> m (Step s b)
fstep s
fs a
x
                case Step s b
res of
                    FL.Done b
b -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s b. b -> Step s b
FL.Done b
b
                    FL.Partial s
fs1 -> SPEC -> s -> s -> m (Step s b)
go SPEC
SPEC s
fs1 s
s
            Skip s
s -> SPEC -> s -> s -> m (Step s b)
go SPEC
SPEC s
fs s
s
            Step s a
Stop -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s b. s -> Step s b
FL.Partial s
fs

------------------------------------------------------------------------------
-- Right Folds
------------------------------------------------------------------------------

-- Adapted from the vector package.
--
-- XXX Use of SPEC constructor in folds causes 2x performance degradation in
-- one shot operations, but helps immensely in operations composed of multiple
-- combinators or the same combinator many times. There seems to be an
-- opportunity to optimize here, can we get both, better perf for single ops
-- as well as composed ops? Without SPEC, all single operation benchmarks
-- become 2x faster.

-- The way we want a left fold to be strict, dually we want the right fold to
-- be lazy.  The correct signature of the fold function to keep it lazy must be
-- (a -> m b -> m b) instead of (a -> b -> m b). We were using the latter
-- earlier, which is incorrect. In the latter signature we have to feed the
-- value to the fold function after evaluating the monadic action, depending on
-- the bind behavior of the monad, the action may get evaluated immediately
-- introducing unnecessary strictness to the fold. If the implementation is
-- lazy the following example, must work:
--
-- S.foldrM (\x t -> if x then return t else return False) (return True)
--  (S.fromList [False,undefined] :: SerialT IO Bool)
--
{-# INLINE_NORMAL foldrM #-}
foldrM :: Monad m => (a -> m b -> m b) -> m b -> Stream m a -> m b
foldrM :: forall (m :: * -> *) a b.
Monad m =>
(a -> m b -> m b) -> m b -> Stream m a -> m b
foldrM a -> m b -> m b
f m b
z (Stream State Stream m a -> s -> m (Step s a)
step s
state) = SPEC -> s -> m b
go SPEC
SPEC s
state
  where
    {-# INLINE_LATE go #-}
    go :: SPEC -> s -> m b
go !SPEC
_ s
st = do
          Step s a
r <- State Stream m a -> s -> m (Step s a)
step forall (t :: (* -> *) -> * -> *) (m :: * -> *) a. State t m a
defState s
st
          case Step s a
r of
            Yield a
x s
s -> a -> m b -> m b
f a
x (SPEC -> s -> m b
go SPEC
SPEC s
s)
            Skip s
s    -> SPEC -> s -> m b
go SPEC
SPEC s
s
            Step s a
Stop      -> m b
z

{-# INLINE_NORMAL foldrMx #-}
foldrMx :: Monad m
    => (a -> m x -> m x) -> m x -> (m x -> m b) -> Stream m a -> m b
foldrMx :: forall (m :: * -> *) a x b.
Monad m =>
(a -> m x -> m x) -> m x -> (m x -> m b) -> Stream m a -> m b
foldrMx a -> m x -> m x
fstep m x
final m x -> m b
convert (Stream State Stream m a -> s -> m (Step s a)
step s
state) = m x -> m b
convert forall a b. (a -> b) -> a -> b
$ SPEC -> s -> m x
go SPEC
SPEC s
state
  where
    {-# INLINE_LATE go #-}
    go :: SPEC -> s -> m x
go !SPEC
_ s
st = do
          Step s a
r <- State Stream m a -> s -> m (Step s a)
step forall (t :: (* -> *) -> * -> *) (m :: * -> *) a. State t m a
defState s
st
          case Step s a
r of
            Yield a
x s
s -> a -> m x -> m x
fstep a
x (SPEC -> s -> m x
go SPEC
SPEC s
s)
            Skip s
s    -> SPEC -> s -> m x
go SPEC
SPEC s
s
            Step s a
Stop      -> m x
final

-- Note that foldr works on pure values, therefore it becomes necessarily
-- strict when the monad m is strict. In that case it cannot terminate early,
-- it would evaluate all of its input.  Though, this should work fine with lazy
-- monads. For example, if "any" is implemented using "foldr" instead of
-- "foldrM" it performs the same with Identity monad but performs 1000x slower
-- with IO monad.
--
{-# INLINE_NORMAL foldr #-}
foldr :: Monad m => (a -> b -> b) -> b -> Stream m a -> m b
foldr :: forall (m :: * -> *) a b.
Monad m =>
(a -> b -> b) -> b -> Stream m a -> m b
foldr a -> b -> b
f b
z = forall (m :: * -> *) a b.
Monad m =>
(a -> m b -> m b) -> m b -> Stream m a -> m b
foldrM (forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 a -> b -> b
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a. Monad m => a -> m a
return) (forall (m :: * -> *) a. Monad m => a -> m a
return b
z)

-- this performs horribly, should not be used
{-# INLINE_NORMAL foldrS #-}
foldrS
    :: Monad m
    => (a -> Stream m b -> Stream m b)
    -> Stream m b
    -> Stream m a
    -> Stream m b
foldrS :: forall (m :: * -> *) a b.
Monad m =>
(a -> Stream m b -> Stream m b)
-> Stream m b -> Stream m a -> Stream m b
foldrS a -> Stream m b -> Stream m b
f Stream m b
final (Stream State Stream m a -> s -> m (Step s a)
step s
state) = SPEC -> s -> Stream m b
go SPEC
SPEC s
state
  where
    {-# INLINE_LATE go #-}
    go :: SPEC -> s -> Stream m b
go !SPEC
_ s
st = do
        -- defState??
        Step s a
r <- forall (m :: * -> *) a. Applicative m => m a -> Stream m a
fromEffect forall a b. (a -> b) -> a -> b
$ State Stream m a -> s -> m (Step s a)
step forall (t :: (* -> *) -> * -> *) (m :: * -> *) a. State t m a
defState s
st
        case Step s a
r of
          Yield a
x s
s -> a -> Stream m b -> Stream m b
f a
x (SPEC -> s -> Stream m b
go SPEC
SPEC s
s)
          Skip s
s    -> SPEC -> s -> Stream m b
go SPEC
SPEC s
s
          Step s a
Stop      -> Stream m b
final

-- Right fold to some transformer (T) monad.  This can be useful to implement
-- stateless combinators like map, filtering, insertions, takeWhile, dropWhile.
--
{-# INLINE_NORMAL foldrT #-}
foldrT :: (Monad m, Monad (t m), MonadTrans t)
    => (a -> t m b -> t m b) -> t m b -> Stream m a -> t m b
foldrT :: forall (m :: * -> *) (t :: (* -> *) -> * -> *) a b.
(Monad m, Monad (t m), MonadTrans t) =>
(a -> t m b -> t m b) -> t m b -> Stream m a -> t m b
foldrT a -> t m b -> t m b
f t m b
final (Stream State Stream m a -> s -> m (Step s a)
step s
state) = SPEC -> s -> t m b
go SPEC
SPEC s
state
  where
    {-# INLINE_LATE go #-}
    go :: SPEC -> s -> t m b
go !SPEC
_ s
st = do
          Step s a
r <- forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall a b. (a -> b) -> a -> b
$ State Stream m a -> s -> m (Step s a)
step forall (t :: (* -> *) -> * -> *) (m :: * -> *) a. State t m a
defState s
st
          case Step s a
r of
            Yield a
x s
s -> a -> t m b -> t m b
f a
x (SPEC -> s -> t m b
go SPEC
SPEC s
s)
            Skip s
s    -> SPEC -> s -> t m b
go SPEC
SPEC s
s
            Step s a
Stop      -> t m b
final

------------------------------------------------------------------------------
-- Left Folds
------------------------------------------------------------------------------

-- XXX run begin action only if the stream is not empty.
{-# INLINE_NORMAL foldlMx' #-}
foldlMx' :: Monad m => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> m b
foldlMx' :: forall (m :: * -> *) x a b.
Monad m =>
(x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> m b
foldlMx' x -> a -> m x
fstep m x
begin x -> m b
done (Stream State Stream m a -> s -> m (Step s a)
step s
state) =
    m x
begin forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \x
x -> SPEC -> x -> s -> m b
go SPEC
SPEC x
x s
state
  where
    -- XXX !acc?
    {-# INLINE_LATE go #-}
    go :: SPEC -> x -> s -> m b
go !SPEC
_ x
acc s
st = x
acc seq :: forall a b. a -> b -> b
`seq` do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step forall (t :: (* -> *) -> * -> *) (m :: * -> *) a. State t m a
defState s
st
        case Step s a
r of
            Yield a
x s
s -> do
                x
acc' <- x -> a -> m x
fstep x
acc a
x
                SPEC -> x -> s -> m b
go SPEC
SPEC x
acc' s
s
            Skip s
s -> SPEC -> x -> s -> m b
go SPEC
SPEC x
acc s
s
            Step s a
Stop   -> x -> m b
done x
acc

{-# INLINE foldlx' #-}
foldlx' :: Monad m => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> m b
foldlx' :: forall (m :: * -> *) x a b.
Monad m =>
(x -> a -> x) -> x -> (x -> b) -> Stream m a -> m b
foldlx' x -> a -> x
fstep x
begin x -> b
done =
    forall (m :: * -> *) x a b.
Monad m =>
(x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> m b
foldlMx' (\x
b a
a -> forall (m :: * -> *) a. Monad m => a -> m a
return (x -> a -> x
fstep x
b a
a)) (forall (m :: * -> *) a. Monad m => a -> m a
return x
begin) (forall (m :: * -> *) a. Monad m => a -> m a
return forall b c a. (b -> c) -> (a -> b) -> a -> c
. x -> b
done)

-- Adapted from the vector package.
-- XXX implement in terms of foldlMx'?
{-# INLINE_NORMAL foldlM' #-}
foldlM' :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> m b
foldlM' :: forall (m :: * -> *) b a.
Monad m =>
(b -> a -> m b) -> m b -> Stream m a -> m b
foldlM' b -> a -> m b
fstep m b
mbegin (Stream State Stream m a -> s -> m (Step s a)
step s
state) = do
    b
begin <- m b
mbegin
    SPEC -> b -> s -> m b
go SPEC
SPEC b
begin s
state
  where
    {-# INLINE_LATE go #-}
    go :: SPEC -> b -> s -> m b
go !SPEC
_ b
acc s
st = b
acc seq :: forall a b. a -> b -> b
`seq` do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step forall (t :: (* -> *) -> * -> *) (m :: * -> *) a. State t m a
defState s
st
        case Step s a
r of
            Yield a
x s
s -> do
                b
acc' <- b -> a -> m b
fstep b
acc a
x
                SPEC -> b -> s -> m b
go SPEC
SPEC b
acc' s
s
            Skip s
s -> SPEC -> b -> s -> m b
go SPEC
SPEC b
acc s
s
            Step s a
Stop   -> forall (m :: * -> *) a. Monad m => a -> m a
return b
acc

{-# INLINE foldl' #-}
foldl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> m b
foldl' :: forall (m :: * -> *) b a.
Monad m =>
(b -> a -> b) -> b -> Stream m a -> m b
foldl' b -> a -> b
fstep b
begin = forall (m :: * -> *) b a.
Monad m =>
(b -> a -> m b) -> m b -> Stream m a -> m b
foldlM' (\b
b a
a -> forall (m :: * -> *) a. Monad m => a -> m a
return (b -> a -> b
fstep b
b a
a)) (forall (m :: * -> *) a. Monad m => a -> m a
return b
begin)

------------------------------------------------------------------------------
-- Special folds
------------------------------------------------------------------------------

-- | Run a streaming composition, discard the results.
{-# INLINE_LATE drain #-}
drain :: Monad m => Stream m a -> m ()
-- drain = foldrM (\_ xs -> xs) (return ())
drain :: forall (m :: * -> *) a. Monad m => Stream m a -> m ()
drain (Stream State Stream m a -> s -> m (Step s a)
step s
state) = SPEC -> s -> m ()
go SPEC
SPEC s
state
  where
    go :: SPEC -> s -> m ()
go !SPEC
_ s
st = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step forall (t :: (* -> *) -> * -> *) (m :: * -> *) a. State t m a
defState s
st
        case Step s a
r of
            Yield a
_ s
s -> SPEC -> s -> m ()
go SPEC
SPEC s
s
            Skip s
s    -> SPEC -> s -> m ()
go SPEC
SPEC s
s
            Step s a
Stop      -> forall (m :: * -> *) a. Monad m => a -> m a
return ()

------------------------------------------------------------------------------
-- To Containers
------------------------------------------------------------------------------

{-# INLINE_NORMAL toList #-}
toList :: Monad m => Stream m a -> m [a]
toList :: forall (m :: * -> *) a. Monad m => Stream m a -> m [a]
toList = forall (m :: * -> *) a b.
Monad m =>
(a -> b -> b) -> b -> Stream m a -> m b
foldr (:) []

-- Use foldr/build fusion to fuse with list consumers
-- This can be useful when using the IsList instance
{-# INLINE_LATE toListFB #-}
toListFB :: (a -> b -> b) -> b -> Stream Identity a -> b
toListFB :: forall a b. (a -> b -> b) -> b -> Stream Identity a -> b
toListFB a -> b -> b
c b
n (Stream State Stream Identity a -> s -> Identity (Step s a)
step s
state) = s -> b
go s
state
  where
    go :: s -> b
go s
st = case forall a. Identity a -> a
runIdentity (State Stream Identity a -> s -> Identity (Step s a)
step forall (t :: (* -> *) -> * -> *) (m :: * -> *) a. State t m a
defState s
st) of
             Yield a
x s
s -> a
x a -> b -> b
`c` s -> b
go s
s
             Skip s
s    -> s -> b
go s
s
             Step s a
Stop      -> b
n

{-# RULES "toList Identity" toList = toListId #-}
{-# INLINE_EARLY toListId #-}
toListId :: Stream Identity a -> Identity [a]
toListId :: forall a. Stream Identity a -> Identity [a]
toListId Stream Identity a
s = forall a. a -> Identity a
Identity forall a b. (a -> b) -> a -> b
$ forall a. (forall b. (a -> b -> b) -> b -> b) -> [a]
build (\a -> b -> b
c b
n -> forall a b. (a -> b -> b) -> b -> Stream Identity a -> b
toListFB a -> b -> b
c b
n Stream Identity a
s)

------------------------------------------------------------------------------
-- Multi-stream folds
------------------------------------------------------------------------------

-- Adapted from the vector package.
{-# INLINE_NORMAL eqBy #-}
eqBy :: Monad m => (a -> b -> Bool) -> Stream m a -> Stream m b -> m Bool
eqBy :: forall (m :: * -> *) a b.
Monad m =>
(a -> b -> Bool) -> Stream m a -> Stream m b -> m Bool
eqBy a -> b -> Bool
eq (Stream State Stream m a -> s -> m (Step s a)
step1 s
t1) (Stream State Stream m b -> s -> m (Step s b)
step2 s
t2) = SPEC -> s -> s -> m Bool
eq_loop0 SPEC
SPEC s
t1 s
t2
  where
    eq_loop0 :: SPEC -> s -> s -> m Bool
eq_loop0 !SPEC
_ s
s1 s
s2 = do
      Step s a
r <- State Stream m a -> s -> m (Step s a)
step1 forall (t :: (* -> *) -> * -> *) (m :: * -> *) a. State t m a
defState s
s1
      case Step s a
r of
        Yield a
x s
s1' -> SPEC -> a -> s -> s -> m Bool
eq_loop1 SPEC
SPEC a
x s
s1' s
s2
        Skip    s
s1' -> SPEC -> s -> s -> m Bool
eq_loop0 SPEC
SPEC   s
s1' s
s2
        Step s a
Stop        -> s -> m Bool
eq_null s
s2

    eq_loop1 :: SPEC -> a -> s -> s -> m Bool
eq_loop1 !SPEC
_ a
x s
s1 s
s2 = do
      Step s b
r <- State Stream m b -> s -> m (Step s b)
step2 forall (t :: (* -> *) -> * -> *) (m :: * -> *) a. State t m a
defState s
s2
      case Step s b
r of
        Yield b
y s
s2'
          | a -> b -> Bool
eq a
x b
y    -> SPEC -> s -> s -> m Bool
eq_loop0 SPEC
SPEC   s
s1 s
s2'
          | Bool
otherwise -> forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False
        Skip    s
s2'   -> SPEC -> a -> s -> s -> m Bool
eq_loop1 SPEC
SPEC a
x s
s1 s
s2'
        Step s b
Stop          -> forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False

    eq_null :: s -> m Bool
eq_null s
s2 = do
      Step s b
r <- State Stream m b -> s -> m (Step s b)
step2 forall (t :: (* -> *) -> * -> *) (m :: * -> *) a. State t m a
defState s
s2
      case Step s b
r of
        Yield b
_ s
_ -> forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False
        Skip s
s2'  -> s -> m Bool
eq_null s
s2'
        Step s b
Stop      -> forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True

-- Adapted from the vector package.
-- | Compare two streams lexicographically
{-# INLINE_NORMAL cmpBy #-}
cmpBy
    :: Monad m
    => (a -> b -> Ordering) -> Stream m a -> Stream m b -> m Ordering
cmpBy :: forall (m :: * -> *) a b.
Monad m =>
(a -> b -> Ordering) -> Stream m a -> Stream m b -> m Ordering
cmpBy a -> b -> Ordering
cmp (Stream State Stream m a -> s -> m (Step s a)
step1 s
t1) (Stream State Stream m b -> s -> m (Step s b)
step2 s
t2) = SPEC -> s -> s -> m Ordering
cmp_loop0 SPEC
SPEC s
t1 s
t2
  where
    cmp_loop0 :: SPEC -> s -> s -> m Ordering
cmp_loop0 !SPEC
_ s
s1 s
s2 = do
      Step s a
r <- State Stream m a -> s -> m (Step s a)
step1 forall (t :: (* -> *) -> * -> *) (m :: * -> *) a. State t m a
defState s
s1
      case Step s a
r of
        Yield a
x s
s1' -> SPEC -> a -> s -> s -> m Ordering
cmp_loop1 SPEC
SPEC a
x s
s1' s
s2
        Skip    s
s1' -> SPEC -> s -> s -> m Ordering
cmp_loop0 SPEC
SPEC   s
s1' s
s2
        Step s a
Stop        -> s -> m Ordering
cmp_null s
s2

    cmp_loop1 :: SPEC -> a -> s -> s -> m Ordering
cmp_loop1 !SPEC
_ a
x s
s1 s
s2 = do
      Step s b
r <- State Stream m b -> s -> m (Step s b)
step2 forall (t :: (* -> *) -> * -> *) (m :: * -> *) a. State t m a
defState s
s2
      case Step s b
r of
        Yield b
y s
s2' -> case a
x a -> b -> Ordering
`cmp` b
y of
                         Ordering
EQ -> SPEC -> s -> s -> m Ordering
cmp_loop0 SPEC
SPEC s
s1 s
s2'
                         Ordering
c  -> forall (m :: * -> *) a. Monad m => a -> m a
return Ordering
c
        Skip    s
s2' -> SPEC -> a -> s -> s -> m Ordering
cmp_loop1 SPEC
SPEC a
x s
s1 s
s2'
        Step s b
Stop        -> forall (m :: * -> *) a. Monad m => a -> m a
return Ordering
GT

    cmp_null :: s -> m Ordering
cmp_null s
s2 = do
      Step s b
r <- State Stream m b -> s -> m (Step s b)
step2 forall (t :: (* -> *) -> * -> *) (m :: * -> *) a. State t m a
defState s
s2
      case Step s b
r of
        Yield b
_ s
_ -> forall (m :: * -> *) a. Monad m => a -> m a
return Ordering
LT
        Skip s
s2'  -> s -> m Ordering
cmp_null s
s2'
        Step s b
Stop      -> forall (m :: * -> *) a. Monad m => a -> m a
return Ordering
EQ

------------------------------------------------------------------------------
-- Transformations
------------------------------------------------------------------------------

-- Adapted from the vector package.
-- | Map a monadic function over a 'Stream'
{-# INLINE_NORMAL mapM #-}
mapM :: Monad m => (a -> m b) -> Stream m a -> Stream m b
mapM :: forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Stream m a -> Stream m b
mapM a -> m b
f (Stream State Stream m a -> s -> m (Step s a)
step s
state) = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream forall {m :: * -> *} {a}. State Stream m a -> s -> m (Step s b)
step' s
state
  where
    {-# INLINE_LATE step' #-}
    step' :: State Stream m a -> s -> m (Step s b)
step' State Stream m a
gst s
st = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
st
        case Step s a
r of
            Yield a
x s
s -> a -> m b
f a
x forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \b
a -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield b
a s
s
            Skip s
s    -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip s
s
            Step s a
Stop      -> forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

{-# INLINE map #-}
map :: Monad m => (a -> b) -> Stream m a -> Stream m b
map :: forall (m :: * -> *) a b.
Monad m =>
(a -> b) -> Stream m a -> Stream m b
map a -> b
f = forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Stream m a -> Stream m b
mapM (forall (m :: * -> *) a. Monad m => a -> m a
return forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> b
f)

instance Functor m => Functor (Stream m) where
    {-# INLINE fmap #-}
    fmap :: forall a b. (a -> b) -> Stream m a -> Stream m b
fmap a -> b
f (Stream State Stream m a -> s -> m (Step s a)
step s
state) = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream forall {m :: * -> *} {a}. State Stream m a -> s -> m (Step s b)
step' s
state
      where
        {-# INLINE_LATE step' #-}
        step' :: State Stream m a -> s -> m (Step s b)
step' State Stream m a
gst s
st = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f) (State Stream m a -> s -> m (Step s a)
step (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
st)

    {-# INLINE (<$) #-}
    <$ :: forall a b. a -> Stream m b -> Stream m a
(<$) = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. a -> b -> a
const

-------------------------------------------------------------------------------
-- Filtering
-------------------------------------------------------------------------------

-- Adapted from the vector package.
{-# INLINE_NORMAL take #-}
take :: Applicative m => Int -> Stream m a -> Stream m a
take :: forall (m :: * -> *) a.
Applicative m =>
Int -> Stream m a -> Stream m a
take Int
n (Stream State Stream m a -> s -> m (Step s a)
step s
state) = Int
n seq :: forall a b. a -> b -> b
`seq` forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m a -> (s, Int) -> m (Step (s, Int) a)
step' (s
state, Int
0)

    where

    {-# INLINE_LATE step' #-}
    step' :: State Stream m a -> (s, Int) -> m (Step (s, Int) a)
step' State Stream m a
gst (s
st, Int
i) | Int
i forall a. Ord a => a -> a -> Bool
< Int
n = do
        (\case
            Yield a
x s
s -> forall s a. a -> s -> Step s a
Yield a
x (s
s, Int
i forall a. Num a => a -> a -> a
+ Int
1)
            Skip s
s    -> forall s a. s -> Step s a
Skip (s
s, Int
i)
            Step s a
Stop      -> forall s a. Step s a
Stop) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> State Stream m a -> s -> m (Step s a)
step State Stream m a
gst s
st
    step' State Stream m a
_ (s
_, Int
_) = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall s a. Step s a
Stop

-- Adapted from the vector package.
{-# INLINE_NORMAL takeWhileM #-}
takeWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
takeWhileM :: forall (m :: * -> *) a.
Monad m =>
(a -> m Bool) -> Stream m a -> Stream m a
takeWhileM a -> m Bool
f (Stream State Stream m a -> s -> m (Step s a)
step s
state) = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m a -> s -> m (Step s a)
step' s
state
  where
    {-# INLINE_LATE step' #-}
    step' :: State Stream m a -> s -> m (Step s a)
step' State Stream m a
gst s
st = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step State Stream m a
gst s
st
        case Step s a
r of
            Yield a
x s
s -> do
                Bool
b <- a -> m Bool
f a
x
                forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ if Bool
b then forall s a. a -> s -> Step s a
Yield a
x s
s else forall s a. Step s a
Stop
            Skip s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip s
s
            Step s a
Stop   -> forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

{-# INLINE takeWhile #-}
takeWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
takeWhile :: forall (m :: * -> *) a.
Monad m =>
(a -> Bool) -> Stream m a -> Stream m a
takeWhile a -> Bool
f = forall (m :: * -> *) a.
Monad m =>
(a -> m Bool) -> Stream m a -> Stream m a
takeWhileM (forall (m :: * -> *) a. Monad m => a -> m a
return forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Bool
f)

-- Like takeWhile but with an inverted condition and also taking
-- the matching element.

{-# INLINE_NORMAL takeEndByM #-}
takeEndByM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
takeEndByM :: forall (m :: * -> *) a.
Monad m =>
(a -> m Bool) -> Stream m a -> Stream m a
takeEndByM a -> m Bool
f (Stream State Stream m a -> s -> m (Step s a)
step s
state) = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m a -> Maybe s -> m (Step (Maybe s) a)
step' (forall a. a -> Maybe a
Just s
state)
  where
    {-# INLINE_LATE step' #-}
    step' :: State Stream m a -> Maybe s -> m (Step (Maybe s) a)
step' State Stream m a
gst (Just s
st) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step State Stream m a
gst s
st
        case Step s a
r of
            Yield a
x s
s -> do
                Bool
b <- a -> m Bool
f a
x
                forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$
                    if Bool -> Bool
not Bool
b
                    then forall s a. a -> s -> Step s a
Yield a
x (forall a. a -> Maybe a
Just s
s)
                    else forall s a. a -> s -> Step s a
Yield a
x forall a. Maybe a
Nothing
            Skip s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall a. a -> Maybe a
Just s
s)
            Step s a
Stop   -> forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

    step' State Stream m a
_ Maybe s
Nothing = forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

{-# INLINE takeEndBy #-}
takeEndBy :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
takeEndBy :: forall (m :: * -> *) a.
Monad m =>
(a -> Bool) -> Stream m a -> Stream m a
takeEndBy a -> Bool
f = forall (m :: * -> *) a.
Monad m =>
(a -> m Bool) -> Stream m a -> Stream m a
takeEndByM (forall (m :: * -> *) a. Monad m => a -> m a
return forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Bool
f)

------------------------------------------------------------------------------
-- Combine N Streams - concatAp
------------------------------------------------------------------------------

{-# INLINE_NORMAL concatAp #-}
concatAp :: Functor f => Stream f (a -> b) -> Stream f a -> Stream f b
concatAp :: forall (f :: * -> *) a b.
Functor f =>
Stream f (a -> b) -> Stream f a -> Stream f b
concatAp (Stream State Stream f (a -> b) -> s -> f (Step s (a -> b))
stepa s
statea) (Stream State Stream f a -> s -> f (Step s a)
stepb s
stateb) =
    forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream forall {m :: * -> *} {a}.
State Stream m a
-> Either s (a -> b, s, s) -> f (Step (Either s (a -> b, s, s)) b)
step' (forall a b. a -> Either a b
Left s
statea)

    where

    {-# INLINE_LATE step' #-}
    step' :: State Stream m a
-> Either s (a -> b, s, s) -> f (Step (Either s (a -> b, s, s)) b)
step' State Stream m a
gst (Left s
st) = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap
        (\case
            Yield a -> b
f s
s -> forall s a. s -> Step s a
Skip (forall a b. b -> Either a b
Right (a -> b
f, s
s, s
stateb))
            Skip    s
s -> forall s a. s -> Step s a
Skip (forall a b. a -> Either a b
Left s
s)
            Step s (a -> b)
Stop      -> forall s a. Step s a
Stop)
        (State Stream f (a -> b) -> s -> f (Step s (a -> b))
stepa (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
st)
    step' State Stream m a
gst (Right (a -> b
f, s
os, s
st)) = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap
        (\case
            Yield a
a s
s -> forall s a. a -> s -> Step s a
Yield (a -> b
f a
a) (forall a b. b -> Either a b
Right (a -> b
f, s
os, s
s))
            Skip s
s    -> forall s a. s -> Step s a
Skip (forall a b. b -> Either a b
Right (a -> b
f,s
os, s
s))
            Step s a
Stop      -> forall s a. s -> Step s a
Skip (forall a b. a -> Either a b
Left s
os))
        (State Stream f a -> s -> f (Step s a)
stepb (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
st)

{-# INLINE_NORMAL apSequence #-}
apSequence :: Functor f => Stream f a -> Stream f b -> Stream f b
apSequence :: forall (f :: * -> *) a b.
Functor f =>
Stream f a -> Stream f b -> Stream f b
apSequence (Stream State Stream f a -> s -> f (Step s a)
stepa s
statea) (Stream State Stream f b -> s -> f (Step s b)
stepb s
stateb) =
    forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream f b -> Either s (s, s) -> f (Step (Either s (s, s)) b)
step (forall a b. a -> Either a b
Left s
statea)

    where

    {-# INLINE_LATE step #-}
    step :: State Stream f b -> Either s (s, s) -> f (Step (Either s (s, s)) b)
step State Stream f b
gst (Left s
st) =
        forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap
            (\case
                 Yield a
_ s
s -> forall s a. s -> Step s a
Skip (forall a b. b -> Either a b
Right (s
s, s
stateb))
                 Skip s
s -> forall s a. s -> Step s a
Skip (forall a b. a -> Either a b
Left s
s)
                 Step s a
Stop -> forall s a. Step s a
Stop)
            (State Stream f a -> s -> f (Step s a)
stepa (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream f b
gst) s
st)
    step State Stream f b
gst (Right (s
ostate, s
st)) =
        forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap
            (\case
                 Yield b
b s
s -> forall s a. a -> s -> Step s a
Yield b
b (forall a b. b -> Either a b
Right (s
ostate, s
s))
                 Skip s
s -> forall s a. s -> Step s a
Skip (forall a b. b -> Either a b
Right (s
ostate, s
s))
                 Step s b
Stop -> forall s a. s -> Step s a
Skip (forall a b. a -> Either a b
Left s
ostate))
            (State Stream f b -> s -> f (Step s b)
stepb State Stream f b
gst s
st)

{-# INLINE_NORMAL apDiscardSnd #-}
apDiscardSnd :: Functor f => Stream f a -> Stream f b -> Stream f a
apDiscardSnd :: forall (f :: * -> *) a b.
Functor f =>
Stream f a -> Stream f b -> Stream f a
apDiscardSnd (Stream State Stream f a -> s -> f (Step s a)
stepa s
statea) (Stream State Stream f b -> s -> f (Step s b)
stepb s
stateb) =
    forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream f a
-> Either s (s, s, a) -> f (Step (Either s (s, s, a)) a)
step (forall a b. a -> Either a b
Left s
statea)

    where

    {-# INLINE_LATE step #-}
    step :: State Stream f a
-> Either s (s, s, a) -> f (Step (Either s (s, s, a)) a)
step State Stream f a
gst (Left s
st) =
        forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap
            (\case
                 Yield a
b s
s -> forall s a. s -> Step s a
Skip (forall a b. b -> Either a b
Right (s
s, s
stateb, a
b))
                 Skip s
s -> forall s a. s -> Step s a
Skip (forall a b. a -> Either a b
Left s
s)
                 Step s a
Stop -> forall s a. Step s a
Stop)
            (State Stream f a -> s -> f (Step s a)
stepa State Stream f a
gst s
st)
    step State Stream f a
gst (Right (s
ostate, s
st, a
b)) =
        forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap
            (\case
                 Yield b
_ s
s -> forall s a. a -> s -> Step s a
Yield a
b (forall a b. b -> Either a b
Right (s
ostate, s
s, a
b))
                 Skip s
s -> forall s a. s -> Step s a
Skip (forall a b. b -> Either a b
Right (s
ostate, s
s, a
b))
                 Step s b
Stop -> forall s a. s -> Step s a
Skip (forall a b. a -> Either a b
Left s
ostate))
            (State Stream f b -> s -> f (Step s b)
stepb (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream f a
gst) s
st)

instance Applicative f => Applicative (Stream f) where
    {-# INLINE pure #-}
    pure :: forall a. a -> Stream f a
pure = forall (m :: * -> *) a. Applicative m => a -> Stream m a
fromPure

    {-# INLINE (<*>) #-}
    <*> :: forall a b. Stream f (a -> b) -> Stream f a -> Stream f b
(<*>) = forall (f :: * -> *) a b.
Functor f =>
Stream f (a -> b) -> Stream f a -> Stream f b
concatAp

#if MIN_VERSION_base(4,10,0)
    {-# INLINE liftA2 #-}
    liftA2 :: forall a b c.
(a -> b -> c) -> Stream f a -> Stream f b -> Stream f c
liftA2 a -> b -> c
f Stream f a
x = forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
(<*>) (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b -> c
f Stream f a
x)
#endif

    {-# INLINE (*>) #-}
    *> :: forall a b. Stream f a -> Stream f b -> Stream f b
(*>) = forall (f :: * -> *) a b.
Functor f =>
Stream f a -> Stream f b -> Stream f b
apSequence

    {-# INLINE (<*) #-}
    <* :: forall a b. Stream f a -> Stream f b -> Stream f a
(<*) = forall (f :: * -> *) a b.
Functor f =>
Stream f a -> Stream f b -> Stream f a
apDiscardSnd

------------------------------------------------------------------------------
-- Combine N Streams - unfoldMany
------------------------------------------------------------------------------

-- Define a unique structure to use in inspection testing
data ConcatMapUState o i =
      ConcatMapUOuter o
    | ConcatMapUInner o i

-- | @unfoldMany unfold stream@ uses @unfold@ to map the input stream elements
-- to streams and then flattens the generated streams into a single output
-- stream.

-- This is like 'concatMap' but uses an unfold with an explicit state to
-- generate the stream instead of a 'Stream' type generator. This allows better
-- optimization via fusion.  This can be many times more efficient than
-- 'concatMap'.

{-# INLINE_NORMAL unfoldMany #-}
unfoldMany :: Monad m => Unfold m a b -> Stream m a -> Stream m b
unfoldMany :: forall (m :: * -> *) a b.
Monad m =>
Unfold m a b -> Stream m a -> Stream m b
unfoldMany (Unfold s -> m (Step s b)
istep a -> m s
inject) (Stream State Stream m a -> s -> m (Step s a)
ostep s
ost) =
    forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream forall {m :: * -> *} {a}.
State Stream m a
-> ConcatMapUState s s -> m (Step (ConcatMapUState s s) b)
step (forall o i. o -> ConcatMapUState o i
ConcatMapUOuter s
ost)
  where
    {-# INLINE_LATE step #-}
    step :: State Stream m a
-> ConcatMapUState s s -> m (Step (ConcatMapUState s s) b)
step State Stream m a
gst (ConcatMapUOuter s
o) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
ostep (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
o
        case Step s a
r of
            Yield a
a s
o' -> do
                s
i <- a -> m s
inject a
a
                s
i seq :: forall a b. a -> b -> b
`seq` forall (m :: * -> *) a. Monad m => a -> m a
return (forall s a. s -> Step s a
Skip (forall o i. o -> i -> ConcatMapUState o i
ConcatMapUInner s
o' s
i))
            Skip s
o' -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall o i. o -> ConcatMapUState o i
ConcatMapUOuter s
o')
            Step s a
Stop -> forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

    step State Stream m a
_ (ConcatMapUInner s
o s
i) = do
        Step s b
r <- s -> m (Step s b)
istep s
i
        forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ case Step s b
r of
            Yield b
x s
i' -> forall s a. a -> s -> Step s a
Yield b
x (forall o i. o -> i -> ConcatMapUState o i
ConcatMapUInner s
o s
i')
            Skip s
i'    -> forall s a. s -> Step s a
Skip (forall o i. o -> i -> ConcatMapUState o i
ConcatMapUInner s
o s
i')
            Step s b
Stop       -> forall s a. s -> Step s a
Skip (forall o i. o -> ConcatMapUState o i
ConcatMapUOuter s
o)

------------------------------------------------------------------------------
-- Combine N Streams - concatMap
------------------------------------------------------------------------------

-- Adapted from the vector package.
{-# INLINE_NORMAL concatMapM #-}
concatMapM :: Monad m => (a -> m (Stream m b)) -> Stream m a -> Stream m b
concatMapM :: forall (m :: * -> *) a b.
Monad m =>
(a -> m (Stream m b)) -> Stream m a -> Stream m b
concatMapM a -> m (Stream m b)
f (Stream State Stream m a -> s -> m (Step s a)
step s
state) = forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream forall {m :: * -> *} {a}.
State Stream m a
-> Either s (Stream m b, s)
-> m (Step (Either s (Stream m b, s)) b)
step' (forall a b. a -> Either a b
Left s
state)
  where
    {-# INLINE_LATE step' #-}
    step' :: State Stream m a
-> Either s (Stream m b, s)
-> m (Step (Either s (Stream m b, s)) b)
step' State Stream m a
gst (Left s
st) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
st
        case Step s a
r of
            Yield a
a s
s -> do
                Stream m b
b_stream <- a -> m (Stream m b)
f a
a
                forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall a b. b -> Either a b
Right (Stream m b
b_stream, s
s))
            Skip s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall a b. a -> Either a b
Left s
s)
            Step s a
Stop -> forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

    -- XXX flattenArrays is 5x faster than "concatMap fromArray". if somehow we
    -- can get inner_step to inline and fuse here we can perhaps get the same
    -- performance using "concatMap fromArray".
    --
    -- XXX using the pattern synonym "Stream" causes a major performance issue
    -- here even if the synonym does not include an adaptState call. Need to
    -- find out why. Is that something to be fixed in GHC?
    step' State Stream m a
gst (Right (UnStream State Stream m b -> s -> m (Step s b)
inner_step s
inner_st, s
st)) = do
        Step s b
r <- State Stream m b -> s -> m (Step s b)
inner_step (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
inner_st
        case Step s b
r of
            Yield b
b s
inner_s ->
                forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield b
b (forall a b. b -> Either a b
Right (forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m b -> s -> m (Step s b)
inner_step s
inner_s, s
st))
            Skip s
inner_s ->
                forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall a b. b -> Either a b
Right (forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream State Stream m b -> s -> m (Step s b)
inner_step s
inner_s, s
st))
            Step s b
Stop -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall a b. a -> Either a b
Left s
st)

{-# INLINE concatMap #-}
concatMap :: Monad m => (a -> Stream m b) -> Stream m a -> Stream m b
concatMap :: forall (m :: * -> *) a b.
Monad m =>
(a -> Stream m b) -> Stream m a -> Stream m b
concatMap a -> Stream m b
f = forall (m :: * -> *) a b.
Monad m =>
(a -> m (Stream m b)) -> Stream m a -> Stream m b
concatMapM (forall (m :: * -> *) a. Monad m => a -> m a
return forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Stream m b
f)

-- XXX The idea behind this rule is to rewrite any calls to "concatMap
-- fromArray" automatically to flattenArrays which is much faster.  However, we
-- need an INLINE_EARLY on concatMap for this rule to fire. But if we use
-- INLINE_EARLY on concatMap or fromArray then direct uses of
-- "concatMap fromArray" (without the RULE) become much slower, this means
-- "concatMap f" in general would become slower. Need to find a solution to
-- this.
--
-- {-# RULES "concatMap Array.toStreamD"
--      concatMap Array.toStreamD = Array.flattenArray #-}

-- NOTE: even though concatMap for StreamD is 4x faster compared to StreamK,
-- the monad instance does not seem to be significantly faster.
instance Monad m => Monad (Stream m) where
    {-# INLINE return #-}
    return :: forall a. a -> Stream m a
return = forall (f :: * -> *) a. Applicative f => a -> f a
pure

    {-# INLINE (>>=) #-}
    >>= :: forall a b. Stream m a -> (a -> Stream m b) -> Stream m b
(>>=) = forall a b c. (a -> b -> c) -> b -> a -> c
flip forall (m :: * -> *) a b.
Monad m =>
(a -> Stream m b) -> Stream m a -> Stream m b
concatMap

    {-# INLINE (>>) #-}
    >> :: forall a b. Stream m a -> Stream m b -> Stream m b
(>>) = forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
(*>)

------------------------------------------------------------------------------
-- Grouping/Splitting
------------------------------------------------------------------------------

-- s = stream state, fs = fold state
{-# ANN type FoldManyPost Fuse #-}
data FoldManyPost s fs b a
    = FoldManyPostStart s
    | FoldManyPostLoop s fs
    | FoldManyPostYield b (FoldManyPost s fs b a)
    | FoldManyPostDone

-- | Like foldMany but with the following differences:
--
-- * If the stream is empty the default value of the fold would still be
-- emitted in the output.
-- * At the end of the stream if the last application of the fold did not
-- receive any input it would still yield the default fold accumulator as the
-- last value.
--
{-# INLINE_NORMAL foldManyPost #-}
foldManyPost :: Monad m => Fold m a b -> Stream m a -> Stream m b
foldManyPost :: forall (m :: * -> *) a b.
Monad m =>
Fold m a b -> Stream m a -> Stream m b
foldManyPost (Fold s -> a -> m (Step s b)
fstep m (Step s b)
initial s -> m b
extract) (Stream State Stream m a -> s -> m (Step s a)
step s
state) =
    forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream forall {m :: * -> *} {a} {a}.
State Stream m a
-> FoldManyPost s s b a -> m (Step (FoldManyPost s s b a) b)
step' (forall s fs b a. s -> FoldManyPost s fs b a
FoldManyPostStart s
state)

    where

    {-# INLINE consume #-}
    consume :: a -> s -> s -> m (Step (FoldManyPost s s b a) a)
consume a
x s
s s
fs = do
        Step s b
res <- s -> a -> m (Step s b)
fstep s
fs a
x
        forall (m :: * -> *) a. Monad m => a -> m a
return
            forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip
            forall a b. (a -> b) -> a -> b
$ case Step s b
res of
                  FL.Done b
b -> forall s fs b a.
b -> FoldManyPost s fs b a -> FoldManyPost s fs b a
FoldManyPostYield b
b (forall s fs b a. s -> FoldManyPost s fs b a
FoldManyPostStart s
s)
                  FL.Partial s
ps -> forall s fs b a. s -> fs -> FoldManyPost s fs b a
FoldManyPostLoop s
s s
ps

    {-# INLINE_LATE step' #-}
    step' :: State Stream m a
-> FoldManyPost s s b a -> m (Step (FoldManyPost s s b a) b)
step' State Stream m a
_ (FoldManyPostStart s
st) = do
        Step s b
r <- m (Step s b)
initial
        forall (m :: * -> *) a. Monad m => a -> m a
return
            forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip
            forall a b. (a -> b) -> a -> b
$ case Step s b
r of
                  FL.Done b
b -> forall s fs b a.
b -> FoldManyPost s fs b a -> FoldManyPost s fs b a
FoldManyPostYield b
b (forall s fs b a. s -> FoldManyPost s fs b a
FoldManyPostStart s
st)
                  FL.Partial s
fs -> forall s fs b a. s -> fs -> FoldManyPost s fs b a
FoldManyPostLoop s
st s
fs
    step' State Stream m a
gst (FoldManyPostLoop s
st s
fs) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
st
        case Step s a
r of
            Yield a
x s
s -> forall {s} {a} {a}.
a -> s -> s -> m (Step (FoldManyPost s s b a) a)
consume a
x s
s s
fs
            Skip s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall s fs b a. s -> fs -> FoldManyPost s fs b a
FoldManyPostLoop s
s s
fs)
            Step s a
Stop -> do
                b
b <- s -> m b
extract s
fs
                forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall s fs b a.
b -> FoldManyPost s fs b a -> FoldManyPost s fs b a
FoldManyPostYield b
b forall s fs b a. FoldManyPost s fs b a
FoldManyPostDone)
    step' State Stream m a
_ (FoldManyPostYield b
b FoldManyPost s s b a
next) = forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield b
b FoldManyPost s s b a
next
    step' State Stream m a
_ FoldManyPost s s b a
FoldManyPostDone = forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

{-# ANN type FoldMany Fuse #-}
data FoldMany s fs b a
    = FoldManyStart s
    | FoldManyFirst fs s
    | FoldManyLoop s fs
    | FoldManyYield b (FoldMany s fs b a)
    | FoldManyDone

-- | Apply a fold multiple times until the stream ends. If the stream is empty
-- the output would be empty.
--
-- @foldMany f = parseMany (fromFold f)@
--
-- A terminating fold may terminate even without accepting a single input. So
-- we run the fold's initial action before evaluating the stream. However, this
-- means that if later the stream does not yield anything we have to discard
-- the fold's initial result which could have generated an effect.
--
{-# INLINE_NORMAL foldMany #-}
foldMany :: Monad m => Fold m a b -> Stream m a -> Stream m b
foldMany :: forall (m :: * -> *) a b.
Monad m =>
Fold m a b -> Stream m a -> Stream m b
foldMany (Fold s -> a -> m (Step s b)
fstep m (Step s b)
initial s -> m b
extract) (Stream State Stream m a -> s -> m (Step s a)
step s
state) =
    forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream forall {m :: * -> *} {a} {a}.
State Stream m a
-> FoldMany s s b a -> m (Step (FoldMany s s b a) b)
step' (forall s fs b a. s -> FoldMany s fs b a
FoldManyStart s
state)

    where

    {-# INLINE consume #-}
    consume :: a -> s -> s -> m (Step (FoldMany s s b a) a)
consume a
x s
s s
fs = do
        Step s b
res <- s -> a -> m (Step s b)
fstep s
fs a
x
        forall (m :: * -> *) a. Monad m => a -> m a
return
            forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip
            forall a b. (a -> b) -> a -> b
$ case Step s b
res of
                  FL.Done b
b -> forall s fs b a. b -> FoldMany s fs b a -> FoldMany s fs b a
FoldManyYield b
b (forall s fs b a. s -> FoldMany s fs b a
FoldManyStart s
s)
                  FL.Partial s
ps -> forall s fs b a. s -> fs -> FoldMany s fs b a
FoldManyLoop s
s s
ps

    {-# INLINE_LATE step' #-}
    step' :: State Stream m a
-> FoldMany s s b a -> m (Step (FoldMany s s b a) b)
step' State Stream m a
_ (FoldManyStart s
st) = do
        Step s b
r <- m (Step s b)
initial
        forall (m :: * -> *) a. Monad m => a -> m a
return
            forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip
            forall a b. (a -> b) -> a -> b
$ case Step s b
r of
                  FL.Done b
b -> forall s fs b a. b -> FoldMany s fs b a -> FoldMany s fs b a
FoldManyYield b
b (forall s fs b a. s -> FoldMany s fs b a
FoldManyStart s
st)
                  FL.Partial s
fs -> forall s fs b a. fs -> s -> FoldMany s fs b a
FoldManyFirst s
fs s
st
    step' State Stream m a
gst (FoldManyFirst s
fs s
st) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
st
        case Step s a
r of
            Yield a
x s
s -> forall {s} {a} {a}. a -> s -> s -> m (Step (FoldMany s s b a) a)
consume a
x s
s s
fs
            Skip s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall s fs b a. fs -> s -> FoldMany s fs b a
FoldManyFirst s
fs s
s)
            Step s a
Stop -> forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop
    step' State Stream m a
gst (FoldManyLoop s
st s
fs) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
st
        case Step s a
r of
            Yield a
x s
s -> forall {s} {a} {a}. a -> s -> s -> m (Step (FoldMany s s b a) a)
consume a
x s
s s
fs
            Skip s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall s fs b a. s -> fs -> FoldMany s fs b a
FoldManyLoop s
s s
fs)
            Step s a
Stop -> do
                b
b <- s -> m b
extract s
fs
                forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall s fs b a. b -> FoldMany s fs b a -> FoldMany s fs b a
FoldManyYield b
b forall s fs b a. FoldMany s fs b a
FoldManyDone)
    step' State Stream m a
_ (FoldManyYield b
b FoldMany s s b a
next) = forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield b
b FoldMany s s b a
next
    step' State Stream m a
_ FoldMany s s b a
FoldManyDone = forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

{-# INLINE chunksOf #-}
chunksOf :: Monad m => Int -> Fold m a b -> Stream m a -> Stream m b
chunksOf :: forall (m :: * -> *) a b.
Monad m =>
Int -> Fold m a b -> Stream m a -> Stream m b
chunksOf Int
n Fold m a b
f = forall (m :: * -> *) a b.
Monad m =>
Fold m a b -> Stream m a -> Stream m b
foldMany (forall (m :: * -> *) a b.
Monad m =>
Int -> Fold m a b -> Fold m a b
FL.take Int
n Fold m a b
f)

-- Keep the argument order consistent with consumeIterateM.
--
-- | Like 'foldMany' but for the 'Refold' type.  The supplied action is used as
-- the initial value for each refold.
--
-- /Internal/
{-# INLINE_NORMAL refoldMany #-}
refoldMany :: Monad m => Refold m x a b -> m x -> Stream m a -> Stream m b
refoldMany :: forall (m :: * -> *) x a b.
Monad m =>
Refold m x a b -> m x -> Stream m a -> Stream m b
refoldMany (Refold s -> a -> m (Step s b)
fstep x -> m (Step s b)
inject s -> m b
extract) m x
action (Stream State Stream m a -> s -> m (Step s a)
step s
state) =
    forall (m :: * -> *) a s.
(State Stream m a -> s -> m (Step s a)) -> s -> Stream m a
Stream forall {m :: * -> *} {a} {a}.
State Stream m a
-> FoldMany s s b a -> m (Step (FoldMany s s b a) b)
step' (forall s fs b a. s -> FoldMany s fs b a
FoldManyStart s
state)

    where

    {-# INLINE consume #-}
    consume :: a -> s -> s -> m (Step (FoldMany s s b a) a)
consume a
x s
s s
fs = do
        Step s b
res <- s -> a -> m (Step s b)
fstep s
fs a
x
        forall (m :: * -> *) a. Monad m => a -> m a
return
            forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip
            forall a b. (a -> b) -> a -> b
$ case Step s b
res of
                  FL.Done b
b -> forall s fs b a. b -> FoldMany s fs b a -> FoldMany s fs b a
FoldManyYield b
b (forall s fs b a. s -> FoldMany s fs b a
FoldManyStart s
s)
                  FL.Partial s
ps -> forall s fs b a. s -> fs -> FoldMany s fs b a
FoldManyLoop s
s s
ps

    {-# INLINE_LATE step' #-}
    step' :: State Stream m a
-> FoldMany s s b a -> m (Step (FoldMany s s b a) b)
step' State Stream m a
_ (FoldManyStart s
st) = do
        Step s b
r <- m x
action forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= x -> m (Step s b)
inject
        forall (m :: * -> *) a. Monad m => a -> m a
return
            forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip
            forall a b. (a -> b) -> a -> b
$ case Step s b
r of
                  FL.Done b
b -> forall s fs b a. b -> FoldMany s fs b a -> FoldMany s fs b a
FoldManyYield b
b (forall s fs b a. s -> FoldMany s fs b a
FoldManyStart s
st)
                  FL.Partial s
fs -> forall s fs b a. fs -> s -> FoldMany s fs b a
FoldManyFirst s
fs s
st
    step' State Stream m a
gst (FoldManyFirst s
fs s
st) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
st
        case Step s a
r of
            Yield a
x s
s -> forall {s} {a} {a}. a -> s -> s -> m (Step (FoldMany s s b a) a)
consume a
x s
s s
fs
            Skip s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall s fs b a. fs -> s -> FoldMany s fs b a
FoldManyFirst s
fs s
s)
            Step s a
Stop -> forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop
    step' State Stream m a
gst (FoldManyLoop s
st s
fs) = do
        Step s a
r <- State Stream m a -> s -> m (Step s a)
step (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
gst) s
st
        case Step s a
r of
            Yield a
x s
s -> forall {s} {a} {a}. a -> s -> s -> m (Step (FoldMany s s b a) a)
consume a
x s
s s
fs
            Skip s
s -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall s fs b a. s -> fs -> FoldMany s fs b a
FoldManyLoop s
s s
fs)
            Step s a
Stop -> do
                b
b <- s -> m b
extract s
fs
                forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. s -> Step s a
Skip (forall s fs b a. b -> FoldMany s fs b a -> FoldMany s fs b a
FoldManyYield b
b forall s fs b a. FoldMany s fs b a
FoldManyDone)
    step' State Stream m a
_ (FoldManyYield b
b FoldMany s s b a
next) = forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall s a. a -> s -> Step s a
Yield b
b FoldMany s s b a
next
    step' State Stream m a
_ FoldMany s s b a
FoldManyDone = forall (m :: * -> *) a. Monad m => a -> m a
return forall s a. Step s a
Stop

------------------------------------------------------------------------------
-- Other instances
------------------------------------------------------------------------------

instance MonadTrans Stream where
    {-# INLINE lift #-}
    lift :: forall (m :: * -> *) a. Monad m => m a -> Stream m a
lift = forall (m :: * -> *) a. Applicative m => m a -> Stream m a
fromEffect

instance (MonadThrow m) => MonadThrow (Stream m) where
    throwM :: forall e a. Exception e => e -> Stream m a
throwM = forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM