Revert "Revert scanl types"

This reverts commit 0a7d7c6.

Author: adithyaov

src/Streamly/Internal/Data/Fold/Channel/Type.hs

@@ -221,6 +221,7 @@ sendExceptionToDriver sv e = do
 data FromSVarState m a b =
       FromSVarRead (Channel m a b)
     | FromSVarLoop (Channel m a b) [ChildEvent a]
+    | FromSVarStop
 
 {-# INLINE_NORMAL fromInputQueue #-}
 fromInputQueue :: MonadIO m => Channel m a b -> D.Stream m a
@@ -246,6 +247,34 @@ fromInputQueue svar = D.Stream step (FromSVarRead svar)
             ChildStopChannel -> return D.Stop
             _ -> undefined
 
+    step _ FromSVarStop = undefined
+
+{-# INLINE_NORMAL fromInputQueueRaw #-}
+fromInputQueueRaw :: MonadIO m => Channel m a b -> D.Stream m (ChildEvent a)
+fromInputQueueRaw svar = D.Stream step (FromSVarRead svar)
+
+    where
+
+    {-# INLINE_LATE step #-}
+    step _ (FromSVarRead sv) = do
+        list <- readInputQ sv
+        -- Reversing the output is important to guarantee that we process the
+        -- outputs in the same order as they were generated by the constituent
+        -- streams.
+        return $ D.Skip $ FromSVarLoop sv (Prelude.reverse list)
+
+    step _ (FromSVarLoop sv []) = return $ D.Skip $ FromSVarRead sv
+    step _ (FromSVarLoop sv (ev : es)) = do
+        case ev of
+            -- XXX Separate input and output events. Input events cannot have
+            -- Stop event and output events cannot have ChildStopChannel
+            -- event.
+            ChildYield a -> return $ D.Yield (ChildYield a) (FromSVarLoop sv es)
+            ChildStopChannel -> return $ D.Yield ChildStopChannel FromSVarStop
+            _ -> undefined
+
+    step _ FromSVarStop = pure $ D.Stop
+
 {-# INLINE readInputQChan #-}
 readInputQChan :: Channel m a b -> IO ([ChildEvent a], Int)
 readInputQChan chan = do
@@ -330,8 +359,14 @@ newChannelWith outq outqDBell modifier f = do
         let f1 = Fold.rmapM (void . sendYieldToDriver chan) f
          in D.fold f1 $ fromInputQueue chan
 
+data SendChannelRaw s a
+    = SCREmptyBuffer s
+    | SCRBuffered s a
+    | SCRDrain
+
 {-# INLINE scanToChannel #-}
-scanToChannel :: MonadIO m => Channel m a b -> Scanl m a b -> Scanl m a ()
+scanToChannel ::
+    MonadIO m => Channel m a b -> Scanl m a b -> Scanl m (ChildEvent a) ()
 scanToChannel chan (Scanl step initial extract final) =
     Scanl step1 initial1 extract1 final1
 
@@ -343,24 +378,49 @@ scanToChannel chan (Scanl step initial extract final) =
             Fold.Partial s -> do
                 b <- extract s
                 void $ sendPartialToDriver chan b
-                return $ Fold.Partial s
+                return $ Fold.Partial (SCREmptyBuffer s)
             Fold.Done b ->
                 Fold.Done <$> void (sendYieldToDriver chan b)
 
-    step1 st x = do
+    step1 (SCREmptyBuffer st) (ChildYield x) =
+        return $ Fold.Partial (SCRBuffered st (ChildYield x))
+
+    step1 (SCRBuffered st (ChildYield x)) (ChildYield x1) = do
         r <- step st x
         case r of
             Fold.Partial s -> do
                 b <- extract s
                 void $ sendPartialToDriver chan b
-                return $ Fold.Partial s
-            Fold.Done b ->
-                Fold.Done <$> void (sendYieldToDriver chan b)
+                return $ Fold.Partial (SCRBuffered s (ChildYield x1))
+            Fold.Done b -> do
+                -- Even if the original scan is Done, we don't end this
+                -- concurrent fold here as we want to drain all the input in the
+                -- input buffer.
+                --
+                -- This scan only ends when it receives a ChildStopChannel
+                -- event.
+                --
+                void (sendYieldToDriver chan b)
+                pure $ Fold.Partial SCRDrain
+    step1 (SCRBuffered st (ChildYield x)) ChildStopChannel = do
+        r <- step st x
+        b <-
+            case r of
+                Fold.Partial s -> extract s
+                Fold.Done b0 -> pure b0
+        Fold.Done <$> void (sendYieldToDriver chan b)
+    step1 SCRDrain ChildStopChannel = pure $ Fold.Done ()
+    step1 SCRDrain _ = pure $ Fold.Partial SCRDrain
+    step1 _ _ = error "scanToChannel: Unsupported constructor"
 
     extract1 _ = return ()
 
     -- XXX Should we not discard the result?
-    final1 st = void (final st)
+    final1 (SCREmptyBuffer st) = void (final st)
+    -- XXX We are losing the input here.
+    -- XXX Should we consume the input and finalize it instead?
+    final1 (SCRBuffered st _val) = void (final st)
+    final1 SCRDrain = pure ()
 
 {-# INLINABLE newChannelWithScan #-}
 {-# SPECIALIZE newChannelWithScan ::
@@ -386,7 +446,7 @@ newChannelWithScan outq outqDBell modifier f = do
     where
 
     {-# NOINLINE work #-}
-    work chan = D.drain $ D.scanl (scanToChannel chan f) $ fromInputQueue chan
+    work chan = D.drain $ D.scanl (scanToChannel chan f) $ fromInputQueueRaw chan
 
 {-# INLINABLE newChannel #-}
 {-# SPECIALIZE newChannel ::