on 8dabc2e help sign in
1931 lines 56 kB history raw
1-- |
2-- Module : Streamly.Internal.Data.Fold
3-- Copyright : (c) 2019 Composewell Technologies
4-- (c) 2013 Gabriel Gonzalez
5-- License : BSD3
6-- Maintainer : streamly@composewell.com
7-- Stability : experimental
8-- Portability : GHC
9--
10-- See "Streamly.Data.Fold" for an overview and
11-- "Streamly.Internal.Data.Fold.Types" for design notes.
12--
13-- IMPORTANT: keep the signatures consistent with the folds in Streamly.Prelude
15module Streamly.Internal.Data.Fold
16 (
17 -- * Fold Type
18 Step (..)
19 , Fold (..)
21 -- * Constructors
22 -- | Which constructor to use?
23 --
24 -- * @foldl*@: If the fold never terminates i.e. does not use the 'Done'
25 -- constructor otherwise use the @mkFold@ variants.
26 -- * @*M@: Use the @M@ suffix variants if any of the step, initial, or
27 -- extract function is monadic, otherwise use the pure variants.
28 -- * @*_@: Use the @_@ suffix variants if the extract function is 'id' or
29 -- 'return'.
30 --
31 , foldl'
32 , foldlM'
33 , foldl1'
34 , foldr
35 , foldrM
36 , mkFold
37 , mkFold_
38 , mkFoldM
39 , mkFoldM_
41 -- * Folds
42 -- ** Identity
43 , fromPure
44 , fromEffect
46 -- ** Accumulators
47 -- *** Semigroups and Monoids
48 , sconcat
49 , mconcat
50 , foldMap
51 , foldMapM
53 -- *** Reducers
54 , drain
55 , drainBy
56 , last
57 , length
58 , genericLength
59 , countDistinct
60 , countDistinctInt
61 , mean
62 , variance
63 , stdDev
64 , rollingHash
65 , rollingHashWithSalt
66 , rollingHashFirstN
67 -- , rollingHashLastN
69 -- *** Saturating Reducers
70 -- | 'product' terminates if it becomes 0. Other folds can theoretically
71 -- saturate on bounded types, and therefore terminate, however, they will
72 -- run forever on unbounded types like Integer/Double.
73 , sum
74 , product
75 , maximumBy
76 , maximum
77 , minimumBy
78 , minimum
80 -- *** Collectors
81 -- | Avoid using these folds in scalable or performance critical
82 -- applications, they buffer all the input in GC memory which can be
83 -- detrimental to performance if the input is large.
84 , toList
85 , toListRev
86 -- $toListRev
87 , toStream
88 , toStreamRev
89 , toMap
91 -- ** Terminating Folds
92 , drainN
93 -- , lastN
94 -- , (!!)
95 , genericIndex
96 , index
97 , head
98 -- , findM
99 , find
100 , lookup
101 , findIndex
102 , elemIndex
103 , null
104 , elem
105 , notElem
106 , all
107 , any
108 , and
109 , or
110 -- , the
112 -- * Combinators
113 -- ** Utilities
114 , with
116 -- ** Transforming the Monad
117 , hoist
118 , generally
120 -- ** Mapping on output
121 , rmapM
123 -- ** Mapping on Input
124 , transform
125 , lmap
126 --, lsequence
127 , lmapM
128 , scan
129 , postscan
130 , indexed
132 -- ** Filtering
133 , filter
134 , filterM
135 , foldFilter
136 , satisfy
137 , sampleFromthen
138 -- , ldeleteBy
139 -- , luniq
141 -- ** Mapping Filters
142 , catMaybes
143 , mapMaybe
144 -- , mapMaybeM
146 -- ** Scanning Filters
147 , findIndices
148 {-
149 , elemIndices
151 -- ** Insertion
152 -- | Insertion adds more elements to the stream.
154 , insertBy
155 , intersperseM
157 -- ** Reordering
158 , reverse
159 -}
161 -- ** Trimming
162 , take
164 -- By elements
165 , takeEndBy
166 , takeEndBy_
167 -- , takeEndBySeq
168 {-
169 , drop
170 , dropWhile
171 , dropWhileM
172 -}
174 -- ** Serial Append
175 , serialWith
176 , serial_
177 -- , tail
178 -- , init
179 , splitAt -- spanN
180 -- , splitIn -- sessionN
182 -- ** Parallel Distribution
183 , teeWith
184 , tee
185 , teeWithFst
186 , teeWithMin
187 , distribute
188 -- , distributeFst
189 -- , distributeMin
191 -- ** Parallel Alternative
192 , shortest
193 , longest
195 -- ** Partitioning
196 , partitionByM
197 , partitionByFstM
198 , partitionByMinM
199 , partitionBy
200 , partition
202 -- ** Demultiplexing
203 -- | Direct values in the input stream to different folds using an n-ary
204 -- fold selector.
205 , demux -- XXX rename this to demux_
206 , demuxWith
207 , demuxDefault -- XXX rename this to demux
208 , demuxDefaultWith
209 -- , demuxWithSel
210 -- , demuxWithMin
212 -- ** Classifying
213 -- | In an input stream of key value pairs fold values for different keys
214 -- in individual output buckets using the given fold.
215 , classify
216 , classifyWith
217 , classifyScanWith
218 -- , classifyWithSel
219 -- , classifyWithMin
221 -- ** Unzipping
222 , unzip
223 -- These two can be expressed using lmap/lmapM and unzip
224 , unzipWith
225 , unzipWithM
226 , unzipWithFstM
227 , unzipWithMinM
229 -- ** Zipping
230 , zipWithM
231 , zip
233 -- ** Splitting
234 , many
235 , chunksOf
236 , chunksBetween
238 -- ** Nesting
239 , concatSequence
240 , concatMap
242 -- * Running A Fold
243 -- | Normally you would run a fold to completion by supplying it a stream,
244 -- e.g. using 'Stream.fold'. However, you could also run a fold partially
245 -- by using 'duplicate' on it and then running it with a stream.
246 -- Alternatively, 'initialize', 'snoc' and 'finish' can be used to run a
247 -- fold incrementally, however, that may not be the most efficient way to
248 -- run a fold.
249 , initialize
250 , snoc
251 , duplicate
252 , finish
254 -- * Deprecated
255 , sequence
256 , mapM
257 )
258where
260import Data.Bifunctor (first)
261import Data.Either (isLeft, isRight)
262import Data.Functor.Identity (Identity(..))
263import Data.Int (Int64)
264import Data.Map.Strict (Map)
265import Data.Maybe (isJust, fromJust)
266#if __GLASGOW_HASKELL__ < 804
267import Data.Semigroup (Semigroup((<>)))
268#endif
269import Streamly.Internal.BaseCompat (fromLeft, fromRight)
270import Streamly.Internal.Data.Either.Strict
271 (Either'(..), fromLeft', fromRight', isLeft', isRight')
272import Streamly.Internal.Data.Pipe.Type (Pipe (..), PipeState(..))
273import Streamly.Internal.Data.Tuple.Strict (Tuple'(..), Tuple3'(..))
274import Streamly.Internal.Data.Stream.Serial (SerialT(..))
276import qualified Data.IntSet as IntSet
277import qualified Data.Map.Strict as Map
278import qualified Data.Set as Set
279import qualified Streamly.Internal.Data.Pipe.Type as Pipe
280-- import qualified Streamly.Internal.Data.Stream.IsStream.Enumeration as Stream
281import qualified Prelude
283import Prelude hiding
284 ( filter, foldl1, drop, dropWhile, take, takeWhile, zipWith
285 , foldl, foldr, map, mapM_, sequence, all, any, sum, product, elem
286 , notElem, maximum, minimum, head, last, tail, length, null
287 , reverse, iterate, init, and, or, lookup, (!!)
288 , scanl, scanl1, replicate, concatMap, mconcat, foldMap, unzip
289 , span, splitAt, break, mapM, zip)
290import Streamly.Internal.Data.Fold.Type
292-- $setup
293-- >>> :m
294-- >>> :set -package streamly
295-- >>> import Prelude hiding (break, map, span, splitAt)
296-- >>> import qualified Streamly.Prelude as Stream
297-- >>> import qualified Streamly.Internal.Data.Stream.IsStream as Stream (parse, foldMany)
298-- >>> import qualified Streamly.Data.Fold as Fold
299-- >>> import qualified Streamly.Internal.Data.Fold as Fold
300-- >>> import qualified Streamly.Internal.Data.Fold.Type as Fold
301-- >>> import qualified Streamly.Internal.Data.Parser as Parser
302-- >>> import Streamly.Internal.Data.Stream.Serial (SerialT(..))
303-- >>> import Data.IORef (newIORef, readIORef, writeIORef)
305------------------------------------------------------------------------------
306-- hoist
307------------------------------------------------------------------------------
309-- | Change the underlying monad of a fold
310--
311-- /Pre-release/
312hoist :: (forall x. m x -> n x) -> Fold m a b -> Fold n a b
313hoist f (Fold step initial extract) =
314 Fold (\x a -> f $ step x a) (f initial) (f . extract)
316-- | Adapt a pure fold to any monad
317--
318-- > generally = Fold.hoist (return . runIdentity)
319--
320-- /Pre-release/
321generally :: Monad m => Fold Identity a b -> Fold m a b
322generally = hoist (return . runIdentity)
324------------------------------------------------------------------------------
325-- Transformations on fold inputs
326------------------------------------------------------------------------------
328-- | Flatten the monadic output of a fold to pure output.
329--
330-- @since 0.7.0
331{-# DEPRECATED sequence "Use \"rmapM id\" instead" #-}
332{-# INLINE sequence #-}
333sequence :: Monad m => Fold m a (m b) -> Fold m a b
334sequence = rmapM id
336-- | Map a monadic function on the output of a fold.
337--
338-- @since 0.7.0
339{-# DEPRECATED mapM "Use rmapM instead" #-}
340{-# INLINE mapM #-}
341mapM :: Monad m => (b -> m c) -> Fold m a b -> Fold m a c
342mapM = rmapM
344-- | @mapMaybe f fold@ maps a 'Maybe' returning function @f@ on the input of
345-- the fold, filters out 'Nothing' elements, and return the values extracted
346-- from 'Just'.
347--
348-- >>> f x = if even x then Just x else Nothing
349-- >>> fld = Fold.mapMaybe f Fold.toList
350-- >>> Stream.fold fld (Stream.enumerateFromTo 1 10)
351-- [2,4,6,8,10]
352--
353-- @since 0.8.0
354{-# INLINE mapMaybe #-}
355mapMaybe :: (Monad m) => (a -> Maybe b) -> Fold m b r -> Fold m a r
356mapMaybe f = lmap f . filter isJust . lmap fromJust
358------------------------------------------------------------------------------
359-- Transformations on fold inputs
360------------------------------------------------------------------------------
362-- rename to lpipe?
363--
364-- | Apply a transformation on a 'Fold' using a 'Pipe'.
365--
366-- /Pre-release/
367{-# INLINE transform #-}
368transform :: Monad m => Pipe m a b -> Fold m b c -> Fold m a c
369transform (Pipe pstep1 pstep2 pinitial) (Fold fstep finitial fextract) =
370 Fold step initial extract
372 where
374 initial = first (Tuple' pinitial) <$> finitial
376 step (Tuple' ps fs) x = do
377 r <- pstep1 ps x
378 go fs r
380 where
382 -- XXX use SPEC?
383 go acc (Pipe.Yield b (Consume ps')) = do
384 acc' <- fstep acc b
385 return
386 $ case acc' of
387 Partial s -> Partial $ Tuple' ps' s
388 Done b2 -> Done b2
389 go acc (Pipe.Yield b (Produce ps')) = do
390 acc' <- fstep acc b
391 r <- pstep2 ps'
392 case acc' of
393 Partial s -> go s r
394 Done b2 -> return $ Done b2
395 go acc (Pipe.Continue (Consume ps')) =
396 return $ Partial $ Tuple' ps' acc
397 go acc (Pipe.Continue (Produce ps')) = do
398 r <- pstep2 ps'
399 go acc r
401 extract (Tuple' _ fs) = fextract fs
403-- | Scan the input of a 'Fold' to change it in a stateful manner using another
404-- 'Fold'.
405-- /Pre-release/
406{-# INLINE scan #-}
407scan :: Monad m => Fold m a b -> Fold m b c -> Fold m a c
408scan (Fold stepL initialL extractL) (Fold stepR initialR extractR) =
409 Fold step initial extract
411 where
413 -- XXX It can be moved out and used for both scan and postscan
414 {-# INLINE runStep #-}
415 runStep actionL sR = do
416 rL <- actionL
417 case rL of
418 Done bL -> do
419 rR <- stepR sR bL
420 case rR of
421 Partial sR1 -> Done <$> extractR sR1
422 Done bR -> return $ Done bR
423 Partial sL -> do
424 !b <- extractL sL
425 rR <- stepR sR b
426 return
427 $ case rR of
428 Partial sR1 -> Partial (sL, sR1)
429 Done bR -> Done bR
431 initial = do
432 r <- initialR
433 case r of
434 Partial sR -> runStep initialL sR
435 Done b -> return $ Done b
437 step (sL, sR) x = runStep (stepL sL x) sR
439 extract = extractR . snd
441-- | Postscan the input of a 'Fold' to change it in a stateful manner using
442-- another 'Fold'.
443-- /Pre-release/
444{-# INLINE postscan #-}
445postscan :: Monad m => Fold m a b -> Fold m b c -> Fold m a c
446postscan (Fold stepL initialL extractL) (Fold stepR initialR extractR) =
447 Fold step initial extract
449 where
451 {-# INLINE runStep #-}
452 runStep actionL sR = do
453 rL <- actionL
454 case rL of
455 Done bL -> do
456 rR <- stepR sR bL
457 case rR of
458 Partial sR1 -> Done <$> extractR sR1
459 Done bR -> return $ Done bR
460 Partial sL -> do
461 !b <- extractL sL
462 rR <- stepR sR b
463 return
464 $ case rR of
465 Partial sR1 -> Partial (sL, sR1)
466 Done bR -> Done bR
468 initial = do
469 r <- initialR
470 rL <- initialL
471 case r of
472 Partial sR ->
473 case rL of
474 Done _ -> Done <$> extractR sR
475 Partial sL -> return $ Partial (sL, sR)
476 Done b -> return $ Done b
478 step (sL, sR) x = runStep (stepL sL x) sR
480 extract = extractR . snd
482------------------------------------------------------------------------------
483-- Filters
484------------------------------------------------------------------------------
486-- | Convert a predicate into a filtering fold.
487--
488-- >>> f = Fold.foldFilter (Fold.satisfy (> 5)) Fold.sum
489-- >>> Stream.fold f $ Stream.fromList [1..10]
490-- 40
491--
492-- /Pre-release/
493{-# INLINE satisfy #-}
494satisfy :: Monad m => (a -> Bool) -> Fold m a (Maybe a)
495satisfy f = Fold step (return $ Partial ()) (const (return Nothing))
497 where
499 step () a = return $ Done $ if f a then Just a else Nothing
501-- | Use a 'Maybe' returning fold as a filtering scan.
502--
503-- >>> f = Fold.foldFilter (Fold.satisfy (> 5)) Fold.sum
504-- >>> Stream.fold f $ Stream.fromList [1..10]
505-- 40
506--
507-- The above snippet is equivalent to:
508--
509-- >>> f = Fold.filter (> 5) Fold.sum
510-- >>> Stream.fold f $ Stream.fromList [1..10]
511-- 40
512--
513-- /Pre-release/
514{-# INLINE foldFilter #-}
515foldFilter :: Monad m => Fold m a (Maybe b) -> Fold m b c -> Fold m a c
516foldFilter f1 f2 = many f1 (catMaybes f2)
518------------------------------------------------------------------------------
519-- Left folds
520------------------------------------------------------------------------------
522------------------------------------------------------------------------------
523-- Run Effects
524------------------------------------------------------------------------------
526-- |
527-- > drainBy f = lmapM f drain
528-- > drainBy = Fold.foldMapM (void . f)
529--
530-- Drain all input after passing it through a monadic function. This is the
531-- dual of mapM_ on stream producers.
532--
533-- See also: 'Streamly.Prelude.mapM_'
534--
535-- @since 0.7.0
536{-# INLINE drainBy #-}
537drainBy :: Monad m => (a -> m b) -> Fold m a ()
538drainBy f = lmapM f drain
540-- | Extract the last element of the input stream, if any.
541--
542-- > last = fmap getLast $ Fold.foldMap (Last . Just)
543--
544-- @since 0.7.0
545{-# INLINE last #-}
546last :: Monad m => Fold m a (Maybe a)
547last = foldl1' (\_ x -> x)
549------------------------------------------------------------------------------
550-- To Summary
551------------------------------------------------------------------------------
553-- | Like 'length', except with a more general 'Num' return value
554--
555-- > genericLength = fmap getSum $ foldMap (Sum . const 1)
556--
557-- /Pre-release/
558{-# INLINE genericLength #-}
559genericLength :: (Monad m, Num b) => Fold m a b
560genericLength = foldl' (\n _ -> n + 1) 0
562-- | Determine the length of the input stream.
563--
564-- > length = fmap getSum $ Fold.foldMap (Sum . const 1)
565--
566-- @since 0.7.0
567{-# INLINE length #-}
568length :: Monad m => Fold m a Int
569length = genericLength
571-- XXX Try Hash set
572-- XXX Add a countDistinct window fold
573-- XXX Add a bloom filter fold
575-- | Count non-duplicate elements in the stream.
576--
577-- The memory used is proportional to the number of distinct elements in the
578-- stream, to guard against using too much memory use it as a scan and
579-- terminate if the count reaches more than a threshold.
580--
581-- /Space/: \(\mathcal{O}(n)\)
582--
583-- /Pre-release/
584--
585{-# INLINE countDistinct #-}
586countDistinct :: (Monad m, Ord a) => Fold m a Int
587countDistinct = fmap (\(Tuple' _ n) -> n) $ foldl' step initial
589 where
591 initial = Tuple' Set.empty 0
593 step (Tuple' set n) x = do
594 if Set.member x set
595 then
596 Tuple' set n
597 else
598 let cnt = n + 1
599 in Tuple' (Set.insert x set) cnt
601-- | Like 'countDistinct' but specialized to a stream of 'Int', for better
602-- performance.
603--
604-- /Pre-release/
605{-# INLINE countDistinctInt #-}
606countDistinctInt :: Monad m => Fold m Int Int
607countDistinctInt = fmap (\(Tuple' _ n) -> n) $ foldl' step initial
609 where
611 initial = Tuple' IntSet.empty 0
613 step (Tuple' set n) x = do
614 if IntSet.member x set
615 then
616 Tuple' set n
617 else
618 let cnt = n + 1
619 in Tuple' (IntSet.insert x set) cnt
621-- | Determine the sum of all elements of a stream of numbers. Returns additive
622-- identity (@0@) when the stream is empty. Note that this is not numerically
623-- stable for floating point numbers.
624--
625-- > sum = fmap getSum $ Fold.foldMap Sum
626--
627-- @since 0.7.0
628{-# INLINE sum #-}
629sum :: (Monad m, Num a) => Fold m a a
630sum = foldl' (+) 0
632-- | Determine the product of all elements of a stream of numbers. Returns
633-- multiplicative identity (@1@) when the stream is empty. The fold terminates
634-- when it encounters (@0@) in its input.
635--
636-- Compare with @Fold.foldMap Product@.
637--
638-- @since 0.7.0
639-- /Since 0.8.0 (Added 'Eq' constraint)/
640{-# INLINE product #-}
641product :: (Monad m, Num a, Eq a) => Fold m a a
642product = mkFold_ step (Partial 1)
644 where
646 step x a =
647 if a == 0
648 then Done 0
649 else Partial $ x * a
651------------------------------------------------------------------------------
652-- To Summary (Maybe)
653------------------------------------------------------------------------------
655-- | Determine the maximum element in a stream using the supplied comparison
656-- function.
657--
658-- @since 0.7.0
659{-# INLINE maximumBy #-}
660maximumBy :: Monad m => (a -> a -> Ordering) -> Fold m a (Maybe a)
661maximumBy cmp = foldl1' max'
663 where
665 max' x y =
666 case cmp x y of
667 GT -> x
668 _ -> y
670-- |
671-- @
672-- maximum = Fold.maximumBy compare
673-- @
674--
675-- Determine the maximum element in a stream.
676--
677-- Compare with @Fold.foldMap Max@.
678--
679-- @since 0.7.0
680{-# INLINE maximum #-}
681maximum :: (Monad m, Ord a) => Fold m a (Maybe a)
682maximum = foldl1' max
684-- | Computes the minimum element with respect to the given comparison function
685--
686-- @since 0.7.0
687{-# INLINE minimumBy #-}
688minimumBy :: Monad m => (a -> a -> Ordering) -> Fold m a (Maybe a)
689minimumBy cmp = foldl1' min'
691 where
693 min' x y =
694 case cmp x y of
695 GT -> y
696 _ -> x
698-- | Determine the minimum element in a stream using the supplied comparison
699-- function.
700--
701-- @
702-- minimum = 'minimumBy' compare
703-- @
704--
705-- Compare with @Fold.foldMap Min@.
706--
707-- @since 0.7.0
708{-# INLINE minimum #-}
709minimum :: (Monad m, Ord a) => Fold m a (Maybe a)
710minimum = foldl1' min
712------------------------------------------------------------------------------
713-- To Summary (Statistical)
714------------------------------------------------------------------------------
716-- | Compute a numerically stable arithmetic mean of all elements in the input
717-- stream.
718--
719-- @since 0.7.0
720{-# INLINE mean #-}
721mean :: (Monad m, Fractional a) => Fold m a a
722mean = fmap done $ foldl' step begin
724 where
726 begin = Tuple' 0 0
728 step (Tuple' x n) y =
729 let n1 = n + 1
730 in Tuple' (x + (y - x) / n1) n1
732 done (Tuple' x _) = x
734-- | Compute a numerically stable (population) variance over all elements in
735-- the input stream.
736--
737-- @since 0.7.0
738{-# INLINE variance #-}
739variance :: (Monad m, Fractional a) => Fold m a a
740variance = fmap done $ foldl' step begin
742 where
744 begin = Tuple3' 0 0 0
746 step (Tuple3' n mean_ m2) x = Tuple3' n' mean' m2'
748 where
750 n' = n + 1
751 mean' = (n * mean_ + x) / (n + 1)
752 delta = x - mean_
753 m2' = m2 + delta * delta * n / (n + 1)
755 done (Tuple3' n _ m2) = m2 / n
757-- | Compute a numerically stable (population) standard deviation over all
758-- elements in the input stream.
759--
760-- @since 0.7.0
761{-# INLINE stdDev #-}
762stdDev :: (Monad m, Floating a) => Fold m a a
763stdDev = sqrt <$> variance
765-- | Compute an 'Int' sized polynomial rolling hash
766--
767-- > H = salt * k ^ n + c1 * k ^ (n - 1) + c2 * k ^ (n - 2) + ... + cn * k ^ 0
768--
769-- Where @c1@, @c2@, @cn@ are the elements in the input stream and @k@ is a
770-- constant.
771--
772-- This hash is often used in Rabin-Karp string search algorithm.
773--
774-- See https://en.wikipedia.org/wiki/Rolling_hash
775--
776-- @since 0.8.0
777{-# INLINE rollingHashWithSalt #-}
778rollingHashWithSalt :: (Monad m, Enum a) => Int64 -> Fold m a Int64
779rollingHashWithSalt = foldl' step
781 where
783 k = 2891336453 :: Int64
785 step cksum a = cksum * k + fromIntegral (fromEnum a)
787-- | A default salt used in the implementation of 'rollingHash'.
788{-# INLINE defaultSalt #-}
789defaultSalt :: Int64
790defaultSalt = -2578643520546668380
792-- | Compute an 'Int' sized polynomial rolling hash of a stream.
793--
794-- > rollingHash = Fold.rollingHashWithSalt defaultSalt
795--
796-- @since 0.8.0
797{-# INLINE rollingHash #-}
798rollingHash :: (Monad m, Enum a) => Fold m a Int64
799rollingHash = rollingHashWithSalt defaultSalt
801-- | Compute an 'Int' sized polynomial rolling hash of the first n elements of
802-- a stream.
803--
804-- > rollingHashFirstN = Fold.take n Fold.rollingHash
805--
806-- /Pre-release/
807{-# INLINE rollingHashFirstN #-}
808rollingHashFirstN :: (Monad m, Enum a) => Int -> Fold m a Int64
809rollingHashFirstN n = take n rollingHash
811------------------------------------------------------------------------------
812-- Monoidal left folds
813------------------------------------------------------------------------------
815-- | Append the elements of an input stream to a provided starting value.
816--
817-- >>> Stream.fold (Fold.sconcat 10) (Stream.map Data.Monoid.Sum $ Stream.enumerateFromTo 1 10)
818-- Sum {getSum = 65}
819--
820-- @
821-- sconcat = Fold.foldl' (<>)
822-- @
823--
824-- @since 0.8.0
825{-# INLINE sconcat #-}
826sconcat :: (Monad m, Semigroup a) => a -> Fold m a a
827sconcat = foldl' (<>)
829-- | Fold an input stream consisting of monoidal elements using 'mappend'
830-- and 'mempty'.
831--
832-- >>> Stream.fold Fold.mconcat (Stream.map Data.Monoid.Sum $ Stream.enumerateFromTo 1 10)
833-- Sum {getSum = 55}
834--
835-- > mconcat = Fold.sconcat mempty
836--
837-- @since 0.7.0
838{-# INLINE mconcat #-}
839mconcat ::
840 ( Monad m
841#if !MIN_VERSION_base(4,11,0)
842 , Semigroup a
843#endif
844 , Monoid a) => Fold m a a
845mconcat = sconcat mempty
847-- |
848-- > foldMap f = Fold.lmap f Fold.mconcat
849--
850-- Make a fold from a pure function that folds the output of the function
851-- using 'mappend' and 'mempty'.
852--
853-- >>> Stream.fold (Fold.foldMap Data.Monoid.Sum) $ Stream.enumerateFromTo 1 10
854-- Sum {getSum = 55}
855--
856-- @since 0.7.0
857{-# INLINE foldMap #-}
858foldMap :: (Monad m, Monoid b
859#if !MIN_VERSION_base(4,11,0)
860 , Semigroup b
861#endif
862 ) => (a -> b) -> Fold m a b
863foldMap f = lmap f mconcat
865-- |
866-- > foldMapM f = Fold.lmapM f Fold.mconcat
867--
868-- Make a fold from a monadic function that folds the output of the function
869-- using 'mappend' and 'mempty'.
870--
871-- >>> Stream.fold (Fold.foldMapM (return . Data.Monoid.Sum)) $ Stream.enumerateFromTo 1 10
872-- Sum {getSum = 55}
873--
874-- @since 0.7.0
875{-# INLINE foldMapM #-}
876foldMapM :: (Monad m, Monoid b) => (a -> m b) -> Fold m a b
877foldMapM act = foldlM' step (pure mempty)
879 where
881 step m a = do
882 m' <- act a
883 return $! mappend m m'
885------------------------------------------------------------------------------
886-- To Containers
887------------------------------------------------------------------------------
889-- $toListRev
890-- This is more efficient than 'Streamly.Internal.Data.Fold.toList'. toList is
891-- exactly the same as reversing the list after 'toListRev'.
893-- | Buffers the input stream to a list in the reverse order of the input.
894--
895-- > toListRev = Fold.foldl' (flip (:)) []
896--
897-- /Warning!/ working on large lists accumulated as buffers in memory could be
898-- very inefficient, consider using "Streamly.Array" instead.
899--
900-- @since 0.8.0
902-- xn : ... : x2 : x1 : []
903{-# INLINE toListRev #-}
904toListRev :: Monad m => Fold m a [a]
905toListRev = foldl' (flip (:)) []
907------------------------------------------------------------------------------
908-- Partial Folds
909------------------------------------------------------------------------------
911-- | A fold that drains the first n elements of its input, running the effects
912-- and discarding the results.
913--
914-- > drainN n = Fold.take n Fold.drain
915--
916-- /Pre-release/
917{-# INLINE drainN #-}
918drainN :: Monad m => Int -> Fold m a ()
919drainN n = take n drain
921------------------------------------------------------------------------------
922-- To Elements
923------------------------------------------------------------------------------
925-- | Like 'index', except with a more general 'Integral' argument
926--
927-- /Pre-release/
928{-# INLINE genericIndex #-}
929genericIndex :: (Integral i, Monad m) => i -> Fold m a (Maybe a)
930genericIndex i = mkFold step (Partial 0) (const Nothing)
932 where
934 step j a =
935 if i == j
936 then Done $ Just a
937 else Partial (j + 1)
939-- | Lookup the element at the given index.
940--
941-- See also: 'Streamly.Prelude.!!'
942--
943-- @since 0.7.0
944{-# INLINE index #-}
945index :: Monad m => Int -> Fold m a (Maybe a)
946index = genericIndex
948-- | Extract the first element of the stream, if any.
949--
950-- @since 0.7.0
951{-# INLINE head #-}
952head :: Monad m => Fold m a (Maybe a)
953head = mkFold_ (const (Done . Just)) (Partial Nothing)
955-- | Returns the first element that satisfies the given predicate.
956--
957-- @since 0.7.0
958{-# INLINE find #-}
959find :: Monad m => (a -> Bool) -> Fold m a (Maybe a)
960find predicate = mkFold step (Partial ()) (const Nothing)
962 where
964 step () a =
965 if predicate a
966 then Done (Just a)
967 else Partial ()
969-- | In a stream of (key-value) pairs @(a, b)@, return the value @b@ of the
970-- first pair where the key equals the given value @a@.
971--
972-- > lookup = snd <$> Fold.find ((==) . fst)
973--
974-- @since 0.7.0
975{-# INLINE lookup #-}
976lookup :: (Eq a, Monad m) => a -> Fold m (a,b) (Maybe b)
977lookup a0 = mkFold step (Partial ()) (const Nothing)
979 where
981 step () (a, b) =
982 if a == a0
983 then Done $ Just b
984 else Partial ()
986-- | Returns the first index that satisfies the given predicate.
987--
988-- @since 0.7.0
989{-# INLINE findIndex #-}
990findIndex :: Monad m => (a -> Bool) -> Fold m a (Maybe Int)
991findIndex predicate = mkFold step (Partial 0) (const Nothing)
993 where
995 step i a =
996 if predicate a
997 then Done $ Just i
998 else Partial (i + 1)
1000-- XXX If we have a Continue contructor we won't need a Maybe type for the
1001-- scan. The fold will have the ability to skip producing output.
1003-- | Returns all indices that satisfy the given predicate.
1005{-# INLINE findIndices #-}
1006findIndices :: Monad m => (a -> Bool) -> Fold m a (Maybe Int)
1007findIndices predicate =
1008 fmap (either (const Nothing) Just) $ foldl' step (Left (-1))
1010 where
1012 step i a =
1013 if predicate a
1014 then Right (either id id i + 1)
1015 else Left (either id id i + 1)
1017-- | Returns the first index where a given value is found in the stream.
1019-- > elemIndex a = Fold.findIndex (== a)
1021-- @since 0.7.0
1022{-# INLINE elemIndex #-}
1023elemIndex :: (Eq a, Monad m) => a -> Fold m a (Maybe Int)
1024elemIndex a = findIndex (a ==)
1026------------------------------------------------------------------------------
1027-- To Boolean
1028------------------------------------------------------------------------------
1030-- | Return 'True' if the input stream is empty.
1032-- > null = fmap isJust Fold.head
1034-- @since 0.7.0
1035{-# INLINE null #-}
1036null :: Monad m => Fold m a Bool
1037null = mkFold (\() _ -> Done False) (Partial ()) (const True)
1039-- | Returns 'True' if any of the elements of a stream satisfies a predicate.
1041-- >>> Stream.fold (Fold.any (== 0)) $ Stream.fromList [1,0,1]
1042-- True
1044-- > any p = Fold.lmap p Fold.or
1046-- @since 0.7.0
1047{-# INLINE any #-}
1048any :: Monad m => (a -> Bool) -> Fold m a Bool
1049any predicate = mkFold_ step initial
1051 where
1053 initial = Partial False
1055 step _ a =
1056 if predicate a
1057 then Done True
1058 else Partial False
1060-- | Return 'True' if the given element is present in the stream.
1062-- > elem a = Fold.any (== a)
1064-- @since 0.7.0
1065{-# INLINE elem #-}
1066elem :: (Eq a, Monad m) => a -> Fold m a Bool
1067elem a = any (a ==)
1069-- | Returns 'True' if all elements of a stream satisfy a predicate.
1071-- >>> Stream.fold (Fold.all (== 0)) $ Stream.fromList [1,0,1]
1072-- False
1074-- > all p = Fold.lmap p Fold.and
1076-- @since 0.7.0
1077{-# INLINE all #-}
1078all :: Monad m => (a -> Bool) -> Fold m a Bool
1079all predicate = mkFold_ step initial
1081 where
1083 initial = Partial True
1085 step _ a =
1086 if predicate a
1087 then Partial True
1088 else Done False
1090-- | Returns 'True' if the given element is not present in the stream.
1092-- > notElem a = Fold.all (/= a)
1094-- @since 0.7.0
1095{-# INLINE notElem #-}
1096notElem :: (Eq a, Monad m) => a -> Fold m a Bool
1097notElem a = all (a /=)
1099-- | Returns 'True' if all elements are 'True', 'False' otherwise
1101-- > and = Fold.all (== True)
1103-- @since 0.7.0
1104{-# INLINE and #-}
1105and :: Monad m => Fold m Bool Bool
1106and = all (== True)
1108-- | Returns 'True' if any element is 'True', 'False' otherwise
1110-- > or = Fold.any (== True)
1112-- @since 0.7.0
1113{-# INLINE or #-}
1114or :: Monad m => Fold m Bool Bool
1115or = any (== True)
1117------------------------------------------------------------------------------
1118-- Grouping/Splitting
1119------------------------------------------------------------------------------
1121------------------------------------------------------------------------------
1122-- Grouping without looking at elements
1123------------------------------------------------------------------------------
1125------------------------------------------------------------------------------
1126-- Binary APIs
1127------------------------------------------------------------------------------
1129-- | @splitAt n f1 f2@ composes folds @f1@ and @f2@ such that first @n@
1130-- elements of its input are consumed by fold @f1@ and the rest of the stream
1131-- is consumed by fold @f2@.
1133-- >>> let splitAt_ n xs = Stream.fold (Fold.splitAt n Fold.toList Fold.toList) $ Stream.fromList xs
1135-- >>> splitAt_ 6 "Hello World!"
1136-- ("Hello ","World!")
1138-- >>> splitAt_ (-1) [1,2,3]
1139-- ([],[1,2,3])
1141-- >>> splitAt_ 0 [1,2,3]
1142-- ([],[1,2,3])
1144-- >>> splitAt_ 1 [1,2,3]
1145-- ([1],[2,3])
1147-- >>> splitAt_ 3 [1,2,3]
1148-- ([1,2,3],[])
1150-- >>> splitAt_ 4 [1,2,3]
1151-- ([1,2,3],[])
1153-- > splitAt n f1 f2 = Fold.serialWith (,) (Fold.take n f1) f2
1155-- /Internal/
1157{-# INLINE splitAt #-}
1158splitAt
1159 :: Monad m
1160 => Int
1161 -> Fold m a b
1162 -> Fold m a c
1163 -> Fold m a (b, c)
1164splitAt n fld = serialWith (,) (take n fld)
1166------------------------------------------------------------------------------
1167-- Element Aware APIs
1168------------------------------------------------------------------------------
1170------------------------------------------------------------------------------
1171-- Binary APIs
1172------------------------------------------------------------------------------
1174-- Note: Keep this consistent with S.splitOn. In fact we should eliminate
1175-- S.splitOn in favor of the fold.
1177-- XXX Use Fold.many instead once it is fixed.
1179-- | Like 'takeEndBy' but drops the element on which the predicate succeeds.
1181-- >>> Stream.fold (Fold.takeEndBy_ (== '\n') Fold.toList) $ Stream.fromList "hello\nthere\n"
1182-- "hello"
1184-- >>> Stream.toList $ Stream.foldMany (Fold.takeEndBy_ (== '\n') Fold.toList) $ Stream.fromList "hello\nthere\n"
1185-- ["hello","there"]
1187-- > Stream.splitOnSuffix p f = Stream.foldMany (Fold.takeEndBy_ p f)
1189-- See 'Streamly.Prelude.splitOnSuffix' for more details on splitting a
1190-- stream using 'takeEndBy_'.
1192-- @since 0.8.0
1193{-# INLINE takeEndBy_ #-}
1194takeEndBy_ :: Monad m => (a -> Bool) -> Fold m a b -> Fold m a b
1195takeEndBy_ predicate (Fold fstep finitial fextract) =
1196 Fold step finitial fextract
1198 where
1200 step s a =
1201 if not (predicate a)
1202 then fstep s a
1203 else Done <$> fextract s
1205-- | Take the input, stop when the predicate succeeds taking the succeeding
1206-- element as well.
1208-- >>> Stream.fold (Fold.takeEndBy (== '\n') Fold.toList) $ Stream.fromList "hello\nthere\n"
1209-- "hello\n"
1211-- >>> Stream.toList $ Stream.foldMany (Fold.takeEndBy (== '\n') Fold.toList) $ Stream.fromList "hello\nthere\n"
1212-- ["hello\n","there\n"]
1214-- > Stream.splitWithSuffix p f = Stream.foldMany (Fold.takeEndBy p f)
1216-- See 'Streamly.Prelude.splitWithSuffix' for more details on splitting a
1217-- stream using 'takeEndBy'.
1219-- @since 0.8.0
1220{-# INLINE takeEndBy #-}
1221takeEndBy :: Monad m => (a -> Bool) -> Fold m a b -> Fold m a b
1222takeEndBy predicate (Fold fstep finitial fextract) =
1223 Fold step finitial fextract
1225 where
1227 step s a = do
1228 res <- fstep s a
1229 if not (predicate a)
1230 then return res
1231 else do
1232 case res of
1233 Partial s1 -> Done <$> fextract s1
1234 Done b -> return $ Done b
1236------------------------------------------------------------------------------
1237-- Binary splitting on a separator
1238------------------------------------------------------------------------------
1241-- | Find the first occurrence of the specified sequence in the input stream
1242-- and break the input stream into two parts, the first part consisting of the
1243-- stream before the sequence and the second part consisting of the sequence
1244-- and the rest of the stream.
1246-- > let breakOn_ pat xs = S.fold (S.breakOn pat FL.toList FL.toList) $ S.fromList xs
1248-- >>> breakOn_ "dear" "Hello dear world!"
1249-- > ("Hello ","dear world!")
1251{-# INLINE breakOn #-}
1252breakOn :: Monad m => Array a -> Fold m a b -> Fold m a c -> Fold m a (b,c)
1253breakOn pat f m = undefined
1256------------------------------------------------------------------------------
1257-- Distributing
1258------------------------------------------------------------------------------
1260-- | Distribute one copy of the stream to each fold and zip the results.
1262-- @
1263-- |-------Fold m a b--------|
1264-- ---stream m a---| |---m (b,c)
1265-- |-------Fold m a c--------|
1266-- @
1267-- >>> Stream.fold (Fold.tee Fold.sum Fold.length) (Stream.enumerateFromTo 1.0 100.0)
1268-- (5050.0,100)
1270-- > tee = teeWith (,)
1272-- @since 0.7.0
1273{-# INLINE tee #-}
1274tee :: Monad m => Fold m a b -> Fold m a c -> Fold m a (b,c)
1275tee = teeWith (,)
1277-- XXX use "List" instead of "[]"?, use Array for output to scale it to a large
1278-- number of consumers? For polymorphic case a vector could be helpful. For
1279-- Storables we can use arrays. Will need separate APIs for those.
1281-- | Distribute one copy of the stream to each fold and collect the results in
1282-- a container.
1284-- @
1286-- |-------Fold m a b--------|
1287-- ---stream m a---| |---m [b]
1288-- |-------Fold m a b--------|
1289-- | |
1290-- ...
1291-- @
1293-- >>> Stream.fold (Fold.distribute [Fold.sum, Fold.length]) (Stream.enumerateFromTo 1 5)
1294-- [15,5]
1296-- > distribute = Prelude.foldr (Fold.teeWith (:)) (Fold.fromPure [])
1298-- This is the consumer side dual of the producer side 'sequence' operation.
1300-- Stops when all the folds stop.
1302-- @since 0.7.0
1303{-# INLINE distribute #-}
1304distribute :: Monad m => [Fold m a b] -> Fold m a [b]
1305distribute = Prelude.foldr (teeWith (:)) (fromPure [])
1307------------------------------------------------------------------------------
1308-- Partitioning
1309------------------------------------------------------------------------------
1311{-# INLINE partitionByMUsing #-}
1312partitionByMUsing :: Monad m =>
1313 ( (x -> y -> (x, y))
1314 -> Fold m (Either b c) x
1315 -> Fold m (Either b c) y
1316 -> Fold m (Either b c) (x, y)
1318 -> (a -> m (Either b c))
1319 -> Fold m b x
1320 -> Fold m c y
1321 -> Fold m a (x, y)
1322partitionByMUsing t f fld1 fld2 =
1323 let l = lmap (fromLeft undefined) fld1 -- :: Fold m (Either b c) x
1324 r = lmap (fromRight undefined) fld2 -- :: Fold m (Either b c) y
1325 in lmapM f (t (,) (filter isLeft l) (filter isRight r))
1327-- | Partition the input over two folds using an 'Either' partitioning
1328-- predicate.
1330-- @
1332-- |-------Fold b x--------|
1333-- -----stream m a --> (Either b c)----| |----(x,y)
1334-- |-------Fold c y--------|
1335-- @
1337-- Send input to either fold randomly:
1339-- @
1340-- > import System.Random (randomIO)
1341-- > randomly a = randomIO >>= \\x -> return $ if x then Left a else Right a
1342-- > Stream.fold (Fold.partitionByM randomly Fold.length Fold.length) (Stream.enumerateFromTo 1 100)
1343-- (59,41)
1344-- @
1346-- Send input to the two folds in a proportion of 2:1:
1348-- >>> :{
1349-- proportionately m n = do
1350-- ref <- newIORef $ cycle $ concat [replicate m Left, replicate n Right]
1351-- return $ \a -> do
1352-- r <- readIORef ref
1353-- writeIORef ref $ tail r
1354-- return $ Prelude.head r a
1355-- :}
1357-- >>> :{
1358-- main = do
1359-- f <- proportionately 2 1
1360-- r <- Stream.fold (Fold.partitionByM f Fold.length Fold.length) (Stream.enumerateFromTo (1 :: Int) 100)
1361-- print r
1362-- :}
1364-- >>> main
1365-- (67,33)
1368-- This is the consumer side dual of the producer side 'mergeBy' operation.
1370-- When one fold is done, any input meant for it is ignored until the other
1371-- fold is also done.
1373-- Stops when both the folds stop.
1375-- /See also: 'partitionByFstM' and 'partitionByMinM'./
1377-- /Pre-release/
1378{-# INLINE partitionByM #-}
1379partitionByM :: Monad m
1380 => (a -> m (Either b c)) -> Fold m b x -> Fold m c y -> Fold m a (x, y)
1381partitionByM = partitionByMUsing teeWith
1383-- | Similar to 'partitionByM' but terminates when the first fold terminates.
1386{-# INLINE partitionByFstM #-}
1387partitionByFstM :: Monad m
1388 => (a -> m (Either b c)) -> Fold m b x -> Fold m c y -> Fold m a (x, y)
1389partitionByFstM = partitionByMUsing teeWithFst
1391-- | Similar to 'partitionByM' but terminates when any fold terminates.
1394{-# INLINE partitionByMinM #-}
1395partitionByMinM :: Monad m =>
1396 (a -> m (Either b c)) -> Fold m b x -> Fold m c y -> Fold m a (x, y)
1397partitionByMinM = partitionByMUsing teeWithMin
1399-- Note: we could use (a -> Bool) instead of (a -> Either b c), but the latter
1400-- makes the signature clearer as to which case belongs to which fold.
1401-- XXX need to check the performance in both cases.
1402-- | Same as 'partitionByM' but with a pure partition function.
1404-- Count even and odd numbers in a stream:
1406-- >>> :{
1407-- let f = Fold.partitionBy (\n -> if even n then Left n else Right n)
1408-- (fmap (("Even " ++) . show) Fold.length)
1409-- (fmap (("Odd " ++) . show) Fold.length)
1410-- in Stream.fold f (Stream.enumerateFromTo 1 100)
1411-- :}
1412-- ("Even 50","Odd 50")
1414-- /Pre-release/
1415{-# INLINE partitionBy #-}
1416partitionBy :: Monad m
1417 => (a -> Either b c) -> Fold m b x -> Fold m c y -> Fold m a (x, y)
1418partitionBy f = partitionByM (return . f)
1420-- | Compose two folds such that the combined fold accepts a stream of 'Either'
1421-- and routes the 'Left' values to the first fold and 'Right' values to the
1422-- second fold.
1424-- > partition = partitionBy id
1426-- @since 0.7.0
1427{-# INLINE partition #-}
1428partition :: Monad m
1429 => Fold m b x -> Fold m c y -> Fold m (Either b c) (x, y)
1430partition = partitionBy id
1433-- | Send one item to each fold in a round-robin fashion. This is the consumer
1434-- side dual of producer side 'mergeN' operation.
1436-- partitionN :: Monad m => [Fold m a b] -> Fold m a [b]
1437-- partitionN fs = Fold step begin done
1440-- TODO Demultiplex an input element into a number of typed variants. We want
1441-- to statically restrict the target values within a set of predefined types,
1442-- an enumeration of a GADT. We also want to make sure that the Map contains
1443-- only those types and the full set of those types.
1445-- TODO Instead of the input Map it should probably be a lookup-table using an
1446-- array and not in GC memory. The same applies to the output Map as well.
1447-- However, that would only be helpful if we have a very large data structure,
1448-- need to measure and see how it scales.
1450-- This is the consumer side dual of the producer side 'mux' operation (XXX to
1451-- be implemented).
1453-- | Split the input stream based on a key field and fold each split using a
1454-- specific fold collecting the results in a map from the keys to the results.
1455-- Useful for cases like protocol handlers to handle different type of packets
1456-- using different handlers.
1458-- @
1460-- |-------Fold m a b
1461-- -----stream m a-----Map-----|
1462-- |-------Fold m a b
1463-- |
1464-- ...
1465-- @
1467-- Any input that does not map to a fold in the input Map is silently ignored.
1469-- > demuxWith f kv = fmap fst $ demuxDefaultWith f kv drain
1471-- /Pre-release/
1473{-# INLINE demuxWith #-}
1474demuxWith :: (Monad m, Ord k)
1475 => (a -> (k, a')) -> Map k (Fold m a' b) -> Fold m a (Map k b)
1476demuxWith f kv = fmap fst $ demuxDefaultWith f kv drain
1478-- | Fold a stream of key value pairs using a map of specific folds for each
1479-- key into a map from keys to the results of fold outputs of the corresponding
1480-- values.
1482-- >>> import qualified Data.Map
1483-- >>> :{
1484-- let table = Data.Map.fromList [("SUM", Fold.sum), ("PRODUCT", Fold.product)]
1485-- input = Stream.fromList [("SUM",1),("PRODUCT",2),("SUM",3),("PRODUCT",4)]
1486-- in Stream.fold (Fold.demux table) input
1487-- :}
1488-- fromList [("PRODUCT",8),("SUM",4)]
1490-- > demux = demuxWith id
1492-- /Pre-release/
1493{-# INLINE demux #-}
1494demux :: (Monad m, Ord k)
1495 => Map k (Fold m a b) -> Fold m (k, a) (Map k b)
1496demux = demuxWith id
1498data DemuxState s b doneMap runMap =
1499 DemuxMapAndDefault !s !doneMap !runMap
1500 | DemuxOnlyMap b !doneMap !runMap
1501 | DemuxOnlyDefault s !doneMap
1503-- | Like 'demuxWith' but uses a default catchall fold to handle inputs which
1504-- do not have a specific fold in the map to handle them.
1506-- If any fold in the map stops, inputs meant for that fold are sent to the
1507-- catchall fold. If the catchall fold stops then inputs that do not match any
1508-- fold are ignored.
1510-- Stops when all the folds, including the catchall fold, stop.
1512-- /Pre-release/
1514{-# INLINE demuxDefaultWith #-}
1515demuxDefaultWith :: (Monad m, Ord k)
1516 => (a -> (k, a'))
1517 -> Map k (Fold m a' b)
1518 -> Fold m (k, a') c
1519 -> Fold m a (Map k b, c)
1520demuxDefaultWith f kv (Fold dstep dinitial dextract) =
1521 Fold step initial extract
1523 where
1525 initial = do
1526 let runInit (Fold step1 initial1 done1) = do
1527 r <- initial1
1528 return
1529 $ case r of
1530 Partial _ -> Right' (Fold step1 (return r) done1)
1531 Done b -> Left' b
1533 -- initialize folds in the kv map and separate the ones that are done
1534 -- from running ones
1535 kv1 <- Prelude.mapM runInit kv
1536 let runMap = Map.map fromRight' $ Map.filter isRight' kv1
1537 doneMap = Map.map fromLeft' $ Map.filter isLeft' kv1
1539 -- Run the default fold, and decide the next state based on its result
1540 dres <- dinitial
1541 return
1542 $ case dres of
1543 Partial s ->
1544 Partial
1545 $ if Map.size runMap > 0
1546 then DemuxMapAndDefault s doneMap runMap
1547 else DemuxOnlyDefault s doneMap
1548 Done b ->
1549 if Map.size runMap > 0
1550 then Partial $ DemuxOnlyMap b doneMap runMap
1551 else Done (doneMap, b)
1553 {-# INLINE runFold #-}
1554 runFold fPartial fDone doneMap runMap (Fold step1 initial1 done1) k a1 = do
1555 resi <- initial1
1556 case resi of
1557 Partial st -> do
1558 res <- step1 st a1
1559 return $ case res of
1560 Partial s ->
1561 let fld = Fold step1 (return $ Partial s) done1
1562 runMap1 = Map.insert k fld runMap
1563 in Partial $ fPartial doneMap runMap1
1564 Done b -> do
1565 let runMap1 = Map.delete k runMap
1566 doneMap1 = Map.insert k b doneMap
1567 if Map.size runMap1 == 0
1568 then fDone doneMap1
1569 else Partial $ fPartial doneMap1 runMap1
1570 Done _ -> error "Bug: demuxDefaultWith: Done fold"
1572 step (DemuxMapAndDefault dacc doneMap runMap) a = do
1573 let (k, a1) = f a
1574 case Map.lookup k runMap of
1575 Nothing -> do
1576 res <- dstep dacc (k, a1)
1577 return
1578 $ Partial
1579 $ case res of
1580 Partial s -> DemuxMapAndDefault s doneMap runMap
1581 Done b -> DemuxOnlyMap b doneMap runMap
1582 Just fld ->
1583 runFold
1584 (DemuxMapAndDefault dacc)
1585 (Partial . DemuxOnlyDefault dacc)
1586 doneMap runMap fld k a1
1588 step (DemuxOnlyMap dval doneMap runMap) a = do
1589 let (k, a1) = f a
1590 case Map.lookup k runMap of
1591 Nothing -> return $ Partial $ DemuxOnlyMap dval doneMap runMap
1592 Just fld ->
1593 runFold
1594 (DemuxOnlyMap dval)
1595 (Done . (, dval))
1596 doneMap runMap fld k a1
1597 step (DemuxOnlyDefault dacc doneMap) a = do
1598 let (k, a1) = f a
1599 res <- dstep dacc (k, a1)
1600 return
1601 $ case res of
1602 Partial s -> Partial $ DemuxOnlyDefault s doneMap
1603 Done b -> Done (doneMap, b)
1605 runExtract (Fold _ initial1 done1) = do
1606 res <- initial1
1607 case res of
1608 Partial s -> done1 s
1609 Done b -> return b
1611 extract (DemuxMapAndDefault dacc doneMap runMap) = do
1612 b <- dextract dacc
1613 runMap1 <- Prelude.mapM runExtract runMap
1614 return (doneMap `Map.union` runMap1, b)
1615 extract (DemuxOnlyMap dval doneMap runMap) = do
1616 runMap1 <- Prelude.mapM runExtract runMap
1617 return (doneMap `Map.union` runMap1, dval)
1618 extract (DemuxOnlyDefault dacc doneMap) = do
1619 b <- dextract dacc
1620 return (doneMap, b)
1622-- |
1623-- > demuxDefault = demuxDefaultWith id
1625-- /Pre-release/
1626{-# INLINE demuxDefault #-}
1627demuxDefault :: (Monad m, Ord k)
1628 => Map k (Fold m a b) -> Fold m (k, a) b -> Fold m (k, a) (Map k b, b)
1629demuxDefault = demuxDefaultWith id
1631{-# INLINE classifyScanWith #-}
1632classifyScanWith :: (Monad m, Ord k) =>
1633 -- Note: we need to return the Map itself to display the in-progress values
1634 -- e.g. to implement top. We could possibly create a separate abstraction
1635 -- for that use case. We return an action because we want it to be lazy so
1636 -- that the downstream consumers can choose to process or discard it.
1637 (a -> k) -> Fold m a b -> Fold m a (m (Map k b), Maybe (k, b))
1638classifyScanWith f (Fold step1 initial1 extract1) =
1639 fmap extract $ foldlM' step initial
1641 where
1643 initial = return $ Tuple3' Map.empty Set.empty Nothing
1645 {-# INLINE initFold #-}
1646 initFold kv set k a = do
1647 x <- initial1
1648 case x of
1649 Partial s -> do
1650 r <- step1 s a
1651 return
1652 $ case r of
1653 Partial s1 ->
1654 Tuple3' (Map.insert k s1 kv) set Nothing
1655 Done b ->
1656 Tuple3' kv set (Just (k, b))
1657 Done b -> return (Tuple3' kv (Set.insert k set) (Just (k, b)))
1659 step (Tuple3' kv set _) a = do
1660 let k = f a
1661 case Map.lookup k kv of
1662 Nothing -> do
1663 if Set.member k set
1664 then return (Tuple3' kv set Nothing)
1665 else initFold kv set k a
1666 Just s -> do
1667 r <- step1 s a
1668 return
1669 $ case r of
1670 Partial s1 ->
1671 Tuple3' (Map.insert k s1 kv) set Nothing
1672 Done b ->
1673 let kv1 = Map.delete k kv
1674 in Tuple3' kv1 (Set.insert k set) (Just (k, b))
1676 extract (Tuple3' kv _ x) = (Prelude.mapM extract1 kv, x)
1678{-# INLINE toMap #-}
1679toMap :: (Monad m, Ord k) => Fold m (k, a) (Map k a)
1680toMap = foldl' (\kv (k, v) -> Map.insert k v kv) Map.empty
1682-- | Split the input stream based on a key field and fold each split using the
1683-- given fold. Useful for map/reduce, bucketizing the input in different bins
1684-- or for generating histograms.
1686-- >>> :{
1687-- let input = Stream.fromList [("ONE",1),("ONE",1.1),("TWO",2), ("TWO",2.2)]
1688-- in Stream.fold (Fold.classifyWith fst (Fold.lmap snd Fold.toList)) input
1689-- :}
1690-- fromList [("ONE",[1.0,1.1]),("TWO",[2.0,2.2])]
1692-- Once the classifier fold terminates for a particular key any further inputs
1693-- in that bucket are ignored.
1695-- Space used is proportional to the number of keys seen till now and
1696-- monotonically increases because it stores whether a key has been seen or
1697-- not.
1699-- /Stops: never/
1701-- /Pre-release/
1703{-# INLINE classifyWith #-}
1704classifyWith :: (Monad m, Ord k) =>
1705 (a -> k) -> Fold m a b -> Fold m a (Map k b)
1706classifyWith f fld =
1707 let
1708 classifier = classifyScanWith f fld
1709 getMap Nothing = pure Map.empty
1710 getMap (Just action) = action
1711 aggregator =
1712 teeWith Map.union
1713 (rmapM getMap $ lmap fst last)
1714 (lmap snd $ catMaybes toMap)
1715 in postscan classifier aggregator
1717-- | Given an input stream of key value pairs and a fold for values, fold all
1718-- the values belonging to each key. Useful for map/reduce, bucketizing the
1719-- input in different bins or for generating histograms.
1721-- >>> :{
1722-- let input = Stream.fromList [("ONE",1),("ONE",1.1),("TWO",2), ("TWO",2.2)]
1723-- in Stream.fold (Fold.classify Fold.toList) input
1724-- :}
1725-- fromList [("ONE",[1.0,1.1]),("TWO",[2.0,2.2])]
1727-- Same as:
1729-- > classify fld = Fold.classifyWith fst (lmap snd fld)
1731-- /Pre-release/
1732{-# INLINE classify #-}
1733classify :: (Monad m, Ord k) => Fold m a b -> Fold m (k, a) (Map k b)
1734classify fld = classifyWith fst (lmap snd fld)
1736------------------------------------------------------------------------------
1737-- Unzipping
1738------------------------------------------------------------------------------
1740{-# INLINE unzipWithMUsing #-}
1741unzipWithMUsing :: Monad m =>
1742 ( (x -> y -> (x, y))
1743 -> Fold m (b, c) x
1744 -> Fold m (b, c) y
1745 -> Fold m (b, c) (x, y)
1747 -> (a -> m (b, c))
1748 -> Fold m b x
1749 -> Fold m c y
1750 -> Fold m a (x, y)
1751unzipWithMUsing t f fld1 fld2 =
1752 let f1 = lmap fst fld1 -- :: Fold m (b, c) b
1753 f2 = lmap snd fld2 -- :: Fold m (b, c) c
1754 in lmapM f (t (,) f1 f2)
1756-- | Like 'unzipWith' but with a monadic splitter function.
1758-- @unzipWithM k f1 f2 = lmapM k (unzip f1 f2)@
1760-- /Pre-release/
1761{-# INLINE unzipWithM #-}
1762unzipWithM :: Monad m
1763 => (a -> m (b,c)) -> Fold m b x -> Fold m c y -> Fold m a (x,y)
1764unzipWithM = unzipWithMUsing teeWith
1766-- | Similar to 'unzipWithM' but terminates when the first fold terminates.
1769{-# INLINE unzipWithFstM #-}
1770unzipWithFstM :: Monad m =>
1771 (a -> m (b, c)) -> Fold m b x -> Fold m c y -> Fold m a (x, y)
1772unzipWithFstM = unzipWithMUsing teeWithFst
1774-- | Similar to 'unzipWithM' but terminates when any fold terminates.
1777{-# INLINE unzipWithMinM #-}
1778unzipWithMinM :: Monad m =>
1779 (a -> m (b,c)) -> Fold m b x -> Fold m c y -> Fold m a (x,y)
1780unzipWithMinM = unzipWithMUsing teeWithMin
1782-- | Split elements in the input stream into two parts using a pure splitter
1783-- function, direct each part to a different fold and zip the results.
1785-- @unzipWith f fld1 fld2 = Fold.lmap f (Fold.unzip fld1 fld2)@
1787-- This fold terminates when both the input folds terminate.
1789-- /Pre-release/
1790{-# INLINE unzipWith #-}
1791unzipWith :: Monad m
1792 => (a -> (b,c)) -> Fold m b x -> Fold m c y -> Fold m a (x,y)
1793unzipWith f = unzipWithM (return . f)
1795-- | Send the elements of tuples in a stream of tuples through two different
1796-- folds.
1798-- @
1800-- |-------Fold m a x--------|
1801-- ---------stream of (a,b)--| |----m (x,y)
1802-- |-------Fold m b y--------|
1804-- @
1806-- > unzip = Fold.unzipWith id
1808-- This is the consumer side dual of the producer side 'zip' operation.
1810-- @since 0.7.0
1811{-# INLINE unzip #-}
1812unzip :: Monad m => Fold m a x -> Fold m b y -> Fold m (a,b) (x,y)
1813unzip = unzipWith id
1815------------------------------------------------------------------------------
1816-- Combining streams and folds - Zipping
1817------------------------------------------------------------------------------
1819-- | Zip a stream with the input of a fold using the supplied function.
1821-- /Unimplemented/
1823{-# INLINE zipWithM #-}
1824zipWithM :: -- Monad m =>
1825 (a -> b -> m c) -> t m a -> Fold m c x -> Fold m b x
1826zipWithM = undefined
1828-- | Zip a stream with the input of a fold.
1830-- /Unimplemented/
1832{-# INLINE zip #-}
1833zip :: Monad m => t m a -> Fold m (a, b) x -> Fold m b x
1834zip = zipWithM (curry return)
1836-- | Pair each element of a fold input with its index, starting from index 0.
1838-- /Unimplemented/
1839{-# INLINE indexed #-}
1840indexed :: -- forall m a b. Monad m =>
1841 Fold m (Int, a) b -> Fold m a b
1842indexed = undefined -- zip (Stream.enumerateFrom 0 :: SerialT m Int)
1844-- | Change the predicate function of a Fold from @a -> b@ to accept an
1845-- additional state input @(s, a) -> b@. Convenient to filter with an
1846-- addiitonal index or time input.
1848-- @
1849-- filterWithIndex = with indexed filter
1850-- filterWithAbsTime = with timestamped filter
1851-- filterWithRelTime = with timeIndexed filter
1852-- @
1854-- /Pre-release/
1855{-# INLINE with #-}
1856with ::
1857 (Fold m (s, a) b -> Fold m a b)
1858 -> (((s, a) -> c) -> Fold m (s, a) b -> Fold m (s, a) b)
1859 -> (((s, a) -> c) -> Fold m a b -> Fold m a b)
1860with f comb g = f . comb g . lmap snd
1862-- | @sampleFromthen offset stride@ samples the element at @offset@ index and
1863-- then every element at strides of @stride@.
1865-- /Unimplemented/
1866{-# INLINE sampleFromthen #-}
1867sampleFromthen :: Monad m => Int -> Int -> Fold m a b -> Fold m a b
1868sampleFromthen offset size =
1869 with indexed filter (\(i, _) -> (i + offset) `mod` size == 0)
1871------------------------------------------------------------------------------
1872-- Nesting
1873------------------------------------------------------------------------------
1875-- | @concatSequence f t@ applies folds from stream @t@ sequentially and
1876-- collects the results using the fold @f@.
1878-- /Unimplemented/
1880{-# INLINE concatSequence #-}
1881concatSequence ::
1882 -- IsStream t =>
1883 Fold m b c -> t (Fold m a b) -> Fold m a c
1884concatSequence _f _p = undefined
1886-- | Group the input stream into groups of elements between @low@ and @high@.
1887-- Collection starts in chunks of @low@ and then keeps doubling until we reach
1888-- @high@. Each chunk is folded using the provided fold function.
1890-- This could be useful, for example, when we are folding a stream of unknown
1891-- size to a stream of arrays and we want to minimize the number of
1892-- allocations.
1894-- NOTE: this would be an application of "many" using a terminating fold.
1896-- /Unimplemented/
1898{-# INLINE chunksBetween #-}
1899chunksBetween :: -- Monad m =>
1900 Int -> Int -> Fold m a b -> Fold m b c -> Fold m a c
1901chunksBetween _low _high _f1 _f2 = undefined
1903-- | A fold that buffers its input to a pure stream.
1905-- /Warning!/ working on large streams accumulated as buffers in memory could
1906-- be very inefficient, consider using "Streamly.Data.Array" instead.
1908-- >>> toStream = fmap SerialT Fold.toStreamK
1910-- /Pre-release/
1911{-# INLINE toStream #-}
1912toStream :: Monad m => Fold m a (SerialT n a)
1913toStream = fmap SerialT toStreamK
1915-- This is more efficient than 'toStream'. toStream is exactly the same as
1916-- reversing the stream after toStreamRev.
1918-- | Buffers the input stream to a pure stream in the reverse order of the
1919-- input.
1921-- >>> toStreamRev = fmap SerialT Fold.toStreamKRev
1923-- /Warning!/ working on large streams accumulated as buffers in memory could
1924-- be very inefficient, consider using "Streamly.Data.Array" instead.
1926-- /Pre-release/
1928-- xn : ... : x2 : x1 : []
1929{-# INLINE toStreamRev #-}
1930toStreamRev :: Monad m => Fold m a (SerialT n a)
1931toStreamRev = fmap SerialT toStreamKRev