@@ -1,5 +1,6 @@
11{-# LANGUAGE BangPatterns #-}
22{-# LANGUAGE CPP #-}
3+{-# LANGUAGE ExistentialQuantification #-}
34{-# LANGUAGE FlexibleContexts #-}
45{-# LANGUAGE RankNTypes #-}
56{-# LANGUAGE RecordWildCards #-}
@@ -284,8 +285,14 @@ module Streamly.Fold
284285 , partitionByM
285286 , partitionBy
286287 , partition
288+
289+ -- * Demultiplexing
290+ , demux
291+ , demuxWith
287292 , demux_
288293 , demuxWith_
294+
295+ -- * Classifying
289296 , classify
290297 , classifyWith
291298
@@ -1825,19 +1832,69 @@ partition = partitionBy id
18251832--
18261833-- partitionN :: Monad m => [Fold m a b] -> Fold m a [b]
18271834-- partitionN fs = Fold step begin done
1835+-}
18281836
1829−-- | Demultiplex an input element into a number of typed variants. We want to
1837+-- Demultiplex an input element into a number of typed variants. We want to
18301838-- statically restrict the target values within a set of predefined types, an
18311839-- enumeration of a GADT. We also want to make sure that the Map contains only
18321840-- those types and the full set of those types. Instead of Map it should
18331841-- probably be a lookup-table using an array and not in GC memory.
18341842--
1835−-- This is the consumer side dual of the producer side 'mux' operation.
1843+-- This is the consumer side dual of the producer side 'mux' operation (XXX to
1844+-- be implemented).
1845+
1846+-- | Split the input stream based on a key field and fold each split using a
1847+-- specific fold collecting the results in a map from the keys to the results.
1848+-- Useful for cases like protocol handlers to handle different type of packets
1849+-- using different handlers.
18361850--
1837−-- demux :: (Monad m, Ord k)
1838−-- => (a -> k) -> Map k (Fold m a b) -> Fold m a (Map k b)
1839−-- demux f kv = Fold step begin done
1840−-}
1851+-- @
1852+--
1853+-- |-------Fold m a b
1854+-- -----stream m a-----Map-----|
1855+-- |-------Fold m a b
1856+-- |
1857+-- ...
1858+-- @
1859+--
1860+-- @since 0.7.0
1861+demuxWith :: (Monad m, Ord k)
1862+ => (a -> k) -> Map k (Fold m a b) -> Fold m a (Map k b)
1863+demuxWith f kv = Fold step initial extract
1864+
1865+ where
1866+
1867+ initial = return kv
1868+ step mp a =
1869+ -- XXX should we raise an exception in Nothing case?
1870+ -- Ideally we should enforce that it is a total map over k so that look
1871+ -- up never fails
1872+ -- XXX we could use a monadic update function for a single lookup and
1873+ -- update in the map.
1874+ let k = f a
1875+ in case Map.lookup k mp of
1876+ Nothing -> return mp
1877+ Just (Fold step' acc extract') -> do
1878+ !r <- acc >>= \x -> step' x a
1879+ return $ Map.insert k (Fold step' (return r) extract') mp
1880+ extract = mapM (\(Fold s acc e) -> acc >>= e)
1881+
1882+-- | Fold a stream of key value pairs using a map of specific folds for each
1883+-- key into a map from keys to the results of fold outputs of the corresponding
1884+-- values.
1885+--
1886+-- @
1887+-- > let table = Data.Map.fromList [(\"SUM", FL.sum), (\"PRODUCT", FL.product)]
1888+-- input = S.fromList [(\"SUM",1),(\"PRODUCT",2),(\"SUM",3),(\"PRODUCT",4)]
1889+-- in FL.foldl' (FL.demux table) input
1890+-- One 1
1891+-- Two 2
1892+-- @
1893+--
1894+-- @since 0.7.0
1895+demux :: (Monad m, Ord k)
1896+ => Map k (Fold m a b) -> Fold m (k, a) (Map k b)
1897+demux fs = demuxWith fst (Map.map (lmap snd) fs)
18411898
18421899-- | Split the input stream based on a key field and fold each split using a
18431900-- specific fold without collecting the results. Useful for cases like protocol
@@ -1852,36 +1909,41 @@ partition = partitionBy id
18521909-- ...
18531910-- @
18541911--
1855−-- @
1856−-- > let prn = FL.drainBy print
1857−-- > let table = Data.Map.fromList [(\"ONE", prn), (\"TWO", prn)]
1858−-- input = S.fromList [(\"ONE",1),(\"TWO",2)]
1859−-- in FL.foldl' (FL.demux_ table) input
1860−-- One 1
1861−-- Two 2
1862−-- @
18631912--
18641913-- @since 0.7.0
1865−demuxWith_ :: (Monad m, Ord k) => (a -> k) -> Map k (Fold m a ()) -> Fold m a ()
1914+
1915+-- demuxWith_ can be slightly faster than demuxWith because we do not need to
1916+-- update the Map in this case. This may be significant only if the map is
1917+-- large.
1918+demuxWith_ :: (Monad m, Ord k)
1919+ => (a -> k) -> Map k (Fold m a b) -> Fold m a ()
18661920demuxWith_ f kv = Fold step initial extract
18671921
18681922 where
18691923
1870− initial = return ()
1871− step () a =
1924+ initial = mapM (\(Fold s i e) -> i >>= \r -> return (Fold s (return r) e)) kv
1925+ step mp a =
18721926 -- XXX should we raise an exception in Nothing case?
18731927 -- Ideally we should enforce that it is a total map over k so that look
18741928 -- up never fails
1875− case Map.lookup (f a) kv of
1876− Nothing -> return ()
1877− Just (Fold step' initial' extract') ->
1878− initial' >>= \x -> step' x a >>= extract'
1879− extract = return
1929+ case Map.lookup (f a) mp of
1930+ Nothing -> return mp
1931+ Just (Fold step' acc extract') -> do
1932+ acc >>= \x -> step' x a
1933+ return mp
1934+ extract mp = mapM (\(Fold s acc e) -> acc >>= e) mp >> return ()
18801935
18811936-- | Given a stream of key value pairs and a map from keys to folds, fold the
1882−-- values for each key using the corresponding folds.
1937+-- values for each key using the corresponding folds, discarding the outputs.
18831938--
1884−-- See 'demuxWith_' for details.
1939+-- @
1940+-- > let prn = FL.drainBy print
1941+-- > let table = Data.Map.fromList [(\"ONE", prn), (\"TWO", prn)]
1942+-- input = S.fromList [(\"ONE",1),(\"TWO",2)]
1943+-- in FL.foldl' (FL.demux_ table) input
1944+-- One 1
1945+-- Two 2
1946+-- @
18851947--
18861948-- @since 0.7.0
18871949demux_ :: (Monad m, Ord k) => Map k (Fold m a ()) -> Fold m (k, a) ()