-- |
-- Module      : Streamly.Internal.Data.Unfold.Type
-- Copyright   : (c) 2019 Composewell Technologies
-- License     : BSD3
-- Maintainer  : streamly@composewell.com
-- Stability   : experimental
-- Portability : GHC
--
-- To run the examples in this module:
--
-- >>> import qualified Streamly.Prelude as Stream
-- >>> import qualified Streamly.Data.Fold as Fold
-- >>> import qualified Streamly.Internal.Data.Unfold as Unfold

module Streamly.Internal.Data.Unfold.Type
    ( Unfold (..)

    -- * Basic Constructors
    , mkUnfoldM
    , mkUnfoldrM
    , unfoldrM
    , unfoldr
    , functionM
    , function
    , identity

    -- * From Values
    , fromEffect
    , fromPure

    -- * Transformations
    , lmap
    , lmapM
    , map
    , supply
    , supplyFirst
    , supplySecond

    -- * Trimming
    , takeWhileMWithInput
    , takeWhileM
    , takeWhile

    -- * Nesting
    , ConcatState (..)
    , many
    , manyInterleave

    -- Applicative
    , apSequence
    , apDiscardSnd
    , crossWithM
    , crossWith
    , cross
    , apply

    -- Monad
    , concatMapM
    , concatMap
    , bind

    , zipWithM
    , zipWith
    )
where

#include "inline.hs"

