@@ -1,5 +1,6 @@
1{-# LANGUAGE BangPatterns #-}
1{-# LANGUAGE BangPatterns #-}
2{-# LANGUAGE CPP #-}
2{-# LANGUAGE CPP #-}
3+{-# LANGUAGE ExistentialQuantification #-}
3{-# LANGUAGE FlexibleContexts #-}
4{-# LANGUAGE FlexibleContexts #-}
4{-# LANGUAGE RankNTypes #-}
5{-# LANGUAGE RankNTypes #-}
5{-# LANGUAGE RecordWildCards #-}
6{-# LANGUAGE RecordWildCards #-}
@@ -284,8 +285,14 @@ module Streamly.Fold
284 , partitionByM
285 , partitionByM
285 , partitionBy
286 , partitionBy
286 , partition
287 , partition
288 , demuxWith_
293 , demuxWith_
289 , classify
296 , classify
290 , classifyWith
297 , classifyWith
@@ -1825,19 +1832,69 @@ partition = partitionBy id
1826-- partitionN :: Monad m => [Fold m a b] -> Fold m a [b]
1833-- partitionN :: Monad m => [Fold m a b] -> Fold m a [b]
1827-- partitionN fs = Fold step begin done
1834-- partitionN fs = Fold step begin done
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
1830-- statically restrict the target values within a set of predefined types, an
1838-- statically restrict the target values within a set of predefined types, an
1831-- enumeration of a GADT. We also want to make sure that the Map contains only
1839-- enumeration of a GADT. We also want to make sure that the Map contains only
1832-- those types and the full set of those types. Instead of Map it should
1840-- those types and the full set of those types. Instead of Map it should
1833-- probably be a lookup-table using an array and not in GC memory.
1841-- probably be a lookup-table using an array and not in GC memory.
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
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.
1853+-- |-------Fold m a b
1854+-- -----stream m a-----Map-----|
1855+-- |-------Fold m a b
1837−-- demux :: (Monad m, Ord k)
1861+demuxWith :: (Monad m, Ord k)
1838−-- => (a -> k) -> Map k (Fold m a b) -> Fold m a (Map k b)
1862+ => (a -> k) -> Map k (Fold m a b) -> Fold m a (Map k b)
1839−-- demux f kv = Fold step begin done
1863+demuxWith f kv = Fold step initial extract
1867+ initial = return kv
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
1872+ -- XXX we could use a monadic update function for a single lookup and
1873+ -- update in the map.
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)
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
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
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)
1842-- | Split the input stream based on a key field and fold each split using a
1899-- | Split the input stream based on a key field and fold each split using a
1843-- specific fold without collecting the results. Useful for cases like protocol
1900-- specific fold without collecting the results. Useful for cases like protocol
@@ -1852,36 +1909,41 @@ partition = partitionBy id
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
1864-- @since 0.7.0
1913-- @since 0.7.0
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
1865−demuxWith_ :: (Monad m, Ord k) => (a -> k) -> Map k (Fold m a ()) -> Fold m a ()
1918+demuxWith_ :: (Monad m, Ord k)
1919+ => (a -> k) -> Map k (Fold m a b) -> Fold m a ()
1866demuxWith_ f kv = Fold step initial extract
1920demuxWith_ f kv = Fold step initial extract
1870− initial = return ()
1924+ initial = mapM (\(Fold s i e) -> i >>= \r -> return (Fold s (return r) e)) kv
1871− step () a =
1925+ step mp a =
1872 -- XXX should we raise an exception in Nothing case?
1926 -- XXX should we raise an exception in Nothing case?
1873 -- Ideally we should enforce that it is a total map over k so that look
1927 -- Ideally we should enforce that it is a total map over k so that look
1874 -- up never fails
1928 -- up never fails
1875− case Map.lookup (f a) kv of
1929+ case Map.lookup (f a) mp of
1876− Nothing -> return ()
1930+ Nothing -> return mp
1877− Just (Fold step' initial' extract') ->
1931+ Just (Fold step' acc extract') -> do
1878− initial' >>= \x -> step' x a >>= extract'
1932+ acc >>= \x -> step' x a
1879− extract = return
1934+ extract mp = mapM (\(Fold s acc e) -> acc >>= e) mp >> return ()
1881-- | Given a stream of key value pairs and a map from keys to folds, fold the
1936-- | 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.
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
1886-- @since 0.7.0
1948-- @since 0.7.0
1887demux_ :: (Monad m, Ord k) => Map k (Fold m a ()) -> Fold m (k, a) ()
1949demux_ :: (Monad m, Ord k) => Map k (Fold m a ()) -> Fold m (k, a) ()