@@ -100,6 +100,8 @@ module Streamly.Internal.Data.Stream.Channel.Types
100100
101101 , withDiagMVar
102102 , printSVar
103+
104+ , concatMapDivK
103105 )
104106where
105107
@@ -114,11 +116,13 @@ import Data.IORef (IORef, newIORef, readIORef, writeIORef)
114116import Data.Set (Set)
115117import Streamly.Internal.Data.Atomics
116118 (atomicModifyIORefCAS, atomicModifyIORefCAS_, storeLoadBarrier)
119+import Streamly.Internal.Data.SVar.Type (adaptState)
117120import Streamly.Internal.Data.Time.Clock (Clock(Monotonic), getTime)
118121import Streamly.Internal.Data.Time.Units (AbsTime, NanoSecond64(..))
119122import System.IO (hPutStrLn, stderr)
120123
121124import qualified Data.Set as Set
125+import qualified Streamly.Internal.Data.Stream.StreamK.Type as K
122126
123127------------------------------------------------------------------------------
124128-- Common types
@@ -708,3 +712,25 @@ cleanupSVar workerSet = do
708712 self <- myThreadId
709713 Prelude.mapM_ (`throwTo` ThreadAbort)
710714 (Prelude.filter (/= self) $ Set.toList workers)
715+
716+-------------------------------------------------------------------------------
717+-- Evaluator
718+-------------------------------------------------------------------------------
719+
720+-- | @concatMapDivK useTail useHead stream@, divides the stream in head and
721+-- tail, maps a stream generator on the head and maps an action on the tail of
722+-- a stream. Returns the stream generated by the head.
723+--
724+-- Used for concurrent evaluation of streams using a Channel.
725+{-# INLINE concatMapDivK #-}
726+concatMapDivK :: Monad m =>
727+ (K.Stream m a -> m ())
728+ -> (a -> K.Stream m b)
729+ -> K.Stream m a
730+ -> K.Stream m b
731+concatMapDivK useTail useHead stream =
732+ K.mkStream $ \st yld sng stp -> do
733+ let foldShared = K.foldStreamShared st yld sng stp
734+ single a = foldShared $ useHead a
735+ yieldk a r = useTail r >> single a
736+ in K.foldStreamShared (adaptState st) yieldk single stp stream