on master help sign in
Use monadic action to generate fold in demux*
800ae8aHarendra Kumar
benchmark/Streamly/Benchmark/Data/Fold.hs+6−6
@@ -193,27 +193,27 @@ partitionByMinM =
193
193
194{-# INLINE demuxWith #-}
194{-# INLINE demuxWith #-}
195demuxWith :: (Monad m, Ord k) =>
195demuxWith :: (Monad m, Ord k) =>
196− (a -> k) -> (k -> Fold m a b) -> SerialT m a -> m (Map k b)
196+ (a -> k) -> (k -> m (Fold m a b)) -> SerialT m a -> m (Map k b)
197demuxWith f g = S.fold (FL.demuxWith f g)
197demuxWith f g = S.fold (FL.demuxWith f g)
198
198
199{-# INLINE demuxWithInt #-}
199{-# INLINE demuxWithInt #-}
200demuxWithInt :: Monad m =>
200demuxWithInt :: Monad m =>
201− (a -> Int) -> (Int -> Fold m a b) -> SerialT m a -> m (IntMap b)
201+ (a -> Int) -> (Int -> m (Fold m a b)) -> SerialT m a -> m (IntMap b)
202demuxWithInt f g = S.fold (FL.demuxWith f g)
202demuxWithInt f g = S.fold (FL.demuxWith f g)
203
203
204{-# INLINE demuxWithHash #-}
204{-# INLINE demuxWithHash #-}
205demuxWithHash :: (Monad m, Ord k, Hashable k) =>
205demuxWithHash :: (Monad m, Ord k, Hashable k) =>
206− (a -> k) -> (k -> Fold m a b) -> SerialT m a -> m (HashMap k b)
206+ (a -> k) -> (k -> m (Fold m a b)) -> SerialT m a -> m (HashMap k b)
207demuxWithHash f g = S.fold (FL.demuxWith f g)
207demuxWithHash f g = S.fold (FL.demuxWith f g)
208
208
209{-# INLINE demuxMutWith #-}
209{-# INLINE demuxMutWith #-}
210demuxMutWith :: (MonadIO m, Ord k) =>
210demuxMutWith :: (MonadIO m, Ord k) =>
211− (a -> k) -> (k -> Fold m a b) -> SerialT m a -> m (Map k b)
211+ (a -> k) -> (k -> m (Fold m a b)) -> SerialT m a -> m (Map k b)
212demuxMutWith f g = S.fold (FL.demuxMutWith f g)
212demuxMutWith f g = S.fold (FL.demuxMutWith f g)
213
213
214{-# INLINE demuxMutWithHash #-}
214{-# INLINE demuxMutWithHash #-}
215demuxMutWithHash :: (MonadIO m, Ord k, Hashable k) =>
215demuxMutWithHash :: (MonadIO m, Ord k, Hashable k) =>
216− (a -> k) -> (k -> Fold m a b) -> SerialT m a -> m (HashMap k b)
216+ (a -> k) -> (k -> m (Fold m a b)) -> SerialT m a -> m (HashMap k b)
217demuxMutWithHash f g = S.fold (FL.demuxMutWith f g)
217demuxMutWithHash f g = S.fold (FL.demuxMutWith f g)
218
218
219{-# INLINE classifyWith #-}
219{-# INLINE classifyWith #-}
@@ -468,7 +468,7 @@ o_n_heap_serial value =
468 getKey buckets = (`mod` buckets)
468 getKey buckets = (`mod` buckets)
469
469
470 getFold k =
470 getFold k =
471− case k of
471+ return $ case k of
472 0 -> FL.sum
472 0 -> FL.sum
473 1 -> FL.length
473 1 -> FL.length
474 _ -> FL.length
474 _ -> FL.length
core/src/Streamly/Internal/Data/Fold.hs+10−10
@@ -1497,7 +1497,7 @@ partition = partitionBy id
1497{-# INLINE demuxScanWith #-}
1497{-# INLINE demuxScanWith #-}
1498demuxScanWith :: (Monad m, IsMap f, Traversable f) =>
1498demuxScanWith :: (Monad m, IsMap f, Traversable f) =>
1499 (a -> Key f)
1499 (a -> Key f)
1500− -> (Key f -> Fold m a b)
1500+ -> (Key f -> m (Fold m a b))
1501 -> Fold m a (m (f b), Maybe (Key f, b))
1501 -> Fold m a (m (f b), Maybe (Key f, b))
1502demuxScanWith getKey getFold = fmap extract $ foldlM' step initial
1502demuxScanWith getKey getFold = fmap extract $ foldlM' step initial
1503
1503
@@ -1521,7 +1521,7 @@ demuxScanWith getKey getFold = fmap extract $ foldlM' step initial
1521
1521
1522 step (Tuple' kv _) a = do
1522 step (Tuple' kv _) a = do
1523 let k = getKey a
1523 let k = getKey a
1524− let fld = getFold k
1524+ fld <- getFold k
1525 case IsMap.mapLookup k kv of
1525 case IsMap.mapLookup k kv of
1526 Nothing -> runFold kv fld (k, a)
1526 Nothing -> runFold kv fld (k, a)
1527 Just f -> runFold kv f (k, a)
1527 Just f -> runFold kv f (k, a)
@@ -1544,7 +1544,7 @@ demuxScanWith getKey getFold = fmap extract $ foldlM' step initial
1544{-# INLINE demuxScanMutWith #-}
1544{-# INLINE demuxScanMutWith #-}
1545demuxScanMutWith :: (MonadIO m, IsMap f, Traversable f) =>
1545demuxScanMutWith :: (MonadIO m, IsMap f, Traversable f) =>
1546 (a -> Key f)
1546 (a -> Key f)
1547− -> (Key f -> Fold m a b)
1547+ -> (Key f -> m (Fold m a b))
1548 -> Fold m a (m (f b), Maybe (Key f, b))
1548 -> Fold m a (m (f b), Maybe (Key f, b))
1549demuxScanMutWith getKey getFold = fmap extract $ foldlM' step initial
1549demuxScanMutWith getKey getFold = fmap extract $ foldlM' step initial
1550
1550
@@ -1586,7 +1586,7 @@ demuxScanMutWith getKey getFold = fmap extract $ foldlM' step initial
1586 let k = getKey a
1586 let k = getKey a
1587 case IsMap.mapLookup k kv of
1587 case IsMap.mapLookup k kv of
1588 Nothing -> do
1588 Nothing -> do
1589− let f = getFold k
1589+ f <- getFold k
1590 initFold kv f (k, a)
1590 initFold kv f (k, a)
1591 Just ref -> do
1591 Just ref -> do
1592 f <- liftIO $ readIORef ref
1592 f <- liftIO $ readIORef ref
@@ -1607,7 +1607,7 @@ demuxScanMutWith getKey getFold = fmap extract $ foldlM' step initial
1607--
1607--
1608{-# INLINE demuxWith #-}
1608{-# INLINE demuxWith #-}
1609demuxWith :: (Monad m, IsMap f, Traversable f) =>
1609demuxWith :: (Monad m, IsMap f, Traversable f) =>
1610− (a -> Key f) -> (Key f -> Fold m a b) -> Fold m a (f b)
1610+ (a -> Key f) -> (Key f -> m (Fold m a b)) -> Fold m a (f b)
1611demuxWith getKey getFold =
1611demuxWith getKey getFold =
1612 let
1612 let
1613 classifier = demuxScanWith getKey getFold
1613 classifier = demuxScanWith getKey getFold
@@ -1625,7 +1625,7 @@ demuxWith getKey getFold =
1625--
1625--
1626{-# INLINE demuxMutWith #-}
1626{-# INLINE demuxMutWith #-}
1627demuxMutWith :: (MonadIO m, IsMap f, Traversable f) =>
1627demuxMutWith :: (MonadIO m, IsMap f, Traversable f) =>
1628− (a -> Key f) -> (Key f -> Fold m a b) -> Fold m a (f b)
1628+ (a -> Key f) -> (Key f -> m (Fold m a b)) -> Fold m a (f b)
1629demuxMutWith getKey getFold =
1629demuxMutWith getKey getFold =
1630 let
1630 let
1631 classifier = demuxScanMutWith getKey getFold
1631 classifier = demuxScanMutWith getKey getFold
@@ -1641,8 +1641,8 @@ demuxMutWith getKey getFold =
1641--
1641--
1642-- >>> import Data.Map (Map)
1642-- >>> import Data.Map (Map)
1643-- >>> :{
1643-- >>> :{
1644−-- let f "SUM" = Fold.sum
1644+-- let f "SUM" = return Fold.sum
1645−-- f _ = Fold.product
1645+-- f _ = return Fold.product
1646-- input = Stream.fromList [("SUM",1),("PRODUCT",2),("SUM",3),("PRODUCT",4)]
1646-- input = Stream.fromList [("SUM",1),("PRODUCT",2),("SUM",3),("PRODUCT",4)]
1647-- in Stream.fold (Fold.demux f) input :: IO (Map String Int)
1647-- in Stream.fold (Fold.demux f) input :: IO (Map String Int)
1648-- :}
1648-- :}
@@ -1653,8 +1653,8 @@ demuxMutWith getKey getFold =
1653-- /Pre-release/
1653-- /Pre-release/
1654{-# INLINE demux #-}
1654{-# INLINE demux #-}
1655demux :: (Monad m, IsMap f, Traversable f) =>
1655demux :: (Monad m, IsMap f, Traversable f) =>
1656− (Key f -> Fold m a b) -> Fold m (Key f, a) (f b)
1656+ (Key f -> m (Fold m a b)) -> Fold m (Key f, a) (f b)
1657−demux f = demuxWith fst (lmap snd . f)
1657+demux f = demuxWith fst (fmap (lmap snd) . f)
1658
1658
1659------------------------------------------------------------------------------
1659------------------------------------------------------------------------------
1660-- Classify: Like demux but uses the same fold for all keys.
1660-- Classify: Like demux but uses the same fold for all keys.
test/Streamly/Test/Data/Fold.hs+5−5
@@ -558,9 +558,9 @@ headAndRest ls = monadicIO $ do
558
558
559demux :: Expectation
559demux :: Expectation
560demux =
560demux =
561− let table "SUM" = FL.sum
561+ let table "SUM" = return FL.sum
562− table "PRODUCT" = FL.product
562+ table "PRODUCT" = return FL.product
563− table _ = FL.length
563+ table _ = return FL.length
564 input = Stream.fromList (
564 input = Stream.fromList (
565 [ ("SUM", 1)
565 [ ("SUM", 1)
566 , ("abc", 1)
566 , ("abc", 1)
@@ -584,8 +584,8 @@ demuxWith =
584 let getKey x | even x = "SUM"
584 let getKey x | even x = "SUM"
585 | otherwise = "PRODUCT"
585 | otherwise = "PRODUCT"
586
586
587− getFold "SUM" = FL.sum
587+ getFold "SUM" = return FL.sum
588− getFold "PRODUCT" = FL.product
588+ getFold "PRODUCT" = return FL.product
589 getFold _ = error "demuxWith: bug"
589 getFold _ = error "demuxWith: bug"
590
590
591 input = Stream.fromList [1, 2, 3, 4 :: Int]
591 input = Stream.fromList [1, 2, 3, 4 :: Int]