@@ -23,6 +23,8 @@ module Streamly.Internal.Data.Stream.Async.Channel
2323 , maxRate
2424 , constRate
2525 , inspectMode
26+ , StopWhen (..)
27+ , stopWhen
2628
2729 -- * Operations
2830 , fromChannel
@@ -129,35 +131,46 @@ evalWithD modifier m = D.Stream step Nothing
129131-- appending two streams
130132-------------------------------------------------------------------------------
131133
132−-- | Create a new channel and add both the streams to it for async evaluation.
133−-- The output stream is the result of the evaluation.
134−{-# INLINE appendWithK #-}
135−appendWithK :: MonadAsync m =>
136− (Config -> Config) -> K.Stream m a -> K.Stream m a -> K.Stream m a
137−appendWithK modifier stream1 stream2 = K.concatEffect action
134+{-# INLINE appendGeneric #-}
135+appendGeneric :: MonadAsync m =>
136+ ((Config -> Config) -> m (Channel m a))
137+ -> (Config -> Config)
138+ -> K.Stream m a
139+ -> K.Stream m a
140+ -> K.Stream m a
141+appendGeneric newChan modifier stream1 stream2 = K.concatEffect action
138142
139143 where
140144
141145 action = do
142− chan <- Append.newChannel modifier
143− toChannelK chan stream2
144− toChannelK chan stream1
146+ chan <- newChan modifier
147+ let cfg = modifier defaultConfig
148+ done = K.nilM (stopChannel chan)
149+ case getStopWhen cfg of
150+ AllStop -> do
151+ toChannelK chan stream2
152+ toChannelK chan stream1
153+ FirstStops -> do
154+ toChannelK chan stream2
155+ toChannelK chan (K.serial stream1 done)
156+ AnyStops -> do
157+ toChannelK chan (K.serial stream2 done)
158+ toChannelK chan (K.serial stream1 done)
145159 return $ Stream.toStreamK $ fromChannel chan
146160
161+-- | Create a new channel and add both the streams to it for async evaluation.
162+-- The output stream is the result of the evaluation.
163+{-# INLINE appendWithK #-}
164+appendWithK :: MonadAsync m =>
165+ (Config -> Config) -> K.Stream m a -> K.Stream m a -> K.Stream m a
166+appendWithK = appendGeneric Append.newChannel
167+
147168-- | Create a new channel and add both the streams to it for async evaluation.
148169-- The output stream is the result of the evaluation.
149170{-# INLINE interleaveWithK #-}
150171interleaveWithK :: MonadAsync m =>
151172 (Config -> Config) -> K.Stream m a -> K.Stream m a -> K.Stream m a
152−interleaveWithK modifier stream1 stream2 = K.concatEffect action
153−
154− where
155−
156− action = do
157− chan <- Interleave.newChannel modifier
158− toChannelK chan stream2
159− toChannelK chan stream1
160− return $ Stream.toStreamK $ fromChannel chan
173+interleaveWithK = appendGeneric Interleave.newChannel
161174
162175-- | Evaluate the first stream in the current thread and add the second stream
163176-- to the supplied channel. This is to be used by a worker thread.
@@ -219,26 +232,62 @@ concatMapWithChanK chan f stream =
219232 in K.concatMapEffect (`run` stream) (mkEnqueue chan run)
220233 -- K.concatMapWith (_appendWithChanK chan) f stream
221234
235+{-# INLINE concatMapWithChanKAny #-}
236+concatMapWithChanKAny :: MonadAsync m =>
237+ Channel m b -> (a -> K.Stream m b) -> K.Stream m a -> K.Stream m b
238+concatMapWithChanKAny chan f stream =
239+ let done = K.nilM (stopChannel chan)
240+ run q = concatMapDivK q (\x -> K.serial (f x) done)
241+ in K.concatMapEffect (`run` stream) (mkEnqueue chan run)
242+
243+{-# INLINE concatMapWithChanKFirst #-}
244+concatMapWithChanKFirst :: MonadAsync m =>
245+ Channel m b -> (a -> K.Stream m b) -> K.Stream m a -> K.Stream m b
246+concatMapWithChanKFirst chan f stream =
247+ let done = K.nilM (stopChannel chan)
248+ run q = concatMapDivK q f
249+ in K.concatEffect $ do
250+ res <- K.uncons stream
251+ case res of
252+ Nothing -> return K.nil
253+ Just (h, t) -> do
254+ q <- mkEnqueue chan run
255+ q t
256+ return $ K.serial (f h) done
257+
258+{-# INLINE concatMapWithChanKGeneric #-}
259+concatMapWithChanKGeneric :: MonadAsync m =>
260+ (Config -> Config)
261+ -> Channel m b
262+ -> (a -> K.Stream m b)
263+ -> K.Stream m a
264+ -> K.Stream m b
265+concatMapWithChanKGeneric modifier chan f stream = do
266+ let cfg = modifier defaultConfig
267+ case getStopWhen cfg of
268+ AllStop -> concatMapWithChanK chan f stream
269+ FirstStops -> concatMapWithChanKFirst chan f stream
270+ AnyStops -> concatMapWithChanKAny chan f stream
271+
222272-- XXX Add a deep evaluation variant that evaluates individual elements in the
223273-- generated streams in parallel.
224274
225275-- | Allocate a channel and use it to concurrently evaluate the streams
226276-- generated by the mapped function.
227277--
228−-- >> concatMapWithK modifier f = withChannelK (`concatMapWithChanK` f) modifier
229−--
230278{-# INLINE concatMapWithK #-}
231279concatMapWithK :: MonadAsync m =>
232280 (Config -> Config) -> (a -> K.Stream m b) -> K.Stream m a -> K.Stream m b
233−concatMapWithK modifier f = withChannelK (`concatMapWithChanK` f) modifier
281+concatMapWithK modifier f =
282+ let g = concatMapWithChanKGeneric modifier
283+ in withChannelK (`g` f) modifier
234284
235285-- | Allocate a channel and use it to concurrently evaluate the streams
236286-- generated by the mapped function.
237287--
238−-- >> concatMapInterleaveWithK modifier f = withInterleaveChannelK (`concatMapWithChanK` f) modifier
239−--
240288{-# INLINE concatMapInterleaveWithK #-}
241289concatMapInterleaveWithK :: MonadAsync m =>
242290 (Config -> Config) -> (a -> K.Stream m b) -> K.Stream m a -> K.Stream m b
243291concatMapInterleaveWithK modifier f =
244− withInterleaveChannelK (`concatMapWithChanK` f) modifier
292+ let g = concatMapWithChanKGeneric modifier
293+ in withInterleaveChannelK (`g` f) modifier