@@ -214,14 +214,18 @@ module Streamly.Internal.Data.Fold
215 -- ** Demultiplexing
215 -- ** Demultiplexing
216 -- | Direct values in the input stream to different folds using an n-ary
216 -- | Direct values in the input stream to different folds using an n-ary
217− -- fold selector. 'demux' is a generalization of 'classify' where each key
217+ -- fold selector. 'demux' is a generalization of 'classify' (and
218− -- of the classifier can use a different fold.
218+ -- 'partition') where each key of the classifier can use a different fold.
219− , demux -- XXX rename this to demux_
219+ , demux
220 , demuxWith
220 , demuxWith
221− , demuxDefault -- XXX rename this to demux
221+ , demuxScanWith
222− , demuxDefaultWith
222+ , demuxScanMutWith
223− -- , demuxWithSel
223+ , demuxMutWith
225+ -- TODO: These can be implemented using the above operations
226+ -- , demuxWithSel -- Stop when the fold for the specified key stops
224− -- , demuxWithMin
227+ -- , demuxWithMin -- Stop when any of the folds stop
228+ -- , demuxWithAll -- Stop when all the folds stop (run once)
226 -- ** Classifying
230 -- ** Classifying
227 -- | In an input stream of key value pairs fold values for different keys
231 -- | In an input stream of key value pairs fold values for different keys
@@ -281,8 +285,6 @@ import Data.Int (Int64)
281import Data.IORef (newIORef, readIORef, writeIORef)
285import Data.IORef (newIORef, readIORef, writeIORef)
282import Data.Map.Strict (Map)
286import Data.Map.Strict (Map)
283import Data.Maybe (isJust, fromJust)
287import Data.Maybe (isJust, fromJust)
284−import Streamly.Internal.Data.Either.Strict
285− (Either'(..), fromLeft', fromRight', isLeft', isRight')
286import Streamly.Internal.Data.IsMap (IsMap(..))
288import Streamly.Internal.Data.IsMap (IsMap(..))
287import Streamly.Internal.Data.Pipe.Type (Pipe (..), PipeState(..))
289import Streamly.Internal.Data.Pipe.Type (Pipe (..), PipeState(..))
288import Streamly.Internal.Data.Unfold.Type (Unfold(..))
290import Streamly.Internal.Data.Unfold.Type (Unfold(..))
@@ -290,7 +292,6 @@ import Streamly.Internal.Data.Tuple.Strict (Tuple'(..), Tuple3'(..))
290import Streamly.Internal.Data.Stream.Serial (SerialT(..))
292import Streamly.Internal.Data.Stream.Serial (SerialT(..))
292import qualified Data.IntSet as IntSet
294import qualified Data.IntSet as IntSet
293−import qualified Data.Map.Strict as Map
294import qualified Data.Set as Set
295import qualified Data.Set as Set
295import qualified Streamly.Internal.Data.IsMap as IsMap
296import qualified Streamly.Internal.Data.IsMap as IsMap
296import qualified Prelude
297import qualified Prelude
@@ -1467,201 +1468,213 @@ partition = partitionBy id
1467-- partitionN fs = Fold step begin done
1468-- partitionN fs = Fold step begin done
1471+------------------------------------------------------------------------------
1472+-- demux: in a key value stream fold each key sub-stream with a different fold
1473+------------------------------------------------------------------------------
1470-- TODO Demultiplex an input element into a number of typed variants. We want
1475-- TODO Demultiplex an input element into a number of typed variants. We want
1471-- to statically restrict the target values within a set of predefined types,
1476-- to statically restrict the target values within a set of predefined types,
1472−-- an enumeration of a GADT. We also want to make sure that the Map contains
1477+-- an enumeration of a GADT.
1473−-- only those types and the full set of those types.
1475−-- TODO Instead of the input Map it should probably be a lookup-table using an
1476−-- array and not in GC memory. The same applies to the output Map as well.
1477−-- However, that would only be helpful if we have a very large data structure,
1478−-- need to measure and see how it scales.
1480-- This is the consumer side dual of the producer side 'mux' operation (XXX to
1479-- This is the consumer side dual of the producer side 'mux' operation (XXX to
1481-- be implemented).
1480-- be implemented).
1482+-- XXX If we use Refold in it, it can perhaps fuse/be more efficient. For
1483+-- example we can store just the result rather than storing the whole fold in
1483−-- | Split the input stream based on a key field and fold each split using a
1486+-- | In a key value stream, fold values corresponding to each key with a key
1484−-- specific fold collecting the results in a map from the keys to the results.
1487+-- specific fold. The fold returns the fold result as the second component of
1485−-- Useful for cases like protocol handlers to handle different type of packets
1488+-- the output tuple whenever a fold terminates. The first component of the
1486−-- using different handlers.
1489+-- tuple is a container of in-progress folds. If a fold terminates, another
1490+-- instance of the fold is started upon receiving an input with that key.
1488−-- @
1492+-- This can be used to scan a stream and collect the results from the scan
1490−-- |-------Fold m a b
1495+-- /Pre-release/
1491−-- -----stream m a-----Map-----|
1492−-- |-------Fold m a b
1497−-- Any input that does not map to a fold in the input Map is silently ignored.
1497+{-# INLINE demuxScanWith #-}
1498+demuxScanWith :: (Monad m, IsMap f, Traversable f) =>
1500+ -> (Key f -> Fold m a b)
1501+ -> Fold m a (m (f b), Maybe (Key f, b))
1502+demuxScanWith getKey getFold = fmap extract $ foldlM' step initial
1506+ initial = return $ Tuple' IsMap.mapEmpty Nothing
1508+ {-# INLINE runFold #-}
1509+ runFold kv (Fold step1 initial1 extract1) (k, a) = do
1517+ let fld = Fold step1 (return res1) extract1
1518+ in Tuple' (IsMap.mapInsert k fld kv) Nothing
1519+ Done b -> Tuple' (IsMap.mapDelete k kv) (Just (k, b))
1520+ Done b -> return $ Tuple' kv (Just (k, b))
1522+ step (Tuple' kv _) a = do
1524+ let fld = getFold k
1525+ case IsMap.mapLookup k kv of
1526+ Nothing -> runFold kv fld (k, a)
1527+ Just f -> runFold kv f (k, a)
1529+ extract (Tuple' kv x) = (Prelude.mapM f kv, x)
1533+ f (Fold _ i e) = do
1539+-- | This is specialized version of 'demuxScanWith' that uses mutable cells as
1540+-- fold accumulators for better performance.
1499−-- >>> demuxWith f kv = fmap fst $ Fold.demuxDefaultWith f kv Fold.drain
1542+-- >>> demuxScanMutWith = Fold.demuxScanWith
1501−-- /Pre-release/
1544+{-# INLINE demuxScanMutWith #-}
1545+demuxScanMutWith :: (MonadIO m, IsMap f, Traversable f) =>
1547+ -> (Key f -> Fold m a b)
1548+ -> Fold m a (m (f b), Maybe (Key f, b))
1549+demuxScanMutWith getKey getFold = fmap extract $ foldlM' step initial
1553+ initial = return $ Tuple' IsMap.mapEmpty Nothing
1555+ {-# INLINE initFold #-}
1556+ initFold kv (Fold step1 initial1 extract1) (k, a) = do
1563+ let fld = Fold step1 (return res1) extract1
1564+ ref <- liftIO $ newIORef fld
1565+ return $ Tuple' (IsMap.mapInsert k ref kv) Nothing
1566+ Done b -> return $ Tuple' kv (Just (k, b))
1567+ Done b -> return $ Tuple' kv (Just (k, b))
1569+ {-# INLINE runFold #-}
1570+ runFold kv ref (Fold step1 initial1 extract1) (k, a) = do
1577+ let fld = Fold step1 (return res1) extract1
1578+ liftIO $ writeIORef ref fld
1579+ return $ Tuple' kv Nothing
1581+ let kv1 = IsMap.mapDelete k kv
1582+ in return $ Tuple' kv1 (Just (k, b))
1583+ Done _ -> error "demuxScanMutWith: unreachable"
1585+ step (Tuple' kv _) a = do
1587+ case IsMap.mapLookup k kv of
1590+ initFold kv f (k, a)
1592+ f <- liftIO $ readIORef ref
1593+ runFold kv ref f (k, a)
1595+ extract (Tuple' kv x) = (Prelude.mapM f kv, x)
1600+ (Fold _ i e) <- liftIO $ readIORef ref
1606+-- | This collects all the results of 'demuxScanWith' in a container.
1503{-# INLINE demuxWith #-}
1608{-# INLINE demuxWith #-}
1504−demuxWith :: (Monad m, Ord k)
1609+demuxWith :: (Monad m, IsMap f, Traversable f) =>
1610+ (a -> Key f) -> (Key f -> Fold m a b) -> Fold m a (f b)
1611+demuxWith getKey getFold =
1613+ classifier = demuxScanWith getKey getFold
1614+ getMap Nothing = pure IsMap.mapEmpty
1615+ getMap (Just action) = action
1617+ teeWith IsMap.mapUnion
1618+ (rmapM getMap $ lmap fst last)
1619+ (lmap snd $ catMaybes toMap)
1620+ in postscan classifier aggregator
1622+-- | Same as 'demuxWith' but uses 'demuxScanMutWith' for better performance.
1624+-- >>> demuxMutWith = Fold.demuxWith
1626+{-# INLINE demuxMutWith #-}
1627+demuxMutWith :: (MonadIO m, IsMap f, Traversable f) =>
1505− => (a -> (k, a')) -> Map k (Fold m a' b) -> Fold m a (Map k b)
1628+ (a -> Key f) -> (Key f -> Fold m a b) -> Fold m a (f b)
1506−demuxWith f kv = fmap fst $ demuxDefaultWith f kv drain
1629+demuxMutWith getKey getFold =
1631+ classifier = demuxScanMutWith getKey getFold
1632+ getMap Nothing = pure IsMap.mapEmpty
1633+ getMap (Just action) = action
1635+ teeWith IsMap.mapUnion
1636+ (rmapM getMap $ lmap fst last)
1637+ (lmap snd $ catMaybes toMap)
1638+ in postscan classifier aggregator
1508−-- | Fold a stream of key value pairs using a map of specific folds for each
1640+-- | Fold a stream of key value pairs using a function that maps keys to folds.
1509−-- key into a map from keys to the results of fold outputs of the corresponding
1512−-- >>> import qualified Data.Map
1642+-- >>> import Data.Map (Map)
1513-- >>> :{
1643-- >>> :{
1514−-- let table = Data.Map.fromList [("SUM", Fold.sum), ("PRODUCT", Fold.product)]
1644+-- let f "SUM" = Fold.sum
1645+-- f _ = Fold.product
1515-- input = Stream.fromList [("SUM",1),("PRODUCT",2),("SUM",3),("PRODUCT",4)]
1646-- input = Stream.fromList [("SUM",1),("PRODUCT",2),("SUM",3),("PRODUCT",4)]
1516−-- in Stream.fold (Fold.demux table) input
1647+-- in Stream.fold (Fold.demux f) input :: IO (Map String Int)
1518-- fromList [("PRODUCT",8),("SUM",4)]
1649-- fromList [("PRODUCT",8),("SUM",4)]
1520−-- >>> demux = Fold.demuxWith id
1651+-- >>> demux f = Fold.demuxWith fst (Fold.lmap snd . f)
1522-- /Pre-release/
1653-- /Pre-release/
1523{-# INLINE demux #-}
1654{-# INLINE demux #-}
1524−demux :: (Monad m, Ord k)
1655+demux :: (Monad m, IsMap f, Traversable f) =>
1525− => Map k (Fold m a b) -> Fold m (k, a) (Map k b)
1526−demux = demuxWith id
1528−data DemuxState s b doneMap runMap =
1529− DemuxMapAndDefault !s !doneMap !runMap
1530− | DemuxOnlyMap b !doneMap !runMap
1531− | DemuxOnlyDefault s !doneMap
1533−-- | Like 'demuxWith' but uses a default catchall fold to handle inputs which
1534−-- do not have a specific fold in the map to handle them.
1536−-- If any fold in the map stops, inputs meant for that fold are sent to the
1537−-- catchall fold. If the catchall fold stops then inputs that do not match any
1538−-- fold are ignored.
1540−-- Stops when all the folds, including the catchall fold, stop.
1544−{-# INLINE demuxDefaultWith #-}
1545−demuxDefaultWith :: (Monad m, Ord k)
1547− -> Map k (Fold m a' b)
1656+ (Key f -> Fold m a b) -> Fold m (Key f, a) (f b)
1548− -> Fold m (k, a') c
1657+demux f = demuxWith fst (lmap snd . f)
1549− -> Fold m a (Map k b, c)
1550−demuxDefaultWith f kv (Fold dstep dinitial dextract) =
1551− Fold step initial extract
1553− where
1659+------------------------------------------------------------------------------
1554−
1660+-- Classify: Like demux but uses the same fold for all keys.
1555− initial = do
1661+------------------------------------------------------------------------------
1556− let runInit (Fold step1 initial1 done1) = do
1560− Partial _ -> Right' (Fold step1 (return r) done1)
1563− -- initialize folds in the kv map and separate the ones that are done
1564− -- from running ones
1565− kv1 <- Prelude.mapM runInit kv
1566− let runMap = Map.map fromRight' $ Map.filter isRight' kv1
1567− doneMap = Map.map fromLeft' $ Map.filter isLeft' kv1
1569− -- Run the default fold, and decide the next state based on its result
1575− $ if Map.size runMap > 0
1576− then DemuxMapAndDefault s doneMap runMap
1577− else DemuxOnlyDefault s doneMap
1579− if Map.size runMap > 0
1580− then Partial $ DemuxOnlyMap b doneMap runMap
1581− else Done (doneMap, b)
1583− {-# INLINE runFold #-}
1663+-- XXX Change these to make the behavior similar to demux* variants. We can
1584− runFold fPartial fDone doneMap runMap (Fold step1 initial1 done1) k a1 = do
1664+-- implement this using classifyScanManyWith. Maintain a set of done folds in
1585− resi <- initial1
1665+-- the underlying monad, and when initial is called look it up, if the fold is
1586− case resi of
1666+-- done then initial would set a flag in the state to ignore the input or
1587− Partial st -> do
1667+-- return an error.
1589− return $ case res of
1591− let fld = Fold step1 (return $ Partial s) done1
1592− runMap1 = Map.insert k fld runMap
1593− in Partial $ fPartial doneMap runMap1
1595− let runMap1 = Map.delete k runMap
1596− doneMap1 = Map.insert k b doneMap
1597− if Map.size runMap1 == 0
1598− then fDone doneMap1
1599− else Partial $ fPartial doneMap1 runMap1
1600− Done _ -> error "Bug: demuxDefaultWith: Done fold"
1602− step (DemuxMapAndDefault dacc doneMap runMap) a = do
1604− case Map.lookup k runMap of
1606− res <- dstep dacc (k, a1)
1610− Partial s -> DemuxMapAndDefault s doneMap runMap
1611− Done b -> DemuxOnlyMap b doneMap runMap
1614− (DemuxMapAndDefault dacc)
1615− (Partial . DemuxOnlyDefault dacc)
1616− doneMap runMap fld k a1
1618− step (DemuxOnlyMap dval doneMap runMap) a = do
1620− case Map.lookup k runMap of
1621− Nothing -> return $ Partial $ DemuxOnlyMap dval doneMap runMap
1624− (DemuxOnlyMap dval)
1626− doneMap runMap fld k a1
1627− step (DemuxOnlyDefault dacc doneMap) a = do
1629− res <- dstep dacc (k, a1)
1632− Partial s -> Partial $ DemuxOnlyDefault s doneMap
1633− Done b -> Done (doneMap, b)
1635− runExtract (Fold _ initial1 done1) = do
1638− Partial s -> done1 s
1641− extract (DemuxMapAndDefault dacc doneMap runMap) = do
1643− runMap1 <- Prelude.mapM runExtract runMap
1644− return (doneMap `Map.union` runMap1, b)
1645− extract (DemuxOnlyMap dval doneMap runMap) = do
1646− runMap1 <- Prelude.mapM runExtract runMap
1647− return (doneMap `Map.union` runMap1, dval)
1648− extract (DemuxOnlyDefault dacc doneMap) = do
1650− return (doneMap, b)
1653−-- >>> demuxDefault = Fold.demuxDefaultWith id
1656−{-# INLINE demuxDefault #-}
1657−demuxDefault :: (Monad m, Ord k)
1658− => Map k (Fold m a b) -> Fold m (k, a) b -> Fold m (k, a) (Map k b, b)
1659−demuxDefault = demuxDefaultWith id
1661-- | Folds the values for each key using the supplied fold. When scanning, as
1669-- | Folds the values for each key using the supplied fold. When scanning, as
1662-- soon as the fold is complete, its result is available in the second
1670-- soon as the fold is complete, its result is available in the second
1663-- component of the tuple. The first component of the tuple is a snapshot of
1671-- component of the tuple. The first component of the tuple is a snapshot of
1664-- the in-progress folds.
1672-- the in-progress folds.
1674+-- Once the fold for a key is done, any future values of the key are ignored.
1676+-- >>> classifyScanWith f fld = Fold.demuxScanWith f (const fld)
1665{-# INLINE classifyScanWith #-}
1678{-# INLINE classifyScanWith #-}
1666classifyScanWith :: (Monad m, IsMap f, Traversable f, Ord (Key f)) =>
1679classifyScanWith :: (Monad m, IsMap f, Traversable f, Ord (Key f)) =>
1667 -- Note: we need to return the Map itself to display the in-progress values
1680 -- Note: we need to return the Map itself to display the in-progress values
@@ -1716,6 +1729,10 @@ classifyScanWith f (Fold step1 initial1 extract1) =
1716-- | Same as classifyScanWith except that it uses mutable IORef cells in the
1729-- | Same as classifyScanWith except that it uses mutable IORef cells in the
1717-- Map providing better performance. Be aware that if this is used as a scan,
1730-- Map providing better performance. Be aware that if this is used as a scan,
1718-- the values in the intermediate Maps would be mutable.
1731-- the values in the intermediate Maps would be mutable.
1733+-- >>> classifyScanMutWith = classifyScanWith
1734+-- >>> classifyScanMutWith f fld = demuxScanMutWith f (const fld)
1719{-# INLINE classifyScanMutWith #-}
1736{-# INLINE classifyScanMutWith #-}
1720classifyScanMutWith :: (MonadIO m, IsMap f, Traversable f, Ord (Key f)) =>
1737classifyScanMutWith :: (MonadIO m, IsMap f, Traversable f, Ord (Key f)) =>
1721 (a -> Key f) -> Fold m a b -> Fold m a (m (f b), Maybe (Key f, b))
1738 (a -> Key f) -> Fold m a b -> Fold m a (m (f b), Maybe (Key f, b))
@@ -1805,8 +1822,11 @@ classifyWith f fld =
1805 (lmap snd $ catMaybes toMap)
1822 (lmap snd $ catMaybes toMap)
1806 in postscan classifier aggregator
1823 in postscan classifier aggregator
1808−-- | Like classifyWith but maybe faster because it uses mutable cells as fold
1825+-- | Same as 'classifyWith' but maybe faster because it uses mutable cells as
1809−-- accumulators in the Map.
1826+-- fold accumulators in the Map.
1828+-- >>> classifyMutWith = Fold.classifyWith
1810{-# INLINE classifyMutWith #-}
1830{-# INLINE classifyMutWith #-}
1811classifyMutWith :: (MonadIO m, IsMap f, Traversable f, Ord (Key f)) =>
1831classifyMutWith :: (MonadIO m, IsMap f, Traversable f, Ord (Key f)) =>
1812 (a -> Key f) -> Fold m a b -> Fold m a (f b)
1832 (a -> Key f) -> Fold m a b -> Fold m a (f b)