@@ -56,6 +56,7 @@ module Streamly.Fold
5656
5757 -- ** Full Folds
5858 , drain
59+ , drainBy
5960 , last
6061 , length
6162 , sum
@@ -275,6 +276,7 @@ module Streamly.Fold
275276 -- select the fold. This is useful to demultiplex the input stream.
276277 , partitionByM
277278 , partitionBy
279+ , demux_
278280
279281 -- ** Unzipping
280282 , unzip
@@ -294,8 +296,10 @@ module Streamly.Fold
294296 )
295297where
296298
299+import Control.Monad (void)
297300import Control.Monad.IO.Class (MonadIO(..))
298301import Data.Functor.Identity (Identity)
302+import Data.Map.Strict (Map)
299303import Foreign.Storable (Storable(..))
300304import Prelude
301305 hiding (filter, drop, dropWhile, take, takeWhile, zipWith, foldr,
@@ -305,6 +309,8 @@ import Prelude
305309 scanl, scanl1, replicate, concatMap, mconcat, foldMap, unzip,
306310 span, splitAt, break)
307311
312+import qualified Data.Map.Strict as Map
313+
308314import Streamly.Array (Array)
309315import Streamly.Fold.Types (Fold(..))
310316import Streamly.Streams.Serial (SerialT)
@@ -439,6 +445,15 @@ drain = Fold step begin done
439445 step _ _ = return ()
440446 done = return
441447
448+-- |
449+-- > drainBy f = lmapM f drain
450+--
451+-- Drain all input after passing it through a monadic function. This is the
452+-- dual of mapM_ on stream producers.
453+{-# INLINABLE drainBy #-}
454+drainBy :: Monad m => (a -> m b) -> Fold m a ()
455+drainBy f = Fold (const (void . f)) (return ()) return
456+
442457-- | Extract the last element of the input stream, if any.
443458--
444459-- @since 0.7.0
@@ -1683,6 +1698,44 @@ partitionBy f = partitionByM (return . f)
16831698-- demux f kv = Fold step begin done
16841699-}
16851700
1701+-- | Demultiplex to multiple consumers without collecting the results. Useful
1702+-- to run different effectful computations depending on the value of the stream
1703+-- elements, for example handling network packets of different types using
1704+-- different handlers.
1705+--
1706+-- @
1707+--
1708+-- |-------Fold m a ()
1709+-- -----stream m a-----Map-----|
1710+-- |-------Fold m a ()
1711+-- |
1712+-- ...
1713+-- @
1714+--
1715+-- @
1716+-- > let prn = FL.drainBy print
1717+-- > let table = Data.Map.fromList [(\"ONE", prn), (\"TWO", prn)]
1718+-- input = S.fromList [(\"ONE",1),(\"TWO",2)]
1719+-- in FL.foldl' (FL.demux_ table) input
1720+-- One 1
1721+-- Two 2
1722+-- @
1723+demux_ :: (Monad m, Ord k) => Map k (Fold m a ()) -> Fold m (k, a) ()
1724+demux_ kv = Fold step initial extract
1725+
1726+ where
1727+
1728+ initial = return ()
1729+ step () (k, a) =
1730+ -- XXX should we raise an exception in Nothing case?
1731+ -- Ideally we should enforce that it is a total map over k so that look
1732+ -- up never fails
1733+ case Map.lookup k kv of
1734+ Nothing -> return ()
1735+ Just (Fold step' initial' extract') ->
1736+ initial' >>= \x -> step' x a >>= extract'
1737+ extract = return
1738+
16861739------------------------------------------------------------------------------
16871740-- Unzipping
16881741------------------------------------------------------------------------------