{-# LANGUAGE UndecidableInstances #-}

#include "inline.hs"

-- |
-- Module      : Streamly.Internal.Data.Stream.StreamK.Type
-- Copyright   : (c) 2017 Composewell Technologies
--
-- License     : BSD3
-- Maintainer  : streamly@composewell.com
-- Stability   : experimental
-- Portability : GHC
--
--
-- Continuation passing style (CPS) stream implementation. The symbol 'K' below
-- denotes a function as well as a Kontinuation.
--
module Streamly.Internal.Data.Stream.StreamK.Type
    (
    -- * The stream type
      Stream (..)
    , toStreamK
    , fromStreamK

    -- * foldr/build
    , mkStream
    , foldStream
    , foldStreamShared
    , foldrM
    , foldrS
    , foldrSShared
    , foldrSM
    , build
    , buildS
    , buildM
    , buildSM
    , augmentS
    , augmentSM

    -- * Construction
    , fromStopK
    , fromYieldK
    , consK
    , cons
    , (.:)
    , consM
    , consMBy
    , nil
    , nilM

    -- * Generation
    , fromEffect
    , fromPure
    , unfoldr
    , unfoldrMWith
    , repeat
    , repeatMWith
    , replicateMWith
    , fromIndicesMWith
    , iterateMWith
    , fromFoldable
    , fromFoldableM
    , mfix

    -- * Elimination
    , uncons
    , foldl'
    , foldlx'
    , drain
    , null
    , tail
    , init

    -- * Transformation
    , conjoin
    , serial
    , map
    , mapMWith
    , mapMSerial
    , unShare
    , concatMapWith
    , concatMap
    , bindWith
    , concatPairsWith
    , apWith
    , apSerial
    , apSerialDiscardFst
    , apSerialDiscardSnd
    , foldlS
    , reverse

    -- * Reader
    , withLocal
    )
where

import Control.Monad (ap, (>=>))
import Control.Monad.Reader.Class (MonadReader(..))
import Control.Monad.Trans.Class (MonadTrans(lift))
import Data.Function (fix)
#if __GLASGOW_HASKELL__ < 804
import Data.Semigroup (Semigroup(..))
#endif
import Streamly.Internal.Data.SVar.Type (State, adaptState, defState)
import Prelude hiding
    (map, mapM, concatMap, foldr, repeat, null, reverse, tail, init)
import qualified Prelude

------------------------------------------------------------------------------
-- Basic stream type
------------------------------------------------------------------------------

-- | The type @Stream m a@ represents a monadic stream of values of type 'a'
-- constructed using actions in monad 'm'. It uses stop, singleton and yield
-- continuations equivalent to the following direct style type:
--
-- @
-- data Stream m a = Stop | Singleton a | Yield a (Stream m a)
-- @
--
-- To facilitate parallel composition we maintain a local state in an 'SVar'
-- that is shared across and is used for synchronization of the streams being
-- composed.
--
-- The singleton case can be expressed in terms of stop and yield but we have
-- it as a separate case to optimize composition operations for streams with
-- single element.  We build singleton streams in the implementation of 'pure'
-- for Applicative and Monad, and in 'lift' for MonadTrans.

-- XXX remove the Stream type parameter from State as it is always constant.
-- We can remove it from SVar as well

newtype Stream m a =
    MkStream (forall r.
               State Stream m a         -- state
            -> (a -> Stream m a -> m r) -- yield
            -> (a -> m r)               -- singleton
            -> m r                      -- stop
            -> m r
            )

{-# INLINE fromStreamK #-}
fromStreamK :: Stream m a -> Stream m a
fromStreamK :: forall (m :: * -> *) a. Stream m a -> Stream m a
fromStreamK = forall a. a -> a
id

{-# INLINE toStreamK #-}
toStreamK :: Stream m a -> Stream m a
toStreamK :: forall (m :: * -> *) a. Stream m a -> Stream m a
toStreamK = forall a. a -> a
id

mkStream
    :: (forall r. State Stream m a
        -> (a -> Stream m a -> m r)
        -> (a -> m r)
        -> m r
        -> m r)
    -> Stream m a
mkStream :: 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
mkStream = 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
MkStream

-- | A terminal function that has no continuation to follow.
type StopK m = forall r. m r -> m r

-- | A monadic continuation, it is a function that yields a value of type "a"
-- and calls the argument (a -> m r) as a continuation with that value. We can
-- also think of it as a callback with a handler (a -> m r).  Category
-- theorists call it a codensity type, a special type of right kan extension.
type YieldK m a = forall r. (a -> m r) -> m r

_wrapM :: Monad m => m a -> YieldK m a
_wrapM :: forall (m :: * -> *) a. Monad m => m a -> YieldK m a
_wrapM m a
m = (m a
m forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>=)

-- | Make an empty stream from a stop function.
fromStopK :: StopK m -> Stream m a
fromStopK :: forall (m :: * -> *) a. StopK m -> Stream m a
fromStopK StopK m
k = 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
mkStream forall a b. (a -> b) -> a -> b
$ \State Stream m a
_ a -> Stream m a -> m r
_ a -> m r
_ m r
stp -> StopK m
k m r
stp

-- | Make a singleton stream from a callback function. The callback function
-- calls the one-shot yield continuation to yield an element.
fromYieldK :: YieldK m a -> Stream m a
fromYieldK :: forall (m :: * -> *) a. YieldK m a -> Stream m a
fromYieldK YieldK m a
k = 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
mkStream forall a b. (a -> b) -> a -> b
$ \State Stream m a
_ a -> Stream m a -> m r
_ a -> m r
sng m r
_ -> YieldK m a
k a -> m r
sng

-- | Add a yield function at the head of the stream.
consK :: YieldK m a -> Stream m a -> Stream m a
consK :: forall (m :: * -> *) a. YieldK m a -> Stream m a -> Stream m a
consK YieldK m a
k Stream m a
r = 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
mkStream forall a b. (a -> b) -> a -> b
$ \State Stream m a
_ a -> Stream m a -> m r
yld a -> m r
_ m r
_ -> YieldK m a
k (a -> Stream m a -> m r
`yld` Stream m a
r)

-- XXX Build a stream from a repeating callback function.

------------------------------------------------------------------------------
-- Construction
------------------------------------------------------------------------------

infixr 5 `cons`

-- faster than consM because there is no bind.
-- | Construct a stream by adding a pure value at the head of an existing
-- stream. For serial streams this is the same as @(return a) \`consM` r@ but
-- more efficient. For concurrent streams this is not concurrent whereas
-- 'consM' is concurrent. For example:
--
-- @
-- > toList $ 1 \`cons` 2 \`cons` 3 \`cons` nil
-- [1,2,3]
-- @
--
-- @since 0.1.0
{-# INLINE_NORMAL cons #-}
cons :: a -> Stream m a -> Stream m a
cons :: forall a (m :: * -> *). a -> Stream m a -> Stream m a
cons a
a Stream m a
r = 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
mkStream forall a b. (a -> b) -> a -> b
$ \State Stream m a
_ a -> Stream m a -> m r
yield a -> m r
_ m r
_ -> a -> Stream m a -> m r
yield a
a Stream m a
r

infixr 5 .:

-- | Operator equivalent of 'cons'.
--
-- @
-- > toList $ 1 .: 2 .: 3 .: nil
-- [1,2,3]
-- @
--
-- @since 0.1.1
{-# INLINE (.:) #-}
(.:) :: a -> Stream m a -> Stream m a
.: :: forall a (m :: * -> *). a -> Stream m a -> Stream m a
(.:) = forall a (m :: * -> *). a -> Stream m a -> Stream m a
cons

-- | An empty stream.
--
-- @
-- > toList nil
-- []
-- @
--
-- @since 0.1.0
{-# INLINE_NORMAL nil #-}
nil :: Stream m a
nil :: forall (m :: * -> *) a. Stream m a
nil = 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
mkStream forall a b. (a -> b) -> a -> b
$ \State Stream m a
_ a -> Stream m a -> m r
_ a -> m r
_ m r
stp -> m r
stp

-- | An empty stream producing a side effect.
--
-- @
-- > toList (nilM (print "nil"))
-- "nil"
-- []
-- @
--
-- /Pre-release/
{-# 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.
(forall r.
 State Stream m a
 -> (a -> Stream m a -> m r) -> (a -> m r) -> m r -> m r)
-> Stream m a
mkStream forall a b. (a -> b) -> a -> b
$ \State Stream m a
_ a -> Stream m a -> m r
_ a -> m r
_ m r
stp -> m b
m forall (f :: * -> *) a b. Applicative f => f a -> f b -> f b
*> m r
stp

{-# INLINE_NORMAL fromPure #-}
fromPure :: a -> Stream m a
fromPure :: forall a (m :: * -> *). a -> Stream m a
fromPure a
a = 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
mkStream forall a b. (a -> b) -> a -> b
$ \State Stream m a
_ a -> Stream m a -> m r
_ a -> m r
single m r
_ -> a -> m r
single a
a

{-# INLINE_NORMAL fromEffect #-}
fromEffect :: Monad m => m a -> Stream m a
fromEffect :: forall (m :: * -> *) a. Monad m => m a -> Stream m a
fromEffect m a
m = 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
mkStream forall a b. (a -> b) -> a -> b
$ \State Stream m a
_ a -> Stream m a -> m r
_ a -> m r
single m r
_ -> m a
m forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= a -> m r
single

infixr 5 `consM`

-- NOTE: specializing the function outside the instance definition seems to
-- improve performance quite a bit at times, even if we have the same
-- SPECIALIZE in the instance definition.
{-# INLINE consM #-}
{-# SPECIALIZE consM :: IO a -> Stream IO a -> Stream IO a #-}
consM :: Monad m => m a -> Stream m a -> Stream m a
consM :: forall (m :: * -> *) a. Monad m => m a -> Stream m a -> Stream m a
consM m a
m Stream m a
r = 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
MkStream forall a b. (a -> b) -> a -> b
$ \State Stream m a
_ a -> Stream m a -> m r
yld a -> m r
_ m r
_ -> m a
m forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (a -> Stream m a -> m r
`yld` Stream m a
r)

-- XXX specialize to IO?
{-# INLINE consMBy #-}
consMBy :: Monad m =>
    (Stream m a -> Stream m a -> Stream m a) -> m a -> Stream m a -> Stream m a
consMBy :: forall (m :: * -> *) a.
Monad m =>
(Stream m a -> Stream m a -> Stream m a)
-> m a -> Stream m a -> Stream m a
consMBy Stream m a -> Stream m a -> Stream m a
f m a
m Stream m a
r = forall (m :: * -> *) a. Monad m => m a -> Stream m a
fromEffect m a
m Stream m a -> Stream m a -> Stream m a
`f` Stream m a
r

------------------------------------------------------------------------------
-- Folding a stream
------------------------------------------------------------------------------

-- | Fold a stream by providing an SVar, a stop continuation, a singleton
-- continuation and a yield continuation. The stream would share the current
-- SVar passed via the State.
{-# INLINE_EARLY foldStreamShared #-}
foldStreamShared
    :: State Stream m a
    -> (a -> Stream m a -> m r)
    -> (a -> m r)
    -> m r
    -> Stream m a
    -> m r
foldStreamShared :: forall (m :: * -> *) a r.
State Stream m a
-> (a -> Stream m a -> m r)
-> (a -> m r)
-> m r
-> Stream m a
-> m r
foldStreamShared State Stream m a
s a -> Stream m a -> m r
yield a -> m r
single m r
stop (MkStream forall r.
State Stream m a
-> (a -> Stream m a -> m r) -> (a -> m r) -> m r -> m r
k) = forall r.
State Stream m a
-> (a -> Stream m a -> m r) -> (a -> m r) -> m r -> m r
k State Stream m a
s a -> Stream m a -> m r
yield a -> m r
single m r
stop

-- | Fold a stream by providing a State, stop continuation, a singleton
-- continuation and a yield continuation. The stream will not use the SVar
-- passed via State.
{-# INLINE foldStream #-}
foldStream
    :: State Stream m a
    -> (a -> Stream m a -> m r)
    -> (a -> m r)
    -> m r
    -> Stream m a
    -> m r
foldStream :: forall (m :: * -> *) a r.
State Stream m a
-> (a -> Stream m a -> m r)
-> (a -> m r)
-> m r
-> Stream m a
-> m r
foldStream State Stream m a
s a -> Stream m a -> m r
yield a -> m r
single m r
stop (MkStream forall r.
State Stream m a
-> (a -> Stream m a -> m r) -> (a -> m r) -> m r -> m r
k) =
    forall r.
State Stream m a
-> (a -> Stream m a -> m r) -> (a -> m r) -> m r -> m r
k (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
s) a -> Stream m a -> m r
yield a -> m r
single m r
stop

-------------------------------------------------------------------------------
-- foldr/build fusion
-------------------------------------------------------------------------------

-- XXX perhaps we can just use foldrSM/buildM everywhere as they are more
-- general and cover foldrS/buildS as well.

-- | The function 'f' decides how to reconstruct the stream. We could
-- reconstruct using a shared state (SVar) or without sharing the state.
--
{-# INLINE foldrSWith #-}
foldrSWith ::
    (forall r. State Stream m b
        -> (b -> Stream m b -> m r)
        -> (b -> m r)
        -> m r
        -> Stream m b
        -> m r)
    -> (a -> Stream m b -> Stream m b)
    -> Stream m b
    -> Stream m a
    -> Stream m b
foldrSWith :: forall (m :: * -> *) b a.
(forall r.
 State Stream m b
 -> (b -> Stream m b -> m r)
 -> (b -> m r)
 -> m r
 -> Stream m b
 -> m r)
-> (a -> Stream m b -> Stream m b)
-> Stream m b
-> Stream m a
-> Stream m b
foldrSWith forall r.
State Stream m b
-> (b -> Stream m b -> m r)
-> (b -> m r)
-> m r
-> Stream m b
-> m r
f a -> Stream m b -> Stream m b
step Stream m b
final Stream m a
m = Stream m a -> Stream m b
go Stream m a
m
    where
    go :: Stream m a -> Stream m b
go Stream m a
m1 = 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
mkStream forall a b. (a -> b) -> a -> b
$ \State Stream m b
st b -> Stream m b -> m r
yld b -> m r
sng m r
stp ->
        let run :: Stream m b -> m r
run Stream m b
x = forall r.
State Stream m b
-> (b -> Stream m b -> m r)
-> (b -> m r)
-> m r
-> Stream m b
-> m r
f State Stream m b
st b -> Stream m b -> m r
yld b -> m r
sng m r
stp Stream m b
x
            stop :: m r
stop = Stream m b -> m r
run Stream m b
final
            single :: a -> m r
single a
a = Stream m b -> m r
run forall a b. (a -> b) -> a -> b
$ a -> Stream m b -> Stream m b
step a
a Stream m b
final
            yieldk :: a -> Stream m a -> m r
yieldk a
a Stream m a
r = Stream m b -> m r
run forall a b. (a -> b) -> a -> b
$ a -> Stream m b -> Stream m b
step a
a (Stream m a -> Stream m b
go Stream m a
r)
         -- XXX if type a and b are the same we do not need adaptState, can we
         -- save some perf with that?
         -- XXX since we are using adaptState anyway here we can use
         -- foldStreamShared instead, will that save some perf?
         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
foldStream (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m b
st) a -> Stream m a -> m r
yieldk a -> m r
single m r
stop Stream m a
m1

-- XXX we can use rewrite rules just for foldrSWith, if the function f is the
-- same we can rewrite it.

-- | Fold sharing the SVar state within the reconstructed stream
{-# INLINE_NORMAL foldrSShared #-}
foldrSShared ::
       (a -> Stream m b -> Stream m b)
    -> Stream m b
    -> Stream m a
    -> Stream m b
foldrSShared :: forall a (m :: * -> *) b.
(a -> Stream m b -> Stream m b)
-> Stream m b -> Stream m a -> Stream m b
foldrSShared = forall (m :: * -> *) b a.
(forall r.
 State Stream m b
 -> (b -> Stream m b -> m r)
 -> (b -> m r)
 -> m r
 -> Stream m b
 -> m r)
-> (a -> Stream m b -> Stream m b)
-> Stream m b
-> Stream m a
-> Stream m b
foldrSWith forall (m :: * -> *) a r.
State Stream m a
-> (a -> Stream m a -> m r)
-> (a -> m r)
-> m r
-> Stream m a
-> m r
foldStreamShared

-- XXX consM is a typeclass method, therefore rewritten already. Instead maybe
-- we can make consM polymorphic using rewrite rules.
-- {-# RULES "foldrSShared/id"     foldrSShared consM nil = \x -> x #-}
{-# RULES "foldrSShared/nil"
    forall k z. foldrSShared k z nil = z #-}
{-# RULES "foldrSShared/single"
    forall k z x. foldrSShared k z (fromPure x) = k x z #-}
-- {-# RULES "foldrSShared/app" [1]
--     forall ys. foldrSShared consM ys = \xs -> xs `conjoin` ys #-}

-- | Lazy right associative fold to a stream.
{-# INLINE_NORMAL foldrS #-}
foldrS ::
       (a -> Stream m b -> Stream m b)
    -> Stream m b
    -> Stream m a
    -> Stream m b
foldrS :: forall a (m :: * -> *) b.
(a -> Stream m b -> Stream m b)
-> Stream m b -> Stream m a -> Stream m b
foldrS = forall (m :: * -> *) b a.
(forall r.
 State Stream m b
 -> (b -> Stream m b -> m r)
 -> (b -> m r)
 -> m r
 -> Stream m b
 -> m r)
-> (a -> Stream m b -> Stream m b)
-> Stream m b
-> Stream m a
-> Stream m b
foldrSWith forall (m :: * -> *) a r.
State Stream m a
-> (a -> Stream m a -> m r)
-> (a -> m r)
-> m r
-> Stream m a
-> m r
foldStream

{-# RULES "foldrS/id"     foldrS cons nil = \x -> x #-}
{-# RULES "foldrS/nil"    forall k z.   foldrS k z nil  = z #-}
-- See notes in GHC.Base about this rule
-- {-# RULES "foldr/cons"
--  forall k z x xs. foldrS k z (x `cons` xs) = k x (foldrS k z xs) #-}
{-# RULES "foldrS/single" forall k z x. foldrS k z (fromPure x) = k x z #-}
-- {-# RULES "foldrS/app" [1]
--  forall ys. foldrS cons ys = \xs -> xs `conjoin` ys #-}

-------------------------------------------------------------------------------
-- foldrS with monadic cons i.e. consM
-------------------------------------------------------------------------------

{-# INLINE foldrSMWith #-}
foldrSMWith :: Monad m
    => (forall r. State Stream m b
        -> (b -> Stream m b -> m r)
        -> (b -> m r)
        -> m r
        -> Stream m b
        -> m r)
    -> (m a -> Stream m b -> Stream m b)
    -> Stream m b
    -> Stream m a
    -> Stream m b
foldrSMWith :: forall (m :: * -> *) b a.
Monad m =>
(forall r.
 State Stream m b
 -> (b -> Stream m b -> m r)
 -> (b -> m r)
 -> m r
 -> Stream m b
 -> m r)
-> (m a -> Stream m b -> Stream m b)
-> Stream m b
-> Stream m a
-> Stream m b
foldrSMWith forall r.
State Stream m b
-> (b -> Stream m b -> m r)
-> (b -> m r)
-> m r
-> Stream m b
-> m r
f m a -> Stream m b -> Stream m b
step Stream m b
final Stream m a
m = Stream m a -> Stream m b
go Stream m a
m
    where
    go :: Stream m a -> Stream m b
go Stream m a
m1 = 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
mkStream forall a b. (a -> b) -> a -> b
$ \State Stream m b
st b -> Stream m b -> m r
yld b -> m r
sng m r
stp ->
        let run :: Stream m b -> m r
run Stream m b
x = forall r.
State Stream m b
-> (b -> Stream m b -> m r)
-> (b -> m r)
-> m r
-> Stream m b
-> m r
f State Stream m b
st b -> Stream m b -> m r
yld b -> m r
sng m r
stp Stream m b
x
            stop :: m r
stop = Stream m b -> m r
run Stream m b
final
            single :: a -> m r
single a
a = Stream m b -> m r
run forall a b. (a -> b) -> a -> b
$ m a -> Stream m b -> Stream m b
step (forall (m :: * -> *) a. Monad m => a -> m a
return a
a) Stream m b
final
            yieldk :: a -> Stream m a -> m r
yieldk a
a Stream m a
r = Stream m b -> m r
run forall a b. (a -> b) -> a -> b
$ m a -> Stream m b -> Stream m b
step (forall (m :: * -> *) a. Monad m => a -> m a
return a
a) (Stream m a -> Stream m b
go Stream m a
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
foldStream (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m b
st) a -> Stream m a -> m r
yieldk a -> m r
single m r
stop Stream m a
m1

{-# INLINE_NORMAL foldrSM #-}
foldrSM :: Monad m
    => (m a -> Stream m b -> Stream m b)
    -> Stream m b
    -> Stream m a
    -> Stream m b
foldrSM :: forall (m :: * -> *) a b.
Monad m =>
(m a -> Stream m b -> Stream m b)
-> Stream m b -> Stream m a -> Stream m b
foldrSM = forall (m :: * -> *) b a.
Monad m =>
(forall r.
 State Stream m b
 -> (b -> Stream m b -> m r)
 -> (b -> m r)
 -> m r
 -> Stream m b
 -> m r)
-> (m a -> Stream m b -> Stream m b)
-> Stream m b
-> Stream m a
-> Stream m b
foldrSMWith forall (m :: * -> *) a r.
State Stream m a
-> (a -> Stream m a -> m r)
-> (a -> m r)
-> m r
-> Stream m a
-> m r
foldStream

-- {-# RULES "foldrSM/id"     foldrSM consM nil = \x -> x #-}
{-# RULES "foldrSM/nil"    forall k z.   foldrSM k z nil  = z #-}
{-# RULES "foldrSM/single" forall k z x. foldrSM k z (fromEffect x) = k x z #-}
-- {-# RULES "foldrSM/app" [1]
--  forall ys. foldrSM consM ys = \xs -> xs `conjoin` ys #-}

-- Like foldrSM but sharing the SVar state within the recostructed stream.
{-# INLINE_NORMAL foldrSMShared #-}
foldrSMShared :: Monad m
    => (m a -> Stream m b -> Stream m b)
    -> Stream m b
    -> Stream m a
    -> Stream m b
foldrSMShared :: forall (m :: * -> *) a b.
Monad m =>
(m a -> Stream m b -> Stream m b)
-> Stream m b -> Stream m a -> Stream m b
foldrSMShared = forall (m :: * -> *) b a.
Monad m =>
(forall r.
 State Stream m b
 -> (b -> Stream m b -> m r)
 -> (b -> m r)
 -> m r
 -> Stream m b
 -> m r)
-> (m a -> Stream m b -> Stream m b)
-> Stream m b
-> Stream m a
-> Stream m b
foldrSMWith forall (m :: * -> *) a r.
State Stream m a
-> (a -> Stream m a -> m r)
-> (a -> m r)
-> m r
-> Stream m a
-> m r
foldStreamShared

-- {-# RULES "foldrSM/id"     foldrSM consM nil = \x -> x #-}
{-# RULES "foldrSMShared/nil"
    forall k z. foldrSMShared k z nil = z #-}
{-# RULES "foldrSMShared/single"
    forall k z x. foldrSMShared k z (fromEffect x) = k x z #-}
-- {-# RULES "foldrSM/app" [1]
--  forall ys. foldrSM consM ys = \xs -> xs `conjoin` ys #-}

-------------------------------------------------------------------------------
-- build
-------------------------------------------------------------------------------

{-# INLINE_NORMAL build #-}
build :: forall m a. (forall b. (a -> b -> b) -> b -> b) -> Stream m a
build :: forall (m :: * -> *) a.
(forall b. (a -> b -> b) -> b -> b) -> Stream m a
build forall b. (a -> b -> b) -> b -> b
g = forall b. (a -> b -> b) -> b -> b
g forall a (m :: * -> *). a -> Stream m a -> Stream m a
cons forall (m :: * -> *) a. Stream m a
nil

{-# RULES "foldrM/build"
    forall k z (g :: forall b. (a -> b -> b) -> b -> b).
    foldrM k z (build g) = g k z #-}

{-# RULES "foldrS/build"
      forall k z (g :: forall b. (a -> b -> b) -> b -> b).
      foldrS k z (build g) = g k z #-}

{-# RULES "foldrS/cons/build"
      forall k z x (g :: forall b. (a -> b -> b) -> b -> b).
      foldrS k z (x `cons` build g) = k x (g k z) #-}

{-# RULES "foldrSShared/build"
      forall k z (g :: forall b. (a -> b -> b) -> b -> b).
      foldrSShared k z (build g) = g k z #-}

{-# RULES "foldrSShared/cons/build"
      forall k z x (g :: forall b. (a -> b -> b) -> b -> b).
      foldrSShared k z (x `cons` build g) = k x (g k z) #-}

-- build a stream by applying cons and nil to a build function
{-# INLINE_NORMAL buildS #-}
buildS ::
       ((a -> Stream m a -> Stream m a) -> Stream m a -> Stream m a)
    -> Stream m a
buildS :: forall a (m :: * -> *).
((a -> Stream m a -> Stream m a) -> Stream m a -> Stream m a)
-> Stream m a
buildS (a -> Stream m a -> Stream m a) -> Stream m a -> Stream m a
g = (a -> Stream m a -> Stream m a) -> Stream m a -> Stream m a
g forall a (m :: * -> *). a -> Stream m a -> Stream m a
cons forall (m :: * -> *) a. Stream m a
nil

{-# RULES "foldrS/buildS"
      forall k z
        (g :: (a -> Stream m a -> Stream m a) -> Stream m a -> Stream m a).
      foldrS k z (buildS g) = g k z #-}

{-# RULES "foldrS/cons/buildS"
      forall k z x
        (g :: (a -> Stream m a -> Stream m a) -> Stream m a -> Stream m a).
      foldrS k z (x `cons` buildS g) = k x (g k z) #-}

{-# RULES "foldrSShared/buildS"
      forall k z
        (g :: (a -> Stream m a -> Stream m a) -> Stream m a -> Stream m a).
      foldrSShared k z (buildS g) = g k z #-}

{-# RULES "foldrSShared/cons/buildS"
      forall k z x
        (g :: (a -> Stream m a -> Stream m a) -> Stream m a -> Stream m a).
      foldrSShared k z (x `cons` buildS g) = k x (g k z) #-}

-- build a stream by applying consM and nil to a build function
{-# INLINE_NORMAL buildSM #-}
buildSM :: Monad m
    => ((m a -> Stream m a -> Stream m a) -> Stream m a -> Stream m a)
    -> Stream m a
buildSM :: forall (m :: * -> *) a.
Monad m =>
((m a -> Stream m a -> Stream m a) -> Stream m a -> Stream m a)
-> Stream m a
buildSM (m a -> Stream m a -> Stream m a) -> Stream m a -> Stream m a
g = (m a -> Stream m a -> Stream m a) -> Stream m a -> Stream m a
g forall (m :: * -> *) a. Monad m => m a -> Stream m a -> Stream m a
consM forall (m :: * -> *) a. Stream m a
nil

{-# RULES "foldrSM/buildSM"
     forall k z
        (g :: (m a -> Stream m a -> Stream m a) -> Stream m a -> Stream m a).
     foldrSM k z (buildSM g) = g k z #-}

{-# RULES "foldrSMShared/buildSM"
     forall k z
        (g :: (m a -> Stream m a -> Stream m a) -> Stream m a -> Stream m a).
     foldrSMShared k z (buildSM g) = g k z #-}

-- Disabled because this may not fire as consM is a class Op
{-
{-# RULES "foldrS/consM/buildSM"
      forall k z x (g :: (m a -> t m a -> t m a) -> t m a -> t m a)
    . foldrSM k z (x `consM` buildSM g)
    = k x (g k z)
#-}
-}

-- Build using monadic build functions (continuations) instead of
-- reconstructing a stream.
{-# INLINE_NORMAL buildM #-}
buildM :: Monad m
    => (forall r. (a -> Stream m a -> m r)
        -> (a -> m r)
        -> m r
        -> m r
       )
    -> Stream m a
buildM :: forall (m :: * -> *) a.
Monad m =>
(forall r. (a -> Stream m a -> m r) -> (a -> m r) -> m r -> m r)
-> Stream m a
buildM forall r. (a -> Stream m a -> m r) -> (a -> m r) -> m r -> m r
g = 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
mkStream forall a b. (a -> b) -> a -> b
$ \State Stream m a
st a -> Stream m a -> m r
yld a -> m r
sng m r
stp ->
    forall r. (a -> Stream m a -> m r) -> (a -> m r) -> m r -> m r
g (\a
a Stream m a
r -> forall (m :: * -> *) a r.
State Stream m a
-> (a -> Stream m a -> m r)
-> (a -> m r)
-> m r
-> Stream m a
-> m r
foldStream State Stream m a
st a -> Stream m a -> m r
yld a -> m r
sng m r
stp (forall (m :: * -> *) a. Monad m => a -> m a
return a
a forall (m :: * -> *) a. Monad m => m a -> Stream m a -> Stream m a
`consM` Stream m a
r)) a -> m r
sng m r
stp

-- | Like 'buildM' but shares the SVar state across computations.
{-# INLINE_NORMAL sharedMWith #-}
sharedMWith :: Monad m
    => (m a -> Stream m a -> Stream m a)
    -> (forall r. (a -> Stream m a -> m r)
        -> (a -> m r)
        -> m r
        -> m r
       )
    -> Stream m a
sharedMWith :: forall (m :: * -> *) a.
Monad m =>
(m a -> Stream m a -> Stream m a)
-> (forall r. (a -> Stream m a -> m r) -> (a -> m r) -> m r -> m r)
-> Stream m a
sharedMWith m a -> Stream m a -> Stream m a
cns forall r. (a -> Stream m a -> m r) -> (a -> m r) -> m r -> m r
g = 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
mkStream forall a b. (a -> b) -> a -> b
$ \State Stream m a
st a -> Stream m a -> m r
yld a -> m r
sng m r
stp ->
    forall r. (a -> Stream m a -> m r) -> (a -> m r) -> m r -> m r
g (\a
a Stream m a
r -> forall (m :: * -> *) a r.
State Stream m a
-> (a -> Stream m a -> m r)
-> (a -> m r)
-> m r
-> Stream m a
-> m r
foldStreamShared State Stream m a
st a -> Stream m a -> m r
yld a -> m r
sng m r
stp (forall (m :: * -> *) a. Monad m => a -> m a
return a
a m a -> Stream m a -> Stream m a
`cns` Stream m a
r)) a -> m r
sng m r
stp

-------------------------------------------------------------------------------
-- augment
-------------------------------------------------------------------------------

{-# INLINE_NORMAL augmentS #-}
augmentS ::
       ((a -> Stream m a -> Stream m a) -> Stream m a -> Stream m a)
    -> Stream m a
    -> Stream m a
augmentS :: forall a (m :: * -> *).
((a -> Stream m a -> Stream m a) -> Stream m a -> Stream m a)
-> Stream m a -> Stream m a
augmentS (a -> Stream m a -> Stream m a) -> Stream m a -> Stream m a
g Stream m a
xs = (a -> Stream m a -> Stream m a) -> Stream m a -> Stream m a
g forall a (m :: * -> *). a -> Stream m a -> Stream m a
cons Stream m a
xs

{-# RULES "augmentS/nil"
    forall (g :: (a -> Stream m a -> Stream m a) -> Stream m a -> Stream m a).
    augmentS g nil = buildS g
    #-}

{-# RULES "foldrS/augmentS"
    forall k z xs
        (g :: (a -> Stream m a -> Stream m a) -> Stream m a -> Stream m a).
    foldrS k z (augmentS g xs) = g k (foldrS k z xs)
    #-}

{-# RULES "augmentS/buildS"
    forall (g :: (a -> Stream m a -> Stream m a) -> Stream m a -> Stream m a)
           (h :: (a -> Stream m a -> Stream m a) -> Stream m a -> Stream m a).
    augmentS g (buildS h) = buildS (\c n -> g c (h c n))
    #-}

{-# INLINE_NORMAL augmentSM #-}
augmentSM :: Monad m =>
       ((m a -> Stream m a -> Stream m a) -> Stream m a -> Stream m a)
    -> Stream m a -> Stream m a
augmentSM :: forall (m :: * -> *) a.
Monad m =>
((m a -> Stream m a -> Stream m a) -> Stream m a -> Stream m a)
-> Stream m a -> Stream m a
augmentSM (m a -> Stream m a -> Stream m a) -> Stream m a -> Stream m a
g Stream m a
xs = (m a -> Stream m a -> Stream m a) -> Stream m a -> Stream m a
g forall (m :: * -> *) a. Monad m => m a -> Stream m a -> Stream m a
consM Stream m a
xs

{-# RULES "augmentSM/nil"
    forall
        (g :: (m a -> Stream m a -> Stream m a) -> Stream m a -> Stream m a).
    augmentSM g nil = buildSM g
    #-}

{-# RULES "foldrSM/augmentSM"
    forall k z xs
        (g :: (m a -> Stream m a -> Stream m a) -> Stream m a -> Stream m a).
    foldrSM k z (augmentSM g xs) = g k (foldrSM k z xs)
    #-}

{-# RULES "augmentSM/buildSM"
    forall
        (g :: (m a -> Stream m a -> Stream m a) -> Stream m a -> Stream m a)
        (h :: (m a -> Stream m a -> Stream m a) -> Stream m a -> Stream m a).
    augmentSM g (buildSM h) = buildSM (\c n -> g c (h c n))
    #-}

-------------------------------------------------------------------------------
-- Experimental foldrM/buildM
-------------------------------------------------------------------------------

-- | Lazy right fold with a monadic step function.
{-# INLINE_NORMAL foldrM #-}
foldrM :: (a -> m b -> m b) -> m b -> Stream m a -> m b
foldrM :: forall a (m :: * -> *) b.
(a -> m b -> m b) -> m b -> Stream m a -> m b
foldrM a -> m b -> m b
step m b
acc Stream m a
m = Stream m a -> m b
go Stream m a
m
    where
    go :: Stream m a -> m b
go Stream m a
m1 =
        let stop :: m b
stop = m b
acc
            single :: a -> m b
single a
a = a -> m b -> m b
step a
a m b
acc
            yieldk :: a -> Stream m a -> m b
yieldk a
a Stream m a
r = a -> m b -> m b
step a
a (Stream m a -> m b
go Stream m a
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
foldStream forall (t :: (* -> *) -> * -> *) (m :: * -> *) a. State t m a
defState a -> Stream m a -> m b
yieldk a -> m b
single m b
stop Stream m a
m1

{-# INLINE_NORMAL foldrMKWith #-}
foldrMKWith
    :: (State Stream m a
        -> (a -> Stream m a -> m b)
        -> (a -> m b)
        -> m b
        -> Stream m a
        -> m b)
    -> (a -> m b -> m b)
    -> m b
    -> ((a -> Stream m a -> m b) -> (a -> m b) -> m b -> m b)
    -> m b
foldrMKWith :: forall (m :: * -> *) a b.
(State Stream m a
 -> (a -> Stream m a -> m b)
 -> (a -> m b)
 -> m b
 -> Stream m a
 -> m b)
-> (a -> m b -> m b)
-> m b
-> ((a -> Stream m a -> m b) -> (a -> m b) -> m b -> m b)
-> m b
foldrMKWith State Stream m a
-> (a -> Stream m a -> m b)
-> (a -> m b)
-> m b
-> Stream m a
-> m b
f a -> m b -> m b
step m b
acc = ((a -> Stream m a -> m b) -> (a -> m b) -> m b -> m b) -> m b
go
    where
    go :: ((a -> Stream m a -> m b) -> (a -> m b) -> m b -> m b) -> m b
go (a -> Stream m a -> m b) -> (a -> m b) -> m b -> m b
k =
        let stop :: m b
stop = m b
acc
            single :: a -> m b
single a
a = a -> m b -> m b
step a
a m b
acc
            yieldk :: a -> Stream m a -> m b
yieldk a
a Stream m a
r = a -> m b -> m b
step a
a (((a -> Stream m a -> m b) -> (a -> m b) -> m b -> m b) -> m b
go (\a -> Stream m a -> m b
yld a -> m b
sng m b
stp -> State Stream m a
-> (a -> Stream m a -> m b)
-> (a -> m b)
-> m b
-> Stream m a
-> m b
f forall (t :: (* -> *) -> * -> *) (m :: * -> *) a. State t m a
defState a -> Stream m a -> m b
yld a -> m b
sng m b
stp Stream m a
r))
        in (a -> Stream m a -> m b) -> (a -> m b) -> m b -> m b
k a -> Stream m a -> m b
yieldk a -> m b
single m b
stop

{-
{-# RULES "foldrM/buildS"
      forall k z (g :: (a -> t m a -> t m a) -> t m a -> t m a)
    . foldrM k z (buildS g)
    = g k z
#-}
-}
-- XXX in which case will foldrM/buildM fusion be useful?
{-# RULES "foldrM/buildM"
    forall step acc (g :: (forall r.
           (a -> Stream m a -> m r)
        -> (a -> m r)
        -> m r
        -> m r
       )).
    foldrM step acc (buildM g) = foldrMKWith foldStream step acc g
    #-}

{-
{-# RULES "foldrM/sharedM"
    forall step acc (g :: (forall r.
           (a -> Stream m a -> m r)
        -> (a -> m r)
        -> m r
        -> m r
       )).
    foldrM step acc (sharedM g) = foldrMKWith foldStreamShared step acc g
    #-}
-}

------------------------------------------------------------------------------
-- Left fold
------------------------------------------------------------------------------

-- | Strict left fold with an extraction function. Like the standard strict
-- left fold, but applies a user supplied extraction function (the third
-- argument) to the folded value at the end. This is designed to work with the
-- @foldl@ library. The suffix @x@ is a mnemonic for extraction.
--
-- Note that the accumulator is always evaluated including the initial value.
{-# INLINE foldlx' #-}
foldlx' :: forall m a b x. Monad m
    => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> m b
foldlx' :: forall (m :: * -> *) a b x.
Monad m =>
(x -> a -> x) -> x -> (x -> b) -> Stream m a -> m b
foldlx' x -> a -> x
step x
begin x -> b
done Stream m a
m = Stream m x -> m b
get forall a b. (a -> b) -> a -> b
$ Stream m a -> x -> Stream m x
go Stream m a
m x
begin
    where
    {-# NOINLINE get #-}
    get :: Stream m x -> m b
    get :: Stream m x -> m b
get Stream m x
m1 =
        -- XXX we are not strictly evaluating the accumulator here. Is this
        -- okay?
        let single :: x -> m b
single = forall (m :: * -> *) a. Monad m => a -> m a
return forall b c a. (b -> c) -> (a -> b) -> a -> c
. x -> b
done
        -- XXX this is foldSingleton. why foldStreamShared?
         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
foldStreamShared forall a. HasCallStack => a
undefined forall a. HasCallStack => a
undefined x -> m b
single forall a. HasCallStack => a
undefined Stream m x
m1

    -- Note, this can be implemented by making a recursive call to "go",
    -- however that is more expensive because of unnecessary recursion
    -- that cannot be tail call optimized. Unfolding recursion explicitly via
    -- continuations is much more efficient.
    go :: Stream m a -> x -> Stream m x
    go :: Stream m a -> x -> Stream m x
go Stream m a
m1 !x
acc = 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
mkStream forall a b. (a -> b) -> a -> b
$ \State Stream m x
_ x -> Stream m x -> m r
yld x -> m r
sng m r
_ ->
        let stop :: m r
stop = x -> m r
sng x
acc
            single :: a -> m r
single a
a = x -> m r
sng forall a b. (a -> b) -> a -> b
$ x -> a -> x
step x
acc a
a
            -- XXX this is foldNonEmptyStream
            yieldk :: a -> Stream m a -> m r
yieldk a
a Stream m a
r = forall (m :: * -> *) a r.
State Stream m a
-> (a -> Stream m a -> m r)
-> (a -> m r)
-> m r
-> Stream m a
-> m r
foldStream forall (t :: (* -> *) -> * -> *) (m :: * -> *) a. State t m a
defState x -> Stream m x -> m r
yld x -> m r
sng forall a. HasCallStack => a
undefined forall a b. (a -> b) -> a -> b
$
                Stream m a -> x -> Stream m x
go Stream m a
r (x -> a -> x
step x
acc a
a)
        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
foldStream forall (t :: (* -> *) -> * -> *) (m :: * -> *) a. State t m a
defState a -> Stream m a -> m r
yieldk a -> m r
single m r
stop Stream m a
m1

-- | Strict left associative fold.
{-# 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
step b
begin = forall (m :: * -> *) a b x.
Monad m =>
(x -> a -> x) -> x -> (x -> b) -> Stream m a -> m b
foldlx' b -> a -> b
step b
begin forall a. a -> a
id

------------------------------------------------------------------------------
-- Specialized folds
------------------------------------------------------------------------------

-- XXX use foldrM to implement folds where possible
-- XXX This (commented) definition of drain and mapM_ perform much better on
-- some benchmarks but worse on others. Need to investigate why, may there is
-- an optimization opportunity that we can exploit.
-- drain = foldrM (\_ xs -> return () >> xs) (return ())

-- |
-- > drain = foldl' (\_ _ -> ()) ()
-- > drain = mapM_ (\_ -> return ())
{-# INLINE drain #-}
drain :: Monad m => Stream m a -> m ()
drain :: forall (m :: * -> *) a. Monad m => Stream m a -> m ()
drain = forall a (m :: * -> *) b.
(a -> m b -> m b) -> m b -> Stream m a -> m b
foldrM (\a
_ m ()
xs -> m ()
xs) (forall (m :: * -> *) a. Monad m => a -> m a
return ())
{-
drain = go
    where
    go m1 =
        let stop = return ()
            single _ = return ()
            yieldk _ r = go r
         in foldStream defState yieldk single stop m1
-}

{-# INLINE null #-}
null :: Monad m => Stream m a -> m Bool
-- null = foldrM (\_ _ -> return True) (return False)
null :: forall (m :: * -> *) a. Monad m => Stream m a -> m Bool
null Stream m a
m =
    let stop :: m Bool
stop      = forall (m :: * -> *) a. Monad m => a -> m a
return Bool
True
        single :: p -> m Bool
single p
_  = forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False
        yieldk :: p -> p -> m Bool
yieldk p
_ p
_ = forall (m :: * -> *) a. Monad m => a -> m a
return Bool
False
    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
foldStream forall (t :: (* -> *) -> * -> *) (m :: * -> *) a. State t m a
defState forall {m :: * -> *} {p} {p}. Monad m => p -> p -> m Bool
yieldk forall {m :: * -> *} {p}. Monad m => p -> m Bool
single m Bool
stop Stream m a
m

------------------------------------------------------------------------------
-- Semigroup
------------------------------------------------------------------------------

infixr 6 `serial`

-- | Appends two streams sequentially, yielding all elements from the first
-- stream, and then all elements from the second stream.
--
{-# INLINE serial #-}
serial :: Stream m a -> Stream m a -> Stream m a
-- XXX This doubles the time of toNullAp benchmark, may not be fusing properly
-- serial xs ys = augmentS (\c n -> foldrS c n xs) ys
serial :: forall (m :: * -> *) a. Stream m a -> Stream m a -> Stream m a
serial Stream m a
m1 Stream m a
m2 = Stream m a -> Stream m a
go Stream m a
m1
    where
    go :: Stream m a -> Stream m a
go Stream m a
m = 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
mkStream forall a b. (a -> b) -> a -> b
$ \State Stream m a
st a -> Stream m a -> m r
yld a -> m r
sng m r
stp ->
               let stop :: m r
stop       = forall (m :: * -> *) a r.
State Stream m a
-> (a -> Stream m a -> m r)
-> (a -> m r)
-> m r
-> Stream m a
-> m r
foldStream State Stream m a
st a -> Stream m a -> m r
yld a -> m r
sng m r
stp Stream m a
m2
                   single :: a -> m r
single a
a   = a -> Stream m a -> m r
yld a
a Stream m a
m2
                   yieldk :: a -> Stream m a -> m r
yieldk a
a Stream m a
r = a -> Stream m a -> m r
yld a
a (Stream m a -> Stream m a
go Stream m a
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
foldStream State Stream m a
st a -> Stream m a -> m r
yieldk a -> m r
single m r
stop Stream m a
m

-- join/merge/append streams depending on consM
{-# INLINE conjoin #-}
conjoin :: Monad m => Stream m a -> Stream m a -> Stream m a
conjoin :: forall (m :: * -> *) a.
Monad m =>
Stream m a -> Stream m a -> Stream m a
conjoin Stream m a
xs = forall (m :: * -> *) a.
Monad m =>
((m a -> Stream m a -> Stream m a) -> Stream m a -> Stream m a)
-> Stream m a -> Stream m a
augmentSM (\m a -> Stream m a -> Stream m a
c Stream m a
n -> forall (m :: * -> *) a b.
Monad m =>
(m a -> Stream m b -> Stream m b)
-> Stream m b -> Stream m a -> Stream m b
foldrSM m a -> Stream m a -> Stream m a
c Stream m a
n Stream m a
xs)

instance Semigroup (Stream m a) where
    <> :: Stream m a -> Stream m a -> Stream m a
(<>) = forall (m :: * -> *) a. Stream m a -> Stream m a -> Stream m a
serial

------------------------------------------------------------------------------
-- Monoid
------------------------------------------------------------------------------

instance Monoid (Stream m a) where
    mempty :: Stream m a
mempty = forall (m :: * -> *) a. Stream m a
nil
    mappend :: Stream m a -> Stream m a -> Stream m a
mappend = forall a. Semigroup a => a -> a -> a
(<>)

-------------------------------------------------------------------------------
-- Functor
-------------------------------------------------------------------------------

-- Note eta expanded
{-# INLINE_LATE mapFB #-}
mapFB :: forall b m a.
       (b -> Stream m b -> Stream m b)
    -> (a -> b)
    -> a
    -> Stream m b
    -> Stream m b
mapFB :: forall b (m :: * -> *) a.
(b -> Stream m b -> Stream m b)
-> (a -> b) -> a -> Stream m b -> Stream m b
mapFB b -> Stream m b -> Stream m b
c a -> b
f = \a
x Stream m b
ys -> b -> Stream m b -> Stream m b
c (a -> b
f a
x) Stream m b
ys

{-# RULES
"mapFB/mapFB" forall c f g. mapFB (mapFB c f) g = mapFB c (f . g)
"mapFB/id"    forall c.     mapFB c (\x -> x)   = c
    #-}

{-# INLINE map #-}
map :: (a -> b) -> Stream m a -> Stream m b
map :: forall a b (m :: * -> *). (a -> b) -> Stream m a -> Stream m b
map a -> b
f Stream m a
xs = forall a (m :: * -> *).
((a -> Stream m a -> Stream m a) -> Stream m a -> Stream m a)
-> Stream m a
buildS (\b -> Stream m b -> Stream m b
c Stream m b
n -> forall a (m :: * -> *) b.
(a -> Stream m b -> Stream m b)
-> Stream m b -> Stream m a -> Stream m b
foldrS (forall b (m :: * -> *) a.
(b -> Stream m b -> Stream m b)
-> (a -> b) -> a -> Stream m b -> Stream m b
mapFB b -> Stream m b -> Stream m b
c a -> b
f) Stream m b
n Stream m a
xs)

-- XXX This definition might potentially be more efficient, but the cost in the
-- benchmark is dominated by unfoldrM cost so we cannot correctly determine
-- differences in the mapping cost. We should perhaps deduct the cost of
-- unfoldrM from the benchmarks and then compare.
{-
map f m = go m
    where
        go m1 =
            mkStream $ \st yld sng stp ->
            let single     = sng . f
                yieldk a r = yld (f a) (go r)
            in foldStream (adaptState st) yieldk single stp m1
-}

{-# INLINE_LATE mapMFB #-}
mapMFB :: Monad m => (m b -> t m b -> t m b) -> (a -> m b) -> m a -> t m b -> t m b
mapMFB :: forall (m :: * -> *) b (t :: (* -> *) -> * -> *) a.
Monad m =>
(m b -> t m b -> t m b) -> (a -> m b) -> m a -> t m b -> t m b
mapMFB m b -> t m b -> t m b
c a -> m b
f m a
x = m b -> t m b -> t m b
c (m a
x forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= a -> m b
f)

{-# RULES
    "mapMFB/mapMFB" forall c f g. mapMFB (mapMFB c f) g = mapMFB c (f >=> g)
    #-}
-- XXX These rules may never fire because pure/return type class rules will
-- fire first.
{-
"mapMFB/pure"    forall c.     mapMFB c (\x -> pure x)   = c
"mapMFB/return"  forall c.     mapMFB c (\x -> return x) = c
-}

-- This is experimental serial version supporting fusion.
--
-- XXX what if we do not want to fuse two concurrent mapMs?
-- XXX we can combine two concurrent mapM only if the SVar is of the same type
-- So for now we use it only for serial streams.
-- XXX fusion would be easier for monomoprhic stream types.
-- {-# RULES "mapM serial" mapM = mapMSerial #-}
{-# INLINE mapMSerial #-}
mapMSerial :: Monad m => (a -> m b) -> Stream m a -> Stream m b
mapMSerial :: forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Stream m a -> Stream m b
mapMSerial a -> m b
f Stream m a
xs = forall (m :: * -> *) a.
Monad m =>
((m a -> Stream m a -> Stream m a) -> Stream m a -> Stream m a)
-> Stream m a
buildSM (\m b -> Stream m b -> Stream m b
c Stream m b
n -> forall (m :: * -> *) a b.
Monad m =>
(m a -> Stream m b -> Stream m b)
-> Stream m b -> Stream m a -> Stream m b
foldrSMShared (forall (m :: * -> *) b (t :: (* -> *) -> * -> *) a.
Monad m =>
(m b -> t m b -> t m b) -> (a -> m b) -> m a -> t m b -> t m b
mapMFB m b -> Stream m b -> Stream m b
c a -> m b
f) Stream m b
n Stream m a
xs)

{-# INLINE mapMWith #-}
mapMWith ::
       (m b -> Stream m b -> Stream m b)
    -> (a -> m b)
    -> Stream m a
    -> Stream m b
mapMWith :: forall (m :: * -> *) b a.
(m b -> Stream m b -> Stream m b)
-> (a -> m b) -> Stream m a -> Stream m b
mapMWith m b -> Stream m b -> Stream m b
cns a -> m b
f = forall a (m :: * -> *) b.
(a -> Stream m b -> Stream m b)
-> Stream m b -> Stream m a -> Stream m b
foldrSShared (\a
x Stream m b
xs -> a -> m b
f a
x m b -> Stream m b -> Stream m b
`cns` Stream m b
xs) forall (m :: * -> *) a. Stream m a
nil

{-
-- See note under map definition above.
mapMWith cns f = go
    where
    go m1 = mkStream $ \st yld sng stp ->
        let single a  = f a >>= sng
            yieldk a r = foldStreamShared st yld sng stp $ f a `cns` go r
         in foldStream (adaptState st) yieldk single stp m1
-}

-- XXX in fact use the Stream type everywhere and only use polymorphism in the
-- high level modules/prelude.
instance Monad m => Functor (Stream m) where
    fmap :: forall a b. (a -> b) -> Stream m a -> Stream m b
fmap = forall a b (m :: * -> *). (a -> b) -> Stream m a -> Stream m b
map

-------------------------------------------------------------------------------
-- Transformers
-------------------------------------------------------------------------------

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

-------------------------------------------------------------------------------
-- Nesting
-------------------------------------------------------------------------------

-- | Detach a stream from an SVar
{-# INLINE unShare #-}
unShare :: Stream m a -> Stream m a
unShare :: forall (m :: * -> *) a. Stream m a -> Stream m a
unShare Stream m a
x = 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
mkStream forall a b. (a -> b) -> a -> b
$ \State Stream m a
st a -> Stream m a -> m r
yld a -> m r
sng m r
stp ->
    forall (m :: * -> *) a r.
State Stream m a
-> (a -> Stream m a -> m r)
-> (a -> m r)
-> m r
-> Stream m a
-> m r
foldStream State Stream m a
st a -> Stream m a -> m r
yld a -> m r
sng m r
stp Stream m a
x

-- XXX the function stream and value stream can run in parallel
{-# INLINE apWith #-}
apWith ::
       (Stream m b -> Stream m b -> Stream m b)
    -> Stream m (a -> b)
    -> Stream m a
    -> Stream m b
apWith :: forall (m :: * -> *) b a.
(Stream m b -> Stream m b -> Stream m b)
-> Stream m (a -> b) -> Stream m a -> Stream m b
apWith Stream m b -> Stream m b -> Stream m b
par Stream m (a -> b)
fstream Stream m a
stream = Stream m (a -> b) -> Stream m b
go1 Stream m (a -> b)
fstream

    where

    go1 :: Stream m (a -> b) -> Stream m b
go1 Stream m (a -> b)
m =
        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
mkStream forall a b. (a -> b) -> a -> b
$ \State Stream m b
st b -> Stream m b -> m r
yld b -> m r
sng m r
stp ->
            let foldShared :: Stream m b -> m r
foldShared = forall (m :: * -> *) a r.
State Stream m a
-> (a -> Stream m a -> m r)
-> (a -> m r)
-> m r
-> Stream m a
-> m r
foldStreamShared State Stream m b
st b -> Stream m b -> m r
yld b -> m r
sng m r
stp
                single :: (a -> b) -> m r
single a -> b
f   = Stream m b -> m r
foldShared forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. Stream m a -> Stream m a
unShare (forall a b (m :: * -> *). (a -> b) -> Stream m a -> Stream m b
go2 a -> b
f Stream m a
stream)
                yieldk :: (a -> b) -> Stream m (a -> b) -> m r
yieldk a -> b
f Stream m (a -> b)
r = Stream m b -> m r
foldShared forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. Stream m a -> Stream m a
unShare (forall a b (m :: * -> *). (a -> b) -> Stream m a -> Stream m b
go2 a -> b
f Stream m a
stream) Stream m b -> Stream m b -> Stream m b
`par` Stream m (a -> b) -> Stream m b
go1 Stream m (a -> b)
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
foldStream (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m b
st) (a -> b) -> Stream m (a -> b) -> m r
yieldk (a -> b) -> m r
single m r
stp Stream m (a -> b)
m

    go2 :: (t -> a) -> Stream m t -> Stream m a
go2 t -> a
f Stream m t
m =
        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
mkStream forall a b. (a -> b) -> a -> b
$ \State Stream m a
st a -> Stream m a -> m r
yld a -> m r
sng m r
stp ->
            let single :: t -> m r
single t
a   = a -> m r
sng (t -> a
f t
a)
                yieldk :: t -> Stream m t -> m r
yieldk t
a Stream m t
r = a -> Stream m a -> m r
yld (t -> a
f t
a) ((t -> a) -> Stream m t -> Stream m a
go2 t -> a
f Stream m t
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
foldStream (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
st) t -> Stream m t -> m r
yieldk t -> m r
single m r
stp Stream m t
m

{-# INLINE apSerial #-}
apSerial ::
       Stream m (a -> b)
    -> Stream m a
    -> Stream m b
apSerial :: forall (m :: * -> *) a b.
Stream m (a -> b) -> Stream m a -> Stream m b
apSerial Stream m (a -> b)
fstream Stream m a
stream = forall {a}. Stream m (a -> a) -> Stream m a
go1 Stream m (a -> b)
fstream

    where

    go1 :: Stream m (a -> a) -> Stream m a
go1 Stream m (a -> a)
m =
        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
mkStream forall a b. (a -> b) -> a -> b
$ \State Stream m a
st a -> Stream m a -> m r
yld a -> m r
sng m r
stp ->
            let foldShared :: Stream m a -> m r
foldShared = forall (m :: * -> *) a r.
State Stream m a
-> (a -> Stream m a -> m r)
-> (a -> m r)
-> m r
-> Stream m a
-> m r
foldStreamShared State Stream m a
st a -> Stream m a -> m r
yld a -> m r
sng m r
stp
                single :: (a -> a) -> m r
single a -> a
f   = Stream m a -> m r
foldShared forall a b. (a -> b) -> a -> b
$ forall a b (m :: * -> *). (a -> b) -> Stream m a -> Stream m b
go3 a -> a
f Stream m a
stream
                yieldk :: (a -> a) -> Stream m (a -> a) -> m r
yieldk a -> a
f Stream m (a -> a)
r = Stream m a -> m r
foldShared forall a b. (a -> b) -> a -> b
$ (a -> a) -> Stream m (a -> a) -> Stream m a -> Stream m a
go2 a -> a
f Stream m (a -> a)
r Stream m a
stream
            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
foldStream (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
st) (a -> a) -> Stream m (a -> a) -> m r
yieldk (a -> a) -> m r
single m r
stp Stream m (a -> a)
m

    go2 :: (a -> a) -> Stream m (a -> a) -> Stream m a -> Stream m a
go2 a -> a
f Stream m (a -> a)
r1 Stream m a
m =
        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
mkStream forall a b. (a -> b) -> a -> b
$ \State Stream m a
st a -> Stream m a -> m r
yld a -> m r
sng m r
stp ->
            let foldShared :: Stream m a -> m r
foldShared = forall (m :: * -> *) a r.
State Stream m a
-> (a -> Stream m a -> m r)
-> (a -> m r)
-> m r
-> Stream m a
-> m r
foldStreamShared State Stream m a
st a -> Stream m a -> m r
yld a -> m r
sng m r
stp
                stop :: m r
stop = Stream m a -> m r
foldShared forall a b. (a -> b) -> a -> b
$ Stream m (a -> a) -> Stream m a
go1 Stream m (a -> a)
r1
                single :: a -> m r
single a
a   = a -> Stream m a -> m r
yld (a -> a
f a
a) (Stream m (a -> a) -> Stream m a
go1 Stream m (a -> a)
r1)
                yieldk :: a -> Stream m a -> m r
yieldk a
a Stream m a
r = a -> Stream m a -> m r
yld (a -> a
f a
a) ((a -> a) -> Stream m (a -> a) -> Stream m a -> Stream m a
go2 a -> a
f Stream m (a -> a)
r1 Stream m a
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
foldStream (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
st) a -> Stream m a -> m r
yieldk a -> m r
single m r
stop Stream m a
m

    go3 :: (t -> a) -> Stream m t -> Stream m a
go3 t -> a
f Stream m t
m =
        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
mkStream forall a b. (a -> b) -> a -> b
$ \State Stream m a
st a -> Stream m a -> m r
yld a -> m r
sng m r
stp ->
            let single :: t -> m r
single t
a   = a -> m r
sng (t -> a
f t
a)
                yieldk :: t -> Stream m t -> m r
yieldk t
a Stream m t
r = a -> Stream m a -> m r
yld (t -> a
f t
a) ((t -> a) -> Stream m t -> Stream m a
go3 t -> a
f Stream m t
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
foldStream (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
st) t -> Stream m t -> m r
yieldk t -> m r
single m r
stp Stream m t
m

{-# INLINE apSerialDiscardFst #-}
apSerialDiscardFst ::
       Stream m a
    -> Stream m b
    -> Stream m b
apSerialDiscardFst :: forall (m :: * -> *) a b. Stream m a -> Stream m b -> Stream m b
apSerialDiscardFst Stream m a
fstream Stream m b
stream = forall {a}. Stream m a -> Stream m b
go1 Stream m a
fstream

    where

    go1 :: Stream m a -> Stream m b
go1 Stream m a
m =
        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
mkStream forall a b. (a -> b) -> a -> b
$ \State Stream m b
st b -> Stream m b -> m r
yld b -> m r
sng m r
stp ->
            let foldShared :: Stream m b -> m r
foldShared = forall (m :: * -> *) a r.
State Stream m a
-> (a -> Stream m a -> m r)
-> (a -> m r)
-> m r
-> Stream m a
-> m r
foldStreamShared State Stream m b
st b -> Stream m b -> m r
yld b -> m r
sng m r
stp
                single :: p -> m r
single p
_   = Stream m b -> m r
foldShared Stream m b
stream
                yieldk :: p -> Stream m a -> m r
yieldk p
_ Stream m a
r = Stream m b -> m r
foldShared forall a b. (a -> b) -> a -> b
$ Stream m a -> Stream m b -> Stream m b
go2 Stream m a
r Stream m b
stream
            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
foldStream (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m b
st) forall {p}. p -> Stream m a -> m r
yieldk forall {p}. p -> m r
single m r
stp Stream m a
m

    go2 :: Stream m a -> Stream m b -> Stream m b
go2 Stream m a
r1 Stream m b
m =
        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
mkStream forall a b. (a -> b) -> a -> b
$ \State Stream m b
st b -> Stream m b -> m r
yld b -> m r
sng m r
stp ->
            let foldShared :: Stream m b -> m r
foldShared = forall (m :: * -> *) a r.
State Stream m a
-> (a -> Stream m a -> m r)
-> (a -> m r)
-> m r
-> Stream m a
-> m r
foldStreamShared State Stream m b
st b -> Stream m b -> m r
yld b -> m r
sng m r
stp
                stop :: m r
stop = Stream m b -> m r
foldShared forall a b. (a -> b) -> a -> b
$ Stream m a -> Stream m b
go1 Stream m a
r1
                single :: b -> m r
single b
a   = b -> Stream m b -> m r
yld b
a (Stream m a -> Stream m b
go1 Stream m a
r1)
                yieldk :: b -> Stream m b -> m r
yieldk b
a Stream m b
r = b -> Stream m b -> m r
yld b
a (Stream m a -> Stream m b -> Stream m b
go2 Stream m a
r1 Stream m b
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
foldStream State Stream m b
st b -> Stream m b -> m r
yieldk b -> m r
single m r
stop Stream m b
m

{-# INLINE apSerialDiscardSnd #-}
apSerialDiscardSnd ::
       Stream m a
    -> Stream m b
    -> Stream m a
apSerialDiscardSnd :: forall (m :: * -> *) a b. Stream m a -> Stream m b -> Stream m a
apSerialDiscardSnd Stream m a
fstream Stream m b
stream = forall {a}. Stream m a -> Stream m a
go1 Stream m a
fstream

    where

    go1 :: Stream m a -> Stream m a
go1 Stream m a
m =
        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
mkStream forall a b. (a -> b) -> a -> b
$ \State Stream m a
st a -> Stream m a -> m r
yld a -> m r
sng m r
stp ->
            let foldShared :: Stream m a -> m r
foldShared = forall (m :: * -> *) a r.
State Stream m a
-> (a -> Stream m a -> m r)
-> (a -> m r)
-> m r
-> Stream m a
-> m r
foldStreamShared State Stream m a
st a -> Stream m a -> m r
yld a -> m r
sng m r
stp
                single :: a -> m r
single a
f   = Stream m a -> m r
foldShared forall a b. (a -> b) -> a -> b
$ forall {a} {m :: * -> *} {a}. a -> Stream m a -> Stream m a
go3 a
f Stream m b
stream
                yieldk :: a -> Stream m a -> m r
yieldk a
f Stream m a
r = Stream m a -> m r
foldShared forall a b. (a -> b) -> a -> b
$ a -> Stream m a -> Stream m b -> Stream m a
go2 a
f Stream m a
r Stream m b
stream
            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
foldStream State Stream m a
st a -> Stream m a -> m r
yieldk a -> m r
single m r
stp Stream m a
m

    go2 :: a -> Stream m a -> Stream m b -> Stream m a
go2 a
f Stream m a
r1 Stream m b
m =
        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
mkStream forall a b. (a -> b) -> a -> b
$ \State Stream m a
st a -> Stream m a -> m r
yld a -> m r
sng m r
stp ->
            let foldShared :: Stream m a -> m r
foldShared = forall (m :: * -> *) a r.
State Stream m a
-> (a -> Stream m a -> m r)
-> (a -> m r)
-> m r
-> Stream m a
-> m r
foldStreamShared State Stream m a
st a -> Stream m a -> m r
yld a -> m r
sng m r
stp
                stop :: m r
stop = Stream m a -> m r
foldShared forall a b. (a -> b) -> a -> b
$ Stream m a -> Stream m a
go1 Stream m a
r1
                single :: p -> m r
single p
_   = a -> Stream m a -> m r
yld a
f (Stream m a -> Stream m a
go1 Stream m a
r1)
                yieldk :: p -> Stream m b -> m r
yieldk p
_ Stream m b
r = a -> Stream m a -> m r
yld a
f (a -> Stream m a -> Stream m b -> Stream m a
go2 a
f Stream m a
r1 Stream m b
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
foldStream (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
st) forall {p}. p -> Stream m b -> m r
yieldk forall {p}. p -> m r
single m r
stop Stream m b
m

    go3 :: a -> Stream m a -> Stream m a
go3 a
f Stream m a
m =
        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
mkStream forall a b. (a -> b) -> a -> b
$ \State Stream m a
st a -> Stream m a -> m r
yld a -> m r
sng m r
stp ->
            let single :: p -> m r
single p
_   = a -> m r
sng a
f
                yieldk :: p -> Stream m a -> m r
yieldk p
_ Stream m a
r = a -> Stream m a -> m r
yld a
f (a -> Stream m a -> Stream m a
go3 a
f Stream m a
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
foldStream (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m a
st) forall {p}. p -> Stream m a -> m r
yieldk forall {p}. p -> m r
single m r
stp Stream m a
m

-- XXX This is just concatMapWith with arguments flipped. We need to keep this
-- instead of using a concatMap style definition because the bind
-- implementation in Async and WAsync streams show significant perf degradation
-- if the argument order is changed.
{-# INLINE bindWith #-}
bindWith ::
       (Stream m b -> Stream m b -> Stream m b)
    -> Stream m a
    -> (a -> Stream m b)
    -> Stream m b
bindWith :: forall (m :: * -> *) b a.
(Stream m b -> Stream m b -> Stream m b)
-> Stream m a -> (a -> Stream m b) -> Stream m b
bindWith Stream m b -> Stream m b -> Stream m b
par Stream m a
m1 a -> Stream m b
f = Stream m a -> Stream m b
go Stream m a
m1
    where
        go :: Stream m a -> Stream m b
go Stream m a
m =
            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
mkStream forall a b. (a -> b) -> a -> b
$ \State Stream m b
st b -> Stream m b -> m r
yld b -> m r
sng m r
stp ->
                let foldShared :: Stream m b -> m r
foldShared = forall (m :: * -> *) a r.
State Stream m a
-> (a -> Stream m a -> m r)
-> (a -> m r)
-> m r
-> Stream m a
-> m r
foldStreamShared State Stream m b
st b -> Stream m b -> m r
yld b -> m r
sng m r
stp
                    single :: a -> m r
single a
a   = Stream m b -> m r
foldShared forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. Stream m a -> Stream m a
unShare (a -> Stream m b
f a
a)
                    yieldk :: a -> Stream m a -> m r
yieldk a
a Stream m a
r = Stream m b -> m r
foldShared forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. Stream m a -> Stream m a
unShare (a -> Stream m b
f a
a) Stream m b -> Stream m b -> Stream m b
`par` Stream m a -> Stream m b
go Stream m a
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
foldStream (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m b
st) a -> Stream m a -> m r
yieldk a -> m r
single m r
stp Stream m a
m

-- XXX express in terms of foldrS?
-- XXX can we use a different stream type for the generated stream being
-- falttened so that we can combine them differently and keep the resulting
-- stream different?
-- XXX do we need specialize to IO?
-- XXX can we optimize when c and a are same, by removing the forall using
-- rewrite rules with type applications?

-- | Perform a 'concatMap' using a specified concat strategy. The first
-- argument specifies a merge or concat function that is used to merge the
-- streams generated by the map function. For example, the concat function
-- could be 'serial', 'parallel', 'async', 'ahead' or any other zip or merge
-- function.
--
-- @since 0.7.0
{-# INLINE concatMapWith #-}
concatMapWith
    ::
       (Stream m b -> Stream m b -> Stream m b)
    -> (a -> Stream m b)
    -> Stream m a
    -> Stream m b
concatMapWith :: forall (m :: * -> *) b a.
(Stream m b -> Stream m b -> Stream m b)
-> (a -> Stream m b) -> Stream m a -> Stream m b
concatMapWith Stream m b -> Stream m b -> Stream m b
par a -> Stream m b
f Stream m a
xs = forall (m :: * -> *) b a.
(Stream m b -> Stream m b -> Stream m b)
-> Stream m a -> (a -> Stream m b) -> Stream m b
bindWith Stream m b -> Stream m b -> Stream m b
par Stream m a
xs a -> Stream m b
f

{-# INLINE concatMap #-}
concatMap :: (a -> Stream m b) -> Stream m a -> Stream m b
concatMap :: forall a (m :: * -> *) b.
(a -> Stream m b) -> Stream m a -> Stream m b
concatMap = forall (m :: * -> *) b a.
(Stream m b -> Stream m b -> Stream m b)
-> (a -> Stream m b) -> Stream m a -> Stream m b
concatMapWith forall (m :: * -> *) a. Stream m a -> Stream m a -> Stream m a
serial

{-
-- Fused version.
-- XXX This fuses but when the stream is nil this performs poorly.
-- The filterAllOut benchmark degrades. Need to investigate and fix that.
{-# INLINE concatMap #-}
concatMap :: IsStream t => (a -> t m b) -> t m a -> t m b
concatMap f xs = buildS
    (\c n -> foldrS (\x b -> foldrS c b (f x)) n xs)

-- Stream polymorphic concatMap implementation
-- XXX need to use buildSM/foldrSMShared for parallel behavior
-- XXX unShare seems to degrade the fused performance
{-# INLINE_EARLY concatMap_ #-}
concatMap_ :: IsStream t => (a -> t m b) -> t m a -> t m b
concatMap_ f xs = buildS
     (\c n -> foldrSShared (\x b -> foldrSShared c b (unShare $ f x)) n xs)
-}

-- | See 'Streamly.Internal.Data.Stream.IsStream.concatPairsWith' for
-- documentation.
--
{-# INLINE concatPairsWith #-}
concatPairsWith
    ::
       (Stream m b -> Stream m b -> Stream m b)
    -> (a -> Stream m b)
    -> Stream m a
    -> Stream m b
concatPairsWith :: forall (m :: * -> *) b a.
(Stream m b -> Stream m b -> Stream m b)
-> (a -> Stream m b) -> Stream m a -> Stream m b
concatPairsWith Stream m b -> Stream m b -> Stream m b
combine a -> Stream m b
f Stream m a
str = Stream m (Stream m b) -> Stream m b
go (forall {m :: * -> *}. Stream m a -> Stream m (Stream m b)
leafPairs Stream m a
str)

    where

    go :: Stream m (Stream m b) -> Stream m b
go Stream m (Stream m b)
stream =
        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
mkStream forall a b. (a -> b) -> a -> b
$ \State Stream m b
st b -> Stream m b -> m r
yld b -> m r
sng m r
stp ->
            let foldShared :: Stream m b -> m r
foldShared = forall (m :: * -> *) a r.
State Stream m a
-> (a -> Stream m a -> m r)
-> (a -> m r)
-> m r
-> Stream m a
-> m r
foldStreamShared State Stream m b
st b -> Stream m b -> m r
yld b -> m r
sng m r
stp
                single :: Stream m b -> m r
single Stream m b
a   = Stream m b -> m r
foldShared forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. Stream m a -> Stream m a
unShare Stream m b
a
                yieldk :: Stream m b -> Stream m (Stream m b) -> m r
yieldk Stream m b
a Stream m (Stream m b)
r = Stream m b -> m r
foldShared forall a b. (a -> b) -> a -> b
$ Stream m b -> Stream m (Stream m b) -> Stream m b
go1 Stream m b
a Stream m (Stream m b)
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
foldStream (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m b
st) Stream m b -> Stream m (Stream m b) -> m r
yieldk Stream m b -> m r
single m r
stp Stream m (Stream m b)
stream

    go1 :: Stream m b -> Stream m (Stream m b) -> Stream m b
go1 Stream m b
a1 Stream m (Stream m b)
stream =
        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
mkStream forall a b. (a -> b) -> a -> b
$ \State Stream m b
st b -> Stream m b -> m r
yld b -> m r
sng m r
stp ->
            let foldShared :: Stream m b -> m r
foldShared = forall (m :: * -> *) a r.
State Stream m a
-> (a -> Stream m a -> m r)
-> (a -> m r)
-> m r
-> Stream m a
-> m r
foldStreamShared State Stream m b
st b -> Stream m b -> m r
yld b -> m r
sng m r
stp
                stop :: m r
stop = Stream m b -> m r
foldShared forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. Stream m a -> Stream m a
unShare Stream m b
a1
                single :: Stream m b -> m r
single Stream m b
a = Stream m b -> m r
foldShared forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. Stream m a -> Stream m a
unShare Stream m b
a1 Stream m b -> Stream m b -> Stream m b
`combine` Stream m b
a
                yieldk :: Stream m b -> Stream m (Stream m b) -> m r
yieldk Stream m b
a Stream m (Stream m b)
r =
                    Stream m b -> m r
foldShared forall a b. (a -> b) -> a -> b
$ Stream m (Stream m b) -> Stream m b
go forall a b. (a -> b) -> a -> b
$ Stream m b -> Stream m b -> Stream m b
combine Stream m b
a1 Stream m b
a forall a (m :: * -> *). a -> Stream m a -> Stream m a
`cons` forall {m :: * -> *}.
Stream m (Stream m b) -> Stream m (Stream m b)
nonLeafPairs Stream m (Stream m b)
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
foldStream (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m b
st) Stream m b -> Stream m (Stream m b) -> m r
yieldk Stream m b -> m r
single m r
stop Stream m (Stream m b)
stream

    -- Exactly the same as "go" except that stop continuation extracts the
    -- stream.
    leafPairs :: Stream m a -> Stream m (Stream m b)
leafPairs Stream m a
stream =
        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
mkStream forall a b. (a -> b) -> a -> b
$ \State Stream m (Stream m b)
st Stream m b -> Stream m (Stream m b) -> m r
yld Stream m b -> m r
sng m r
stp ->
            let foldShared :: Stream m (Stream m b) -> m r
foldShared = forall (m :: * -> *) a r.
State Stream m a
-> (a -> Stream m a -> m r)
-> (a -> m r)
-> m r
-> Stream m a
-> m r
foldStreamShared State Stream m (Stream m b)
st Stream m b -> Stream m (Stream m b) -> m r
yld Stream m b -> m r
sng m r
stp
                single :: a -> m r
single a
a   = Stream m b -> m r
sng (a -> Stream m b
f a
a)
                yieldk :: a -> Stream m a -> m r
yieldk a
a Stream m a
r = Stream m (Stream m b) -> m r
foldShared forall a b. (a -> b) -> a -> b
$ a -> Stream m a -> Stream m (Stream m b)
leafPairs1 a
a Stream m a
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
foldStream (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m (Stream m b)
st) a -> Stream m a -> m r
yieldk a -> m r
single m r
stp Stream m a
stream

    leafPairs1 :: a -> Stream m a -> Stream m (Stream m b)
leafPairs1 a
a1 Stream m a
stream =
        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
mkStream forall a b. (a -> b) -> a -> b
$ \State Stream m (Stream m b)
st Stream m b -> Stream m (Stream m b) -> m r
yld Stream m b -> m r
sng m r
_ ->
            let stop :: m r
stop = Stream m b -> m r
sng (a -> Stream m b
f a
a1)
                single :: a -> m r
single a
a = Stream m b -> m r
sng (a -> Stream m b
f a
a1 Stream m b -> Stream m b -> Stream m b
`combine` a -> Stream m b
f a
a)
                yieldk :: a -> Stream m a -> m r
yieldk a
a Stream m a
r = Stream m b -> Stream m (Stream m b) -> m r
yld (a -> Stream m b
f a
a1 Stream m b -> Stream m b -> Stream m b
`combine` a -> Stream m b
f a
a) forall a b. (a -> b) -> a -> b
$ Stream m a -> Stream m (Stream m b)
leafPairs Stream m a
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
foldStream (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m (Stream m b)
st) a -> Stream m a -> m r
yieldk a -> m r
single m r
stop Stream m a
stream

    -- Exactly the same as "leafPairs" except that it does not map "f"
    nonLeafPairs :: Stream m (Stream m b) -> Stream m (Stream m b)
nonLeafPairs Stream m (Stream m b)
stream =
        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
mkStream forall a b. (a -> b) -> a -> b
$ \State Stream m (Stream m b)
st Stream m b -> Stream m (Stream m b) -> m r
yld Stream m b -> m r
sng m r
stp ->
            let foldShared :: Stream m (Stream m b) -> m r
foldShared = forall (m :: * -> *) a r.
State Stream m a
-> (a -> Stream m a -> m r)
-> (a -> m r)
-> m r
-> Stream m a
-> m r
foldStreamShared State Stream m (Stream m b)
st Stream m b -> Stream m (Stream m b) -> m r
yld Stream m b -> m r
sng m r
stp
                single :: Stream m b -> m r
single Stream m b
a   = Stream m b -> m r
sng Stream m b
a
                yieldk :: Stream m b -> Stream m (Stream m b) -> m r
yieldk Stream m b
a Stream m (Stream m b)
r = Stream m (Stream m b) -> m r
foldShared forall a b. (a -> b) -> a -> b
$ Stream m b -> Stream m (Stream m b) -> Stream m (Stream m b)
nonLeafPairs1 Stream m b
a Stream m (Stream m b)
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
foldStream (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m (Stream m b)
st) Stream m b -> Stream m (Stream m b) -> m r
yieldk Stream m b -> m r
single m r
stp Stream m (Stream m b)
stream

    nonLeafPairs1 :: Stream m b -> Stream m (Stream m b) -> Stream m (Stream m b)
nonLeafPairs1 Stream m b
a1 Stream m (Stream m b)
stream =
        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
mkStream forall a b. (a -> b) -> a -> b
$ \State Stream m (Stream m b)
st Stream m b -> Stream m (Stream m b) -> m r
yld Stream m b -> m r
sng m r
_ ->
            let stop :: m r
stop = Stream m b -> m r
sng Stream m b
a1
                single :: Stream m b -> m r
single Stream m b
a = Stream m b -> m r
sng (Stream m b
a1 Stream m b -> Stream m b -> Stream m b
`combine` Stream m b
a)
                yieldk :: Stream m b -> Stream m (Stream m b) -> m r
yieldk Stream m b
a Stream m (Stream m b)
r = Stream m b -> Stream m (Stream m b) -> m r
yld (Stream m b
a1 Stream m b -> Stream m b -> Stream m b
`combine` Stream m b
a) forall a b. (a -> b) -> a -> b
$ Stream m (Stream m b) -> Stream m (Stream m b)
nonLeafPairs Stream m (Stream m b)
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
foldStream (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m (Stream m b)
st) Stream m b -> Stream m (Stream m b) -> m r
yieldk Stream m b -> m r
single m r
stop Stream m (Stream m b)
stream

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

-- NOTE: even though concatMap for StreamD is 3x faster compared to StreamK,
-- the monad instance of StreamD is slower than StreamK after foldr/build
-- fusion.
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 a (m :: * -> *) b.
(a -> Stream m b) -> Stream m a -> Stream m b
concatMap

{-
-- Like concatMap but generates stream using an unfold function. Similar to
-- unfoldMany but for StreamK.
concatUnfoldr :: IsStream t
    => (b -> t m (Maybe (a, b))) -> t m b -> t m a
concatUnfoldr = undefined
-}

------------------------------------------------------------------------------
-- MonadReader
------------------------------------------------------------------------------

{-# INLINABLE withLocal #-}
withLocal :: MonadReader r m => (r -> r) -> Stream m a -> Stream m a
withLocal :: forall r (m :: * -> *) a.
MonadReader r m =>
(r -> r) -> Stream m a -> Stream m a
withLocal r -> r
f Stream m a
m =
    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
mkStream forall a b. (a -> b) -> a -> b
$ \State Stream m a
st a -> Stream m a -> m r
yld a -> m r
sng m r
stp ->
        let single :: a -> m r
single = forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local r -> r
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> m r
sng
            yieldk :: a -> Stream m a -> m r
yieldk a
a Stream m a
r = forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local r -> r
f forall a b. (a -> b) -> a -> b
$ a -> Stream m a -> m r
yld a
a (forall r (m :: * -> *) a.
MonadReader r m =>
(r -> r) -> Stream m a -> Stream m a
withLocal r -> r
f Stream m a
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
foldStream State Stream m a
st a -> Stream m a -> m r
yieldk a -> m r
single (forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local r -> r
f m r
stp) Stream m a
m

-------------------------------------------------------------------------------
-- Generation
-------------------------------------------------------------------------------

{-# INLINE unfoldr #-}
unfoldr :: (b -> Maybe (a, b)) -> b -> Stream m a
unfoldr :: forall b a (m :: * -> *). (b -> Maybe (a, b)) -> b -> Stream m a
unfoldr b -> Maybe (a, b)
next b
s0 = forall (m :: * -> *) a.
(forall b. (a -> b -> b) -> b -> b) -> Stream m a
build forall a b. (a -> b) -> a -> b
$ \a -> b -> b
yld b
stp ->
    let go :: b -> b
go b
s =
            case b -> Maybe (a, b)
next b
s of
                Just (a
a, b
b) -> a -> b -> b
yld a
a (b -> b
go b
b)
                Maybe (a, b)
Nothing -> b
stp
    in b -> b
go b
s0

{-# INLINE unfoldrMWith #-}
unfoldrMWith :: Monad m =>
       (m a -> Stream m a -> Stream m a)
    -> (b -> m (Maybe (a, b)))
    -> b
    -> Stream m a
unfoldrMWith :: forall (m :: * -> *) a b.
Monad m =>
(m a -> Stream m a -> Stream m a)
-> (b -> m (Maybe (a, b))) -> b -> Stream m a
unfoldrMWith m a -> Stream m a -> Stream m a
cns b -> m (Maybe (a, b))
step = b -> Stream m a
go

    where

    go :: b -> Stream m a
go b
s = forall (m :: * -> *) a.
Monad m =>
(m a -> Stream m a -> Stream m a)
-> (forall r. (a -> Stream m a -> m r) -> (a -> m r) -> m r -> m r)
-> Stream m a
sharedMWith m a -> Stream m a -> Stream m a
cns forall a b. (a -> b) -> a -> b
$ \a -> Stream m a -> m r
yld a -> m r
_ m r
stp -> do
                Maybe (a, b)
r <- b -> m (Maybe (a, b))
step b
s
                case Maybe (a, b)
r of
                    Just (a
a, b
b) -> a -> Stream m a -> m r
yld a
a (b -> Stream m a
go b
b)
                    Maybe (a, b)
Nothing -> m r
stp

-- | Generate an infinite stream by repeating a pure value.
--
-- /Pre-release/
{-# INLINE repeat #-}
repeat :: a -> Stream m a
repeat :: forall a (m :: * -> *). a -> Stream m a
repeat a
a = let x :: Stream m a
x = forall a (m :: * -> *). a -> Stream m a -> Stream m a
cons a
a Stream m a
x in forall {m :: * -> *}. Stream m a
x

-- | Like 'repeatM' but takes a stream 'cons' operation to combine the actions
-- in a stream specific manner. A serial cons would repeat the values serially
-- while an async cons would repeat concurrently.
--
-- /Pre-release/
repeatMWith :: (m a -> t m a -> t m a) -> m a -> t m a
repeatMWith :: forall (m :: * -> *) a (t :: (* -> *) -> * -> *).
(m a -> t m a -> t m a) -> m a -> t m a
repeatMWith m a -> t m a -> t m a
cns = m a -> t m a
go

    where

    go :: m a -> t m a
go m a
m = m a
m m a -> t m a -> t m a
`cns` m a -> t m a
go m a
m

{-# INLINE replicateMWith #-}
replicateMWith :: (m a -> Stream m a -> Stream m a) -> Int -> m a -> Stream m a
replicateMWith :: forall (m :: * -> *) a.
(m a -> Stream m a -> Stream m a) -> Int -> m a -> Stream m a
replicateMWith m a -> Stream m a -> Stream m a
cns Int
n m a
m = forall {t}. (Ord t, Num t) => t -> Stream m a
go Int
n

    where

    go :: t -> Stream m a
go t
cnt = if t
cnt forall a. Ord a => a -> a -> Bool
<= t
0 then forall (m :: * -> *) a. Stream m a
nil else m a
m m a -> Stream m a -> Stream m a
`cns` t -> Stream m a
go (t
cnt forall a. Num a => a -> a -> a
- t
1)

{-# INLINE fromIndicesMWith #-}
fromIndicesMWith ::
    (m a -> Stream m a -> Stream m a) -> (Int -> m a) -> Stream m a
fromIndicesMWith :: forall (m :: * -> *) a.
(m a -> Stream m a -> Stream m a) -> (Int -> m a) -> Stream m a
fromIndicesMWith m a -> Stream m a -> Stream m a
cns Int -> m a
gen = Int -> Stream m a
go Int
0

    where

    go :: Int -> Stream m a
go Int
i = 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
mkStream forall a b. (a -> b) -> a -> b
$ \State Stream m a
st a -> Stream m a -> m r
stp a -> m r
sng m r
yld -> do
        forall (m :: * -> *) a r.
State Stream m a
-> (a -> Stream m a -> m r)
-> (a -> m r)
-> m r
-> Stream m a
-> m r
foldStreamShared State Stream m a
st a -> Stream m a -> m r
stp a -> m r
sng m r
yld (Int -> m a
gen Int
i m a -> Stream m a -> Stream m a
`cns` Int -> Stream m a
go (Int
i forall a. Num a => a -> a -> a
+ Int
1))

{-# INLINE iterateMWith #-}
iterateMWith :: Monad m =>
    (m a -> Stream m a -> Stream m a) -> (a -> m a) -> m a -> Stream m a
iterateMWith :: forall (m :: * -> *) a.
Monad m =>
(m a -> Stream m a -> Stream m a)
-> (a -> m a) -> m a -> Stream m a
iterateMWith m a -> Stream m a -> Stream m a
cns a -> m a
step = m a -> Stream m a
go

    where

    go :: m a -> Stream m a
go m a
s = 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
mkStream forall a b. (a -> b) -> a -> b
$ \State Stream m a
st a -> Stream m a -> m r
stp a -> m r
sng m r
yld -> do
        !a
next <- m a
s
        forall (m :: * -> *) a r.
State Stream m a
-> (a -> Stream m a -> m r)
-> (a -> m r)
-> m r
-> Stream m a
-> m r
foldStreamShared State Stream m a
st a -> Stream m a -> m r
stp a -> m r
sng m r
yld (forall (m :: * -> *) a. Monad m => a -> m a
return a
next m a -> Stream m a -> Stream m a
`cns` m a -> Stream m a
go (a -> m a
step a
next))

{-# INLINE headPartial #-}
headPartial :: Monad m => Stream m a -> m a
headPartial :: forall (m :: * -> *) a. Monad m => Stream m a -> m a
headPartial = forall a (m :: * -> *) b.
(a -> m b -> m b) -> m b -> Stream m a -> m b
foldrM (\a
x m a
_ -> forall (m :: * -> *) a. Monad m => a -> m a
return a
x) (forall a. HasCallStack => [Char] -> a
error [Char]
"head of nil")

{-# INLINE tailPartial #-}
tailPartial :: Stream m a -> Stream m a
tailPartial :: forall (m :: * -> *) a. Stream m a -> Stream m a
tailPartial Stream m a
m = 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
mkStream forall a b. (a -> b) -> a -> b
$ \State Stream m a
st a -> Stream m a -> m r
yld a -> m r
sng m r
stp ->
    let stop :: a
stop      = forall a. HasCallStack => [Char] -> a
error [Char]
"tail of nil"
        single :: p -> m r
single p
_  = m r
stp
        yieldk :: p -> Stream m a -> m r
yieldk p
_ Stream m a
r = forall (m :: * -> *) a r.
State Stream m a
-> (a -> Stream m a -> m r)
-> (a -> m r)
-> m r
-> Stream m a
-> m r
foldStream State Stream m a
st a -> Stream m a -> m r
yld a -> m r
sng m r
stp Stream m a
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
foldStream State Stream m a
st forall {p}. p -> Stream m a -> m r
yieldk forall {p}. p -> m r
single forall {a}. a
stop Stream m a
m

{-# INLINE mfix #-}
mfix :: Monad m => (m a -> Stream m a) -> Stream m a
mfix :: forall (m :: * -> *) a.
Monad m =>
(m a -> Stream m a) -> Stream m a
mfix m a -> Stream m a
f = 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
mkStream forall a b. (a -> b) -> a -> b
$ \State Stream m a
st a -> Stream m a -> m r
yld a -> m r
sng m r
stp ->
    let single :: a -> m r
single a
a  = forall (m :: * -> *) a r.
State Stream m a
-> (a -> Stream m a -> m r)
-> (a -> m r)
-> m r
-> Stream m a
-> m r
foldStream State Stream m a
st a -> Stream m a -> m r
yld a -> m r
sng m r
stp forall a b. (a -> b) -> a -> b
$ a
a forall a (m :: * -> *). a -> Stream m a -> Stream m a
`cons` Stream m a
ys
        yieldk :: a -> p -> m r
yieldk a
a p
_ = forall (m :: * -> *) a r.
State Stream m a
-> (a -> Stream m a -> m r)
-> (a -> m r)
-> m r
-> Stream m a
-> m r
foldStream State Stream m a
st a -> Stream m a -> m r
yld a -> m r
sng m r
stp forall a b. (a -> b) -> a -> b
$ a
a forall a (m :: * -> *). a -> Stream m a -> Stream m a
`cons` Stream m a
ys
    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
foldStream State Stream m a
st forall {p}. a -> p -> m r
yieldk a -> m r
single m r
stp Stream m a
xs

    where

    -- fix the head element of the stream
    xs :: Stream m a
xs = forall a. (a -> a) -> a
fix  (m a -> Stream m a
f forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (m :: * -> *) a. Monad m => Stream m a -> m a
headPartial)

    -- now fix the tail recursively
    ys :: Stream m a
ys = forall (m :: * -> *) a.
Monad m =>
(m a -> Stream m a) -> Stream m a
mfix (forall (m :: * -> *) a. Stream m a -> Stream m a
tailPartial forall b c a. (b -> c) -> (a -> b) -> a -> c
. m a -> Stream m a
f)

-------------------------------------------------------------------------------
-- Conversions
-------------------------------------------------------------------------------

-- |
-- @
-- fromFoldable = 'Prelude.foldr' 'cons' 'nil'
-- @
--
-- Construct a stream from a 'Foldable' containing pure values:
--
-- @since 0.2.0
{-# INLINE fromFoldable #-}
fromFoldable :: Foldable f => f a -> Stream m a
fromFoldable :: forall (f :: * -> *) a (m :: * -> *).
Foldable f =>
f a -> Stream m a
fromFoldable = forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
Prelude.foldr forall a (m :: * -> *). a -> Stream m a -> Stream m a
cons forall (m :: * -> *) a. Stream m a
nil

{-# INLINE fromFoldableM #-}
fromFoldableM :: (Foldable f, Monad m) => f (m a) -> Stream m a
fromFoldableM :: forall (f :: * -> *) (m :: * -> *) a.
(Foldable f, Monad m) =>
f (m a) -> Stream m a
fromFoldableM = forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
Prelude.foldr forall (m :: * -> *) a. Monad m => m a -> Stream m a -> Stream m a
consM forall (m :: * -> *) a. Stream m a
nil

-------------------------------------------------------------------------------
-- Deconstruction
-------------------------------------------------------------------------------

{-# INLINE uncons #-}
uncons :: Applicative m => Stream m a -> m (Maybe (a, Stream m a))
uncons :: forall (m :: * -> *) a.
Applicative m =>
Stream m a -> m (Maybe (a, Stream m a))
uncons Stream m a
m =
    let stop :: m (Maybe a)
stop = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Maybe a
Nothing
        single :: a -> f (Maybe (a, Stream m a))
single a
a = forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall a. a -> Maybe a
Just (a
a, forall (m :: * -> *) a. Stream m a
nil))
        yieldk :: a -> b -> f (Maybe (a, b))
yieldk a
a b
r = forall (f :: * -> *) a. Applicative f => a -> f a
pure (forall a. a -> Maybe a
Just (a
a, b
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
foldStream forall (t :: (* -> *) -> * -> *) (m :: * -> *) a. State t m a
defState forall {f :: * -> *} {a} {b}.
Applicative f =>
a -> b -> f (Maybe (a, b))
yieldk forall {f :: * -> *} {a} {m :: * -> *} {a}.
Applicative f =>
a -> f (Maybe (a, Stream m a))
single forall {a}. m (Maybe a)
stop Stream m a
m

{-# INLINE tail #-}
tail :: Applicative m => Stream m a -> m (Maybe (Stream m a))
tail :: forall (m :: * -> *) a.
Applicative m =>
Stream m a -> m (Maybe (Stream m a))
tail =
    let stop :: m (Maybe a)
stop      = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Maybe a
Nothing
        single :: p -> f (Maybe (Stream m a))
single p
_  = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just forall (m :: * -> *) a. Stream m a
nil
        yieldk :: p -> a -> f (Maybe a)
yieldk p
_ a
r = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just a
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
foldStream forall (t :: (* -> *) -> * -> *) (m :: * -> *) a. State t m a
defState forall {f :: * -> *} {p} {a}.
Applicative f =>
p -> a -> f (Maybe a)
yieldk forall {f :: * -> *} {p} {m :: * -> *} {a}.
Applicative f =>
p -> f (Maybe (Stream m a))
single forall {a}. m (Maybe a)
stop

{-# INLINE init #-}
init :: Applicative m => Stream m a -> m (Maybe (Stream m a))
init :: forall (m :: * -> *) a.
Applicative m =>
Stream m a -> m (Maybe (Stream m a))
init = forall (m :: * -> *) a.
Applicative m =>
Stream m a -> m (Maybe (Stream m a))
go1
    where
    go1 :: Stream m t -> m (Maybe (Stream m t))
go1 Stream m t
m1 = do
        (\case
            Maybe (t, Stream m t)
Nothing -> forall a. Maybe a
Nothing
            Just (t
h, Stream m t
t) -> forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall a (m :: * -> *). a -> Stream m a -> Stream m a
go t
h Stream m t
t) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) a.
Applicative m =>
Stream m a -> m (Maybe (a, Stream m a))
uncons Stream m t
m1
    go :: t -> Stream m t -> Stream m t
go t
p Stream m t
m1 = 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
mkStream forall a b. (a -> b) -> a -> b
$ \State Stream m t
_ t -> Stream m t -> m r
yld t -> m r
sng m r
stp ->
        let single :: p -> m r
single p
_ = t -> m r
sng t
p
            yieldk :: t -> Stream m t -> m r
yieldk t
a Stream m t
x = t -> Stream m t -> m r
yld t
p forall a b. (a -> b) -> a -> b
$ t -> Stream m t -> Stream m t
go t
a Stream m t
x
         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
foldStream forall (t :: (* -> *) -> * -> *) (m :: * -> *) a. State t m a
defState t -> Stream m t -> m r
yieldk forall {p}. p -> m r
single m r
stp Stream m t
m1

------------------------------------------------------------------------------
-- Reordering
------------------------------------------------------------------------------

-- | Lazy left fold to a stream.
{-# INLINE foldlS #-}
foldlS ::
    (Stream m b -> a -> Stream m b) -> Stream m b -> Stream m a -> Stream m b
foldlS :: forall (m :: * -> *) b a.
(Stream m b -> a -> Stream m b)
-> Stream m b -> Stream m a -> Stream m b
foldlS Stream m b -> a -> Stream m b
step = Stream m b -> Stream m a -> Stream m b
go
    where
    go :: Stream m b -> Stream m a -> Stream m b
go Stream m b
acc Stream m a
rest = 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
mkStream forall a b. (a -> b) -> a -> b
$ \State Stream m b
st b -> Stream m b -> m r
yld b -> m r
sng m r
stp ->
        let run :: Stream m b -> m r
run Stream m b
x = forall (m :: * -> *) a r.
State Stream m a
-> (a -> Stream m a -> m r)
-> (a -> m r)
-> m r
-> Stream m a
-> m r
foldStream State Stream m b
st b -> Stream m b -> m r
yld b -> m r
sng m r
stp Stream m b
x
            stop :: m r
stop = Stream m b -> m r
run Stream m b
acc
            single :: a -> m r
single a
a = Stream m b -> m r
run forall a b. (a -> b) -> a -> b
$ Stream m b -> a -> Stream m b
step Stream m b
acc a
a
            yieldk :: a -> Stream m a -> m r
yieldk a
a Stream m a
r = Stream m b -> m r
run forall a b. (a -> b) -> a -> b
$ Stream m b -> Stream m a -> Stream m b
go (Stream m b -> a -> Stream m b
step Stream m b
acc a
a) Stream m a
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
foldStream (forall (t :: (* -> *) -> * -> *) (m :: * -> *) a (n :: * -> *) b.
State t m a -> State t n b
adaptState State Stream m b
st) a -> Stream m a -> m r
yieldk a -> m r
single m r
stop Stream m a
rest

{-# INLINE reverse #-}
reverse :: Stream m a -> Stream m a
reverse :: forall (m :: * -> *) a. Stream m a -> Stream m a
reverse = forall (m :: * -> *) b a.
(Stream m b -> a -> Stream m b)
-> Stream m b -> Stream m a -> Stream m b
foldlS (forall a b c. (a -> b -> c) -> b -> a -> c
flip forall a (m :: * -> *). a -> Stream m a -> Stream m a
cons) forall (m :: * -> *) a. Stream m a
nil