@@ -925,24 +925,24 @@ partition = partitionBy id
925925-- @since 0.7.0
926926{-# INLINE demuxWith #-}
927927demuxWith :: (Monad m, Ord k)
928− => (a -> k) -> Map k (Fold m a b) -> Fold m a (Map k b)
928+ => (a -> (k, a')) -> Map k (Fold m a' b) -> Fold m a (Map k b)
929929demuxWith f kv = Fold step initial extract
930930
931931 where
932932
933933 initial = return kv
934− step mp a =
934+ step mp a = case f a of
935+ (k, a') -> Map.alterF twiddle k mp
935936 -- XXX should we raise an exception in Nothing case?
936937 -- Ideally we should enforce that it is a total map over k so that look
937938 -- up never fails
938939 -- XXX we could use a monadic update function for a single lookup and
939940 -- update in the map.
940− let k = f a
941− in case Map.lookup k mp of
942− Nothing -> return mp
943− Just (Fold step' acc extract') -> do
944− !r <- acc >>= \x -> step' x a
945− return $ Map.insert k (Fold step' (return r) extract') mp
941+ where
942+ twiddle Nothing = pure Nothing
943+ twiddle (Just (Fold step' acc extract')) = do
944+ !r <- acc >>= \x -> step' x a'
945+ pure . Just $ Fold step' (return r) extract'
946946 extract = Prelude.mapM (\(Fold _ acc e) -> acc >>= e)
947947
948948-- | Fold a stream of key value pairs using a map of specific folds for each
@@ -961,7 +961,7 @@ demuxWith f kv = Fold step initial extract
961961{-# INLINE demux #-}
962962demux :: (Monad m, Ord k)
963963 => Map k (Fold m a b) -> Fold m (k, a) (Map k b)
964−demux fs = demuxWith fst (Map.map (lmap snd) fs)
964+demux = demuxWith id
965965
966966-- | Split the input stream based on a key field and fold each split using a
967967-- specific fold without collecting the results. Useful for cases like protocol
@@ -984,7 +984,7 @@ demux fs = demuxWith fst (Map.map (lmap snd) fs)
984984-- large.
985985{-# INLINE demuxWith_ #-}
986986demuxWith_ :: (Monad m, Ord k)
987− => (a -> k) -> Map k (Fold m a b) -> Fold m a ()
987+ => (a -> (k, a')) -> Map k (Fold m a' b) -> Fold m a ()
988988demuxWith_ f kv = Fold step initial extract
989989
990990 where
@@ -992,16 +992,17 @@ demuxWith_ f kv = Fold step initial extract
992992 initial = do
993993 Prelude.mapM (\(Fold s i e) ->
994994 i >>= \r -> return (Fold s (return r) e)) kv
995− step mp a =
995+ step mp a
996996 -- XXX should we raise an exception in Nothing case?
997997 -- Ideally we should enforce that it is a total map over k so that look
998998 -- up never fails
999− case Map.lookup (f a) mp of
999+ | (k, a') <- f a
1000+ = case Map.lookup k mp of
10001001 Nothing -> return mp
10011002 Just (Fold step' acc _) -> do
1002− _ <- acc >>= \x -> step' x a
1003+ _ <- acc >>= \x -> step' x a'
10031004 return mp
1004− extract mp = Prelude.mapM (\(Fold _ acc e) -> acc >>= e) mp >> return ()
1005+ extract mp = Prelude.mapM_ (\(Fold _ acc e) -> acc >>= e) mp
10051006
10061007-- | Given a stream of key value pairs and a map from keys to folds, fold the
10071008-- values for each key using the corresponding folds, discarding the outputs.
@@ -1018,7 +1019,7 @@ demuxWith_ f kv = Fold step initial extract
10181019-- @since 0.7.0
10191020{-# INLINE demux_ #-}
10201021demux_ :: (Monad m, Ord k) => Map k (Fold m a ()) -> Fold m (k, a) ()
1021−demux_ fs = demuxWith_ fst (Map.map (lmap snd) fs)
1022+demux_ = demuxWith_ id
10221023
10231024-- TODO If the data is large we may need a map/hashmap in pinned memory instead
10241025-- of a regular Map. That may require a serializable constraint though. We can