on master help sign in
Move concatMapDivK to common channel module
commit770887d62fd63a9852bd642660ee63b20144c691 authorHarendra Kumar <harendra@composewell.com> dateTue 27 Sep 2022, 19:33 (4 years ago) parent1f844ea
2 files changed+28−17
src/Streamly/Internal/Data/Stream/Async/Channel.hs+2−17 src/Streamly/Internal/Data/Stream/Channel/Types.hs+26
src/Streamly/Internal/Data/Stream/Async/Channel.hs+2−17
@@ -44,7 +44,6 @@ import Control.Monad.IO.Class (MonadIO(liftIO))
4444import Streamly.Internal.Control.Concurrent (MonadAsync, askRunInIO)
4545import Streamly.Internal.Data.Stream.Async.Channel.Operations
4646 (fromChannel, fromChannelK, toChannel, toChannelK)
47−import Streamly.Internal.Data.SVar.Type (adaptState)
4847
4948import qualified Streamly.Internal.Data.Stream as Stream
5049import qualified Streamly.Internal.Data.Stream.Async.Channel.Append as Append
@@ -72,6 +71,8 @@ withChannelK evaluator modifier input = K.concatEffect action
7271 toChannelK chan (evaluator chan input)
7372 return $ fromChannelK chan
7473
74+-- XXX Dedup with withChannelK
75+
7576{-# INLINE withInterleaveChannelK #-}
7677withInterleaveChannelK :: MonadAsync m =>
7778 (Channel m b -> K.Stream m a -> K.Stream m b)
@@ -175,22 +176,6 @@ _appendWithChanK chan stream1 stream2 =
175176-- concat
176177-------------------------------------------------------------------------------
177178
178−-- | @concatMapDivK useTail useHead stream@, divides the stream in head and
179−-- tail, maps a stream generator on the head and maps an action on the tail of
180−-- a stream. Returns the stream generated by the head.
181−{-# INLINE concatMapDivK #-}
182−concatMapDivK :: Monad m =>
183− (K.Stream m a -> m ())
184− -> (a -> K.Stream m b)
185− -> K.Stream m a
186− -> K.Stream m b
187−concatMapDivK useTail useHead stream =
188− K.mkStream $ \st yld sng stp -> do
189− let foldShared = K.foldStreamShared st yld sng stp
190− single a = foldShared $ useHead a
191− yieldk a r = useTail r >> single a
192− in K.foldStreamShared (adaptState st) yieldk single stp stream
193−
194179-- | A runner function takes a queuing function @q@ and a stream, it splits the
195180-- input stream, queuing the tail and using the head to generate a stream.
196181-- 'mkEnqueue' takes a runner function and generates the queuing function @q@.
src/Streamly/Internal/Data/Stream/Channel/Types.hs+26
@@ -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