@@ -214,6 +214,7 @@ module Streamly.Internal.Data.Fold
214214 -- in individual output buckets using the given fold.
215215 , classify
216216 , classifyWith
217+ , classifyMutWith
217218 , classifyScanWith
218219 -- , classifyWithSel
219220 -- , classifyWithMin
@@ -260,14 +261,17 @@ where
260261
261262#include "inline.hs"
262263
264+import Control.Monad.IO.Class (MonadIO(..))
263265import Data.Bifunctor (first)
264266import Data.Either (isLeft, isRight, fromLeft, fromRight)
265267import Data.Functor.Identity (Identity(..))
266268import Data.Int (Int64)
269+import Data.IORef (newIORef, readIORef, writeIORef)
267270import Data.Map.Strict (Map)
268271import Data.Maybe (isJust, fromJust)
269272import Streamly.Internal.Data.Either.Strict
270273 (Either'(..), fromLeft', fromRight', isLeft', isRight')
274+import Streamly.Internal.Data.IsMap (IsMap(..))
271275import Streamly.Internal.Data.Pipe.Type (Pipe (..), PipeState(..))
272276import Streamly.Internal.Data.Unfold.Type (Unfold(..))
273277import Streamly.Internal.Data.Tuple.Strict (Tuple'(..), Tuple3'(..))
@@ -276,6 +280,7 @@ import Streamly.Internal.Data.Stream.Serial (SerialT(..))
276280import qualified Data.IntSet as IntSet
277281import qualified Data.Map.Strict as Map
278282import qualified Data.Set as Set
283+import qualified Streamly.Internal.Data.IsMap as IsMap
279284import qualified Streamly.Internal.Data.Pipe.Type as Pipe
280285import qualified Streamly.Internal.Data.Stream.StreamD.Type as StreamD
281286-- import qualified Streamly.Internal.Data.Stream.IsStream.Enumeration as Stream
@@ -1625,19 +1630,23 @@ demuxDefault :: (Monad m, Ord k)
16251630 => Map k (Fold m a b) -> Fold m (k, a) b -> Fold m (k, a) (Map k b, b)
16261631demuxDefault = demuxDefaultWith id
16271632
1633+-- | Folds the values for each key using the supplied fold. When scanning, as
1634+-- soon as the fold is complete, its result is available in the second
1635+-- component of the tuple. The first component of the tuple is a snapshot of
1636+-- the in-progress folds.
16281637{-# INLINE classifyScanWith #-}
1629−classifyScanWith :: (Monad m, Ord k) =>
1638+classifyScanWith :: (Monad m, IsMap f, Traversable f, Ord (Key f)) =>
16301639 -- Note: we need to return the Map itself to display the in-progress values
16311640 -- e.g. to implement top. We could possibly create a separate abstraction
16321641 -- for that use case. We return an action because we want it to be lazy so
16331642 -- that the downstream consumers can choose to process or discard it.
1634− (a -> k) -> Fold m a b -> Fold m a (m (Map k b), Maybe (k, b))
1643+ (a -> Key f) -> Fold m a b -> Fold m a (m (f b), Maybe (Key f, b))
16351644classifyScanWith f (Fold step1 initial1 extract1) =
16361645 fmap extract $ foldlM' step initial
16371646
16381647 where
16391648
1640− initial = return $ Tuple3' Map.empty Set.empty Nothing
1649+ initial = return $ Tuple3' IsMap.mapEmpty Set.empty Nothing
16411650
16421651 {-# INLINE initFold #-}
16431652 initFold kv set k a = do
@@ -1648,14 +1657,14 @@ classifyScanWith f (Fold step1 initial1 extract1) =
16481657 return
16491658 $ case r of
16501659 Partial s1 ->
1651− Tuple3' (Map.insert k s1 kv) set Nothing
1660+ Tuple3' (IsMap.mapInsert k s1 kv) set Nothing
16521661 Done b ->
16531662 Tuple3' kv set (Just (k, b))
16541663 Done b -> return (Tuple3' kv (Set.insert k set) (Just (k, b)))
16551664
16561665 step (Tuple3' kv set _) a = do
16571666 let k = f a
1658− case Map.lookup k kv of
1667+ case IsMap.mapLookup k kv of
16591668 Nothing -> do
16601669 if Set.member k set
16611670 then return (Tuple3' kv set Nothing)
@@ -1665,16 +1674,69 @@ classifyScanWith f (Fold step1 initial1 extract1) =
16651674 return
16661675 $ case r of
16671676 Partial s1 ->
1668− Tuple3' (Map.insert k s1 kv) set Nothing
1677+ Tuple3' (IsMap.mapInsert k s1 kv) set Nothing
16691678 Done b ->
1670− let kv1 = Map.delete k kv
1679+ let kv1 = IsMap.mapDelete k kv
16711680 in Tuple3' kv1 (Set.insert k set) (Just (k, b))
16721681
16731682 extract (Tuple3' kv _ x) = (Prelude.mapM extract1 kv, x)
16741683
1684+-- The code is almost the same as classifyScanWith except the IORef operations.
1685+--
1686+-- | Same as classifyScanWith except that it uses mutable IORef cells in the
1687+-- Map providing better performance. Be aware that if this is used as a scan,
1688+-- the values in the intermediate Maps would be mutable.
1689+{-# INLINE classifyScanMutWith #-}
1690+classifyScanMutWith :: (MonadIO m, IsMap f, Traversable f, Ord (Key f)) =>
1691+ (a -> Key f) -> Fold m a b -> Fold m a (m (f b), Maybe (Key f, b))
1692+classifyScanMutWith f (Fold step1 initial1 extract1) =
1693+ fmap extract $ foldlM' step initial
1694+
1695+ where
1696+
1697+ initial = return $ Tuple3' IsMap.mapEmpty Set.empty Nothing
1698+
1699+ {-# INLINE initFold #-}
1700+ initFold kv set k a = do
1701+ x <- initial1
1702+ case x of
1703+ Partial s -> do
1704+ r <- step1 s a
1705+ case r of
1706+ Partial s1 -> do
1707+ ref <- liftIO $ newIORef s1
1708+ return $ Tuple3' (IsMap.mapInsert k ref kv) set Nothing
1709+ Done b ->
1710+ return $ Tuple3' kv set (Just (k, b))
1711+ Done b -> return (Tuple3' kv (Set.insert k set) (Just (k, b)))
1712+
1713+ step (Tuple3' kv set _) a = do
1714+ let k = f a
1715+ case IsMap.mapLookup k kv of
1716+ Nothing -> do
1717+ if Set.member k set
1718+ then return (Tuple3' kv set Nothing)
1719+ else initFold kv set k a
1720+ Just ref -> do
1721+ s <- liftIO $ readIORef ref
1722+ r <- step1 s a
1723+ case r of
1724+ Partial s1 -> do
1725+ liftIO $ writeIORef ref s1
1726+ return $ Tuple3' kv set Nothing
1727+ Done b ->
1728+ let kv1 = IsMap.mapDelete k kv
1729+ in return
1730+ $ Tuple3' kv1 (Set.insert k set) (Just (k, b))
1731+
1732+ extract (Tuple3' kv _ x) =
1733+ (Prelude.mapM (\ref -> liftIO (readIORef ref) >>= extract1) kv, x)
1734+
1735+-- | Fold a key value stream to a key-value container. If the same key appears
1736+-- multiple times, only the last value is retained.
16751737{-# INLINE toMap #-}
1676−toMap :: (Monad m, Ord k) => Fold m (k, a) (Map k a)
1677−toMap = foldl' (\kv (k, v) -> Map.insert k v kv) Map.empty
1738+toMap :: (Monad m, IsMap f) => Fold m (Key f, a) (f a)
1739+toMap = foldl' (\kv (k, v) -> IsMap.mapInsert k v kv) IsMap.mapEmpty
16781740
16791741-- | Split the input stream based on a key field and fold each split using the
16801742-- given fold. Useful for map/reduce, bucketizing the input in different bins
@@ -1711,6 +1773,22 @@ classifyWith f fld =
17111773 (lmap snd $ catMaybes toMap)
17121774 in postscan classifier aggregator
17131775
1776+-- | Like classifyWith but maybe faster because it uses mutable cells as fold
1777+-- accumulators in the Map.
1778+{-# INLINE classifyMutWith #-}
1779+classifyMutWith :: (MonadIO m, IsMap f, Traversable f, Ord (Key f)) =>
1780+ (a -> Key f) -> Fold m a b -> Fold m a (f b)
1781+classifyMutWith f fld =
1782+ let
1783+ classifier = classifyScanMutWith f fld
1784+ getMap Nothing = pure IsMap.mapEmpty
1785+ getMap (Just action) = action
1786+ aggregator =
1787+ teeWith IsMap.mapUnion
1788+ (rmapM getMap $ lmap fst last)
1789+ (lmap snd $ catMaybes toMap)
1790+ in postscan classifier aggregator
1791+
17141792-- | Given an input stream of key value pairs and a fold for values, fold all
17151793-- the values belonging to each key. Useful for map/reduce, bucketizing the
17161794-- input in different bins or for generating histograms.