3853579 · streamly · git.iko.soy

git.iko.soy/streamly

on master help sign in

Add demux_ and drainBy

3853579Harendra Kumar7 years ago

src/Streamly/Fold.hs+53

compare with master

unifiedsplit

@@ -56,6 +56,7 @@ module Streamly.Fold

57 -- ** Full Folds

58 , drain

59+ , drainBy

59 , last

60 , last

60 , length

61 , length

61 , sum

62 , sum

@@ -275,6 +276,7 @@ module Streamly.Fold

275 -- select the fold. This is useful to demultiplex the input stream.

276 -- select the fold. This is useful to demultiplex the input stream.

276 , partitionByM

277 , partitionByM

277 , partitionBy

278 , partitionBy

279+ , demux_

278

280

279 -- ** Unzipping

281 -- ** Unzipping

280 , unzip

282 , unzip

@@ -294,8 +296,10 @@ module Streamly.Fold

294 )

296 )

295where

297where

296

298

299+import Control.Monad (void)

297import Control.Monad.IO.Class (MonadIO(..))

300import Control.Monad.IO.Class (MonadIO(..))

298import Data.Functor.Identity (Identity)

301import Data.Functor.Identity (Identity)

302+import Data.Map.Strict (Map)

299import Foreign.Storable (Storable(..))

303import Foreign.Storable (Storable(..))

300import Prelude

304import Prelude

301 hiding (filter, drop, dropWhile, take, takeWhile, zipWith, foldr,

305 hiding (filter, drop, dropWhile, take, takeWhile, zipWith, foldr,

@@ -305,6 +309,8 @@ import Prelude

305 scanl, scanl1, replicate, concatMap, mconcat, foldMap, unzip,

309 scanl, scanl1, replicate, concatMap, mconcat, foldMap, unzip,

306 span, splitAt, break)

310 span, splitAt, break)

307

311

312+import qualified Data.Map.Strict as Map

313+

308import Streamly.Array (Array)

314import Streamly.Array (Array)

309import Streamly.Fold.Types (Fold(..))

315import Streamly.Fold.Types (Fold(..))

310import Streamly.Streams.Serial (SerialT)

316import Streamly.Streams.Serial (SerialT)

@@ -439,6 +445,15 @@ drain = Fold step begin done

439 step _ _ = return ()

445 step _ _ = return ()

440 done = return

446 done = return

441

447

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+

442-- | Extract the last element of the input stream, if any.

457-- | Extract the last element of the input stream, if any.

443--

458--

444-- @since 0.7.0

459-- @since 0.7.0

@@ -1683,6 +1698,44 @@ partitionBy f = partitionByM (return . f)

1683-- demux f kv = Fold step begin done

1698-- demux f kv = Fold step begin done

1684-}

1699-}

1685

1700

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+

1686------------------------------------------------------------------------------

1739------------------------------------------------------------------------------

1687-- Unzipping

1740-- Unzipping

1688------------------------------------------------------------------------------

1741------------------------------------------------------------------------------