@@ -125,6 +125,7 @@ module Streamly.Internal.Data.Stream.StreamD
125125
126126 -- ** Specialized Folds
127127 , tap
128+ , tapAsync
128129 , drain
129130 , null
130131 , head
@@ -299,6 +300,7 @@ module Streamly.Internal.Data.Stream.StreamD
299300 )
300301where
301302
303+import Control.Concurrent (myThreadId, takeMVar)
302304import Control.Exception (Exception, SomeException, fromException)
303305import Control.Monad (void, when)
304306import Control.Monad.Catch (MonadCatch, throwM)
@@ -336,6 +338,7 @@ import Streamly.Internal.Data.Unfold.Types (Unfold(..))
336338
337339import Streamly.Internal.Data.Stream.StreamD.Type
338340import Streamly.Internal.Data.SVar
341+import Streamly.Internal.Data.Stream.SVar (fromConsumer, pushToFold)
339342
340343import qualified Streamly.Internal.Data.Pipe.Types as Pipe
341344import qualified Streamly.Internal.Memory.Array.Types as A
@@ -732,6 +735,44 @@ fromSVar svar = Stream step FromSVarInit
732735 liftIO $ printSVar sv "SVar Done"
733736 return Stop
734737
738+-------------------------------------------------------------------------------
739+-- Process events received by a fold consumer from a stream producer
740+-------------------------------------------------------------------------------
741+
742+{-# INLINE_NORMAL fromProducer #-}
743+fromProducer :: (MonadAsync m) => SVar t m a -> Stream m a
744+fromProducer svar = Stream step (FromSVarRead svar)
745+ where
746+
747+ {-# INLINE_LATE step #-}
748+ step _ (FromSVarRead sv) = do
749+ list <- readOutputQ sv
750+ -- Reversing the output is important to guarantee that we process the
751+ -- outputs in the same order as they were generated by the constituent
752+ -- streams.
753+ return $ Skip $ FromSVarLoop sv (Prelude.reverse list)
754+
755+ step _ (FromSVarLoop sv []) = return $ Skip $ FromSVarRead sv
756+ step _ (FromSVarLoop sv (ev : es)) = do
757+ case ev of
758+ ChildYield a -> return $ Yield a (FromSVarLoop sv es)
759+ ChildStop tid e -> do
760+ accountThread sv tid
761+ case e of
762+ Nothing -> do
763+ sendStopToProducer sv
764+ return $ Skip (FromSVarDone sv)
765+ Just _ -> error "Bug: fromProducer: received exception"
766+
767+ step _ (FromSVarDone sv) = do
768+ when (svarInspectMode sv) $ do
769+ t <- liftIO $ getTime Monotonic
770+ liftIO $ writeIORef (svarStopTime (svarStats sv)) (Just t)
771+ liftIO $ printSVar sv "SVar Done"
772+ return Stop
773+
774+ step _ FromSVarInit = undefined
775+
735776-------------------------------------------------------------------------------
736777-- Hoisting the inner monad
737778-------------------------------------------------------------------------------
@@ -3732,3 +3773,74 @@ foldParallel f m = do
37323773 sv <- newParallelVar StopNone defState
37333774 toSVarParallel defState sv (toStreamD m)
37343775 f $ fromStreamD $ fromSVar sv
3776+
3777+-------------------------------------------------------------------------------
3778+-- Concurrent tap
3779+-------------------------------------------------------------------------------
3780+
3781+-- | Create an SVar with a fold consumer that will fold any elements sent to it
3782+-- using the supplied fold function.
3783+{-# INLINE newFoldSVar #-}
3784+newFoldSVar :: MonadAsync m => State t m a -> Fold m a b -> m (SVar t m a)
3785+newFoldSVar stt f = do
3786+ -- Buffer size for the SVar is derived from the current state
3787+ sv <- newParallelVar StopAny (adaptState stt)
3788+ -- Add the producer thread-id to the SVar.
3789+ liftIO myThreadId >>= modifyThread sv
3790+ void $ doFork (work sv) (svarMrun sv) (handleFoldException sv)
3791+ return sv
3792+
3793+ where
3794+
3795+ {-# NOINLINE work #-}
3796+ work sv = void $ runFold f $ fromProducer sv
3797+
3798+data TapState sv st = TapInit | Tapping sv st | TapDone st
3799+
3800+{-# INLINE_NORMAL tapAsync #-}
3801+tapAsync :: MonadAsync m => Fold m a b -> Stream m a -> Stream m a
3802+tapAsync f (Stream step1 state1) = Stream step TapInit
3803+ where
3804+
3805+ drainFold svr = do
3806+ -- In general, a Stop event would come equipped with the result
3807+ -- of the fold. It is not used here but it would be useful in
3808+ -- applicative and distribute.
3809+ done <- fromConsumer svr
3810+ when (not done) $ do
3811+ liftIO $ withDiagMVar svr "teeToSVar: waiting to drain"
3812+ $ takeMVar (outputDoorBellFromConsumer svr)
3813+ drainFold svr
3814+
3815+ stopFold svr = do
3816+ liftIO $ sendStop svr Nothing
3817+ -- drain/wait until a stop event arrives from the fold.
3818+ drainFold svr
3819+
3820+ {-# INLINE_LATE step #-}
3821+ step gst TapInit = do
3822+ sv <- newFoldSVar gst f
3823+ return $ Skip (Tapping sv state1)
3824+
3825+ step gst (Tapping sv st) = do
3826+ r <- step1 gst st
3827+ case r of
3828+ Yield a s -> do
3829+ done <- pushToFold sv a
3830+ if done
3831+ then do
3832+ -- XXX we do not need to wait synchronously here
3833+ stopFold sv
3834+ return $ Yield a (TapDone s)
3835+ else return $ Yield a (Tapping sv s)
3836+ Skip s -> return $ Skip (Tapping sv s)
3837+ Stop -> do
3838+ stopFold sv
3839+ return $ Stop
3840+
3841+ step gst (TapDone st) = do
3842+ r <- step1 gst st
3843+ return $ case r of
3844+ Yield a s -> Yield a (TapDone s)
3845+ Skip s -> Skip (TapDone s)
3846+ Stop -> Stop