@@ -214,14 +214,18 @@ module Streamly.Internal.Data.Fold
214214
215215 -- ** Demultiplexing
216216 -- | 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
218− -- of the classifier can use a different fold.
219− , demux -- XXX rename this to demux_
217+ -- fold selector. 'demux' is a generalization of 'classify' (and
218+ -- 'partition') where each key of the classifier can use a different fold.
219+ , demux
220220 , demuxWith
221− , demuxDefault -- XXX rename this to demux
222− , demuxDefaultWith
223− -- , demuxWithSel
224− -- , demuxWithMin
221+ , demuxScanWith
222+ , demuxScanMutWith
223+ , demuxMutWith
224+
225+ -- TODO: These can be implemented using the above operations
226+ -- , demuxWithSel -- Stop when the fold for the specified key stops
227+ -- , demuxWithMin -- Stop when any of the folds stop
228+ -- , demuxWithAll -- Stop when all the folds stop (run once)
225229
226230 -- ** Classifying
227231 -- | In an input stream of key value pairs fold values for different keys
@@ -281,8 +285,6 @@ import Data.Int (Int64)
281285import Data.IORef (newIORef, readIORef, writeIORef)
282286import Data.Map.Strict (Map)
283287import Data.Maybe (isJust, fromJust)
284−import Streamly.Internal.Data.Either.Strict
285− (Either'(..), fromLeft', fromRight', isLeft', isRight')
286288import Streamly.Internal.Data.IsMap (IsMap(..))
287289import Streamly.Internal.Data.Pipe.Type (Pipe (..), PipeState(..))
288290import Streamly.Internal.Data.Unfold.Type (Unfold(..))
@@ -290,7 +292,6 @@ import Streamly.Internal.Data.Tuple.Strict (Tuple'(..), Tuple3'(..))
290292import Streamly.Internal.Data.Stream.Serial (SerialT(..))
291293
292294import qualified Data.IntSet as IntSet
293−import qualified Data.Map.Strict as Map
294295import qualified Data.Set as Set
295296import qualified Streamly.Internal.Data.IsMap as IsMap
296297import qualified Prelude
@@ -1467,201 +1468,213 @@ partition = partitionBy id
14671468-- partitionN fs = Fold step begin done
14681469-}
14691470
1471+------------------------------------------------------------------------------
1472+-- demux: in a key value stream fold each key sub-stream with a different fold
1473+------------------------------------------------------------------------------
1474+
14701475-- TODO Demultiplex an input element into a number of typed variants. We want
14711476-- 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
1473−-- only those types and the full set of those types.
1474−--
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.
1477+-- an enumeration of a GADT.
14791478--
14801479-- This is the consumer side dual of the producer side 'mux' operation (XXX to
14811480-- be implemented).
1481+--
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
1484+-- the Map.
14821485
1483−-- | Split the input stream based on a key field and fold each split using a
1484−-- specific fold collecting the results in a map from the keys to the results.
1485−-- Useful for cases like protocol handlers to handle different type of packets
1486−-- using different handlers.
1486+-- | In a key value stream, fold values corresponding to each key with a key
1487+-- specific fold. The fold returns the fold result as the second component of
1488+-- the output tuple whenever a fold terminates. The first component of the
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.
14871491--
1488−-- @
1492+-- This can be used to scan a stream and collect the results from the scan
1493+-- output.
14891494--
1490−-- |-------Fold m a b
1491−-- -----stream m a-----Map-----|
1492−-- |-------Fold m a b
1493−-- |
1494−-- ...
1495−-- @
1495+-- /Pre-release/
14961496--
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) =>
1499+ (a -> Key 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
1503+
1504+ where
1505+
1506+ initial = return $ Tuple' IsMap.mapEmpty Nothing
1507+
1508+ {-# INLINE runFold #-}
1509+ runFold kv (Fold step1 initial1 extract1) (k, a) = do
1510+ res <- initial1
1511+ case res of
1512+ Partial s -> do
1513+ res1 <- step1 s a
1514+ return
1515+ $ case res1 of
1516+ Partial _ ->
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))
1521+
1522+ step (Tuple' kv _) a = do
1523+ let k = getKey a
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)
1528+
1529+ extract (Tuple' kv x) = (Prelude.mapM f kv, x)
1530+
1531+ where
1532+
1533+ f (Fold _ i e) = do
1534+ r <- i
1535+ case r of
1536+ Partial s -> e s
1537+ Done b -> return b
1538+
1539+-- | This is specialized version of 'demuxScanWith' that uses mutable cells as
1540+-- fold accumulators for better performance.
14981541--
1499−-- >>> demuxWith f kv = fmap fst $ Fold.demuxDefaultWith f kv Fold.drain
1542+-- >>> demuxScanMutWith = Fold.demuxScanWith
15001543--
1501−-- /Pre-release/
1544+{-# INLINE demuxScanMutWith #-}
1545+demuxScanMutWith :: (MonadIO m, IsMap f, Traversable f) =>
1546+ (a -> Key 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
1550+
1551+ where
1552+
1553+ initial = return $ Tuple' IsMap.mapEmpty Nothing
1554+
1555+ {-# INLINE initFold #-}
1556+ initFold kv (Fold step1 initial1 extract1) (k, a) = do
1557+ res <- initial1
1558+ case res of
1559+ Partial s -> do
1560+ res1 <- step1 s a
1561+ case res1 of
1562+ Partial _ -> 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))
1568+
1569+ {-# INLINE runFold #-}
1570+ runFold kv ref (Fold step1 initial1 extract1) (k, a) = do
1571+ res <- initial1
1572+ case res of
1573+ Partial s -> do
1574+ res1 <- step1 s a
1575+ case res1 of
1576+ Partial _ -> do
1577+ let fld = Fold step1 (return res1) extract1
1578+ liftIO $ writeIORef ref fld
1579+ return $ Tuple' kv Nothing
1580+ Done b ->
1581+ let kv1 = IsMap.mapDelete k kv
1582+ in return $ Tuple' kv1 (Just (k, b))
1583+ Done _ -> error "demuxScanMutWith: unreachable"
1584+
1585+ step (Tuple' kv _) a = do
1586+ let k = getKey a
1587+ case IsMap.mapLookup k kv of
1588+ Nothing -> do
1589+ let f = getFold k
1590+ initFold kv f (k, a)
1591+ Just ref -> do
1592+ f <- liftIO $ readIORef ref
1593+ runFold kv ref f (k, a)
1594+
1595+ extract (Tuple' kv x) = (Prelude.mapM f kv, x)
1596+
1597+ where
1598+
1599+ f ref = do
1600+ (Fold _ i e) <- liftIO $ readIORef ref
1601+ r <- i
1602+ case r of
1603+ Partial s -> e s
1604+ Done b -> return b
1605+
1606+-- | This collects all the results of 'demuxScanWith' in a container.
15021607--
15031608{-# INLINE demuxWith #-}
1504−demuxWith :: (Monad m, Ord k)
1505− => (a -> (k, a')) -> Map k (Fold m a' b) -> Fold m a (Map k b)
1506−demuxWith f kv = fmap fst $ demuxDefaultWith f kv drain
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 =
1612+ let
1613+ classifier = demuxScanWith getKey getFold
1614+ getMap Nothing = pure IsMap.mapEmpty
1615+ getMap (Just action) = action
1616+ aggregator =
1617+ teeWith IsMap.mapUnion
1618+ (rmapM getMap $ lmap fst last)
1619+ (lmap snd $ catMaybes toMap)
1620+ in postscan classifier aggregator
1621+
1622+-- | Same as 'demuxWith' but uses 'demuxScanMutWith' for better performance.
1623+--
1624+-- >>> demuxMutWith = Fold.demuxWith
1625+--
1626+{-# INLINE demuxMutWith #-}
1627+demuxMutWith :: (MonadIO m, IsMap f, Traversable f) =>
1628+ (a -> Key f) -> (Key f -> Fold m a b) -> Fold m a (f b)
1629+demuxMutWith getKey getFold =
1630+ let
1631+ classifier = demuxScanMutWith getKey getFold
1632+ getMap Nothing = pure IsMap.mapEmpty
1633+ getMap (Just action) = action
1634+ aggregator =
1635+ teeWith IsMap.mapUnion
1636+ (rmapM getMap $ lmap fst last)
1637+ (lmap snd $ catMaybes toMap)
1638+ in postscan classifier aggregator
15071639
1508−-- | Fold a stream of key value pairs using a map of specific folds for each
1509−-- key into a map from keys to the results of fold outputs of the corresponding
1510−-- values.
1640+-- | Fold a stream of key value pairs using a function that maps keys to folds.
15111641--
1512−-- >>> import qualified Data.Map
1642+-- >>> import Data.Map (Map)
15131643-- >>> :{
1514−-- let table = Data.Map.fromList [("SUM", Fold.sum), ("PRODUCT", Fold.product)]
1644+-- let f "SUM" = Fold.sum
1645+-- f _ = Fold.product
15151646-- 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)
15171648-- :}
15181649-- fromList [("PRODUCT",8),("SUM",4)]
15191650--
1520−-- >>> demux = Fold.demuxWith id
1651+-- >>> demux f = Fold.demuxWith fst (Fold.lmap snd . f)
15211652--
15221653-- /Pre-release/
15231654{-# INLINE demux #-}
1524−demux :: (Monad m, Ord k)
1525− => Map k (Fold m a b) -> Fold m (k, a) (Map k b)
1526−demux = demuxWith id
1527−
1528−data DemuxState s b doneMap runMap =
1529− DemuxMapAndDefault !s !doneMap !runMap
1530− | DemuxOnlyMap b !doneMap !runMap
1531− | DemuxOnlyDefault s !doneMap
1532−
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.
1535−--
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.
1539−--
1540−-- Stops when all the folds, including the catchall fold, stop.
1541−--
1542−-- /Pre-release/
1543−--
1544−{-# INLINE demuxDefaultWith #-}
1545−demuxDefaultWith :: (Monad m, Ord k)
1546− => (a -> (k, a'))
1547− -> Map k (Fold m a' b)
1548− -> Fold m (k, a') c
1549− -> Fold m a (Map k b, c)
1550−demuxDefaultWith f kv (Fold dstep dinitial dextract) =
1551− Fold step initial extract
1655+demux :: (Monad m, IsMap f, Traversable f) =>
1656+ (Key f -> Fold m a b) -> Fold m (Key f, a) (f b)
1657+demux f = demuxWith fst (lmap snd . f)
15521658
1553− where
1554−
1555− initial = do
1556− let runInit (Fold step1 initial1 done1) = do
1557− r <- initial1
1558− return
1559− $ case r of
1560− Partial _ -> Right' (Fold step1 (return r) done1)
1561− Done b -> Left' b
1562−
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
1568−
1569− -- Run the default fold, and decide the next state based on its result
1570− dres <- dinitial
1571− return
1572− $ case dres of
1573− Partial s ->
1574− Partial
1575− $ if Map.size runMap > 0
1576− then DemuxMapAndDefault s doneMap runMap
1577− else DemuxOnlyDefault s doneMap
1578− Done b ->
1579− if Map.size runMap > 0
1580− then Partial $ DemuxOnlyMap b doneMap runMap
1581− else Done (doneMap, b)
1659+------------------------------------------------------------------------------
1660+-- Classify: Like demux but uses the same fold for all keys.
1661+------------------------------------------------------------------------------
15821662
1583− {-# INLINE runFold #-}
1584− runFold fPartial fDone doneMap runMap (Fold step1 initial1 done1) k a1 = do
1585− resi <- initial1
1586− case resi of
1587− Partial st -> do
1588− res <- step1 st a1
1589− return $ case res of
1590− Partial s ->
1591− let fld = Fold step1 (return $ Partial s) done1
1592− runMap1 = Map.insert k fld runMap
1593− in Partial $ fPartial doneMap runMap1
1594− Done b -> do
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"
1601−
1602− step (DemuxMapAndDefault dacc doneMap runMap) a = do
1603− let (k, a1) = f a
1604− case Map.lookup k runMap of
1605− Nothing -> do
1606− res <- dstep dacc (k, a1)
1607− return
1608− $ Partial
1609− $ case res of
1610− Partial s -> DemuxMapAndDefault s doneMap runMap
1611− Done b -> DemuxOnlyMap b doneMap runMap
1612− Just fld ->
1613− runFold
1614− (DemuxMapAndDefault dacc)
1615− (Partial . DemuxOnlyDefault dacc)
1616− doneMap runMap fld k a1
1617−
1618− step (DemuxOnlyMap dval doneMap runMap) a = do
1619− let (k, a1) = f a
1620− case Map.lookup k runMap of
1621− Nothing -> return $ Partial $ DemuxOnlyMap dval doneMap runMap
1622− Just fld ->
1623− runFold
1624− (DemuxOnlyMap dval)
1625− (Done . (, dval))
1626− doneMap runMap fld k a1
1627− step (DemuxOnlyDefault dacc doneMap) a = do
1628− let (k, a1) = f a
1629− res <- dstep dacc (k, a1)
1630− return
1631− $ case res of
1632− Partial s -> Partial $ DemuxOnlyDefault s doneMap
1633− Done b -> Done (doneMap, b)
1634−
1635− runExtract (Fold _ initial1 done1) = do
1636− res <- initial1
1637− case res of
1638− Partial s -> done1 s
1639− Done b -> return b
1640−
1641− extract (DemuxMapAndDefault dacc doneMap runMap) = do
1642− b <- dextract dacc
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
1649− b <- dextract dacc
1650− return (doneMap, b)
1651−
1652−-- |
1653−-- >>> demuxDefault = Fold.demuxDefaultWith id
1654−--
1655−-- /Pre-release/
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
1663+-- XXX Change these to make the behavior similar to demux* variants. We can
1664+-- implement this using classifyScanManyWith. Maintain a set of done folds in
1665+-- the underlying monad, and when initial is called look it up, if the fold is
1666+-- done then initial would set a flag in the state to ignore the input or
1667+-- return an error.
16601668
16611669-- | Folds the values for each key using the supplied fold. When scanning, as
16621670-- soon as the fold is complete, its result is available in the second
16631671-- component of the tuple. The first component of the tuple is a snapshot of
16641672-- the in-progress folds.
1673+--
1674+-- Once the fold for a key is done, any future values of the key are ignored.
1675+--
1676+-- >>> classifyScanWith f fld = Fold.demuxScanWith f (const fld)
1677+--
16651678{-# INLINE classifyScanWith #-}
16661679classifyScanWith :: (Monad m, IsMap f, Traversable f, Ord (Key f)) =>
16671680 -- Note: we need to return the Map itself to display the in-progress values
@@ -1716,6 +1729,10 @@ classifyScanWith f (Fold step1 initial1 extract1) =
17161729-- | Same as classifyScanWith except that it uses mutable IORef cells in the
17171730-- Map providing better performance. Be aware that if this is used as a scan,
17181731-- the values in the intermediate Maps would be mutable.
1732+--
1733+-- >>> classifyScanMutWith = classifyScanWith
1734+-- >>> classifyScanMutWith f fld = demuxScanMutWith f (const fld)
1735+--
17191736{-# INLINE classifyScanMutWith #-}
17201737classifyScanMutWith :: (MonadIO m, IsMap f, Traversable f, Ord (Key f)) =>
17211738 (a -> Key f) -> Fold m a b -> Fold m a (m (f b), Maybe (Key f, b))
@@ -1805,8 +1822,11 @@ classifyWith f fld =
18051822 (lmap snd $ catMaybes toMap)
18061823 in postscan classifier aggregator
18071824
1808−-- | Like classifyWith but maybe faster because it uses mutable cells as fold
1809−-- accumulators in the Map.
1825+-- | Same as 'classifyWith' but maybe faster because it uses mutable cells as
1826+-- fold accumulators in the Map.
1827+--
1828+-- >>> classifyMutWith = Fold.classifyWith
1829+--
18101830{-# INLINE classifyMutWith #-}
18111831classifyMutWith :: (MonadIO m, IsMap f, Traversable f, Ord (Key f)) =>
18121832 (a -> Key f) -> Fold m a b -> Fold m a (f b)