-- import Control.Arrow (Arrow(..))
-- import Control.Category (Category(..))
import Control.Monad ((>=>))
import Data.Void (Void)
import Fusion.Plugin.Types (Fuse(..))
import Streamly.Internal.Data.Tuple.Strict (Tuple'(..))
import Streamly.Internal.Data.Stream.StreamD.Step (Step(..))

import Prelude hiding (map, concatMap, zipWith, takeWhile)

-- $setup
-- >>> import qualified Streamly.Prelude as Stream
-- >>> import qualified Streamly.Data.Fold as Fold
-- >>> import qualified Streamly.Internal.Data.Unfold as Unfold

------------------------------------------------------------------------------
-- Monadic Unfolds
------------------------------------------------------------------------------

-- The order of arguments allows 'Category' and 'Arrow' instances but precludes
-- contravariant and contra-applicative.
--
-- An unfold is akin to a reader. It is the streaming equivalent of a reader.
-- The argument @a@ is the environment of the reader. That's the reason the
-- default unfolds in various modules are called "read".
--
-- | An @Unfold m a b@ is a generator of a stream of values of type @b@ from a
-- seed of type 'a' in 'Monad' @m@.
--
-- @since 0.7.0

data Unfold m a b =
    -- | @Unfold step inject@
    forall s. Unfold (s -> m (Step s b)) (a -> m s)

------------------------------------------------------------------------------
-- Basic constructors
------------------------------------------------------------------------------

-- | Make an unfold from @step@ and @inject@ functions.
--
-- /Pre-release/
{-# INLINE mkUnfoldM #-}
mkUnfoldM :: (s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
mkUnfoldM :: (s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
mkUnfoldM = (s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
Unfold

-- | Make an unfold from a step function.
--
-- See also: 'unfoldrM'
--
-- /Pre-release/
{-# INLINE mkUnfoldrM #-}
mkUnfoldrM :: Applicative m => (a -> m (Step a b)) -> Unfold m a b
mkUnfoldrM :: (a -> m (Step a b)) -> Unfold m a b
mkUnfoldrM a -> m (Step a b)
step = (a -> m (Step a b)) -> (a -> m a) -> Unfold m a b
forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
Unfold a -> m (Step a b)
step a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure

-- The type 'Step' is isomorphic to 'Maybe'. Ideally unfoldrM should be the
-- same as mkUnfoldrM, this is for compatibility with traditional Maybe based
-- unfold step functions.
--
-- | Build a stream by unfolding a /monadic/ step function starting from a seed.
-- The step function returns the next element in the stream and the next seed
-- value. When it is done it returns 'Nothing' and the stream ends.
--
-- /Since: 0.8.0/
--
{-# INLINE unfoldrM #-}
unfoldrM :: Applicative m => (a -> m (Maybe (b, a))) -> Unfold m a b
unfoldrM :: (a -> m (Maybe (b, a))) -> Unfold m a b
unfoldrM a -> m (Maybe (b, a))
next = (a -> m (Step a b)) -> (a -> m a) -> Unfold m a b
forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
Unfold a -> m (Step a b)
step a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
  where
    {-# INLINE_LATE step #-}
    step :: a -> m (Step a b)
step a
st =
        (\case
            Just (b
x, a
s) -> b -> a -> Step a b
forall s a. a -> s -> Step s a
Yield b
x a
s
            Maybe (b, a)
Nothing     -> Step a b
forall s a. Step s a
Stop) (Maybe (b, a) -> Step a b) -> m (Maybe (b, a)) -> m (Step a b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> m (Maybe (b, a))
next a
st

-- | Like 'unfoldrM' but uses a pure step function.
--
-- >>> :{
--  f [] = Nothing
--  f (x:xs) = Just (x, xs)
-- :}
--
-- >>> Unfold.fold Fold.toList (Unfold.unfoldr f) [1,2,3]
-- [1,2,3]
--
-- /Since: 0.8.0/
--
{-# INLINE unfoldr #-}
unfoldr :: Applicative m => (a -> Maybe (b, a)) -> Unfold m a b
unfoldr :: (a -> Maybe (b, a)) -> Unfold m a b
unfoldr a -> Maybe (b, a)
step = (a -> m (Maybe (b, a))) -> Unfold m a b
forall (m :: * -> *) a b.
Applicative m =>
(a -> m (Maybe (b, a))) -> Unfold m a b
unfoldrM (Maybe (b, a) -> m (Maybe (b, a))
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe (b, a) -> m (Maybe (b, a)))
-> (a -> Maybe (b, a)) -> a -> m (Maybe (b, a))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Maybe (b, a)
step)

------------------------------------------------------------------------------
-- Map input
------------------------------------------------------------------------------

-- | Map a function on the input argument of the 'Unfold'.
--
-- >>> u = Unfold.lmap (fmap (+1)) Unfold.fromList
-- >>> Unfold.fold Fold.toList u [1..5]
-- [2,3,4,5,6]
--
-- @
-- lmap f = Unfold.many (Unfold.function f)
-- @
--
-- /Since: 0.8.0/
{-# INLINE_NORMAL lmap #-}
lmap :: (a -> c) -> Unfold m c b -> Unfold m a b
lmap :: (a -> c) -> Unfold m c b -> Unfold m a b
lmap a -> c
f (Unfold s -> m (Step s b)
ustep c -> m s
uinject) = (s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
Unfold s -> m (Step s b)
ustep (c -> m s
uinject (c -> m s) -> (a -> c) -> a -> m s
forall b c a. (b -> c) -> (a -> b) -> a -> c
Prelude.. a -> c
f)

-- | Map an action on the input argument of the 'Unfold'.
--
-- @
-- lmapM f = Unfold.many (Unfold.functionM f)
-- @
--
-- /Since: 0.8.0/
{-# INLINE_NORMAL lmapM #-}
lmapM :: Monad m => (a -> m c) -> Unfold m c b -> Unfold m a b
lmapM :: (a -> m c) -> Unfold m c b -> Unfold m a b
lmapM a -> m c
f (Unfold s -> m (Step s b)
ustep c -> m s
uinject) = (s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
Unfold s -> m (Step s b)
ustep (a -> m c
f (a -> m c) -> (c -> m s) -> a -> m s
forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> c -> m s
uinject)

-- | Supply the seed to an unfold closing the input end of the unfold.
--
-- @
-- supply a = Unfold.lmap (Prelude.const a)
-- @
--
-- /Pre-release/
--
{-# INLINE_NORMAL supply #-}
supply :: a -> Unfold m a b -> Unfold m Void b
supply :: a -> Unfold m a b -> Unfold m Void b
supply a
a = (Void -> a) -> Unfold m a b -> Unfold m Void b
forall a c (m :: * -> *) b.
(a -> c) -> Unfold m c b -> Unfold m a b
lmap (a -> Void -> a
forall a b. a -> b -> a
Prelude.const a
a)

-- | Supply the first component of the tuple to an unfold that accepts a tuple
-- as a seed resulting in a fold that accepts the second component of the tuple
-- as a seed.
--
-- @
-- supplyFirst a = Unfold.lmap (a, )
-- @
--
-- /Pre-release/
--
{-# INLINE_NORMAL supplyFirst #-}
supplyFirst :: a -> Unfold m (a, b) c -> Unfold m b c
supplyFirst :: a -> Unfold m (a, b) c -> Unfold m b c
supplyFirst a
a = (b -> (a, b)) -> Unfold m (a, b) c -> Unfold m b c
forall a c (m :: * -> *) b.
(a -> c) -> Unfold m c b -> Unfold m a b
lmap (a
a, )

-- | Supply the second component of the tuple to an unfold that accepts a tuple
-- as a seed resulting in a fold that accepts the first component of the tuple
-- as a seed.
--
-- @
-- supplySecond b = Unfold.lmap (, b)
-- @
--
-- /Pre-release/
--
{-# INLINE_NORMAL supplySecond #-}
supplySecond :: b -> Unfold m (a, b) c -> Unfold m a c
supplySecond :: b -> Unfold m (a, b) c -> Unfold m a c
supplySecond b
b = (a -> (a, b)) -> Unfold m (a, b) c -> Unfold m a c
forall a c (m :: * -> *) b.
(a -> c) -> Unfold m c b -> Unfold m a b
lmap (, b
b)

------------------------------------------------------------------------------
-- Filter input
------------------------------------------------------------------------------

{-# INLINE_NORMAL takeWhileMWithInput #-}
takeWhileMWithInput :: Monad m =>
    (a -> b -> m Bool) -> Unfold m a b -> Unfold m a b
takeWhileMWithInput :: (a -> b -> m Bool) -> Unfold m a b -> Unfold m a b
takeWhileMWithInput a -> b -> m Bool
f (Unfold s -> m (Step s b)
step1 a -> m s
inject1) = (Tuple' a s -> m (Step (Tuple' a s) b))
-> (a -> m (Tuple' a s)) -> Unfold m a b
forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
Unfold Tuple' a s -> m (Step (Tuple' a s) b)
step a -> m (Tuple' a s)
inject

    where

    inject :: a -> m (Tuple' a s)
inject a
a = do
        s
s <- a -> m s
inject1 a
a
        Tuple' a s -> m (Tuple' a s)
forall (m :: * -> *) a. Monad m => a -> m a
return (Tuple' a s -> m (Tuple' a s)) -> Tuple' a s -> m (Tuple' a s)
forall a b. (a -> b) -> a -> b
$ a -> s -> Tuple' a s
forall a b. a -> b -> Tuple' a b
Tuple' a
a s
s

    {-# INLINE_LATE step #-}
    step :: Tuple' a s -> m (Step (Tuple' a s) b)
step (Tuple' a
a s
st) = do
        Step s b
r <- s -> m (Step s b)
step1 s
st
        case Step s b
r of
            Yield b
x s
s -> do
                Bool
b <- a -> b -> m Bool
f a
a b
x
                Step (Tuple' a s) b -> m (Step (Tuple' a s) b)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (Tuple' a s) b -> m (Step (Tuple' a s) b))
-> Step (Tuple' a s) b -> m (Step (Tuple' a s) b)
forall a b. (a -> b) -> a -> b
$ if Bool
b then b -> Tuple' a s -> Step (Tuple' a s) b
forall s a. a -> s -> Step s a
Yield b
x (a -> s -> Tuple' a s
forall a b. a -> b -> Tuple' a b
Tuple' a
a s
s) else Step (Tuple' a s) b
forall s a. Step s a
Stop
            Skip s
s -> Step (Tuple' a s) b -> m (Step (Tuple' a s) b)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (Tuple' a s) b -> m (Step (Tuple' a s) b))
-> Step (Tuple' a s) b -> m (Step (Tuple' a s) b)
forall a b. (a -> b) -> a -> b
$ Tuple' a s -> Step (Tuple' a s) b
forall s a. s -> Step s a
Skip (a -> s -> Tuple' a s
forall a b. a -> b -> Tuple' a b
Tuple' a
a s
s)
            Step s b
Stop   -> Step (Tuple' a s) b -> m (Step (Tuple' a s) b)
forall (m :: * -> *) a. Monad m => a -> m a
return Step (Tuple' a s) b
forall s a. Step s a
Stop

-- | Same as 'takeWhile' but with a monadic predicate.
--
-- /Since: 0.8.0/
--
{-# INLINE_NORMAL takeWhileM #-}
takeWhileM :: Monad m => (b -> m Bool) -> Unfold m a b -> Unfold m a b
-- XXX Check if the compiler simplifies the following to the same as the custom
-- implementation below (the Tuple' should help eliminate the unused param):
--
-- takeWhileM f = takeWhileMWithInput (\_ b -> f b)
takeWhileM :: (b -> m Bool) -> Unfold m a b -> Unfold m a b
takeWhileM b -> m Bool
f (Unfold s -> m (Step s b)
step1 a -> m s
inject1) = (s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
Unfold s -> m (Step s b)
step a -> m s
inject1
  where
    {-# INLINE_LATE step #-}
    step :: s -> m (Step s b)
step s
st = do
        Step s b
r <- s -> m (Step s b)
step1 s
st
        case Step s b
r of
            Yield b
x s
s -> do
                Bool
b <- b -> m Bool
f b
x
                Step s b -> m (Step s b)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step s b -> m (Step s b)) -> Step s b -> m (Step s b)
forall a b. (a -> b) -> a -> b
$ if Bool
b then b -> s -> Step s b
forall s a. a -> s -> Step s a
Yield b
x s
s else Step s b
forall s a. Step s a
Stop
            Skip s
s -> Step s b -> m (Step s b)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step s b -> m (Step s b)) -> Step s b -> m (Step s b)
forall a b. (a -> b) -> a -> b
$ s -> Step s b
forall s a. s -> Step s a
Skip s
s
            Step s b
Stop   -> Step s b -> m (Step s b)
forall (m :: * -> *) a. Monad m => a -> m a
return Step s b
forall s a. Step s a
Stop

-- | End the stream generated by the 'Unfold' as soon as the predicate fails
-- on an element.
--
-- /Since: 0.8.0/
--
{-# INLINE takeWhile #-}
takeWhile :: Monad m => (b -> Bool) -> Unfold m a b -> Unfold m a b
takeWhile :: (b -> Bool) -> Unfold m a b -> Unfold m a b
takeWhile b -> Bool
f = (b -> m Bool) -> Unfold m a b -> Unfold m a b
forall (m :: * -> *) b a.
Monad m =>
(b -> m Bool) -> Unfold m a b -> Unfold m a b
takeWhileM (Bool -> m Bool
forall (m :: * -> *) a. Monad m => a -> m a
return (Bool -> m Bool) -> (b -> Bool) -> b -> m Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. b -> Bool
f)

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

-- | Map a function on the output of the unfold (the type @b@).
--
-- /Pre-release/
{-# INLINE_NORMAL map #-}
map :: Functor m => (b -> c) -> Unfold m a b -> Unfold m a c
map :: (b -> c) -> Unfold m a b -> Unfold m a c
map b -> c
f (Unfold s -> m (Step s b)
ustep a -> m s
uinject) = (s -> m (Step s c)) -> (a -> m s) -> Unfold m a c
forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
Unfold s -> m (Step s c)
step a -> m s
uinject

    where

    {-# INLINE_LATE step #-}
    step :: s -> m (Step s c)
step s
st = (Step s b -> Step s c) -> m (Step s b) -> m (Step s c)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((b -> c) -> Step s b -> Step s c
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap b -> c
f) (s -> m (Step s b)
ustep s
st)

-- | Maps a function on the output of the unfold (the type @b@).
instance Functor m => Functor (Unfold m a) where
    {-# INLINE fmap #-}
    fmap :: (a -> b) -> Unfold m a a -> Unfold m a b
fmap = (a -> b) -> Unfold m a a -> Unfold m a b
forall (m :: * -> *) b c a.
Functor m =>
(b -> c) -> Unfold m a b -> Unfold m a c
map

------------------------------------------------------------------------------
-- Applicative
------------------------------------------------------------------------------

-- | The unfold discards its input and generates a function stream using the
-- supplied monadic action.
--
-- /Pre-release/
{-# INLINE fromEffect #-}
fromEffect :: Applicative m => m b -> Unfold m a b
fromEffect :: m b -> Unfold m a b
fromEffect m b
m = (Bool -> m (Step Bool b)) -> (a -> m Bool) -> Unfold m a b
forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
Unfold Bool -> m (Step Bool b)
step a -> m Bool
forall (f :: * -> *) p. Applicative f => p -> f Bool
inject

    where

    inject :: p -> f Bool
inject p
_ = Bool -> f Bool
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
False

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

-- | Discards the unfold input and always returns the argument of 'fromPure'.
--
-- > fromPure = fromEffect . pure
--
-- /Pre-release/
fromPure :: Applicative m => b -> Unfold m a b
fromPure :: b -> Unfold m a b
fromPure = m b -> Unfold m a b
forall (m :: * -> *) b a. Applicative m => m b -> Unfold m a b
fromEffect (m b -> Unfold m a b) -> (b -> m b) -> b -> Unfold m a b
forall b c a. (b -> c) -> (a -> b) -> a -> c
Prelude.. b -> m b
forall (f :: * -> *) a. Applicative f => a -> f a
pure

-- | Outer product discarding the first element.
--
-- /Unimplemented/
--
{-# INLINE_NORMAL apSequence #-}
apSequence :: -- Monad m =>
    Unfold m a b -> Unfold m a c -> Unfold m a c
apSequence :: Unfold m a b -> Unfold m a c -> Unfold m a c
apSequence (Unfold s -> m (Step s b)
_step1 a -> m s
_inject1) (Unfold s -> m (Step s c)
_step2 a -> m s
_inject2) = Unfold m a c
forall a. HasCallStack => a
undefined

-- | Outer product discarding the second element.
--
-- /Unimplemented/
--
{-# INLINE_NORMAL apDiscardSnd #-}
apDiscardSnd :: -- Monad m =>
    Unfold m a b -> Unfold m a c -> Unfold m a b
apDiscardSnd :: Unfold m a b -> Unfold m a c -> Unfold m a b
apDiscardSnd (Unfold s -> m (Step s b)
_step1 a -> m s
_inject1) (Unfold s -> m (Step s c)
_step2 a -> m s
_inject2) = Unfold m a b
forall a. HasCallStack => a
undefined

data Cross a s1 b s2 = CrossOuter a s1 | CrossInner a s1 b s2

-- | Create a cross product (vector product or cartesian product) of the
-- output streams of two unfolds using a monadic combining function.
--
-- /Pre-release/
{-# INLINE_NORMAL crossWithM #-}
crossWithM :: Monad m =>
    (b -> c -> m d) -> Unfold m a b -> Unfold m a c -> Unfold m a d
crossWithM :: (b -> c -> m d) -> Unfold m a b -> Unfold m a c -> Unfold m a d
crossWithM b -> c -> m d
f (Unfold s -> m (Step s b)
step1 a -> m s
inject1) (Unfold s -> m (Step s c)
step2 a -> m s
inject2) = (Cross a s b s -> m (Step (Cross a s b s) d))
-> (a -> m (Cross a s b s)) -> Unfold m a d
forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
Unfold Cross a s b s -> m (Step (Cross a s b s) d)
step a -> m (Cross a s b s)
forall b s2. a -> m (Cross a s b s2)
inject

    where

    inject :: a -> m (Cross a s b s2)
inject a
a = do
        s
s1 <- a -> m s
inject1 a
a
        Cross a s b s2 -> m (Cross a s b s2)
forall (m :: * -> *) a. Monad m => a -> m a
return (Cross a s b s2 -> m (Cross a s b s2))
-> Cross a s b s2 -> m (Cross a s b s2)
forall a b. (a -> b) -> a -> b
$ a -> s -> Cross a s b s2
forall a s1 b s2. a -> s1 -> Cross a s1 b s2
CrossOuter a
a s
s1

    {-# INLINE_LATE step #-}
    step :: Cross a s b s -> m (Step (Cross a s b s) d)
step (CrossOuter a
a s
s1) = do
        Step s b
r <- s -> m (Step s b)
step1 s
s1
        case Step s b
r of
            Yield b
b s
s -> do
                s
s2 <- a -> m s
inject2 a
a
                Step (Cross a s b s) d -> m (Step (Cross a s b s) d)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (Cross a s b s) d -> m (Step (Cross a s b s) d))
-> Step (Cross a s b s) d -> m (Step (Cross a s b s) d)
forall a b. (a -> b) -> a -> b
$ Cross a s b s -> Step (Cross a s b s) d
forall s a. s -> Step s a
Skip (a -> s -> b -> s -> Cross a s b s
forall a s1 b s2. a -> s1 -> b -> s2 -> Cross a s1 b s2
CrossInner a
a s
s b
b s
s2)
            Skip s
s    -> Step (Cross a s b s) d -> m (Step (Cross a s b s) d)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (Cross a s b s) d -> m (Step (Cross a s b s) d))
-> Step (Cross a s b s) d -> m (Step (Cross a s b s) d)
forall a b. (a -> b) -> a -> b
$ Cross a s b s -> Step (Cross a s b s) d
forall s a. s -> Step s a
Skip (a -> s -> Cross a s b s
forall a s1 b s2. a -> s1 -> Cross a s1 b s2
CrossOuter a
a s
s)
            Step s b
Stop      -> Step (Cross a s b s) d -> m (Step (Cross a s b s) d)
forall (m :: * -> *) a. Monad m => a -> m a
return Step (Cross a s b s) d
forall s a. Step s a
Stop

    step (CrossInner a
a s
s1 b
b s
s2) = do
        Step s c
r <- s -> m (Step s c)
step2 s
s2
        case Step s c
r of
            Yield c
c s
s -> b -> c -> m d
f b
b c
c m d
-> (d -> m (Step (Cross a s b s) d)) -> m (Step (Cross a s b s) d)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \d
d -> Step (Cross a s b s) d -> m (Step (Cross a s b s) d)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (Cross a s b s) d -> m (Step (Cross a s b s) d))
-> Step (Cross a s b s) d -> m (Step (Cross a s b s) d)
forall a b. (a -> b) -> a -> b
$ d -> Cross a s b s -> Step (Cross a s b s) d
forall s a. a -> s -> Step s a
Yield d
d (a -> s -> b -> s -> Cross a s b s
forall a s1 b s2. a -> s1 -> b -> s2 -> Cross a s1 b s2
CrossInner a
a s
s1 b
b s
s)
            Skip s
s    -> Step (Cross a s b s) d -> m (Step (Cross a s b s) d)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (Cross a s b s) d -> m (Step (Cross a s b s) d))
-> Step (Cross a s b s) d -> m (Step (Cross a s b s) d)
forall a b. (a -> b) -> a -> b
$ Cross a s b s -> Step (Cross a s b s) d
forall s a. s -> Step s a
Skip (a -> s -> b -> s -> Cross a s b s
forall a s1 b s2. a -> s1 -> b -> s2 -> Cross a s1 b s2
CrossInner a
a s
s1 b
b s
s)
            Step s c
Stop      -> Step (Cross a s b s) d -> m (Step (Cross a s b s) d)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (Cross a s b s) d -> m (Step (Cross a s b s) d))
-> Step (Cross a s b s) d -> m (Step (Cross a s b s) d)
forall a b. (a -> b) -> a -> b
$ Cross a s b s -> Step (Cross a s b s) d
forall s a. s -> Step s a
Skip (a -> s -> Cross a s b s
forall a s1 b s2. a -> s1 -> Cross a s1 b s2
CrossOuter a
a s
s1)

-- | Like 'crossWithM' but uses a pure combining function.
--
-- > crossWith f = crossWithM (\b c -> return $ f b c)
--
-- >>> u1 = Unfold.lmap fst Unfold.fromList
-- >>> u2 = Unfold.lmap snd Unfold.fromList
-- >>> u = Unfold.crossWith (,) u1 u2
-- >>> Unfold.fold Fold.toList u ([1,2,3], [4,5,6])
-- [(1,4),(1,5),(1,6),(2,4),(2,5),(2,6),(3,4),(3,5),(3,6)]
--
-- /Since: 0.8.0/
--
{-# INLINE crossWith #-}
crossWith :: Monad m =>
    (b -> c -> d) -> Unfold m a b -> Unfold m a c -> Unfold m a d
crossWith :: (b -> c -> d) -> Unfold m a b -> Unfold m a c -> Unfold m a d
crossWith b -> c -> d
f = (b -> c -> m d) -> Unfold m a b -> Unfold m a c -> Unfold m a d
forall (m :: * -> *) b c d a.
Monad m =>
(b -> c -> m d) -> Unfold m a b -> Unfold m a c -> Unfold m a d
crossWithM (\b
b c
c -> d -> m d
forall (m :: * -> *) a. Monad m => a -> m a
return (d -> m d) -> d -> m d
forall a b. (a -> b) -> a -> b
$ b -> c -> d
f b
b c
c)

-- | See 'crossWith'.
--
-- > cross = crossWith (,)
--
-- To cross the streams from a tuple we can write:
--
-- @
-- crossProduct :: Monad m => Unfold m a b -> Unfold m c d -> Unfold m (a, c) (b, d)
-- crossProduct u1 u2 = cross (lmap fst u1) (lmap snd u2)
-- @
--
-- /Pre-release/
{-# INLINE_NORMAL cross #-}
cross :: Monad m => Unfold m a b -> Unfold m a c -> Unfold m a (b, c)
cross :: Unfold m a b -> Unfold m a c -> Unfold m a (b, c)
cross = (b -> c -> (b, c))
-> Unfold m a b -> Unfold m a c -> Unfold m a (b, c)
forall (m :: * -> *) b c d a.
Monad m =>
(b -> c -> d) -> Unfold m a b -> Unfold m a c -> Unfold m a d
crossWith (,)

apply :: Monad m => Unfold m a (b -> c) -> Unfold m a b -> Unfold m a c
apply :: Unfold m a (b -> c) -> Unfold m a b -> Unfold m a c
apply Unfold m a (b -> c)
u1 Unfold m a b
u2 = ((b -> c, b) -> c) -> Unfold m a (b -> c, b) -> Unfold m a c
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\(b -> c
a, b
b) -> b -> c
a b
b) (Unfold m a (b -> c) -> Unfold m a b -> Unfold m a (b -> c, b)
forall (m :: * -> *) a b c.
Monad m =>
Unfold m a b -> Unfold m a c -> Unfold m a (b, c)
cross Unfold m a (b -> c)
u1 Unfold m a b
u2)

{-
-- | Example:
--
-- >>> rlist = Unfold.lmap fst Unfold.fromList
-- >>> llist = Unfold.lmap snd Unfold.fromList
-- >>> Stream.toList $ Stream.unfold ((,) <$> rlist <*> llist) ([1,2],[3,4])
-- [(1,3),(1,4),(2,3),(2,4)]
--
instance Monad m => Applicative (Unfold m a) where
    {-# INLINE pure #-}
    pure = fromPure

    {-# INLINE (<*>) #-}
    (<*>) = apply

    -- {-# INLINE (*>) #-}
    -- (*>) = apSequence

    -- {-# INLINE (<*) #-}
    -- (<*) = apDiscardSnd
-}

------------------------------------------------------------------------------
-- Monad
------------------------------------------------------------------------------

data ConcatMapState m b s1 x =
      ConcatMapOuter x s1
    | forall s2. ConcatMapInner x s1 s2 (s2 -> m (Step s2 b))

-- | Map an unfold generating action to each element of an unfold and
-- flatten the results into a single stream.
--
{-# INLINE_NORMAL concatMapM #-}
concatMapM :: Monad m
    => (b -> m (Unfold m a c)) -> Unfold m a b -> Unfold m a c
concatMapM :: (b -> m (Unfold m a c)) -> Unfold m a b -> Unfold m a c
concatMapM b -> m (Unfold m a c)
f (Unfold s -> m (Step s b)
step1 a -> m s
inject1) = (ConcatMapState m c s a -> m (Step (ConcatMapState m c s a) c))
-> (a -> m (ConcatMapState m c s a)) -> Unfold m a c
forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
Unfold ConcatMapState m c s a -> m (Step (ConcatMapState m c s a) c)
step a -> m (ConcatMapState m c s a)
forall (m :: * -> *) b. a -> m (ConcatMapState m b s a)
inject

    where

    inject :: a -> m (ConcatMapState m b s a)
inject a
x = do
        s
s <- a -> m s
inject1 a
x
        ConcatMapState m b s a -> m (ConcatMapState m b s a)
forall (m :: * -> *) a. Monad m => a -> m a
return (ConcatMapState m b s a -> m (ConcatMapState m b s a))
-> ConcatMapState m b s a -> m (ConcatMapState m b s a)
forall a b. (a -> b) -> a -> b
$ a -> s -> ConcatMapState m b s a
forall (m :: * -> *) b s1 x. x -> s1 -> ConcatMapState m b s1 x
ConcatMapOuter a
x s
s

    {-# INLINE_LATE step #-}
    step :: ConcatMapState m c s a -> m (Step (ConcatMapState m c s a) c)
step (ConcatMapOuter a
seed s
st) = do
        Step s b
r <- s -> m (Step s b)
step1 s
st
        case Step s b
r of
            Yield b
x s
s -> do
                Unfold s -> m (Step s c)
step2 a -> m s
inject2 <- b -> m (Unfold m a c)
f b
x
                s
innerSt <- a -> m s
inject2 a
seed
                Step (ConcatMapState m c s a) c
-> m (Step (ConcatMapState m c s a) c)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (ConcatMapState m c s a) c
 -> m (Step (ConcatMapState m c s a) c))
-> Step (ConcatMapState m c s a) c
-> m (Step (ConcatMapState m c s a) c)
forall a b. (a -> b) -> a -> b
$ ConcatMapState m c s a -> Step (ConcatMapState m c s a) c
forall s a. s -> Step s a
Skip (a -> s -> s -> (s -> m (Step s c)) -> ConcatMapState m c s a
forall (m :: * -> *) b s1 x s2.
x -> s1 -> s2 -> (s2 -> m (Step s2 b)) -> ConcatMapState m b s1 x
ConcatMapInner a
seed s
s s
innerSt s -> m (Step s c)
step2)
            Skip s
s    -> Step (ConcatMapState m c s a) c
-> m (Step (ConcatMapState m c s a) c)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (ConcatMapState m c s a) c
 -> m (Step (ConcatMapState m c s a) c))
-> Step (ConcatMapState m c s a) c
-> m (Step (ConcatMapState m c s a) c)
forall a b. (a -> b) -> a -> b
$ ConcatMapState m c s a -> Step (ConcatMapState m c s a) c
forall s a. s -> Step s a
Skip (a -> s -> ConcatMapState m c s a
forall (m :: * -> *) b s1 x. x -> s1 -> ConcatMapState m b s1 x
ConcatMapOuter a
seed s
s)
            Step s b
Stop      -> Step (ConcatMapState m c s a) c
-> m (Step (ConcatMapState m c s a) c)
forall (m :: * -> *) a. Monad m => a -> m a
return Step (ConcatMapState m c s a) c
forall s a. Step s a
Stop

    step (ConcatMapInner a
seed s
ost s2
ist s2 -> m (Step s2 c)
istep) = do
        Step s2 c
r <- s2 -> m (Step s2 c)
istep s2
ist
        Step (ConcatMapState m c s a) c
-> m (Step (ConcatMapState m c s a) c)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (ConcatMapState m c s a) c
 -> m (Step (ConcatMapState m c s a) c))
-> Step (ConcatMapState m c s a) c
-> m (Step (ConcatMapState m c s a) c)
forall a b. (a -> b) -> a -> b
$ case Step s2 c
r of
            Yield c
x s2
s -> c -> ConcatMapState m c s a -> Step (ConcatMapState m c s a) c
forall s a. a -> s -> Step s a
Yield c
x (a -> s -> s2 -> (s2 -> m (Step s2 c)) -> ConcatMapState m c s a
forall (m :: * -> *) b s1 x s2.
x -> s1 -> s2 -> (s2 -> m (Step s2 b)) -> ConcatMapState m b s1 x
ConcatMapInner a
seed s
ost s2
s s2 -> m (Step s2 c)
istep)
            Skip s2
s    -> ConcatMapState m c s a -> Step (ConcatMapState m c s a) c
forall s a. s -> Step s a
Skip (a -> s -> s2 -> (s2 -> m (Step s2 c)) -> ConcatMapState m c s a
forall (m :: * -> *) b s1 x s2.
x -> s1 -> s2 -> (s2 -> m (Step s2 b)) -> ConcatMapState m b s1 x
ConcatMapInner a
seed s
ost s2
s s2 -> m (Step s2 c)
istep)
            Step s2 c
Stop      -> ConcatMapState m c s a -> Step (ConcatMapState m c s a) c
forall s a. s -> Step s a
Skip (a -> s -> ConcatMapState m c s a
forall (m :: * -> *) b s1 x. x -> s1 -> ConcatMapState m b s1 x
ConcatMapOuter a
seed s
ost)

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

infixl 1 `bind`

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

{-
-- Note: concatMap and Monad instance for unfolds have performance comparable
-- to Stream. In fact, concatMap is slower than Stream, that may be some
-- optimization issue though.
--
-- Monad allows an unfold to depend on the output of a previous unfold.
-- However, it is probably easier to use streams in such situations.
--
-- | Example:
--
-- >>> :{
--  u = do
--   x <- Unfold.enumerateFromToIntegral 4
--   y <- Unfold.enumerateFromToIntegral x
--   return (x, y)
-- :}
-- >>> Stream.toList $ Stream.unfold u 1
-- [(1,1),(2,1),(2,2),(3,1),(3,2),(3,3),(4,1),(4,2),(4,3),(4,4)]
--
instance Monad m => Monad (Unfold m a) where
    {-# INLINE return #-}
    return = pure

    {-# INLINE (>>=) #-}
    (>>=) = flip concatMap

    -- {-# INLINE (>>) #-}
    -- (>>) = (*>)
-}

-------------------------------------------------------------------------------
-- Category
-------------------------------------------------------------------------------

-- | Lift a monadic function into an unfold. The unfold generates a singleton
-- stream.
--
-- /Since: 0.8.0/

{-# INLINE functionM #-}
functionM :: Applicative m => (a -> m b) -> Unfold m a b
functionM :: (a -> m b) -> Unfold m a b
functionM a -> m b
f = (Maybe a -> m (Step (Maybe a) b))
-> (a -> m (Maybe a)) -> Unfold m a b
forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
Unfold Maybe a -> m (Step (Maybe a) b)
forall a. Maybe a -> m (Step (Maybe a) b)
step a -> m (Maybe a)
forall (f :: * -> *) a. Applicative f => a -> f (Maybe a)
inject

    where

    inject :: a -> f (Maybe a)
inject a
x = Maybe a -> f (Maybe a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe a -> f (Maybe a)) -> Maybe a -> f (Maybe a)
forall a b. (a -> b) -> a -> b
$ a -> Maybe a
forall a. a -> Maybe a
Just a
x

    {-# INLINE_LATE step #-}
    step :: Maybe a -> m (Step (Maybe a) b)
step (Just a
x) = (b -> Maybe a -> Step (Maybe a) b
forall s a. a -> s -> Step s a
`Yield` Maybe a
forall a. Maybe a
Nothing) (b -> Step (Maybe a) b) -> m b -> m (Step (Maybe a) b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> m b
f a
x
    step Maybe a
Nothing = Step (Maybe a) b -> m (Step (Maybe a) b)
forall (f :: * -> *) a. Applicative f => a -> f a
pure Step (Maybe a) b
forall s a. Step s a
Stop

-- | Lift a pure function into an unfold. The unfold generates a singleton
-- stream.
--
-- > function f = functionM $ return . f
--
-- /Since: 0.8.0/

{-# INLINE function #-}
function :: Applicative m => (a -> b) -> Unfold m a b
function :: (a -> b) -> Unfold m a b
function a -> b
f = (a -> m b) -> Unfold m a b
forall (m :: * -> *) a b.
Applicative m =>
(a -> m b) -> Unfold m a b
functionM ((a -> m b) -> Unfold m a b) -> (a -> m b) -> Unfold m a b
forall a b. (a -> b) -> a -> b
$ b -> m b
forall (f :: * -> *) a. Applicative f => a -> f a
pure (b -> m b) -> (a -> b) -> a -> m b
forall b c a. (b -> c) -> (a -> b) -> a -> c
Prelude.. a -> b
f

-- | Identity unfold. The unfold generates a singleton stream having the input
-- as the only element.
--
-- > identity = function Prelude.id
--
-- /Pre-release/
{-# INLINE identity #-}
identity :: Applicative m => Unfold m a a
identity :: Unfold m a a
identity = (a -> a) -> Unfold m a a
forall (m :: * -> *) a b. Applicative m => (a -> b) -> Unfold m a b
function a -> a
forall a. a -> a
Prelude.id

{-# ANN type ConcatState Fuse #-}
data ConcatState s1 s2 = ConcatOuter s1 | ConcatInner s1 s2

-- | Apply the second unfold to each output element of the first unfold and
-- flatten the output in a single stream.
--
-- /Since: 0.8.0/
--
{-# INLINE_NORMAL many #-}
many :: Monad m => Unfold m a b -> Unfold m b c -> Unfold m a c
many :: Unfold m a b -> Unfold m b c -> Unfold m a c
many (Unfold s -> m (Step s b)
step1 a -> m s
inject1) (Unfold s -> m (Step s c)
step2 b -> m s
inject2) = (ConcatState s s -> m (Step (ConcatState s s) c))
-> (a -> m (ConcatState s s)) -> Unfold m a c
forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
Unfold ConcatState s s -> m (Step (ConcatState s s) c)
step a -> m (ConcatState s s)
forall s2. a -> m (ConcatState s s2)
inject

    where

    inject :: a -> m (ConcatState s s2)
inject a
x = do
        s
s <- a -> m s
inject1 a
x
        ConcatState s s2 -> m (ConcatState s s2)
forall (m :: * -> *) a. Monad m => a -> m a
return (ConcatState s s2 -> m (ConcatState s s2))
-> ConcatState s s2 -> m (ConcatState s s2)
forall a b. (a -> b) -> a -> b
$ s -> ConcatState s s2
forall s1 s2. s1 -> ConcatState s1 s2
ConcatOuter s
s

    {-# INLINE_LATE step #-}
    step :: ConcatState s s -> m (Step (ConcatState s s) c)
step (ConcatOuter s
st) = do
        Step s b
r <- s -> m (Step s b)
step1 s
st
        case Step s b
r of
            Yield b
x s
s -> do
                s
innerSt <- b -> m s
inject2 b
x
                Step (ConcatState s s) c -> m (Step (ConcatState s s) c)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (ConcatState s s) c -> m (Step (ConcatState s s) c))
-> Step (ConcatState s s) c -> m (Step (ConcatState s s) c)
forall a b. (a -> b) -> a -> b
$ ConcatState s s -> Step (ConcatState s s) c
forall s a. s -> Step s a
Skip (s -> s -> ConcatState s s
forall s1 s2. s1 -> s2 -> ConcatState s1 s2
ConcatInner s
s s
innerSt)
            Skip s
s    -> Step (ConcatState s s) c -> m (Step (ConcatState s s) c)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (ConcatState s s) c -> m (Step (ConcatState s s) c))
-> Step (ConcatState s s) c -> m (Step (ConcatState s s) c)
forall a b. (a -> b) -> a -> b
$ ConcatState s s -> Step (ConcatState s s) c
forall s a. s -> Step s a
Skip (s -> ConcatState s s
forall s1 s2. s1 -> ConcatState s1 s2
ConcatOuter s
s)
            Step s b
Stop      -> Step (ConcatState s s) c -> m (Step (ConcatState s s) c)
forall (m :: * -> *) a. Monad m => a -> m a
return Step (ConcatState s s) c
forall s a. Step s a
Stop

    step (ConcatInner s
ost s
ist) = do
        Step s c
r <- s -> m (Step s c)
step2 s
ist
        Step (ConcatState s s) c -> m (Step (ConcatState s s) c)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (ConcatState s s) c -> m (Step (ConcatState s s) c))
-> Step (ConcatState s s) c -> m (Step (ConcatState s s) c)
forall a b. (a -> b) -> a -> b
$ case Step s c
r of
            Yield c
x s
s -> c -> ConcatState s s -> Step (ConcatState s s) c
forall s a. a -> s -> Step s a
Yield c
x (s -> s -> ConcatState s s
forall s1 s2. s1 -> s2 -> ConcatState s1 s2
ConcatInner s
ost s
s)
            Skip s
s    -> ConcatState s s -> Step (ConcatState s s) c
forall s a. s -> Step s a
Skip (s -> s -> ConcatState s s
forall s1 s2. s1 -> s2 -> ConcatState s1 s2
ConcatInner s
ost s
s)
            Step s c
Stop      -> ConcatState s s -> Step (ConcatState s s) c
forall s a. s -> Step s a
Skip (s -> ConcatState s s
forall s1 s2. s1 -> ConcatState s1 s2
ConcatOuter s
ost)

{-
-- XXX There are multiple possible ways to combine the unfolds, "many" appends
-- them, we could also have other variants of "many" e.g. manyInterleave.
-- Should we even have a category instance or just use these functions
-- directly?
--
instance Monad m => Category (Unfold m) where
    {-# INLINE id #-}
    id = identity

    {-# INLINE (.) #-}
    (.) = flip many
-}

-------------------------------------------------------------------------------
-- Zipping
-------------------------------------------------------------------------------

-- | Distribute the input to two unfolds and then zip the outputs to a single
-- stream using a monadic zip function.
--
-- Stops as soon as any of the unfolds stops.
--
-- /Pre-release/
{-# INLINE_NORMAL zipWithM #-}
zipWithM :: Monad m
    => (b -> c -> m d) -> Unfold m a b -> Unfold m a c -> Unfold m a d
zipWithM :: (b -> c -> m d) -> Unfold m a b -> Unfold m a c -> Unfold m a d
zipWithM b -> c -> m d
f (Unfold s -> m (Step s b)
step1 a -> m s
inject1) (Unfold s -> m (Step s c)
step2 a -> m s
inject2) = ((s, s, Maybe b) -> m (Step (s, s, Maybe b) d))
-> (a -> m (s, s, Maybe b)) -> Unfold m a d
forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
Unfold (s, s, Maybe b) -> m (Step (s, s, Maybe b) d)
step a -> m (s, s, Maybe b)
forall a. a -> m (s, s, Maybe a)
inject

    where

    inject :: a -> m (s, s, Maybe a)
inject a
x = do
        s
s1 <- a -> m s
inject1 a
x
        s
s2 <- a -> m s
inject2 a
x
        (s, s, Maybe a) -> m (s, s, Maybe a)
forall (m :: * -> *) a. Monad m => a -> m a
return (s
s1, s
s2, Maybe a
forall a. Maybe a
Nothing)

    {-# INLINE_LATE step #-}
    step :: (s, s, Maybe b) -> m (Step (s, s, Maybe b) d)
step (s
s1, s
s2, Maybe b
Nothing) = do
        Step s b
r <- s -> m (Step s b)
step1 s
s1
        Step (s, s, Maybe b) d -> m (Step (s, s, Maybe b) d)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (s, s, Maybe b) d -> m (Step (s, s, Maybe b) d))
-> Step (s, s, Maybe b) d -> m (Step (s, s, Maybe b) d)
forall a b. (a -> b) -> a -> b
$
          case Step s b
r of
            Yield b
x s
s -> (s, s, Maybe b) -> Step (s, s, Maybe b) d
forall s a. s -> Step s a
Skip (s
s, s
s2, b -> Maybe b
forall a. a -> Maybe a
Just b
x)
            Skip s
s    -> (s, s, Maybe b) -> Step (s, s, Maybe b) d
forall s a. s -> Step s a
Skip (s
s, s
s2, Maybe b
forall a. Maybe a
Nothing)
            Step s b
Stop      -> Step (s, s, Maybe b) d
forall s a. Step s a
Stop

    step (s
s1, s
s2, Just b
x) = do
        Step s c
r <- s -> m (Step s c)
step2 s
s2
        case Step s c
r of
            Yield c
y s
s -> do
                d
z <- b -> c -> m d
f b
x c
y
                Step (s, s, Maybe b) d -> m (Step (s, s, Maybe b) d)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (s, s, Maybe b) d -> m (Step (s, s, Maybe b) d))
-> Step (s, s, Maybe b) d -> m (Step (s, s, Maybe b) d)
forall a b. (a -> b) -> a -> b
$ d -> (s, s, Maybe b) -> Step (s, s, Maybe b) d
forall s a. a -> s -> Step s a
Yield d
z (s
s1, s
s, Maybe b
forall a. Maybe a
Nothing)
            Skip s
s -> Step (s, s, Maybe b) d -> m (Step (s, s, Maybe b) d)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (s, s, Maybe b) d -> m (Step (s, s, Maybe b) d))
-> Step (s, s, Maybe b) d -> m (Step (s, s, Maybe b) d)
forall a b. (a -> b) -> a -> b
$ (s, s, Maybe b) -> Step (s, s, Maybe b) d
forall s a. s -> Step s a
Skip (s
s1, s
s, b -> Maybe b
forall a. a -> Maybe a
Just b
x)
            Step s c
Stop   -> Step (s, s, Maybe b) d -> m (Step (s, s, Maybe b) d)
forall (m :: * -> *) a. Monad m => a -> m a
return Step (s, s, Maybe b) d
forall s a. Step s a
Stop

-- | Like 'zipWithM' but with a pure zip function.
--
-- >>> square = fmap (\x -> x * x) Unfold.fromList
-- >>> cube = fmap (\x -> x * x * x) Unfold.fromList
-- >>> u = Unfold.zipWith (,) square cube
-- >>> Unfold.fold Fold.toList u [1..5]
-- [(1,1),(4,8),(9,27),(16,64),(25,125)]
--
-- > zipWith f = zipWithM (\a b -> return $ f a b)
--
-- /Since: 0.8.0/
--
{-# INLINE zipWith #-}
zipWith :: Monad m
    => (b -> c -> d) -> Unfold m a b -> Unfold m a c -> Unfold m a d
zipWith :: (b -> c -> d) -> Unfold m a b -> Unfold m a c -> Unfold m a d
zipWith b -> c -> d
f = (b -> c -> m d) -> Unfold m a b -> Unfold m a c -> Unfold m a d
forall (m :: * -> *) b c d a.
Monad m =>
(b -> c -> m d) -> Unfold m a b -> Unfold m a c -> Unfold m a d
zipWithM (\b
a c
b -> d -> m d
forall (m :: * -> *) a. Monad m => a -> m a
return (b -> c -> d
f b
a c
b))

-------------------------------------------------------------------------------
-- Arrow
-------------------------------------------------------------------------------

{-
-- XXX There are multiple ways of combining the outputs of two unfolds, we
-- could zip, merge, append and more. What is the preferred way for Arrow
-- instance? Should we even have an arrow instance or just use these functions
-- directly?
--
-- | '***' is a zip like operation, in fact it is the same as @Unfold.zipWith
-- (,)@, '&&&' is a tee like operation  i.e. distributes the input to both the
-- unfolds and then zips the output.
--
{-# ANN module "HLint: ignore Use zip" #-}
instance Monad m => Arrow (Unfold m) where
    {-# INLINE arr #-}
    arr = function

    {-# INLINE (***) #-}
    u1 *** u2 = zipWith (,) (lmap fst u1) (lmap snd u2)
-}

------------------------------------------------------------------------------
-- Interleaving
------------------------------------------------------------------------------

-- We can possibly have an "interleave" operation to interleave the streams
-- from two seeds:
--
-- interleave :: Unfold m x a -> Unfold m x a -> Unfold m (x, x) a
--
-- Alternatively, we can use a signature like zipWith:
-- interleave :: Unfold m x a -> Unfold m x a -> Unfold m x a
--
-- We can implement this in terms of manyInterleave, but that may
-- not be as efficient as a custom implementation.
--
-- Similarly we can also have other binary combining ops like append, mergeBy.
-- We already have zipWith.
--

data ManyInterleaveState o i =
      ManyInterleaveOuter o [i]
    | ManyInterleaveInner o [i]
    | ManyInterleaveInnerL [i] [i]
    | ManyInterleaveInnerR [i] [i]

-- | 'Streamly.Internal.Data.Stream.IsStream.unfoldManyInterleave' for
-- documentation and notes.
--
-- This is almost identical to unfoldManyInterleave in StreamD module.
--
-- The 'many' combinator is in fact 'manyAppend' to be more explicit in naming.
--
-- /Internal/
{-# INLINE_NORMAL manyInterleave #-}
manyInterleave :: Monad m => Unfold m a b -> Unfold m c a -> Unfold m c b
manyInterleave :: Unfold m a b -> Unfold m c a -> Unfold m c b
manyInterleave (Unfold s -> m (Step s b)
istep a -> m s
iinject) (Unfold s -> m (Step s a)
ostep c -> m s
oinject) =
    (ManyInterleaveState s s -> m (Step (ManyInterleaveState s s) b))
-> (c -> m (ManyInterleaveState s s)) -> Unfold m c b
forall (m :: * -> *) a b s.
(s -> m (Step s b)) -> (a -> m s) -> Unfold m a b
Unfold ManyInterleaveState s s -> m (Step (ManyInterleaveState s s) b)
step c -> m (ManyInterleaveState s s)
forall i. c -> m (ManyInterleaveState s i)
inject

    where

    inject :: c -> m (ManyInterleaveState s i)
inject c
x = do
        s
ost <- c -> m s
oinject c
x
        ManyInterleaveState s i -> m (ManyInterleaveState s i)
forall (m :: * -> *) a. Monad m => a -> m a
return (s -> [i] -> ManyInterleaveState s i
forall o i. o -> [i] -> ManyInterleaveState o i
ManyInterleaveOuter s
ost [])

    {-# INLINE_LATE step #-}
    step :: ManyInterleaveState s s -> m (Step (ManyInterleaveState s s) b)
step (ManyInterleaveOuter s
o [s]
ls) = do
        Step s a
r <- s -> m (Step s a)
ostep s
o
        case Step s a
r of
            Yield a
a s
o' -> do
                s
i <- a -> m s
iinject a
a
                s
i s
-> m (Step (ManyInterleaveState s s) b)
-> m (Step (ManyInterleaveState s s) b)
`seq` Step (ManyInterleaveState s s) b
-> m (Step (ManyInterleaveState s s) b)
forall (m :: * -> *) a. Monad m => a -> m a
return (ManyInterleaveState s s -> Step (ManyInterleaveState s s) b
forall s a. s -> Step s a
Skip (s -> [s] -> ManyInterleaveState s s
forall o i. o -> [i] -> ManyInterleaveState o i
ManyInterleaveInner s
o' (s
i s -> [s] -> [s]
forall a. a -> [a] -> [a]
: [s]
ls)))
            Skip s
o' -> Step (ManyInterleaveState s s) b
-> m (Step (ManyInterleaveState s s) b)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (ManyInterleaveState s s) b
 -> m (Step (ManyInterleaveState s s) b))
-> Step (ManyInterleaveState s s) b
-> m (Step (ManyInterleaveState s s) b)
forall a b. (a -> b) -> a -> b
$ ManyInterleaveState s s -> Step (ManyInterleaveState s s) b
forall s a. s -> Step s a
Skip (s -> [s] -> ManyInterleaveState s s
forall o i. o -> [i] -> ManyInterleaveState o i
ManyInterleaveOuter s
o' [s]
ls)
            Step s a
Stop -> Step (ManyInterleaveState s s) b
-> m (Step (ManyInterleaveState s s) b)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (ManyInterleaveState s s) b
 -> m (Step (ManyInterleaveState s s) b))
-> Step (ManyInterleaveState s s) b
-> m (Step (ManyInterleaveState s s) b)
forall a b. (a -> b) -> a -> b
$ ManyInterleaveState s s -> Step (ManyInterleaveState s s) b
forall s a. s -> Step s a
Skip ([s] -> [s] -> ManyInterleaveState s s
forall o i. [i] -> [i] -> ManyInterleaveState o i
ManyInterleaveInnerL [s]
ls [])

    step (ManyInterleaveInner s
_ []) = m (Step (ManyInterleaveState s s) b)
forall a. HasCallStack => a
undefined
    step (ManyInterleaveInner s
o (s
st:[s]
ls)) = do
        Step s b
r <- s -> m (Step s b)
istep s
st
        Step (ManyInterleaveState s s) b
-> m (Step (ManyInterleaveState s s) b)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (ManyInterleaveState s s) b
 -> m (Step (ManyInterleaveState s s) b))
-> Step (ManyInterleaveState s s) b
-> m (Step (ManyInterleaveState s s) b)
forall a b. (a -> b) -> a -> b
$ case Step s b
r of
            Yield b
x s
s -> b -> ManyInterleaveState s s -> Step (ManyInterleaveState s s) b
forall s a. a -> s -> Step s a
Yield b
x (s -> [s] -> ManyInterleaveState s s
forall o i. o -> [i] -> ManyInterleaveState o i
ManyInterleaveOuter s
o (s
ss -> [s] -> [s]
forall a. a -> [a] -> [a]
:[s]
ls))
            Skip s
s    -> ManyInterleaveState s s -> Step (ManyInterleaveState s s) b
forall s a. s -> Step s a
Skip (s -> [s] -> ManyInterleaveState s s
forall o i. o -> [i] -> ManyInterleaveState o i
ManyInterleaveInner s
o (s
ss -> [s] -> [s]
forall a. a -> [a] -> [a]
:[s]
ls))
            Step s b
Stop      -> ManyInterleaveState s s -> Step (ManyInterleaveState s s) b
forall s a. s -> Step s a
Skip (s -> [s] -> ManyInterleaveState s s
forall o i. o -> [i] -> ManyInterleaveState o i
ManyInterleaveOuter s
o [s]
ls)

    step (ManyInterleaveInnerL [] []) = Step (ManyInterleaveState s s) b
-> m (Step (ManyInterleaveState s s) b)
forall (m :: * -> *) a. Monad m => a -> m a
return Step (ManyInterleaveState s s) b
forall s a. Step s a
Stop
    step (ManyInterleaveInnerL [] [s]
rs) =
        Step (ManyInterleaveState s s) b
-> m (Step (ManyInterleaveState s s) b)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (ManyInterleaveState s s) b
 -> m (Step (ManyInterleaveState s s) b))
-> Step (ManyInterleaveState s s) b
-> m (Step (ManyInterleaveState s s) b)
forall a b. (a -> b) -> a -> b
$ ManyInterleaveState s s -> Step (ManyInterleaveState s s) b
forall s a. s -> Step s a
Skip ([s] -> [s] -> ManyInterleaveState s s
forall o i. [i] -> [i] -> ManyInterleaveState o i
ManyInterleaveInnerR [] [s]
rs)

    step (ManyInterleaveInnerL (s
st:[s]
ls) [s]
rs) = do
        Step s b
r <- s -> m (Step s b)
istep s
st
        Step (ManyInterleaveState s s) b
-> m (Step (ManyInterleaveState s s) b)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (ManyInterleaveState s s) b
 -> m (Step (ManyInterleaveState s s) b))
-> Step (ManyInterleaveState s s) b
-> m (Step (ManyInterleaveState s s) b)
forall a b. (a -> b) -> a -> b
$ case Step s b
r of
            Yield b
x s
s -> b -> ManyInterleaveState s s -> Step (ManyInterleaveState s s) b
forall s a. a -> s -> Step s a
Yield b
x ([s] -> [s] -> ManyInterleaveState s s
forall o i. [i] -> [i] -> ManyInterleaveState o i
ManyInterleaveInnerL [s]
ls (s
ss -> [s] -> [s]
forall a. a -> [a] -> [a]
:[s]
rs))
            Skip s
s    -> ManyInterleaveState s s -> Step (ManyInterleaveState s s) b
forall s a. s -> Step s a
Skip ([s] -> [s] -> ManyInterleaveState s s
forall o i. [i] -> [i] -> ManyInterleaveState o i
ManyInterleaveInnerL (s
ss -> [s] -> [s]
forall a. a -> [a] -> [a]
:[s]
ls) [s]
rs)
            Step s b
Stop      -> ManyInterleaveState s s -> Step (ManyInterleaveState s s) b
forall s a. s -> Step s a
Skip ([s] -> [s] -> ManyInterleaveState s s
forall o i. [i] -> [i] -> ManyInterleaveState o i
ManyInterleaveInnerL [s]
ls [s]
rs)

    step (ManyInterleaveInnerR [] []) = Step (ManyInterleaveState s s) b
-> m (Step (ManyInterleaveState s s) b)
forall (m :: * -> *) a. Monad m => a -> m a
return Step (ManyInterleaveState s s) b
forall s a. Step s a
Stop
    step (ManyInterleaveInnerR [s]
ls []) =
        Step (ManyInterleaveState s s) b
-> m (Step (ManyInterleaveState s s) b)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (ManyInterleaveState s s) b
 -> m (Step (ManyInterleaveState s s) b))
-> Step (ManyInterleaveState s s) b
-> m (Step (ManyInterleaveState s s) b)
forall a b. (a -> b) -> a -> b
$ ManyInterleaveState s s -> Step (ManyInterleaveState s s) b
forall s a. s -> Step s a
Skip ([s] -> [s] -> ManyInterleaveState s s
forall o i. [i] -> [i] -> ManyInterleaveState o i
ManyInterleaveInnerL [s]
ls [])

    step (ManyInterleaveInnerR [s]
ls (s
st:[s]
rs)) = do
        Step s b
r <- s -> m (Step s b)
istep s
st
        Step (ManyInterleaveState s s) b
-> m (Step (ManyInterleaveState s s) b)
forall (m :: * -> *) a. Monad m => a -> m a
return (Step (ManyInterleaveState s s) b
 -> m (Step (ManyInterleaveState s s) b))
-> Step (ManyInterleaveState s s) b
-> m (Step (ManyInterleaveState s s) b)
forall a b. (a -> b) -> a -> b
$ case Step s b
r of
            Yield b
x s
s -> b -> ManyInterleaveState s s -> Step (ManyInterleaveState s s) b
forall s a. a -> s -> Step s a
Yield b
x ([s] -> [s] -> ManyInterleaveState s s
forall o i. [i] -> [i] -> ManyInterleaveState o i
ManyInterleaveInnerR (s
ss -> [s] -> [s]
forall a. a -> [a] -> [a]
:[s]
ls) [s]
rs)
            Skip s
s    -> ManyInterleaveState s s -> Step (ManyInterleaveState s s) b
forall s a. s -> Step s a
Skip ([s] -> [s] -> ManyInterleaveState s s
forall o i. [i] -> [i] -> ManyInterleaveState o i
ManyInterleaveInnerR [s]
ls (s
ss -> [s] -> [s]
forall a. a -> [a] -> [a]
:[s]
rs))
            Step s b
Stop      -> ManyInterleaveState s s -> Step (ManyInterleaveState s s) b
forall s a. s -> Step s a
Skip ([s] -> [s] -> ManyInterleaveState s s
forall o i. [i] -> [i] -> ManyInterleaveState o i
ManyInterleaveInnerR [s]
ls [s]
rs)