800ae8a · streamly · git.iko.soy

git.iko.soy/streamly

Use monadic action to generate fold in demux*

commit800ae8a8f89bf6fdc469cbaee393612a52eb7b7a authorHarendra Kumar <harendra@composewell.com> dateWed 6 Apr 2022, 00:04 (5 years ago) parentd31f332

3 files changed+21−21

compare with master

unifiedsplit

benchmark/Streamly/Benchmark/Data/Fold.hs+6−6 core/src/Streamly/Internal/Data/Fold.hs+10−10 test/Streamly/Test/Data/Fold.hs+5−5

benchmark/Streamly/Benchmark/Data/Fold.hs+6−6

@@ -193,27 +193,27 @@ partitionByMinM =

193193

194194{-# INLINE demuxWith #-}

195195demuxWith :: (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)

197197demuxWith f g = S.fold (FL.demuxWith f g)

198198

199199{-# INLINE demuxWithInt #-}

200200demuxWithInt :: 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)

202202demuxWithInt f g = S.fold (FL.demuxWith f g)

203203

204204{-# INLINE demuxWithHash #-}

205205demuxWithHash :: (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)

207207demuxWithHash f g = S.fold (FL.demuxWith f g)

208208

209209{-# INLINE demuxMutWith #-}

210210demuxMutWith :: (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)

212212demuxMutWith f g = S.fold (FL.demuxMutWith f g)

213213

214214{-# INLINE demuxMutWithHash #-}

215215demuxMutWithHash :: (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)

217217demuxMutWithHash f g = S.fold (FL.demuxMutWith f g)

218218

219219{-# INLINE classifyWith #-}

@@ -468,7 +468,7 @@ o_n_heap_serial value =

468468 getKey buckets = (`mod` buckets)

469469

470470 getFold k =

471− case k of

471+ return $ case k of

472472 0 -> FL.sum

473473 1 -> FL.length

474474 _ -> FL.length

core/src/Streamly/Internal/Data/Fold.hs+10−10

@@ -1497,7 +1497,7 @@ partition = partitionBy id

14971497{-# INLINE demuxScanWith #-}

14981498demuxScanWith :: (Monad m, IsMap f, Traversable f) =>

14991499 (a -> Key f)

1500− -> (Key f -> Fold m a b)

1500+ -> (Key f -> m (Fold m a b))

15011501 -> Fold m a (m (f b), Maybe (Key f, b))

15021502demuxScanWith getKey getFold = fmap extract $ foldlM' step initial

15031503

@@ -1521,7 +1521,7 @@ demuxScanWith getKey getFold = fmap extract $ foldlM' step initial

15211521

15221522 step (Tuple' kv _) a = do

15231523 let k = getKey a

1524− let fld = getFold k

1524+ fld <- getFold k

15251525 case IsMap.mapLookup k kv of

15261526 Nothing -> runFold kv fld (k, a)

15271527 Just f -> runFold kv f (k, a)

@@ -1544,7 +1544,7 @@ demuxScanWith getKey getFold = fmap extract $ foldlM' step initial

15441544{-# INLINE demuxScanMutWith #-}

15451545demuxScanMutWith :: (MonadIO m, IsMap f, Traversable f) =>

15461546 (a -> Key f)

1547− -> (Key f -> Fold m a b)

1547+ -> (Key f -> m (Fold m a b))

15481548 -> Fold m a (m (f b), Maybe (Key f, b))

15491549demuxScanMutWith getKey getFold = fmap extract $ foldlM' step initial

15501550

@@ -1586,7 +1586,7 @@ demuxScanMutWith getKey getFold = fmap extract $ foldlM' step initial

15861586 let k = getKey a

15871587 case IsMap.mapLookup k kv of

15881588 Nothing -> do

1589− let f = getFold k

1589+ f <- getFold k

15901590 initFold kv f (k, a)

15911591 Just ref -> do

15921592 f <- liftIO $ readIORef ref

@@ -1607,7 +1607,7 @@ demuxScanMutWith getKey getFold = fmap extract $ foldlM' step initial

16071607--

16081608{-# INLINE demuxWith #-}

16091609demuxWith :: (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)

16111611demuxWith getKey getFold =

16121612 let

16131613 classifier = demuxScanWith getKey getFold

@@ -1625,7 +1625,7 @@ demuxWith getKey getFold =

16251625--

16261626{-# INLINE demuxMutWith #-}

16271627demuxMutWith :: (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)

16291629demuxMutWith getKey getFold =

16301630 let

16311631 classifier = demuxScanMutWith getKey getFold

@@ -1641,8 +1641,8 @@ demuxMutWith getKey getFold =

16411641--

16421642-- >>> import Data.Map (Map)

16431643-- >>> :{

1644−-- let f "SUM" = Fold.sum

1645−-- f _ = Fold.product

1644+-- let f "SUM" = return Fold.sum

1645+-- f _ = return Fold.product

16461646-- input = Stream.fromList [("SUM",1),("PRODUCT",2),("SUM",3),("PRODUCT",4)]

16471647-- in Stream.fold (Fold.demux f) input :: IO (Map String Int)

16481648-- :}

@@ -1653,8 +1653,8 @@ demuxMutWith getKey getFold =

16531653-- /Pre-release/

16541654{-# INLINE demux #-}

16551655demux :: (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)

1656+ (Key f -> m (Fold m a b)) -> Fold m (Key f, a) (f b)

1657+demux f = demuxWith fst (fmap (lmap snd) . f)

16581658

16591659------------------------------------------------------------------------------

16601660-- 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

558558

559559demux :: Expectation

560560demux =

561− let table "SUM" = FL.sum

562− table "PRODUCT" = FL.product

563− table _ = FL.length

561+ let table "SUM" = return FL.sum

562+ table "PRODUCT" = return FL.product

563+ table _ = return FL.length

564564 input = Stream.fromList (

565565 [ ("SUM", 1)

566566 , ("abc", 1)

@@ -584,8 +584,8 @@ demuxWith =

584584 let getKey x | even x = "SUM"

585585 | otherwise = "PRODUCT"

586586

587− getFold "SUM" = FL.sum

588− getFold "PRODUCT" = FL.product

587+ getFold "SUM" = return FL.sum

588+ getFold "PRODUCT" = return FL.product

589589 getFold _ = error "demuxWith: bug"

590590

591591 input = Stream.fromList [1, 2, 3, 4 :: Int]