@@ -925,24 +925,24 @@ partition = partitionBy id
925-- @since 0.7.0
925-- @since 0.7.0
926{-# INLINE demuxWith #-}
926{-# INLINE demuxWith #-}
927demuxWith :: (Monad m, Ord k)
927demuxWith :: (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)
929demuxWith f kv = Fold step initial extract
929demuxWith f kv = Fold step initial extract
933 initial = return kv
933 initial = return kv
934− step mp a =
934+ step mp a = case f a of
935+ (k, a') -> Map.alterF twiddle k mp
935 -- XXX should we raise an exception in Nothing case?
936 -- XXX should we raise an exception in Nothing case?
936 -- Ideally we should enforce that it is a total map over k so that look
937 -- Ideally we should enforce that it is a total map over k so that look
937 -- up never fails
938 -- up never fails
938 -- XXX we could use a monadic update function for a single lookup and
939 -- XXX we could use a monadic update function for a single lookup and
939 -- update in the map.
940 -- update in the map.
940− let k = f a
941+ where
941− in case Map.lookup k mp of
942+ twiddle Nothing = pure Nothing
942− Nothing -> return mp
943− Just (Fold step' acc extract') -> do
943+ twiddle (Just (Fold step' acc extract')) = do
944− !r <- acc >>= \x -> step' x a
944+ !r <- acc >>= \x -> step' x a'
945− return $ Map.insert k (Fold step' (return r) extract') mp
945+ pure . Just $ Fold step' (return r) extract'
946 extract = Prelude.mapM (\(Fold _ acc e) -> acc >>= e)
946 extract = Prelude.mapM (\(Fold _ acc e) -> acc >>= e)
948-- | Fold a stream of key value pairs using a map of specific folds for each
948-- | 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
961{-# INLINE demux #-}
961{-# INLINE demux #-}
962demux :: (Monad m, Ord k)
962demux :: (Monad m, Ord k)
963 => Map k (Fold m a b) -> Fold m (k, a) (Map k b)
963 => 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
966-- | Split the input stream based on a key field and fold each split using a
966-- | Split the input stream based on a key field and fold each split using a
967-- specific fold without collecting the results. Useful for cases like protocol
967-- specific fold without collecting the results. Useful for cases like protocol
@@ -984,7 +984,7 @@ demux fs = demuxWith fst (Map.map (lmap snd) fs)
985{-# INLINE demuxWith_ #-}
985{-# INLINE demuxWith_ #-}
986demuxWith_ :: (Monad m, Ord k)
986demuxWith_ :: (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 ()
988demuxWith_ f kv = Fold step initial extract
988demuxWith_ f kv = Fold step initial extract
@@ -992,16 +992,17 @@ demuxWith_ f kv = Fold step initial extract
992 initial = do
992 initial = do
993 Prelude.mapM (\(Fold s i e) ->
993 Prelude.mapM (\(Fold s i e) ->
994 i >>= \r -> return (Fold s (return r) e)) kv
994 i >>= \r -> return (Fold s (return r) e)) kv
995− step mp a =
995+ step mp a
996 -- XXX should we raise an exception in Nothing case?
996 -- XXX should we raise an exception in Nothing case?
997 -- Ideally we should enforce that it is a total map over k so that look
997 -- Ideally we should enforce that it is a total map over k so that look
998 -- up never fails
998 -- up never fails
999− case Map.lookup (f a) mp of
1000+ = case Map.lookup k mp of
1000 Nothing -> return mp
1001 Nothing -> return mp
1001 Just (Fold step' acc _) -> do
1002 Just (Fold step' acc _) -> do
1002− _ <- acc >>= \x -> step' x a
1003+ _ <- acc >>= \x -> step' x a'
1003 return mp
1004 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
1006-- | Given a stream of key value pairs and a map from keys to folds, fold the
1007-- | Given a stream of key value pairs and a map from keys to folds, fold the
1007-- values for each key using the corresponding folds, discarding the outputs.
1008-- values for each key using the corresponding folds, discarding the outputs.
@@ -1018,7 +1019,7 @@ demuxWith_ f kv = Fold step initial extract
1018-- @since 0.7.0
1019-- @since 0.7.0
1019{-# INLINE demux_ #-}
1020{-# INLINE demux_ #-}
1020demux_ :: (Monad m, Ord k) => Map k (Fold m a ()) -> Fold m (k, a) ()
1021demux_ :: (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
1023-- TODO If the data is large we may need a map/hashmap in pinned memory instead
1024-- TODO If the data is large we may need a map/hashmap in pinned memory instead
1024-- of a regular Map. That may require a serializable constraint though. We can
1025-- of a regular Map. That may require a serializable constraint though. We can