content: {-# 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 = loop
    where
      loop ::
        (Control.Monad m) =>
        Stream (Of (a, b)) m r %1 ->
        Stream (Of a) (Stream (Of b) m) r
      loop stream =
        case stream of
          Return r -> Return r
          Effect m -> Effect $ Control.fmap loop $ Control.lift m
          Step ((a, b) :> rest) -> Step (a :> Effect (Step (b :> Return (loop 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 = 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 f st1 st2 =
        case st1 of
          Effect ms -> Effect $ Control.fmap (\s -> loop f s st2) ms
          Return r1 ->
            case st2 of
              Return r2 -> Return $ Left3 (r1, r2)
              st2' -> Return $ Middle3 (r1, st2')
          Step (a :> as) ->
            case st2 of
              Effect ms ->
                Effect $ Control.fmap (\s -> loop f (Step (a :> as)) s) ms
              Return r2 -> Return $ Right3 (Step (a :> as), r2)
              Step (b :> bs) -> Step $ (f a b) :> loop f as 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 f s1 s2 = Control.do
    result <- zipWithR f s1 s2
    case result of
      Left3 rets -> Control.return rets
      Middle3 (r1, s2') -> Control.do
        r2 <- Control.lift $ effects s2'
        Control.return (r1, r2)
      Right3 (s1', r2) -> Control.do
        r1 <- Control.lift $ effects s1'
        Control.return (r1, 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 = 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 = 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 = 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 f s1 s2 s3 =
        case s1 of
          Effect ms -> Effect $ Control.fmap (\s -> loop f s s2 s3) ms
          Return r1 ->
            case (s2, s3) of
              (Return r2, Return r3) -> Return (Left r1, Left r2, Left r3)
              (s2', s3') -> Return (Left r1, Right s2', Right s3')
          Step (a :> as) ->
            case s2 of
              Effect ms ->
                Effect $
                  Control.fmap (\s -> loop f (Step $ a :> as) s s3) ms
              Return r2 -> Return (Right (Step $ a :> as), Left r2, Right s3)
              Step (b :> bs) ->
                case s3 of
                  Effect ms ->
                    Effect $
                      Control.fmap (\s -> loop f (Step $ a :> as) (Step $ b :> bs) s) ms
                  Return r3 ->
                    Return (Right (Step $ a :> as), Right (Step $ b :> bs), Left r3)
                  Step (c :> cs) -> Step $ (f a b c) :> loop f as bs 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 f s1 s2 s3 = Control.do
    result <- zipWith3R f s1 s2 s3
    case result of
      (res1, res2, res3) -> Control.do
        r1 <- Control.lift $ extractResult res1
        r2 <- Control.lift $ extractResult res2
        r3 <- Control.lift $ extractResult res3
        Control.return (r1, r2, 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 = 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 = 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 (Left r) = Control.return r
  extractResult (Right s) = effects 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 = mergeBy 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 f = mergeBy (\x y -> compare (f x) (f 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 comp s1 s2 = loop s1 s2
    where
      loop :: Stream (Of a) m r %1 -> Stream (Of a) m s %1 -> Stream (Of a) m (r, s)
      loop s1 s2 =
        case s1 of
          Return r ->
            Effect $ effects s2 Control.>>= \s -> Control.return $ Return (r, s)
          Effect ms ->
            Effect $
              ms Control.>>= \s1' -> Control.return $ mergeBy comp s1' s2
          Step (a :> as) ->
            case s2 of
              Return s ->
                Effect $ effects as Control.>>= \r -> Control.return $ Return (r, s)
              Effect ms ->
                Effect $
                  ms Control.>>= \s2' ->
                    Control.return $ mergeBy comp (Step (a :> as)) s2'
              Step (b :> bs) -> case comp a b of
                LT -> Step (a :> Step (b :> mergeBy comp as bs))
                _ -> Step (b :> Step (a :> mergeBy comp as bs))
  {-# INLINEABLE mergeBy #-}
  