{-# LANGUAGE GADTs #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE LinearTypes #-}
{-# LANGUAGE QualifiedDo #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# OPTIONS_GHC -Wno-name-shadowing #-}
{-# OPTIONS_HADDOCK hide #-}

-- | This module contains all functions that do something with
-- multiple streams as input or output. This includes combining
-- streams, splitting a stream, etc.
module Streaming.Linear.Internal.Many
  ( -- * Operations that use or return multiple 'Stream's

    -- ** Zips and Unzip
    unzip,
    ZipResidual,
    ZipResidual3,
    zip,
    zipR,
    zipWith,
    zipWithR,
    zip3,
    zip3R,
    zipWith3,
    zipWith3R,
    Either3 (..),

    -- ** Merging
    -- $
    merge,
    mergeOn,
    mergeBy,
  )
where

import qualified Control.Functor.Linear as Control
import Prelude.Linear (($), (&))
import Streaming.Linear.Internal.Consume
import Streaming.Linear.Internal.Type
import Prelude (Either (..), Ord (..), Ordering (..))

-- # Zips and Unzip
-------------------------------------------------------------------------------

-- | The type
--
-- > Data.List.unzip     :: [(a,b)] -> ([a],[b])
--
--   might lead us to expect
--
-- > Streaming.unzip :: Stream (Of (a,b)) m r -> Stream (Of a) m (Stream (Of b) m r)
--
--   which would not stream, since it would have to accumulate the second stream (of @b@s).
--   Of course, @Data.List@ 'Data.List.unzip' doesn't stream either.
--
--   This @unzip@ does
--   stream, though of course you can spoil this by using e.g. 'toList':
--
-- @
-- \>\>\> let xs = Prelude.map (\x -> (x, Prelude.show x)) [1..5 :: Int]
--
-- \>\>\> S.toList $ S.toList $ S.unzip (S.each' xs)
-- ["1","2","3","4","5"] :> ([1,2,3,4,5] :> ())
--
-- \>\>\> Prelude.unzip xs
-- ([1,2,3,4,5],["1","2","3","4","5"])
-- @
--
--    Note the difference of order in the results. It may be of some use to think why.
--    The first application of 'toList' was applied to a stream of integers:
--
-- @
-- \>\>\> :t S.unzip $ S.each' xs
-- S.unzip $ S.each' xs :: Control.Monad m => Stream (Of Int) (Stream (Of String) m) ()
-- @
--
--    Like any fold, 'toList' takes no notice of the monad of effects.
--
-- > toList :: Control.Monad m => Stream (Of a) m r %1-> m (Of [a] r)
--
--    In the case at hand (since I am in @ghci@) @m = Stream (Of String) IO@.
--    So when I apply 'toList', I exhaust that stream of integers, folding
--    it into a list:
--
-- @
-- \>\>\> :t S.toList $ S.unzip $ S.each' xs
-- S.toList $ S.unzip $ S.each' xs
--  :: Control.Monad m => Stream (Of String) m (Of [Int] ())
-- @
--
--    When I apply 'toList' to /this/, I reduce everything to an ordinary action in @IO@,
--    and return a list of strings:
--
-- @
-- \>\>\> S.toList $ S.toList $ S.unzip (S.each' xs)
-- ["1","2","3","4","5"] :> ([1,2,3,4,5] :> ())
-- @
--
-- 'unzip' can be considered a special case of either 'unzips' or 'expand':
--
-- @
--  unzip = 'unzips' . 'maps' (\((a,b) :> x) -> Compose (a :> b :> x))
--  unzip = 'expand' $ \p ((a,b) :> abs) -> b :> p (a :> abs)
-- @
unzip ::
  (Control.Monad m) =>
  Stream (Of (a, b)) m r %1 ->
  Stream (Of a) (Stream (Of b) m) r
unzip :: forall (m :: * -> *) a b r.
Monad m =>
Stream (Of (a, b)) m r %1 -> Stream (Of a) (Stream (Of b) m) r
unzip = forall (m :: * -> *) a b r.
Monad m =>
Stream (Of (a, b)) m r %1 -> Stream (Of a) (Stream (Of b) m) r
loop
  where
    loop ::
      (Control.Monad m) =>
      Stream (Of (a, b)) m r %1 ->
      Stream (Of a) (Stream (Of b) m) r
    loop :: forall (m :: * -> *) a b r.
Monad m =>
Stream (Of (a, b)) m r %1 -> Stream (Of a) (Stream (Of b) m) r
loop Stream (Of (a, b)) m r
stream =
      Stream (Of (a, b)) m r
stream forall a b (p :: Multiplicity) (q :: Multiplicity).
a %p -> (a %p -> b) %q -> b
& \case
        Return r
r -> forall r (f :: * -> *) (m :: * -> *). r -> Stream f m r
Return r
r
        Effect m (Stream (Of (a, b)) m r)
m -> forall (m :: * -> *) (f :: * -> *) r.
m (Stream f m r) -> Stream f m r
Effect forall a b (p :: Multiplicity) (q :: Multiplicity).
(a %p -> b) %q -> a %p -> b
$ forall (f :: * -> *) a b.
Functor f =>
(a %1 -> b) %1 -> f a %1 -> f b
Control.fmap forall (m :: * -> *) a b r.
Monad m =>
Stream (Of (a, b)) m r %1 -> Stream (Of a) (Stream (Of b) m) r
loop forall a b (p :: Multiplicity) (q :: Multiplicity).
(a %p -> b) %q -> a %p -> b
$ forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a %1 -> t m a
Control.lift m (Stream (Of (a, b)) m r)
m
        Step ((a
a, b
b) :> Stream (Of (a, b)) m r
rest) -> forall (f :: * -> *) (m :: * -> *) r.
f (Stream f m r) -> Stream f m r
Step (a
a forall a b. a -> b -> Of a b
:> forall (m :: * -> *) (f :: * -> *) r.
m (Stream f m r) -> Stream f m r
Effect (forall (f :: * -> *) (m :: * -> *) r.
f (Stream f m r) -> Stream f m r
Step (b
b forall a b. a -> b -> Of a b
:> forall r (f :: * -> *) (m :: * -> *). r -> Stream f m r
Return (forall (m :: * -> *) a b r.
Monad m =>
Stream (Of (a, b)) m r %1 -> Stream (Of a) (Stream (Of b) m) r
loop Stream (Of (a, b)) m r
rest))))
{-# INLINEABLE unzip #-}

{- Remarks on the design of zip functions

Zip functions have two design choices:
(1) What do we do with the end-of-stream values of both streams?
(2) If the streams are of different length, do we keep or throw out the
remainder of the longer stream?

\* We are assuming not to take infinite streams as input and instead deal with
reasonably small finite streams.
\* To avoid making choices for the user, we keep both end-of-stream payloads
\* The default zips (ones without a prime in the name) use @effects@ to consume
the remainder stream after zipping. We include zip function variants that
return no remainder (for equal length streams), or the remainder of the
longer stream.

-}

data Either3 a b c where
  Left3 :: a %1 -> Either3 a b c
  Middle3 :: b %1 -> Either3 a b c
  Right3 :: c %1 -> Either3 a b c

-- | The remainder of zipping two streams
type ZipResidual a b m r1 r2 =
  Either3
    (r1, r2)
    (r1, Stream (Of b) m r2)
    (Stream (Of a) m r1, r2)

-- | @zipWithR@ zips two streams applying a function along the way,
-- keeping the remainder of zipping if there is one.  Note. If two streams have
-- the same length, but one needs to perform some effects to obtain the
-- end-of-stream result, that stream is treated as a residual.
zipWithR ::
  (Control.Monad m) =>
  (a -> b -> c) ->
  Stream (Of a) m r1 %1 ->
  Stream (Of b) m r2 %1 ->
  Stream (Of c) m (ZipResidual a b m r1 r2)
zipWithR :: forall (m :: * -> *) a b c r1 r2.
Monad m =>
(a -> b -> c)
-> Stream (Of a) m r1
%1 -> Stream (Of b) m r2
%1 -> Stream (Of c) m (ZipResidual a b m r1 r2)
zipWithR = forall (m :: * -> *) a b c r1 r2.
Monad m =>
(a -> b -> c)
-> Stream (Of a) m r1
%1 -> Stream (Of b) m r2
%1 -> Stream (Of c) m (ZipResidual a b m r1 r2)
loop
  where
    loop ::
      (Control.Monad m) =>
      (a -> b -> c) ->
      Stream (Of a) m r1 %1 ->
      Stream (Of b) m r2 %1 ->
      Stream (Of c) m (ZipResidual a b m r1 r2)
    loop :: forall (m :: * -> *) a b c r1 r2.
Monad m =>
(a -> b -> c)
-> Stream (Of a) m r1
%1 -> Stream (Of b) m r2
%1 -> Stream (Of c) m (ZipResidual a b m r1 r2)
loop a -> b -> c
f Stream (Of a) m r1
st1 Stream (Of b) m r2
st2 =
      Stream (Of a) m r1
st1 forall a b (p :: Multiplicity) (q :: Multiplicity).
a %p -> (a %p -> b) %q -> b
& \case
        Effect m (Stream (Of a) m r1)
ms -> forall (m :: * -> *) (f :: * -> *) r.
m (Stream f m r) -> Stream f m r
Effect forall a b (p :: Multiplicity) (q :: Multiplicity).
(a %p -> b) %q -> a %p -> b
$ forall (f :: * -> *) a b.
Functor f =>
(a %1 -> b) %1 -> f a %1 -> f b
Control.fmap (\Stream (Of a) m r1
s -> forall (m :: * -> *) a b c r1 r2.
Monad m =>
(a -> b -> c)
-> Stream (Of a) m r1
%1 -> Stream (Of b) m r2
%1 -> Stream (Of c) m (ZipResidual a b m r1 r2)
loop a -> b -> c
f Stream (Of a) m r1
s Stream (Of b) m r2
st2) m (Stream (Of a) m r1)
ms
        Return r1
r1 ->
          Stream (Of b) m r2
st2 forall a b (p :: Multiplicity) (q :: Multiplicity).
a %p -> (a %p -> b) %q -> b
& \case
            Return r2
r2 -> forall r (f :: * -> *) (m :: * -> *). r -> Stream f m r
Return forall a b (p :: Multiplicity) (q :: Multiplicity).
(a %p -> b) %q -> a %p -> b
$ forall a b c. a -> Either3 a b c
Left3 (r1
r1, r2
r2)
            Stream (Of b) m r2
st2' -> forall r (f :: * -> *) (m :: * -> *). r -> Stream f m r
Return forall a b (p :: Multiplicity) (q :: Multiplicity).
(a %p -> b) %q -> a %p -> b
$ forall b a c. b -> Either3 a b c
Middle3 (r1
r1, Stream (Of b) m r2
st2')
        Step (a
a :> Stream (Of a) m r1
as) ->
          Stream (Of b) m r2
st2 forall a b (p :: Multiplicity) (q :: Multiplicity).
a %p -> (a %p -> b) %q -> b
& \case
            Effect m (Stream (Of b) m r2)
ms ->
              forall (m :: * -> *) (f :: * -> *) r.
m (Stream f m r) -> Stream f m r
Effect forall a b (p :: Multiplicity) (q :: Multiplicity).
(a %p -> b) %q -> a %p -> b
$ forall (f :: * -> *) a b.
Functor f =>
(a %1 -> b) %1 -> f a %1 -> f b
Control.fmap (\Stream (Of b) m r2
s -> forall (m :: * -> *) a b c r1 r2.
Monad m =>
(a -> b -> c)
-> Stream (Of a) m r1
%1 -> Stream (Of b) m r2
%1 -> Stream (Of c) m (ZipResidual a b m r1 r2)
loop a -> b -> c
f (forall (f :: * -> *) (m :: * -> *) r.
f (Stream f m r) -> Stream f m r
Step (a
a forall a b. a -> b -> Of a b
:> Stream (Of a) m r1
as)) Stream (Of b) m r2
s) m (Stream (Of b) m r2)
ms
            Return r2
r2 -> forall r (f :: * -> *) (m :: * -> *). r -> Stream f m r
Return forall a b (p :: Multiplicity) (q :: Multiplicity).
(a %p -> b) %q -> a %p -> b
$ forall c a b. c -> Either3 a b c
Right3 (forall (f :: * -> *) (m :: * -> *) r.
f (Stream f m r) -> Stream f m r
Step (a
a forall a b. a -> b -> Of a b
:> Stream (Of a) m r1
as), r2
r2)
            Step (b
b :> Stream (Of b) m r2
bs) -> forall (f :: * -> *) (m :: * -> *) r.
f (Stream f m r) -> Stream f m r
Step forall a b (p :: Multiplicity) (q :: Multiplicity).
(a %p -> b) %q -> a %p -> b
$ (a -> b -> c
f a
a b
b) forall a b. a -> b -> Of a b
:> forall (m :: * -> *) a b c r1 r2.
Monad m =>
(a -> b -> c)
-> Stream (Of a) m r1
%1 -> Stream (Of b) m r2
%1 -> Stream (Of c) m (ZipResidual a b m r1 r2)
loop a -> b -> c
f Stream (Of a) m r1
as Stream (Of b) m r2
bs
{-# INLINEABLE zipWithR #-}

zipWith ::
  (Control.Monad m) =>
  (a -> b -> c) ->
  Stream (Of a) m r1 %1 ->
  Stream (Of b) m r2 %1 ->
  Stream (Of c) m (r1, r2)
zipWith :: forall (m :: * -> *) a b c r1 r2.
Monad m =>
(a -> b -> c)
-> Stream (Of a) m r1
%1 -> Stream (Of b) m r2
%1 -> Stream (Of c) m (r1, r2)
zipWith a -> b -> c
f Stream (Of a) m r1
s1 Stream (Of b) m r2
s2 = Control.do
  ZipResidual a b m r1 r2
result <- forall (m :: * -> *) a b c r1 r2.
Monad m =>
(a -> b -> c)
-> Stream (Of a) m r1
%1 -> Stream (Of b) m r2
%1 -> Stream (Of c) m (ZipResidual a b m r1 r2)
zipWithR a -> b -> c
f Stream (Of a) m r1
s1 Stream (Of b) m r2
s2
  ZipResidual a b m r1 r2
result forall a b (p :: Multiplicity) (q :: Multiplicity).
a %p -> (a %p -> b) %q -> b
& \case
    Left3 (r1, r2)
rets -> forall (m :: * -> *) a. Monad m => a %1 -> m a
Control.return (r1, r2)
rets
    Middle3 (r1
r1, Stream (Of b) m r2
s2') -> Control.do
      r2
r2 <- forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a %1 -> t m a
Control.lift forall a b (p :: Multiplicity) (q :: Multiplicity).
(a %p -> b) %q -> a %p -> b
$ forall a (m :: * -> *) r. Monad m => Stream (Of a) m r %1 -> m r
effects Stream (Of b) m r2
s2'
      forall (m :: * -> *) a. Monad m => a %1 -> m a
Control.return (r1
r1, r2
r2)
    Right3 (Stream (Of a) m r1
s1', r2
r2) -> Control.do
      r1
r1 <- forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a %1 -> t m a
Control.lift forall a b (p :: Multiplicity) (q :: Multiplicity).
(a %p -> b) %q -> a %p -> b
$ forall a (m :: * -> *) r. Monad m => Stream (Of a) m r %1 -> m r
effects Stream (Of a) m r1
s1'
      forall (m :: * -> *) a. Monad m => a %1 -> m a
Control.return (r1
r1, r2
r2)
{-# INLINEABLE zipWith #-}

-- | @zip@ zips two streams exhausing the remainder of the longer
-- stream and consuming its effects.
zip ::
  (Control.Monad m) =>
  Stream (Of a) m r1 %1 ->
  Stream (Of b) m r2 %1 ->
  Stream (Of (a, b)) m (r1, r2)
zip :: forall (m :: * -> *) a r1 b r2.
Monad m =>
Stream (Of a) m r1
%1 -> Stream (Of b) m r2 %1 -> Stream (Of (a, b)) m (r1, r2)
zip = forall (m :: * -> *) a b c r1 r2.
Monad m =>
(a -> b -> c)
-> Stream (Of a) m r1
%1 -> Stream (Of b) m r2
%1 -> Stream (Of c) m (r1, r2)
zipWith (,)
{-# INLINE zip #-}

-- | @zipR@ zips two streams keeping the remainder if there is one.
zipR ::
  (Control.Monad m) =>
  Stream (Of a) m r1 %1 ->
  Stream (Of b) m r2 %1 ->
  Stream (Of (a, b)) m (ZipResidual a b m r1 r2)
zipR :: forall (m :: * -> *) a r1 b r2.
Monad m =>
Stream (Of a) m r1
%1 -> Stream (Of b) m r2
%1 -> Stream (Of (a, b)) m (ZipResidual a b m r1 r2)
zipR = forall (m :: * -> *) a b c r1 r2.
Monad m =>
(a -> b -> c)
-> Stream (Of a) m r1
%1 -> Stream (Of b) m r2
%1 -> Stream (Of c) m (ZipResidual a b m r1 r2)
zipWithR (,)
{-# INLINE zipR #-}

-- Remark. For simplicity, we do not create an @Either7@ which is the
-- proper remainder type for 'zip3R'. Our type simply has one impossible
-- case which is when all three streams have a remainder.

-- | The (liberal) remainder of zipping three streams.
-- This has the downside that the possibility of three remainders
-- is allowed, though it will never occur.
type ZipResidual3 a b c m r1 r2 r3 =
  ( Either r1 (Stream (Of a) m r1),
    Either r2 (Stream (Of b) m r2),
    Either r3 (Stream (Of c) m r3)
  )

-- | Like @zipWithR@ but with three streams.
zipWith3R ::
  (Control.Monad m) =>
  (a -> b -> c -> d) ->
  Stream (Of a) m r1 %1 ->
  Stream (Of b) m r2 %1 ->
  Stream (Of c) m r3 %1 ->
  Stream (Of d) m (ZipResidual3 a b c m r1 r2 r3)
zipWith3R :: forall (m :: * -> *) a b c d r1 r2 r3.
Monad m =>
(a -> b -> c -> d)
-> Stream (Of a) m r1
%1 -> Stream (Of b) m r2
%1 -> Stream (Of c) m r3
%1 -> Stream (Of d) m (ZipResidual3 a b c m r1 r2 r3)
zipWith3R = forall (m :: * -> *) a b c d r1 r2 r3.
Monad m =>
(a -> b -> c -> d)
-> Stream (Of a) m r1
%1 -> Stream (Of b) m r2
%1 -> Stream (Of c) m r3
%1 -> Stream (Of d) m (ZipResidual3 a b c m r1 r2 r3)
loop
  where
    loop ::
      (Control.Monad m) =>
      (a -> b -> c -> d) ->
      Stream (Of a) m r1 %1 ->
      Stream (Of b) m r2 %1 ->
      Stream (Of c) m r3 %1 ->
      Stream (Of d) m (ZipResidual3 a b c m r1 r2 r3)
    loop :: forall (m :: * -> *) a b c d r1 r2 r3.
Monad m =>
(a -> b -> c -> d)
-> Stream (Of a) m r1
%1 -> Stream (Of b) m r2
%1 -> Stream (Of c) m r3
%1 -> Stream (Of d) m (ZipResidual3 a b c m r1 r2 r3)
loop a -> b -> c -> d
f Stream (Of a) m r1
s1 Stream (Of b) m r2
s2 Stream (Of c) m r3
s3 =
      Stream (Of a) m r1
s1 forall a b (p :: Multiplicity) (q :: Multiplicity).
a %p -> (a %p -> b) %q -> b
& \case
        Effect m (Stream (Of a) m r1)
ms -> forall (m :: * -> *) (f :: * -> *) r.
m (Stream f m r) -> Stream f m r
Effect forall a b (p :: Multiplicity) (q :: Multiplicity).
(a %p -> b) %q -> a %p -> b
$ forall (f :: * -> *) a b.
Functor f =>
(a %1 -> b) %1 -> f a %1 -> f b
Control.fmap (\Stream (Of a) m r1
s -> forall (m :: * -> *) a b c d r1 r2 r3.
Monad m =>
(a -> b -> c -> d)
-> Stream (Of a) m r1
%1 -> Stream (Of b) m r2
%1 -> Stream (Of c) m r3
%1 -> Stream (Of d) m (ZipResidual3 a b c m r1 r2 r3)
loop a -> b -> c -> d
f Stream (Of a) m r1
s Stream (Of b) m r2
s2 Stream (Of c) m r3
s3) m (Stream (Of a) m r1)
ms
        Return r1
r1 ->
          (Stream (Of b) m r2
s2, Stream (Of c) m r3
s3) forall a b (p :: Multiplicity) (q :: Multiplicity).
a %p -> (a %p -> b) %q -> b
& \case
            (Return r2
r2, Return r3
r3) -> forall r (f :: * -> *) (m :: * -> *). r -> Stream f m r
Return (forall a b. a -> Either a b
Left r1
r1, forall a b. a -> Either a b
Left r2
r2, forall a b. a -> Either a b
Left r3
r3)
            (Stream (Of b) m r2
s2', Stream (Of c) m r3
s3') -> forall r (f :: * -> *) (m :: * -> *). r -> Stream f m r
Return (forall a b. a -> Either a b
Left r1
r1, forall a b. b -> Either a b
Right Stream (Of b) m r2
s2', forall a b. b -> Either a b
Right Stream (Of c) m r3
s3')
        Step (a
a :> Stream (Of a) m r1
as) ->
          Stream (Of b) m r2
s2 forall a b (p :: Multiplicity) (q :: Multiplicity).
a %p -> (a %p -> b) %q -> b
& \case
            Effect m (Stream (Of b) m r2)
ms ->
              forall (m :: * -> *) (f :: * -> *) r.
m (Stream f m r) -> Stream f m r
Effect forall a b (p :: Multiplicity) (q :: Multiplicity).
(a %p -> b) %q -> a %p -> b
$
                forall (f :: * -> *) a b.
Functor f =>
(a %1 -> b) %1 -> f a %1 -> f b
Control.fmap (\Stream (Of b) m r2
s -> forall (m :: * -> *) a b c d r1 r2 r3.
Monad m =>
(a -> b -> c -> d)
-> Stream (Of a) m r1
%1 -> Stream (Of b) m r2
%1 -> Stream (Of c) m r3
%1 -> Stream (Of d) m (ZipResidual3 a b c m r1 r2 r3)
loop a -> b -> c -> d
f (forall (f :: * -> *) (m :: * -> *) r.
f (Stream f m r) -> Stream f m r
Step forall a b (p :: Multiplicity) (q :: Multiplicity).
(a %p -> b) %q -> a %p -> b
$ a
a forall a b. a -> b -> Of a b
:> Stream (Of a) m r1
as) Stream (Of b) m r2
s Stream (Of c) m r3
s3) m (Stream (Of b) m r2)
ms
            Return r2
r2 -> forall r (f :: * -> *) (m :: * -> *). r -> Stream f m r
Return (forall a b. b -> Either a b
Right (forall (f :: * -> *) (m :: * -> *) r.
f (Stream f m r) -> Stream f m r
Step forall a b (p :: Multiplicity) (q :: Multiplicity).
(a %p -> b) %q -> a %p -> b
$ a
a forall a b. a -> b -> Of a b
:> Stream (Of a) m r1
as), forall a b. a -> Either a b
Left r2
r2, forall a b. b -> Either a b
Right Stream (Of c) m r3
s3)
            Step (b
b :> Stream (Of b) m r2
bs) ->
              Stream (Of c) m r3
s3 forall a b (p :: Multiplicity) (q :: Multiplicity).
a %p -> (a %p -> b) %q -> b
& \case
                Effect m (Stream (Of c) m r3)
ms ->
                  forall (m :: * -> *) (f :: * -> *) r.
m (Stream f m r) -> Stream f m r
Effect forall a b (p :: Multiplicity) (q :: Multiplicity).
(a %p -> b) %q -> a %p -> b
$
                    forall (f :: * -> *) a b.
Functor f =>
(a %1 -> b) %1 -> f a %1 -> f b
Control.fmap (\Stream (Of c) m r3
s -> forall (m :: * -> *) a b c d r1 r2 r3.
Monad m =>
(a -> b -> c -> d)
-> Stream (Of a) m r1
%1 -> Stream (Of b) m r2
%1 -> Stream (Of c) m r3
%1 -> Stream (Of d) m (ZipResidual3 a b c m r1 r2 r3)
loop a -> b -> c -> d
f (forall (f :: * -> *) (m :: * -> *) r.
f (Stream f m r) -> Stream f m r
Step forall a b (p :: Multiplicity) (q :: Multiplicity).
(a %p -> b) %q -> a %p -> b
$ a
a forall a b. a -> b -> Of a b
:> Stream (Of a) m r1
as) (forall (f :: * -> *) (m :: * -> *) r.
f (Stream f m r) -> Stream f m r
Step forall a b (p :: Multiplicity) (q :: Multiplicity).
(a %p -> b) %q -> a %p -> b
$ b
b forall a b. a -> b -> Of a b
:> Stream (Of b) m r2
bs) Stream (Of c) m r3
s) m (Stream (Of c) m r3)
ms
                Return r3
r3 ->
                  forall r (f :: * -> *) (m :: * -> *). r -> Stream f m r
Return (forall a b. b -> Either a b
Right (forall (f :: * -> *) (m :: * -> *) r.
f (Stream f m r) -> Stream f m r
Step forall a b (p :: Multiplicity) (q :: Multiplicity).
(a %p -> b) %q -> a %p -> b
$ a
a forall a b. a -> b -> Of a b
:> Stream (Of a) m r1
as), forall a b. b -> Either a b
Right (forall (f :: * -> *) (m :: * -> *) r.
f (Stream f m r) -> Stream f m r
Step forall a b (p :: Multiplicity) (q :: Multiplicity).
(a %p -> b) %q -> a %p -> b
$ b
b forall a b. a -> b -> Of a b
:> Stream (Of b) m r2
bs), forall a b. a -> Either a b
Left r3
r3)
                Step (c
c :> Stream (Of c) m r3
cs) -> forall (f :: * -> *) (m :: * -> *) r.
f (Stream f m r) -> Stream f m r
Step forall a b (p :: Multiplicity) (q :: Multiplicity).
(a %p -> b) %q -> a %p -> b
$ (a -> b -> c -> d
f a
a b
b c
c) forall a b. a -> b -> Of a b
:> forall (m :: * -> *) a b c d r1 r2 r3.
Monad m =>
(a -> b -> c -> d)
-> Stream (Of a) m r1
%1 -> Stream (Of b) m r2
%1 -> Stream (Of c) m r3
%1 -> Stream (Of d) m (ZipResidual3 a b c m r1 r2 r3)
loop a -> b -> c -> d
f Stream (Of a) m r1
as Stream (Of b) m r2
bs Stream (Of c) m r3
cs
{-# INLINEABLE zipWith3R #-}

-- | Like @zipWith@ but with three streams
zipWith3 ::
  (Control.Monad m) =>
  (a -> b -> c -> d) ->
  Stream (Of a) m r1 %1 ->
  Stream (Of b) m r2 %1 ->
  Stream (Of c) m r3 %1 ->
  Stream (Of d) m (r1, r2, r3)
zipWith3 :: forall (m :: * -> *) a b c d r1 r2 r3.
Monad m =>
(a -> b -> c -> d)
-> Stream (Of a) m r1
%1 -> Stream (Of b) m r2
%1 -> Stream (Of c) m r3
%1 -> Stream (Of d) m (r1, r2, r3)
zipWith3 a -> b -> c -> d
f Stream (Of a) m r1
s1 Stream (Of b) m r2
s2 Stream (Of c) m r3
s3 = Control.do
  ZipResidual3 a b c m r1 r2 r3
result <- forall (m :: * -> *) a b c d r1 r2 r3.
Monad m =>
(a -> b -> c -> d)
-> Stream (Of a) m r1
%1 -> Stream (Of b) m r2
%1 -> Stream (Of c) m r3
%1 -> Stream (Of d) m (ZipResidual3 a b c m r1 r2 r3)
zipWith3R a -> b -> c -> d
f Stream (Of a) m r1
s1 Stream (Of b) m r2
s2 Stream (Of c) m r3
s3
  ZipResidual3 a b c m r1 r2 r3
result forall a b (p :: Multiplicity) (q :: Multiplicity).
a %p -> (a %p -> b) %q -> b
& \case
    (Either r1 (Stream (Of a) m r1)
res1, Either r2 (Stream (Of b) m r2)
res2, Either r3 (Stream (Of c) m r3)
res3) -> Control.do
      r1
r1 <- forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a %1 -> t m a
Control.lift forall a b (p :: Multiplicity) (q :: Multiplicity).
(a %p -> b) %q -> a %p -> b
$ forall (m :: * -> *) r a.
Monad m =>
Either r (Stream (Of a) m r) %1 -> m r
extractResult Either r1 (Stream (Of a) m r1)
res1
      r2
r2 <- forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a %1 -> t m a
Control.lift forall a b (p :: Multiplicity) (q :: Multiplicity).
(a %p -> b) %q -> a %p -> b
$ forall (m :: * -> *) r a.
Monad m =>
Either r (Stream (Of a) m r) %1 -> m r
extractResult Either r2 (Stream (Of b) m r2)
res2
      r3
r3 <- forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a %1 -> t m a
Control.lift forall a b (p :: Multiplicity) (q :: Multiplicity).
(a %p -> b) %q -> a %p -> b
$ forall (m :: * -> *) r a.
Monad m =>
Either r (Stream (Of a) m r) %1 -> m r
extractResult Either r3 (Stream (Of c) m r3)
res3
      forall (m :: * -> *) a. Monad m => a %1 -> m a
Control.return (r1
r1, r2
r2, r3
r3)
{-# INLINEABLE zipWith3 #-}

-- | Like @zipR@ but with three streams.
zip3 ::
  (Control.Monad m) =>
  Stream (Of a) m r1 %1 ->
  Stream (Of b) m r2 %1 ->
  Stream (Of c) m r3 %1 ->
  Stream (Of (a, b, c)) m (r1, r2, r3)
zip3 :: forall (m :: * -> *) a r1 b r2 c r3.
Monad m =>
Stream (Of a) m r1
%1 -> Stream (Of b) m r2
%1 -> Stream (Of c) m r3
%1 -> Stream (Of (a, b, c)) m (r1, r2, r3)
zip3 = forall (m :: * -> *) a b c d r1 r2 r3.
Monad m =>
(a -> b -> c -> d)
-> Stream (Of a) m r1
%1 -> Stream (Of b) m r2
%1 -> Stream (Of c) m r3
%1 -> Stream (Of d) m (r1, r2, r3)
zipWith3 (,,)
{-# INLINEABLE zip3 #-}

-- | Like @zipR@ but with three streams.
zip3R ::
  (Control.Monad m) =>
  Stream (Of a) m r1 %1 ->
  Stream (Of b) m r2 %1 ->
  Stream (Of c) m r3 %1 ->
  Stream (Of (a, b, c)) m (ZipResidual3 a b c m r1 r2 r3)
zip3R :: forall (m :: * -> *) a r1 b r2 c r3.
Monad m =>
Stream (Of a) m r1
%1 -> Stream (Of b) m r2
%1 -> Stream (Of c) m r3
%1 -> Stream (Of (a, b, c)) m (ZipResidual3 a b c m r1 r2 r3)
zip3R = forall (m :: * -> *) a b c d r1 r2 r3.
Monad m =>
(a -> b -> c -> d)
-> Stream (Of a) m r1
%1 -> Stream (Of b) m r2
%1 -> Stream (Of c) m r3
%1 -> Stream (Of d) m (ZipResidual3 a b c m r1 r2 r3)
zipWith3R (,,)
{-# INLINEABLE zip3R #-}

-- | Internal function to consume a stream remainder to
-- get the payload
extractResult :: (Control.Monad m) => Either r (Stream (Of a) m r) %1 -> m r
extractResult :: forall (m :: * -> *) r a.
Monad m =>
Either r (Stream (Of a) m r) %1 -> m r
extractResult (Left r
r) = forall (m :: * -> *) a. Monad m => a %1 -> m a
Control.return r
r
extractResult (Right Stream (Of a) m r
s) = forall a (m :: * -> *) r. Monad m => Stream (Of a) m r %1 -> m r
effects Stream (Of a) m r
s

-- # Merging
-------------------------------------------------------------------------------

-- $merging
--   These functions combine two sorted streams of orderable elements
--   into one sorted stream. The elements of the merged stream are
--   guaranteed to be in a sorted order if the two input streams are
--   also sorted.
--
--   The merge operation is /left-biased/: when merging two elements
--   that compare as equal, the left element is chosen first.

-- | Merge two streams of elements ordered with their 'Ord' instance.
--
--   The return values of both streams are returned.
--
-- @
-- \>\>\> S.print $ merge (each [1,3,5]) (each [2,4])
-- 1
-- 2
-- 3
-- 4
-- 5
-- ((), ())
-- @
merge ::
  (Control.Monad m, Ord a) =>
  Stream (Of a) m r %1 ->
  Stream (Of a) m s %1 ->
  Stream (Of a) m (r, s)
merge :: forall (m :: * -> *) a r s.
(Monad m, Ord a) =>
Stream (Of a) m r
%1 -> Stream (Of a) m s %1 -> Stream (Of a) m (r, s)
merge = forall (m :: * -> *) a r s.
Monad m =>
(a -> a -> Ordering)
-> Stream (Of a) m r
%1 -> Stream (Of a) m s
%1 -> Stream (Of a) m (r, s)
mergeBy forall a. Ord a => a -> a -> Ordering
compare
{-# INLINE merge #-}

-- | Merge two streams, ordering them by applying the given function to
--   each element before comparing.
--
--   The return values of both streams are returned.
mergeOn ::
  (Control.Monad m, Ord b) =>
  (a -> b) ->
  Stream (Of a) m r %1 ->
  Stream (Of a) m s %1 ->
  Stream (Of a) m (r, s)
mergeOn :: forall (m :: * -> *) b a r s.
(Monad m, Ord b) =>
(a -> b)
-> Stream (Of a) m r
%1 -> Stream (Of a) m s
%1 -> Stream (Of a) m (r, s)
mergeOn a -> b
f = forall (m :: * -> *) a r s.
Monad m =>
(a -> a -> Ordering)
-> Stream (Of a) m r
%1 -> Stream (Of a) m s
%1 -> Stream (Of a) m (r, s)
mergeBy (\a
x a
y -> forall a. Ord a => a -> a -> Ordering
compare (a -> b
f a
x) (a -> b
f a
y))
{-# INLINE mergeOn #-}

-- | Merge two streams, ordering the elements using the given comparison function.
--
--   The return values of both streams are returned.
mergeBy ::
  forall m a r s.
  (Control.Monad m) =>
  (a -> a -> Ordering) ->
  Stream (Of a) m r %1 ->
  Stream (Of a) m s %1 ->
  Stream (Of a) m (r, s)
mergeBy :: forall (m :: * -> *) a r s.
Monad m =>
(a -> a -> Ordering)
-> Stream (Of a) m r
%1 -> Stream (Of a) m s
%1 -> Stream (Of a) m (r, s)
mergeBy a -> a -> Ordering
comp Stream (Of a) m r
s1 Stream (Of a) m s
s2 = Stream (Of a) m r
%1 -> Stream (Of a) m s %1 -> Stream (Of a) m (r, s)
loop Stream (Of a) m r
s1 Stream (Of a) m s
s2
  where
    loop :: Stream (Of a) m r %1 -> Stream (Of a) m s %1 -> Stream (Of a) m (r, s)
    loop :: Stream (Of a) m r
%1 -> Stream (Of a) m s %1 -> Stream (Of a) m (r, s)
loop Stream (Of a) m r
s1 Stream (Of a) m s
s2 =
      Stream (Of a) m r
s1 forall a b (p :: Multiplicity) (q :: Multiplicity).
a %p -> (a %p -> b) %q -> b
& \case
        Return r
r ->
          forall (m :: * -> *) (f :: * -> *) r.
m (Stream f m r) -> Stream f m r
Effect forall a b (p :: Multiplicity) (q :: Multiplicity).
(a %p -> b) %q -> a %p -> b
$ forall a (m :: * -> *) r. Monad m => Stream (Of a) m r %1 -> m r
effects Stream (Of a) m s
s2 forall (m :: * -> *) a b.
Monad m =>
m a %1 -> (a %1 -> m b) %1 -> m b
Control.>>= \s
s -> forall (m :: * -> *) a. Monad m => a %1 -> m a
Control.return forall a b (p :: Multiplicity) (q :: Multiplicity).
(a %p -> b) %q -> a %p -> b
$ forall r (f :: * -> *) (m :: * -> *). r -> Stream f m r
Return (r
r, s
s)
        Effect m (Stream (Of a) m r)
ms ->
          forall (m :: * -> *) (f :: * -> *) r.
m (Stream f m r) -> Stream f m r
Effect forall a b (p :: Multiplicity) (q :: Multiplicity).
(a %p -> b) %q -> a %p -> b
$
            m (Stream (Of a) m r)
ms forall (m :: * -> *) a b.
Monad m =>
m a %1 -> (a %1 -> m b) %1 -> m b
Control.>>= \Stream (Of a) m r
s1' -> forall (m :: * -> *) a. Monad m => a %1 -> m a
Control.return forall a b (p :: Multiplicity) (q :: Multiplicity).
(a %p -> b) %q -> a %p -> b
$ forall (m :: * -> *) a r s.
Monad m =>
(a -> a -> Ordering)
-> Stream (Of a) m r
%1 -> Stream (Of a) m s
%1 -> Stream (Of a) m (r, s)
mergeBy a -> a -> Ordering
comp Stream (Of a) m r
s1' Stream (Of a) m s
s2
        Step (a
a :> Stream (Of a) m r
as) ->
          Stream (Of a) m s
s2 forall a b (p :: Multiplicity) (q :: Multiplicity).
a %p -> (a %p -> b) %q -> b
& \case
            Return s
s ->
              forall (m :: * -> *) (f :: * -> *) r.
m (Stream f m r) -> Stream f m r
Effect forall a b (p :: Multiplicity) (q :: Multiplicity).
(a %p -> b) %q -> a %p -> b
$ forall a (m :: * -> *) r. Monad m => Stream (Of a) m r %1 -> m r
effects Stream (Of a) m r
as forall (m :: * -> *) a b.
Monad m =>
m a %1 -> (a %1 -> m b) %1 -> m b
Control.>>= \r
r -> forall (m :: * -> *) a. Monad m => a %1 -> m a
Control.return forall a b (p :: Multiplicity) (q :: Multiplicity).
(a %p -> b) %q -> a %p -> b
$ forall r (f :: * -> *) (m :: * -> *). r -> Stream f m r
Return (r
r, s
s)
            Effect m (Stream (Of a) m s)
ms ->
              forall (m :: * -> *) (f :: * -> *) r.
m (Stream f m r) -> Stream f m r
Effect forall a b (p :: Multiplicity) (q :: Multiplicity).
(a %p -> b) %q -> a %p -> b
$
                m (Stream (Of a) m s)
ms forall (m :: * -> *) a b.
Monad m =>
m a %1 -> (a %1 -> m b) %1 -> m b
Control.>>= \Stream (Of a) m s
s2' ->
                  forall (m :: * -> *) a. Monad m => a %1 -> m a
Control.return forall a b (p :: Multiplicity) (q :: Multiplicity).
(a %p -> b) %q -> a %p -> b
$ forall (m :: * -> *) a r s.
Monad m =>
(a -> a -> Ordering)
-> Stream (Of a) m r
%1 -> Stream (Of a) m s
%1 -> Stream (Of a) m (r, s)
mergeBy a -> a -> Ordering
comp (forall (f :: * -> *) (m :: * -> *) r.
f (Stream f m r) -> Stream f m r
Step (a
a forall a b. a -> b -> Of a b
:> Stream (Of a) m r
as)) Stream (Of a) m s
s2'
            Step (a
b :> Stream (Of a) m s
bs) -> case a -> a -> Ordering
comp a
a a
b of
              Ordering
LT -> forall (f :: * -> *) (m :: * -> *) r.
f (Stream f m r) -> Stream f m r
Step (a
a forall a b. a -> b -> Of a b
:> forall (f :: * -> *) (m :: * -> *) r.
f (Stream f m r) -> Stream f m r
Step (a
b forall a b. a -> b -> Of a b
:> forall (m :: * -> *) a r s.
Monad m =>
(a -> a -> Ordering)
-> Stream (Of a) m r
%1 -> Stream (Of a) m s
%1 -> Stream (Of a) m (r, s)
mergeBy a -> a -> Ordering
comp Stream (Of a) m r
as Stream (Of a) m s
bs))
              Ordering
_ -> forall (f :: * -> *) (m :: * -> *) r.
f (Stream f m r) -> Stream f m r
Step (a
b forall a b. a -> b -> Of a b
:> forall (f :: * -> *) (m :: * -> *) r.
f (Stream f m r) -> Stream f m r
Step (a
a forall a b. a -> b -> Of a b
:> forall (m :: * -> *) a r s.
Monad m =>
(a -> a -> Ordering)
-> Stream (Of a) m r
%1 -> Stream (Of a) m s
%1 -> Stream (Of a) m (r, s)
mergeBy a -> a -> Ordering
comp Stream (Of a) m r
as Stream (Of a) m s
bs))
{-# INLINEABLE mergeBy #-}