1{-# LANGUAGE BangPatterns #-}
3{-# LANGUAGE ExistentialQuantification #-}
4{-# LANGUAGE FlexibleContexts #-}
5{-# LANGUAGE RankNTypes #-}
6{-# LANGUAGE RecordWildCards #-}
7{-# LANGUAGE ScopedTypeVariables #-}
10-- Module : Streamly.Internal.Data.Fold
11-- Copyright : (c) 2019 Composewell Technologies
12-- (c) 2013 Gabriel Gonzalez
14-- Maintainer : streamly@composewell.com
15-- Stability : experimental
18-- Also see the "Streamly.Internal.Data.Sink" module that provides specialized left folds
19-- that discard the outputs.
21-- IMPORTANT: keep the signatures consistent with the folds in Streamly.Prelude
23module Streamly.Internal.Data.Fold
34 -- * Fold Creation Utilities
61 -- ** Full Folds (Monoidal)
66 -- ** Full Folds (To Containers)
69 , toListRevF -- experimental
87 -- XXX these are slower than right folds even when full input is used
95 -- ** Covariant Operations
112 -- ** Mapping Filters
116 -- ** Scanning Filters
121 -- | Insertion adds more elements to the stream.
173 -- These can be expressed using lmap/lmapM and unzip
184 , duplicate -- experimental
188 , toParallelSVarLimited
192import Control.Monad (void)
193import Control.Monad.IO.Class (MonadIO(..))
194import Data.Functor.Identity (Identity(..))
195import Data.Map.Strict (Map)
198 hiding (filter, drop, dropWhile, take, takeWhile, zipWith, foldr,
199 foldl, map, mapM_, sequence, all, any, sum, product, elem,
200 notElem, maximum, minimum, head, last, tail, length, null,
201 reverse, iterate, init, and, or, lookup, foldr1, (!!),
202 scanl, scanl1, replicate, concatMap, mconcat, foldMap, unzip,
203 span, splitAt, break, mapM)
205import qualified Data.Map.Strict as Map
206import qualified Prelude
208import Streamly.Internal.Data.Pipe.Types (Pipe (..), PipeState(..))
209import Streamly.Internal.Data.Fold.Types
210import Streamly.Internal.Data.Strict
211import Streamly.Internal.Data.SVar
213import qualified Streamly.Internal.Data.Pipe.Types as Pipe
215------------------------------------------------------------------------------
217------------------------------------------------------------------------------
219-- | Make a fold using a pure step function, a pure initial state and
220-- a pure state extraction function.
225mkPure :: Monad m => (s -> a -> s) -> s -> (s -> b) -> Fold m a b
226mkPure step initial extract =
227 Fold (\s a -> return $ step s a) (return initial) (return . extract)
229-- | Make a fold using a pure step function and a pure initial state. The
230-- final state extracted is identical to the intermediate state.
234{-# INLINE mkPureId #-}
235mkPureId :: Monad m => (b -> a -> b) -> b -> Fold m a b
236mkPureId step initial = mkPure step initial id
238-- | Make a fold with an effectful step function and initial state, and a state
239-- extraction function.
243-- We can just use 'Fold' but it is provided for completeness.
248mkFold :: (s -> a -> m s) -> m s -> (s -> m b) -> Fold m a b
251-- | Make a fold with an effectful step function and initial state. The final
252-- state extracted is identical to the intermediate state.
256{-# INLINE mkFoldId #-}
257mkFoldId :: Monad m => (b -> a -> m b) -> m b -> Fold m a b
258mkFoldId step initial = Fold step initial return
260------------------------------------------------------------------------------
262------------------------------------------------------------------------------
264-- | Change the underlying monad of a fold
267hoist :: (forall x. m x -> n x) -> Fold m a b -> Fold n a b
268hoist f (Fold step initial extract) =
269 Fold (\x a -> f $ step x a) (f initial) (f . extract)
271-- | Adapt a pure fold to any monad
273-- > generally = hoist (return . runIdentity)
276generally :: Monad m => Fold Identity a b -> Fold m a b
277generally = hoist (return . runIdentity)
279------------------------------------------------------------------------------
280-- Transformations on fold inputs
281------------------------------------------------------------------------------
283-- | Flatten the monadic output of a fold to pure output.
286{-# INLINE sequence #-}
287sequence :: Monad m => Fold m a (m b) -> Fold m a b
288sequence (Fold step initial extract) = Fold step initial extract'
294-- | Map a monadic function on the output of a fold.
298mapM :: Monad m => (b -> m c) -> Fold m a b -> Fold m a c
299mapM f = sequence . fmap f
301------------------------------------------------------------------------------
302-- Transformations on fold inputs
303------------------------------------------------------------------------------
307-- | Apply a transformation on a 'Fold' using a 'Pipe'.
310{-# INLINE transform #-}
311transform :: Monad m => Pipe m a b -> Fold m b c -> Fold m a c
312transform (Pipe pstep1 pstep2 pinitial) (Fold fstep finitial fextract) =
313 Fold step initial extract
317 initial = Tuple' <$> return pinitial <*> finitial
318 step (Tuple' ps fs) x = do
324 go acc (Pipe.Yield b (Consume ps')) = do
326 return (Tuple' ps' acc')
328 go acc (Pipe.Yield b (Produce ps')) = do
333 go acc (Pipe.Continue (Consume ps')) = return (Tuple' ps' acc)
335 go acc (Pipe.Continue (Produce ps')) = do
339 extract (Tuple' _ fs) = fextract fs
341------------------------------------------------------------------------------
343------------------------------------------------------------------------------
345-- | @_Fold1 step@ returns a new 'Fold' using just a step function that has the
346-- same type for the accumulator and the element. The result type is the
347-- accumulator type wrapped in 'Maybe'. The initial accumulator is retrieved
348-- from the 'Foldable', the result is 'None' for empty containers.
349{-# INLINABLE _Fold1 #-}
350_Fold1 :: Monad m => (a -> a -> a) -> Fold m a (Maybe a)
351_Fold1 step = Fold step_ (return Nothing') fromStrictMaybe
353 step_ mx a = return $ Just' $
358------------------------------------------------------------------------------
360------------------------------------------------------------------------------
362------------------------------------------------------------------------------
364------------------------------------------------------------------------------
366-- | A fold that drains all its input, running the effects and discarding the
370{-# INLINABLE drain #-}
371drain :: Monad m => Fold m a ()
372drain = Fold step begin done
379-- > drainBy f = lmapM f drain
381-- Drain all input after passing it through a monadic function. This is the
382-- dual of mapM_ on stream producers.
385{-# INLINABLE drainBy #-}
386drainBy :: Monad m => (a -> m b) -> Fold m a ()
387drainBy f = Fold (const (void . f)) (return ()) return
389{-# INLINABLE drainBy2 #-}
390drainBy2 :: Monad m => (a -> m b) -> Fold2 m c a ()
391drainBy2 f = Fold2 (const (void . f)) (\_ -> return ()) return
393-- | Extract the last element of the input stream, if any.
396{-# INLINABLE last #-}
397last :: Monad m => Fold m a (Maybe a)
398last = _Fold1 (flip const)
400------------------------------------------------------------------------------
402------------------------------------------------------------------------------
404-- | Like 'length', except with a more general 'Num' return value
407{-# INLINABLE genericLength #-}
408genericLength :: (Monad m, Num b) => Fold m a b
409genericLength = Fold (\n _ -> return $ n + 1) (return 0) return
411-- | Determine the length of the input stream.
414{-# INLINABLE length #-}
415length :: Monad m => Fold m a Int
416length = genericLength
418-- | Determine the sum of all elements of a stream of numbers. Returns additive
419-- identity (@0@) when the stream is empty. Note that this is not numerically
420-- stable for floating point numbers.
424sum :: (Monad m, Num a) => Fold m a a
425sum = Fold (\x a -> return $ x + a) (return 0) return
427-- | Determine the product of all elements of a stream of numbers. Returns
428-- multiplicative identity (@1@) when the stream is empty.
431{-# INLINABLE product #-}
432product :: (Monad m, Num a) => Fold m a a
433product = Fold (\x a -> return $ x * a) (return 1) return
435------------------------------------------------------------------------------
437------------------------------------------------------------------------------
439-- | Determine the maximum element in a stream using the supplied comparison
443{-# INLINABLE maximumBy #-}
444maximumBy :: Monad m => (a -> a -> Ordering) -> Fold m a (Maybe a)
445maximumBy cmp = _Fold1 max'
447 max' x y = case cmp x y of
453-- maximum = 'maximumBy' compare
456-- Determine the maximum element in a stream.
459{-# INLINABLE maximum #-}
460maximum :: (Monad m, Ord a) => Fold m a (Maybe a)
463-- | Computes the minimum element with respect to the given comparison function
466{-# INLINABLE minimumBy #-}
467minimumBy :: Monad m => (a -> a -> Ordering) -> Fold m a (Maybe a)
468minimumBy cmp = _Fold1 min'
470 min' x y = case cmp x y of
474-- | Determine the minimum element in a stream using the supplied comparison
478{-# INLINABLE minimum #-}
479minimum :: (Monad m, Ord a) => Fold m a (Maybe a)
482------------------------------------------------------------------------------
483-- To Summary (Statistical)
484------------------------------------------------------------------------------
486-- | Compute a numerically stable arithmetic mean of all elements in the input
490{-# INLINABLE mean #-}
491mean :: (Monad m, Fractional a) => Fold m a a
492mean = Fold step (return begin) (return . done)
495 step (Tuple' x n) y = return $
497 in Tuple' (x + (y - x) / n') n'
498 done (Tuple' x _) = x
500-- | Compute a numerically stable (population) variance over all elements in
504{-# INLINABLE variance #-}
505variance :: (Monad m, Fractional a) => Fold m a a
506variance = Fold step (return begin) (return . done)
508 begin = Tuple3' 0 0 0
510 step (Tuple3' n mean_ m2) x = return $ Tuple3' n' mean' m2'
513 mean' = (n * mean_ + x) / (n + 1)
515 m2' = m2 + delta * delta * n / (n + 1)
517 done (Tuple3' n _ m2) = m2 / n
519-- | Compute a numerically stable (population) standard deviation over all
520-- elements in the input stream.
523{-# INLINABLE stdDev #-}
524stdDev :: (Monad m, Floating a) => Fold m a a
525stdDev = sqrt variance
527-- | Compute an 'Int' sized polynomial rolling hash
529-- > H = salt * k ^ n + c1 * k ^ (n - 1) + c2 * k ^ (n - 2) + ... + cn * k ^ 0
531-- Where @c1@, @c2@, @cn@ are the elements in the input stream and @k@ is a
534-- This hash is often used in Rabin-Karp string search algorithm.
536-- See https://en.wikipedia.org/wiki/Rolling_hash
539{-# INLINABLE rollingHashWithSalt #-}
540rollingHashWithSalt :: (Monad m, Enum a) => Int -> Fold m a Int
541rollingHashWithSalt salt = Fold step initial extract
544 initial = return salt
545 step cksum a = return $ cksum * k + fromEnum a
548-- | A default salt used in the implementation of 'rollingHash'.
549{-# INLINE defaultSalt #-}
551#if WORD_SIZE_IN_BITS == 64
552defaultSalt = 0xdc36d1615b7400a4
554defaultSalt = 0x087fc72c
557-- | Compute an 'Int' sized polynomial rolling hash of a stream.
559-- > rollingHash = rollingHashWithSalt defaultSalt
562{-# INLINABLE rollingHash #-}
563rollingHash :: (Monad m, Enum a) => Fold m a Int
564rollingHash = rollingHashWithSalt defaultSalt
566-- | Compute an 'Int' sized polynomial rolling hash of the first n elements of
569-- > rollingHashFirstN = ltake n rollingHash
570{-# INLINABLE rollingHashFirstN #-}
571rollingHashFirstN :: (Monad m, Enum a) => Int -> Fold m a Int
572rollingHashFirstN n = ltake n rollingHash
574------------------------------------------------------------------------------
575-- Monoidal left folds
576------------------------------------------------------------------------------
578-- | Fold an input stream consisting of monoidal elements using 'mappend'
581-- > S.fold FL.mconcat (S.map Sum $ S.enumerateFromTo 1 10)
584{-# INLINABLE mconcat #-}
585mconcat :: (Monad m, Monoid a) => Fold m a a
586mconcat = Fold (\x a -> return $ mappend x a) (return mempty) return
589-- > foldMap f = map f mconcat
591-- Make a fold from a pure function that folds the output of the function
592-- using 'mappend' and 'mempty'.
594-- > S.fold (FL.foldMap Sum) $ S.enumerateFromTo 1 10
597{-# INLINABLE foldMap #-}
598foldMap :: (Monad m, Monoid b) => (a -> b) -> Fold m a b
599foldMap f = lmap f mconcat
602-- > foldMapM f = mapM f mconcat
604-- Make a fold from a monadic function that folds the output of the function
605-- using 'mappend' and 'mempty'.
607-- > S.fold (FL.foldMapM (return . Sum)) $ S.enumerateFromTo 1 10
610{-# INLINABLE foldMapM #-}
611foldMapM :: (Monad m, Monoid b) => (a -> m b) -> Fold m a b
612foldMapM act = Fold step begin done
615 begin = return mempty
618 return $! mappend m m'
620------------------------------------------------------------------------------
622------------------------------------------------------------------------------
624-- | Folds the input stream to a list.
626-- /Warning!/ working on large lists accumulated as buffers in memory could be
627-- very inefficient, consider using "Streamly.Memory.Array" instead.
631-- id . (x1 :) . (x2 :) . (x3 :) . ... . (xn :) $ []
632{-# INLINABLE toList #-}
633toList :: Monad m => Fold m a [a]
634toList = Fold (\f x -> return $ f . (x :))
638------------------------------------------------------------------------------
640------------------------------------------------------------------------------
642-- | A fold that drains the first n elements of its input, running the effects
643-- and discarding the results.
644{-# INLINABLE drainN #-}
645drainN :: Monad m => Int -> Fold m a ()
646drainN n = ltake n drain
648-- | A fold that drains elements of its input as long as the predicate succeeds,
649-- running the effects and discarding the results.
650{-# INLINABLE drainWhile #-}
651drainWhile :: Monad m => (a -> Bool) -> Fold m a ()
652drainWhile p = ltakeWhile p drain
654------------------------------------------------------------------------------
656------------------------------------------------------------------------------
658-- | Like 'index', except with a more general 'Integral' argument
661{-# INLINABLE genericIndex #-}
662genericIndex :: (Integral i, Monad m) => i -> Fold m a (Maybe a)
663genericIndex i = Fold step (return $ Left' 0) done
676-- | Lookup the element at the given index.
679{-# INLINABLE index #-}
680index :: Monad m => Int -> Fold m a (Maybe a)
683-- | Extract the first element of the stream, if any.
686{-# INLINABLE head #-}
687head :: Monad m => Fold m a (Maybe a)
690-- | Returns the first element that satisfies the given predicate.
693{-# INLINABLE find #-}
694find :: Monad m => (a -> Bool) -> Fold m a (Maybe a)
695find predicate = Fold step (return Nothing') fromStrictMaybe
699 Nothing' -> if predicate a
704-- | In a stream of (key-value) pairs @(a, b)@, return the value @b@ of the
705-- first pair where the key equals the given value @a@.
708{-# INLINABLE lookup #-}
709lookup :: (Eq a, Monad m) => a -> Fold m (a,b) (Maybe b)
710lookup a0 = Fold step (return Nothing') fromStrictMaybe
712 step x (a,b) = return $
714 Nothing' -> if a == a0
719-- | Convert strict 'Either'' to lazy 'Maybe'
720{-# INLINABLE hush #-}
721hush :: Either' a b -> Maybe b
722hush (Left' _) = Nothing
723hush (Right' b) = Just b
725-- | Returns the first index that satisfies the given predicate.
728{-# INLINABLE findIndex #-}
729findIndex :: Monad m => (a -> Bool) -> Fold m a (Maybe Int)
730findIndex predicate = Fold step (return $ Left' 0) (return . hush)
740-- | Returns the first index where a given value is found in the stream.
743{-# INLINABLE elemIndex #-}
744elemIndex :: (Eq a, Monad m) => a -> Fold m a (Maybe Int)
745elemIndex a = findIndex (a ==)
747------------------------------------------------------------------------------
749------------------------------------------------------------------------------
751-- | Return 'True' if the input stream is empty.
754{-# INLINABLE null #-}
755null :: Monad m => Fold m a Bool
756null = Fold (\_ _ -> return False) (return True) return
759-- > any p = lmap p or
761-- | Returns 'True' if any of the elements of a stream satisfies a predicate.
765any :: Monad m => (a -> Bool) -> Fold m a Bool
766any predicate = Fold (\x a -> return $ x || predicate a) (return False) return
768-- | Return 'True' if the given element is present in the stream.
771{-# INLINABLE elem #-}
772elem :: (Eq a, Monad m) => a -> Fold m a Bool
776-- > all p = lmap p and
778-- | Returns 'True' if all elements of a stream satisfy a predicate.
782all :: Monad m => (a -> Bool) -> Fold m a Bool
783all predicate = Fold (\x a -> return $ x && predicate a) (return True) return
785-- | Returns 'True' if the given element is not present in the stream.
788{-# INLINABLE notElem #-}
789notElem :: (Eq a, Monad m) => a -> Fold m a Bool
790notElem a = all (a /=)
792-- | Returns 'True' if all elements are 'True', 'False' otherwise
796and :: Monad m => Fold m Bool Bool
797and = Fold (\x a -> return $ x && a) (return True) return
799-- | Returns 'True' if any element is 'True', 'False' otherwise
803or :: Monad m => Fold m Bool Bool
804or = Fold (\x a -> return $ x || a) (return False) return
806------------------------------------------------------------------------------
808------------------------------------------------------------------------------
810-- | Distribute one copy of the stream to each fold and zip the results.
813-- |-------Fold m a b--------|
814-- ---stream m a---| |---m (b,c)
815-- |-------Fold m a c--------|
817-- >>> S.fold (FL.tee FL.sum FL.length) (S.enumerateFromTo 1.0 100.0)
822tee :: Monad m => Fold m a b -> Fold m a c -> Fold m a (b,c)
823tee f1 f2 = (,) <$> f1 <*> f2
825{-# INLINE foldNil #-}
826foldNil :: Monad m => Fold m a [b]
827foldNil = Fold step begin done where
832{-# INLINE foldCons #-}
833foldCons :: Monad m => Fold m a b -> Fold m a [b] -> Fold m a [b]
834foldCons (Fold stepL beginL doneL) (Fold stepR beginR doneR) =
839 begin = Tuple' <$> beginL <*> beginR
840 step (Tuple' xL xR) a = Tuple' <$> stepL xL a <*> stepR xR a
841 done (Tuple' xL xR) = (:) <$> (doneL xL) <*> (doneR xR)
843-- XXX use "List" instead of "[]"?, use Array for output to scale it to a large
844-- number of consumers? For polymorphic case a vector could be helpful. For
845-- Storables we can use arrays. Will need separate APIs for those.
847-- | Distribute one copy of the stream to each fold and collect the results in
852-- |-------Fold m a b--------|
853-- ---stream m a---| |---m [b]
854-- |-------Fold m a b--------|
859-- >>> S.fold (FL.distribute [FL.sum, FL.length]) (S.enumerateFromTo 1 5)
862-- This is the consumer side dual of the producer side 'sequence' operation.
865{-# INLINE distribute #-}
866distribute :: Monad m => [Fold m a b] -> Fold m a [b]
867distribute [] = foldNil
868distribute (x:xs) = foldCons x (distribute xs)
870-- | Like 'distribute' but for folds that return (), this can be more efficient
871-- than 'distribute' as it does not need to maintain state.
873{-# INLINE distribute_ #-}
874distribute_ :: Monad m => [Fold m a ()] -> Fold m a ()
875distribute_ fs = Fold step initial extract
877 initial = Prelude.mapM (\(Fold s i e) ->
878 i >>= \r -> return (Fold s (return r) e)) fs
880 Prelude.mapM_ (\(Fold s i _) -> i >>= \r -> s r a >> return ()) ss
883 Prelude.mapM_ (\(Fold _ i e) -> i >>= \r -> e r) ss
886------------------------------------------------------------------------------
888------------------------------------------------------------------------------
890-- | Partition the input over two folds using an 'Either' partitioning
895-- |-------Fold b x--------|
896-- -----stream m a --> (Either b c)----| |----(x,y)
897-- |-------Fold c y--------|
900-- Send input to either fold randomly:
902-- >>> import System.Random (randomIO)
903-- >>> randomly a = randomIO >>= \x -> return $ if x then Left a else Right a
904-- >>> S.fold (FL.partitionByM randomly FL.length FL.length) (S.enumerateFromTo 1 100)
907-- Send input to the two folds in a proportion of 2:1:
910-- import Data.IORef (newIORef, readIORef, writeIORef)
911-- proportionately m n = do
912-- ref <- newIORef $ cycle $ concat [replicate m Left, replicate n Right]
915-- writeIORef ref $ tail r
919-- f <- proportionately 2 1
920-- r <- S.fold (FL.partitionByM f FL.length FL.length) (S.enumerateFromTo (1 :: Int) 100)
927-- This is the consumer side dual of the producer side 'mergeBy' operation.
930{-# INLINE partitionByM #-}
931partitionByM :: Monad m
932 => (a -> m (Either b c)) -> Fold m b x -> Fold m c y -> Fold m a (x, y)
933partitionByM f (Fold stepL beginL doneL) (Fold stepR beginR doneR) =
939 begin = Tuple' <$> beginL <*> beginR
940 step (Tuple' xL xR) a = do
943 Left b -> Tuple' <$> stepL xL b <*> return xR
944 Right c -> Tuple' <$> return xL <*> stepR xR c
945 done (Tuple' xL xR) = (,) <$> doneL xL <*> doneR xR
947-- Note: we could use (a -> Bool) instead of (a -> Either b c), but the latter
948-- makes the signature clearer as to which case belongs to which fold.
949-- XXX need to check the performance in both cases.
951-- | Same as 'partitionByM' but with a pure partition function.
953-- Count even and odd numbers in a stream:
956-- >>> let f = FL.partitionBy (\\n -> if even n then Left n else Right n)
957-- (fmap (("Even " ++) . show) FL.length)
958-- (fmap (("Odd " ++) . show) FL.length)
959-- in S.fold f (S.enumerateFromTo 1 100)
960-- ("Even 50","Odd 50")
964{-# INLINE partitionBy #-}
965partitionBy :: Monad m
966 => (a -> Either b c) -> Fold m b x -> Fold m c y -> Fold m a (x, y)
967partitionBy f = partitionByM (return . f)
969-- | Compose two folds such that the combined fold accepts a stream of 'Either'
970-- and routes the 'Left' values to the first fold and 'Right' values to the
973-- > partition = partitionBy id
976{-# INLINE partition #-}
978 => Fold m b x -> Fold m c y -> Fold m (Either b c) (x, y)
979partition = partitionBy id
982-- | Send one item to each fold in a round-robin fashion. This is the consumer
983-- side dual of producer side 'mergeN' operation.
985-- partitionN :: Monad m => [Fold m a b] -> Fold m a [b]
986-- partitionN fs = Fold step begin done
989-- TODO Demultiplex an input element into a number of typed variants. We want
990-- to statically restrict the target values within a set of predefined types,
991-- an enumeration of a GADT. We also want to make sure that the Map contains
992-- only those types and the full set of those types.
994-- TODO Instead of the input Map it should probably be a lookup-table using an
995-- array and not in GC memory. The same applies to the output Map as well.
996-- However, that would only be helpful if we have a very large data structure,
997-- need to measure and see how it scales.
999-- This is the consumer side dual of the producer side 'mux' operation (XXX to
1002-- | Split the input stream based on a key field and fold each split using a
1003-- specific fold collecting the results in a map from the keys to the results.
1004-- Useful for cases like protocol handlers to handle different type of packets
1005-- using different handlers.
1009-- |-------Fold m a b
1010-- -----stream m a-----Map-----|
1011-- |-------Fold m a b
1017{-# INLINE demuxWith #-}
1018demuxWith :: (Monad m, Ord k)
1019 => (a -> (k, a')) -> Map k (Fold m a' b) -> Fold m a (Map k b)
1020demuxWith f kv = Fold step initial extract
1025-- alterF is available only since containers version 0.5.8.2
1026#if MIN_VERSION_containers(0,5,8)
1027 step mp a = case f a of
1028 (k, a') -> Map.alterF twiddle k mp
1029 -- XXX should we raise an exception in Nothing case?
1030 -- Ideally we should enforce that it is a total map over k so that look
1032 -- XXX we could use a monadic update function for a single lookup and
1033 -- update in the map.
1035 twiddle Nothing = pure Nothing
1036 twiddle (Just (Fold step' acc extract')) = do
1037 !r <- acc >>= \x -> step' x a'
1038 pure . Just $ Fold step' (return r) extract'
1042 in case Map.lookup k mp of
1043 Nothing -> return mp
1044 Just (Fold step' acc extract') -> do
1045 !r <- acc >>= \x -> step' x a'
1046 return $ Map.insert k (Fold step' (return r) extract') mp
1048 extract = Prelude.mapM (\(Fold _ acc e) -> acc >>= e)
1050-- | Fold a stream of key value pairs using a map of specific folds for each
1051-- key into a map from keys to the results of fold outputs of the corresponding
1055-- > let table = Data.Map.fromList [(\"SUM", FL.sum), (\"PRODUCT", FL.product)]
1056-- input = S.fromList [(\"SUM",1),(\"PRODUCT",2),(\"SUM",3),(\"PRODUCT",4)]
1057-- in S.fold (FL.demux table) input
1064demux :: (Monad m, Ord k)
1065 => Map k (Fold m a b) -> Fold m (k, a) (Map k b)
1068{-# INLINE demuxWithDefault_ #-}
1069demuxWithDefault_ :: (Monad m, Ord k)
1070 => (a -> (k, a')) -> Map k (Fold m a' b) -> Fold m (k, a') b -> Fold m a ()
1071demuxWithDefault_ f kv (Fold dstep dinitial dextract) =
1072 Fold step initial extract
1076 initFold (Fold s i e) = i >>= \r -> return (Fold s (return r) e)
1078 mp <- Prelude.mapM initFold kv
1080 return (Tuple' mp dacc)
1081 step (Tuple' mp dacc) a
1083 = case Map.lookup k mp of
1085 acc <- dstep dacc (k, a')
1086 return (Tuple' mp acc)
1087 Just (Fold step' acc _) -> do
1088 _ <- acc >>= \x -> step' x a'
1089 return (Tuple' mp dacc)
1090 extract (Tuple' mp dacc) = do
1091 void $ dextract dacc
1092 Prelude.mapM_ (\(Fold _ acc e) -> acc >>= e) mp
1094-- | Split the input stream based on a key field and fold each split using a
1095-- specific fold without collecting the results. Useful for cases like protocol
1096-- handlers to handle different type of packets.
1100-- |-------Fold m a ()
1101-- -----stream m a-----Map-----|
1102-- |-------Fold m a ()
1110-- demuxWith_ can be slightly faster than demuxWith because we do not need to
1111-- update the Map in this case. This may be significant only if the map is
1113{-# INLINE demuxWith_ #-}
1114demuxWith_ :: (Monad m, Ord k)
1115 => (a -> (k, a')) -> Map k (Fold m a' b) -> Fold m a ()
1116demuxWith_ f kv = Fold step initial extract
1121 Prelude.mapM (\(Fold s i e) ->
1122 i >>= \r -> return (Fold s (return r) e)) kv
1124 -- XXX should we raise an exception in Nothing case?
1125 -- Ideally we should enforce that it is a total map over k so that look
1128 = case Map.lookup k mp of
1129 Nothing -> return mp
1130 Just (Fold step' acc _) -> do
1131 _ <- acc >>= \x -> step' x a'
1133 extract mp = Prelude.mapM_ (\(Fold _ acc e) -> acc >>= e) mp
1135-- | Given a stream of key value pairs and a map from keys to folds, fold the
1136-- values for each key using the corresponding folds, discarding the outputs.
1139-- > let prn = FL.drainBy print
1140-- > let table = Data.Map.fromList [(\"ONE", prn), (\"TWO", prn)]
1141-- input = S.fromList [(\"ONE",1),(\"TWO",2)]
1142-- in S.fold (FL.demux_ table) input
1148{-# INLINE demux_ #-}
1149demux_ :: (Monad m, Ord k) => Map k (Fold m a ()) -> Fold m (k, a) ()
1150demux_ = demuxWith_ id
1152{-# INLINE demuxDefault_ #-}
1153demuxDefault_ :: (Monad m, Ord k)
1154 => Map k (Fold m a ()) -> Fold m (k, a) () -> Fold m (k, a) ()
1155demuxDefault_ = demuxWithDefault_ id
1157-- TODO If the data is large we may need a map/hashmap in pinned memory instead
1158-- of a regular Map. That may require a serializable constraint though. We can
1159-- have another API for that.
1161-- | Split the input stream based on a key field and fold each split using the
1162-- given fold. Useful for map/reduce, bucketizing the input in different bins
1163-- or for generating histograms.
1166-- > let input = S.fromList [(\"ONE",1),(\"ONE",1.1),(\"TWO",2), (\"TWO",2.2)]
1167-- in S.fold (FL.classify FL.toList) input
1168-- fromList [(\"ONE",[1.1,1.0]),(\"TWO",[2.2,2.0])]
1172{-# INLINE classifyWith #-}
1173classifyWith :: (Monad m, Ord k) => (a -> k) -> Fold m a b -> Fold m a (Map k b)
1174classifyWith f (Fold step initial extract) = Fold step' initial' extract'
1178 initial' = return Map.empty
1181 in case Map.lookup k kv of
1185 return $ Map.insert k r kv
1188 return $ Map.insert k r kv
1189 extract' = Prelude.mapM extract
1191-- | Given an input stream of key value pairs and a fold for values, fold all
1192-- the values belonging to each key. Useful for map/reduce, bucketizing the
1193-- input in different bins or for generating histograms.
1196-- > let input = S.fromList [(\"ONE",1),(\"ONE",1.1),(\"TWO",2), (\"TWO",2.2)]
1197-- in S.fold (FL.classify FL.toList) input
1198-- fromList [(\"ONE",[1.1,1.0]),(\"TWO",[2.2,2.0])]
1205-- > classify fld = classifyWith fst (lmap snd fld)
1207{-# INLINE classify #-}
1208classify :: (Monad m, Ord k) => Fold m a b -> Fold m (k, a) (Map k b)
1209classify fld = classifyWith fst (lmap snd fld)
1211------------------------------------------------------------------------------
1213------------------------------------------------------------------------------
1215-- | Like 'unzipWith' but with a monadic splitter function.
1218{-# INLINE unzipWithM #-}
1219unzipWithM :: Monad m
1220 => (a -> m (b,c)) -> Fold m b x -> Fold m c y -> Fold m a (x,y)
1221unzipWithM f (Fold stepL beginL doneL) (Fold stepR beginR doneR) =
1222 Fold step begin done
1226 step (Tuple' xL xR) a = do
1228 Tuple' <$> stepL xL b <*> stepR xR c
1229 begin = Tuple' <$> beginL <*> beginR
1230 done (Tuple' xL xR) = (,) <$> doneL xL <*> doneR xR
1232-- | Split elements in the input stream into two parts using a pure splitter
1233-- function, direct each part to a different fold and zip the results.
1236{-# INLINE unzipWith #-}
1238 => (a -> (b,c)) -> Fold m b x -> Fold m c y -> Fold m a (x,y)
1239unzipWith f = unzipWithM (return . f)
1241-- | Send the elements of tuples in a stream of tuples through two different
1246-- |-------Fold m a x--------|
1247-- ---------stream of (a,b)--| |----m (x,y)
1248-- |-------Fold m b y--------|
1252-- This is the consumer side dual of the producer side 'zip' operation.
1256unzip :: Monad m => Fold m a x -> Fold m b y -> Fold m (a,b) (x,y)
1259------------------------------------------------------------------------------
1261------------------------------------------------------------------------------
1264-- All the stream flattening transformations can also be applied to a fold
1267-- | This can be used to apply all the stream generation operations on folds.
1268lconcatMap ::(IsStream t, Monad m) => (a -> t m b)
1271lconcatMap s f1 f2 = undefined
1274-- All the grouping transformation that we apply to a stream can also be
1275-- applied to a fold input stream.
1278-- | Group the input stream into groups of elements between @low@ and @high@.
1279-- Collection starts in chunks of @low@ and then keeps doubling until we reach
1280-- @high@. Each chunk is folded using the provided fold function.
1282-- This could be useful, for example, when we are folding a stream of unknown
1283-- size to a stream of arrays and we want to minimize the number of
1288-- XXX we should be able to implement it with parsers/terminating folds.
1290{-# INLINE lchunksInRange #-}
1291lchunksInRange :: Monad m
1292 => Int -> Int -> Fold m a b -> Fold m b c -> Fold m a c
1293lchunksInRange low high (Fold step1 initial1 extract1)
1294 (Fold step2 initial2 extract2) = undefined
1297------------------------------------------------------------------------------
1298-- Fold to a Parallel SVar
1299------------------------------------------------------------------------------
1301{-# INLINE toParallelSVar #-}
1302toParallelSVar :: MonadIO m => SVar t m a -> Maybe WorkerInfo -> Fold m a ()
1303toParallelSVar svar winfo = Fold step initial extract
1308 step () x = liftIO $ do
1309 -- XXX we can have a separate fold for unlimited buffer case to avoid a
1310 -- branch in the step here.
1311 decrementBufferLimit svar
1312 void $ send svar (ChildYield x)
1314 extract () = liftIO $ do
1317{-# INLINE toParallelSVarLimited #-}
1318toParallelSVarLimited :: MonadIO m
1319 => SVar t m a -> Maybe WorkerInfo -> Fold m a ()
1320toParallelSVarLimited svar winfo = Fold step initial extract
1323 initial = return True
1325 step True x = liftIO $ do
1326 yieldLimitOk <- decrementYieldLimit svar
1329 decrementBufferLimit svar
1330 void $ send svar (ChildYield x)
1333 cleanupSVarFromWorker svar
1336 step False _ = return False
1338 extract True = liftIO $ sendStop svar winfo
1339 extract False = return ()
keysesc to close
- j k
- move down, up
- enter
- open
- esc
- back, or up a level
- 1 2 3 4
- readme, files, commits, refs
- /
- filter
- y
- copy the clone command
- c
- compare the chosen commit with another (or shift-click both)
- ?
- show or hide this list