@@ -159,7 +159,9 @@ module Streamly.Internal.Data.Fold
159159 , demux
160160 -- , demuxWith
161161 , demux_
162+ , demuxDefault_
162163 -- , demuxWith_
164+ , demuxWithDefault_
163165
164166 -- * Classifying
165167
@@ -1016,6 +1018,32 @@ demux :: (Monad m, Ord k)
10161018 => Map k (Fold m a b) -> Fold m (k, a) (Map k b)
10171019demux = demuxWith id
10181020
1021+{-# INLINE demuxWithDefault_ #-}
1022+demuxWithDefault_ :: (Monad m, Ord k)
1023+ => (a -> (k, a')) -> Map k (Fold m a' b) -> Fold m (k, a') b -> Fold m a ()
1024+demuxWithDefault_ f kv (Fold dstep dinitial dextract) =
1025+ Fold step initial extract
1026+
1027+ where
1028+
1029+ initFold (Fold s i e) = i >>= \r -> return (Fold s (return r) e)
1030+ initial = do
1031+ mp <- Prelude.mapM initFold kv
1032+ dacc <- dinitial
1033+ return (Tuple' mp dacc)
1034+ step (Tuple' mp dacc) a
1035+ | (k, a') <- f a
1036+ = case Map.lookup k mp of
1037+ Nothing -> do
1038+ acc <- dstep dacc (k, a')
1039+ return (Tuple' mp acc)
1040+ Just (Fold step' acc _) -> do
1041+ _ <- acc >>= \x -> step' x a'
1042+ return (Tuple' mp dacc)
1043+ extract (Tuple' mp dacc) = do
1044+ void $ dextract dacc
1045+ Prelude.mapM_ (\(Fold _ acc e) -> acc >>= e) mp
1046+
10191047-- | Split the input stream based on a key field and fold each split using a
10201048-- specific fold without collecting the results. Useful for cases like protocol
10211049-- handlers to handle different type of packets.
@@ -1074,6 +1102,11 @@ demuxWith_ f kv = Fold step initial extract
10741102demux_ :: (Monad m, Ord k) => Map k (Fold m a ()) -> Fold m (k, a) ()
10751103demux_ = demuxWith_ id
10761104
1105+{-# INLINE demuxDefault_ #-}
1106+demuxDefault_ :: (Monad m, Ord k)
1107+ => Map k (Fold m a ()) -> Fold m (k, a) () -> Fold m (k, a) ()
1108+demuxDefault_ = demuxWithDefault_ id
1109+
10771110-- TODO If the data is large we may need a map/hashmap in pinned memory instead
10781111-- of a regular Map. That may require a serializable constraint though. We can
10791112-- have another API for that.