@@ -7,15 +7,16 @@
77-- Stability : experimental
88-- Portability : GHC
99
10−-- Also see the "Streamly.Internal.Data.Sink" module that provides specialized left folds
11−-- that discard the outputs.
10+-- Also see the "Streamly.Internal.Data.Sink" module that provides specialized
11+-- left folds that discard the outputs.
1212--
1313-- IMPORTANT: keep the signatures consistent with the folds in Streamly.Prelude
1414
1515module Streamly.Internal.Data.Fold
1616 (
1717 -- * Fold Type
18− Fold (..)
18+ Step (..)
19+ , Fold (..)
1920
2021 , hoist
2122 , generally
@@ -24,10 +25,14 @@ module Streamly.Internal.Data.Fold
2425 -- , init
2526
2627 -- * Fold Creation Utilities
27− , mkPure
28− , mkPureId
28+ , mkAccum
29+ , mkAccum_
30+ , mkAccumM
31+ , mkAccumM_
2932 , mkFold
30− , mkFoldId
33+ , mkFold_
34+ , mkFoldM
35+ , mkFoldM_
3136
3237 -- ** Full Folds
3338 , drain
@@ -65,7 +70,7 @@ module Streamly.Internal.Data.Fold
6570
6671 -- ** Partial Folds
6772 , drainN
68− , drainWhile
73+ , drainSepBy
6974 -- , lastN
7075 -- , (!!)
7176 -- , genericIndex
@@ -126,7 +131,7 @@ module Streamly.Internal.Data.Fold
126131 -- ** Trimming
127132 , ltake
128133 -- , lrunFor -- time
129− , ltakeWhile
134+ , takeSepBy
130135 {-
131136 , ltakeWhileM
132137 , ldrop
@@ -151,6 +156,8 @@ module Streamly.Internal.Data.Fold
151156 -- , breakAround
152157 , spanBy
153158 , spanByRolling
159+ , sliceSepBy
160+ , sliceSepWith
154161
155162 -- By sequences
156163 -- , breakOnSeq
@@ -164,33 +171,36 @@ module Streamly.Internal.Data.Fold
164171
165172 -- * Partitioning
166173
167− -- , partitionByM
168− -- , partitionBy
174+ , partitionByM
175+ , partitionByFstM
176+ , partitionByMinM
177+ , partitionBy
169178 , partition
170179
171180 -- * Demultiplexing
172181
173182 , demux
174− -- , demuxWith
183+ , demuxWith
175184 , demux_
176185 , demuxDefault_
177− -- , demuxWith_
186+ , demuxWith_
178187 , demuxWithDefault_
179188
180189 -- * Classifying
181190
182191 , classify
183− -- , classifyWith
192+ , classifyWith
184193
185194 -- * Unzipping
186195 , unzip
187− -- These can be expressed using lmap/lmapM and unzip
188− -- , unzipWith
189− -- , unzipWithM
196+ , unzipWith
197+ , unzipWithM
198+ , unzipWithFstM
199+ , unzipWithMinM
190200
191201 -- * Nested Folds
192202 -- , concatMap
193− , foldChunks
203+ , many
194204 , duplicate
195205
196206 -- * Running Folds
@@ -203,7 +213,7 @@ module Streamly.Internal.Data.Fold
203213 )
204214where
205215
206−import Control.Monad (join, void)
216+import Control.Monad (void, join)
207217import Control.Monad.IO.Class (MonadIO(..))
208218import Data.Functor.Identity (Identity(..))
209219import Data.Int (Int64)
@@ -214,11 +224,13 @@ import Data.Semigroup (Semigroup((<>)))
214224#endif
215225import Streamly.Internal.Data.Pipe.Types (Pipe (..), PipeState(..))
216226import Streamly.Internal.Data.Tuple.Strict (Tuple'(..), Tuple3'(..))
217−import Streamly.Internal.Data.Either.Strict (Either'(..))
218227import Streamly.Internal.Data.Maybe.Strict (Maybe'(..), toMaybe)
228+import Streamly.Internal.Data.Either.Strict (Either'(..))
219229
230+import qualified Streamly.Internal.Data.Fold.Types as FL
220231import qualified Streamly.Internal.Data.Pipe.Types as Pipe
221232import qualified Data.Map.Strict as Map
233+import qualified Data.Set as Set
222234import qualified Prelude
223235
224236import Prelude hiding
@@ -235,46 +247,86 @@ import Streamly.Internal.Data.Fold.Types
235247-- Smart constructors
236248------------------------------------------------------------------------------
237249
238−-- | Make a fold using a pure step function, a pure initial state and
239−-- a pure state extraction function.
250+-- | Make an accumulating (non-terminating) fold using a pure step function, a
251+-- pure initial state and a pure state extraction function.
240252--
241253-- /Internal/
242254--
243−{-# INLINE mkPure #-}
244−mkPure :: Monad m => (s -> a -> s) -> s -> (s -> b) -> Fold m a b
245−mkPure step initial extract =
255+{-# INLINE mkAccum #-}
256+mkAccum :: Monad m => (s -> a -> s) -> s -> (s -> b) -> Fold m a b
257+mkAccum step initial extract =
258+ Fold
259+ (\s a -> return $ Partial $ step s a)
260+ (return initial)
261+ (return . extract)
262+
263+-- | Similar to 'mkAccum' but the final state extracted is identical to the
264+-- intermediate state.
265+--
266+-- /Internal/
267+--
268+{-# INLINE mkAccum_ #-}
269+mkAccum_ :: Monad m => (b -> a -> b) -> b -> Fold m a b
270+mkAccum_ step initial = mkAccum step initial id
271+
272+-- | Make an accumulating (non-terminating) fold with an effectful step
273+-- function, an initial state, and a state extraction function.
274+--
275+-- /Internal/
276+--
277+{-# INLINE mkAccumM #-}
278+mkAccumM :: Functor m => (s -> a -> m s) -> m s -> (s -> m b) -> Fold m a b
279+mkAccumM step = Fold (\s a -> Partial <$> step s a)
280+
281+-- | Similar to 'mkAccumM' but the final state extracted is identical to the
282+-- intermediate state.
283+--
284+-- /Internal/
285+--
286+{-# INLINE mkAccumM_ #-}
287+mkAccumM_ :: Monad m => (b -> a -> m b) -> m b -> Fold m a b
288+mkAccumM_ step initial = mkAccumM step initial return
289+
290+-- | Make a terminating fold using a pure step function, a pure initial state
291+-- and a pure state extraction function.
292+--
293+-- /Internal/
294+--
295+{-# INLINE mkFold #-}
296+mkFold :: Monad m => (s -> a -> Step s b) -> s -> (s -> b) -> Fold m a b
297+mkFold step initial extract =
246298 Fold (\s a -> return $ step s a) (return initial) (return . extract)
247299
248−-- | Make a fold using a pure step function and a pure initial state. The
249−-- final state extracted is identical to the intermediate state.
300+-- | Similar to 'mkFold' but the final state extracted is identical to the
301+-- intermediate state.
250302--
251303-- /Internal/
252304--
253−{-# INLINE mkPureId #-}
254−mkPureId :: Monad m => (b -> a -> b) -> b -> Fold m a b
255−mkPureId step initial = mkPure step initial id
305+{-# INLINE mkFold_ #-}
306+mkFold_ :: Monad m => (b -> a -> Step b b) -> b -> Fold m a b
307+mkFold_ step initial = mkFold step initial id
256308
257−-- | Make a fold with an effectful step function and initial state, and a state
258−-- extraction function.
309+-- | Make a terminating fold with an effectful step function and initial state,
310+-- and a state extraction function.
259311--
260−-- > mkFold = Fold
312+-- > mkFoldM = Fold
261313--
262314-- We can just use 'Fold' but it is provided for completeness.
263315--
264316-- /Internal/
265317--
266−{-# INLINE mkFold #-}
267−mkFold :: (s -> a -> m s) -> m s -> (s -> m b) -> Fold m a b
268−mkFold = Fold
318+{-# INLINE mkFoldM #-}
319+mkFoldM :: (s -> a -> m (Step s b)) -> m s -> (s -> m b) -> Fold m a b
320+mkFoldM = Fold
269321
270−-- | Make a fold with an effectful step function and initial state. The final
271−-- state extracted is identical to the intermediate state.
322+-- | Similar to 'mkFoldM' but the final state extracted is identical to the
323+-- intermediate state.
272324--
273325-- /Internal/
274326--
275−{-# INLINE mkFoldId #-}
276−mkFoldId :: Monad m => (b -> a -> m b) -> m b -> Fold m a b
277−mkFoldId step initial = Fold step initial return
327+{-# INLINE mkFoldM_ #-}
328+mkFoldM_ :: Monad m => (b -> a -> m (Step b b)) -> m b -> Fold m a b
329+mkFoldM_ step initial = Fold step initial return
278330
279331------------------------------------------------------------------------------
280332-- hoist
@@ -304,9 +356,17 @@ generally = hoist (return . runIdentity)
304356-- @since 0.7.0
305357{-# INLINE sequence #-}
306358sequence :: Monad m => Fold m a (m b) -> Fold m a b
307−sequence (Fold step initial extract) = Fold step initial extract'
308− where
309− extract' x = join (extract x)
359+sequence (Fold step initial extract) = Fold step' initial extract'
360+
361+ where
362+
363+ step' s a = do
364+ res <- step s a
365+ case res of
366+ Partial x -> return $ Partial x
367+ Done b -> Done <$> b
368+
369+ extract' = join . extract
310370
311371-- | Map a monadic function on the output of a fold.
312372--
@@ -346,23 +406,28 @@ transform (Pipe pstep1 pstep2 pinitial) (Fold fstep finitial fextract) =
346406 where
347407
348408 initial = Tuple' pinitial <$> finitial
409+
349410 step (Tuple' ps fs) x = do
350411 r <- pstep1 ps x
351412 go fs r
352413
353414 where
415+
354416 -- XXX use SPEC?
355417 go acc (Pipe.Yield b (Consume ps')) = do
356418 acc' <- fstep acc b
357− return (Tuple' ps' acc')
358−
419+ return
420+ $ case acc' of
421+ Partial s -> Partial $ Tuple' ps' s
422+ Done b2 -> Done b2
359423 go acc (Pipe.Yield b (Produce ps')) = do
360424 acc' <- fstep acc b
361425 r <- pstep2 ps'
362− go acc' r
363−
364− go acc (Pipe.Continue (Consume ps')) = return (Tuple' ps' acc)
365−
426+ case acc' of
427+ Partial s -> go s r
428+ Done b2 -> return $ Done b2
429+ go acc (Pipe.Continue (Consume ps')) =
430+ return $ Partial $ Tuple' ps' acc
366431 go acc (Pipe.Continue (Produce ps')) = do
367432 r <- pstep2 ps'
368433 go acc r
@@ -380,11 +445,11 @@ transform (Pipe pstep1 pstep2 pinitial) (Fold fstep finitial fextract) =
380445{-# INLINABLE _Fold1 #-}
381446_Fold1 :: Monad m => (a -> a -> a) -> Fold m a (Maybe a)
382447_Fold1 step = Fold step_ (return Nothing') (return . toMaybe)
383− where
384− step_ mx a = return $ Just' $
385− case mx of
386− Nothing' -> a
387− Just' x -> step x a
448+
449+ where
450+
451+ step_ Nothing' a = return $ Partial $ Just' a
452+ step_ (Just' x) a = return $ Partial $ Just' $ step x a
388453
389454------------------------------------------------------------------------------
390455-- Left folds
@@ -403,9 +468,13 @@ _Fold1 step = Fold step_ (return Nothing') (return . toMaybe)
403468{-# INLINABLE drain #-}
404469drain :: Monad m => Fold m a ()
405470drain = Fold step begin done
471+
406472 where
473+
407474 begin = return ()
408− step _ _ = return ()
475+
476+ step _ _ = return $ FL.Partial ()
477+
409478 done = return
410479
411480-- |
@@ -418,7 +487,7 @@ drain = Fold step begin done
418487-- @since 0.7.0
419488{-# INLINABLE drainBy #-}
420489drainBy :: Monad m => (a -> m b) -> Fold m a ()
421−drainBy f = Fold (const (void . f)) (return ()) return
490+drainBy f = Fold (const (fmap FL.Partial . void . f)) (return ()) return
422491
423492{-# INLINABLE drainBy2 #-}
424493drainBy2 :: Monad m => (a -> m b) -> Fold2 m c a ()
@@ -442,16 +511,16 @@ last = _Fold1 (\_ x -> x)
442511-- > genericLength = fmap getSum $ foldMap (Sum . const 1)
443512--
444513-- @since 0.7.0
445−{-# INLINABLE genericLength #-}
514+{-# INLINE genericLength #-}
446515genericLength :: (Monad m, Num b) => Fold m a b
447−genericLength = Fold (\n _ -> return $ n + 1) (return 0) return
516+genericLength = Fold (\n _ -> return $ Partial $ n + 1) (return 0) return
448517
449518-- | Determine the length of the input stream.
450519--
451520-- > length = fmap getSum $ foldMap (Sum . const 1)
452521--
453522-- @since 0.7.0
454−{-# INLINABLE length #-}
523+{-# INLINE length #-}
455524length :: Monad m => Fold m a Int
456525length = genericLength
457526
@@ -464,17 +533,27 @@ length = genericLength
464533-- @since 0.7.0
465534{-# INLINE sum #-}
466535sum :: (Monad m, Num a) => Fold m a a
467−sum = Fold (\x a -> return $ x + a) (return 0) return
536+sum = Fold (\x a -> return $ Partial $ x + a) (return 0) return
468537
469538-- | Determine the product of all elements of a stream of numbers. Returns
470−-- multiplicative identity (@1@) when the stream is empty.
539+-- multiplicative identity (@1@) when the stream is empty. The fold terminates
540+-- when it encounters (@0@) in its input.
471541--
472542-- > product = fmap getProduct $ FL.foldMap Product
473543--
474544-- @since 0.7.0
475−{-# INLINABLE product #-}
476−product :: (Monad m, Num a) => Fold m a a
477−product = Fold (\x a -> return $ x * a) (return 1) return
545+-- /since 0.8.0 (Added 'Eq' constraint)/
546+{-# INLINE product #-}
547+product :: (Monad m, Num a, Eq a) => Fold m a a
548+product = Fold step (return 1) return
549+
550+ where
551+
552+ step x a =
553+ return
554+ $ if a == 0
555+ then Done 0
556+ else Partial $ x * a
478557
479558------------------------------------------------------------------------------
480559-- To Summary (Maybe)
@@ -484,13 +563,16 @@ product = Fold (\x a -> return $ x * a) (return 1) return
484563-- function.
485564--
486565-- @since 0.7.0
487−{-# INLINABLE maximumBy #-}
566+{-# INLINE maximumBy #-}
488567maximumBy :: Monad m => (a -> a -> Ordering) -> Fold m a (Maybe a)
489568maximumBy cmp = _Fold1 max'
490− where
491− max' x y = case cmp x y of
492− GT -> x
493− _ -> y
569+
570+ where
571+
572+ max' x y =
573+ case cmp x y of
574+ GT -> x
575+ _ -> y
494576
495577-- |
496578-- @
@@ -502,20 +584,23 @@ maximumBy cmp = _Fold1 max'
502584-- Compare with @FL.foldMap Max@.
503585--
504586-- @since 0.7.0
505−{-# INLINABLE maximum #-}
587+{-# INLINE maximum #-}
506588maximum :: (Monad m, Ord a) => Fold m a (Maybe a)
507589maximum = _Fold1 max
508590
509591-- | Computes the minimum element with respect to the given comparison function
510592--
511593-- @since 0.7.0
512−{-# INLINABLE minimumBy #-}
594+{-# INLINE minimumBy #-}
513595minimumBy :: Monad m => (a -> a -> Ordering) -> Fold m a (Maybe a)
514596minimumBy cmp = _Fold1 min'
515− where
516− min' x y = case cmp x y of
517− GT -> y
518− _ -> x
597+
598+ where
599+
600+ min' x y =
601+ case cmp x y of
602+ GT -> y
603+ _ -> x
519604
520605-- | Determine the minimum element in a stream using the supplied comparison
521606-- function.
@@ -527,7 +612,7 @@ minimumBy cmp = _Fold1 min'
527612-- Compare with @FL.foldMap Min@.
528613--
529614-- @since 0.7.0
530−{-# INLINABLE minimum #-}
615+{-# INLINE minimum #-}
531616minimum :: (Monad m, Ord a) => Fold m a (Maybe a)
532617minimum = _Fold1 min
533618
@@ -542,11 +627,16 @@ minimum = _Fold1 min
542627{-# INLINABLE mean #-}
543628mean :: (Monad m, Fractional a) => Fold m a a
544629mean = Fold step (return begin) (return . done)
545− where
630+
631+ where
632+
546633 begin = Tuple' 0 0
547− step (Tuple' x n) y = return $
548− let n' = n + 1
549− in Tuple' (x + (y - x) / n') n'
634+
635+ step (Tuple' x n) y =
636+ return
637+ $ let n' = n + 1
638+ in Partial $ Tuple' (x + (y - x) / n') n'
639+
550640 done (Tuple' x _) = x
551641
552642-- | Compute a numerically stable (population) variance over all elements in
@@ -556,15 +646,19 @@ mean = Fold step (return begin) (return . done)
556646{-# INLINABLE variance #-}
557647variance :: (Monad m, Fractional a) => Fold m a a
558648variance = Fold step (return begin) (return . done)
559− where
649+
650+ where
651+
560652 begin = Tuple3' 0 0 0
561653
562− step (Tuple3' n mean_ m2) x = return $ Tuple3' n' mean' m2'
563− where
564− n' = n + 1
565− mean' = (n * mean_ + x) / (n + 1)
566− delta = x - mean_
567− m2' = m2 + delta * delta * n / (n + 1)
654+ step (Tuple3' n mean_ m2) x = return $ Partial $ Tuple3' n' mean' m2'
655+
656+ where
657+
658+ n' = n + 1
659+ mean' = (n * mean_ + x) / (n + 1)
660+ delta = x - mean_
661+ m2' = m2 + delta * delta * n / (n + 1)
568662
569663 done (Tuple3' n _ m2) = m2 / n
570664
@@ -591,10 +685,15 @@ stdDev = sqrt variance
591685{-# INLINABLE rollingHashWithSalt #-}
592686rollingHashWithSalt :: (Monad m, Enum a) => Int64 -> Fold m a Int64
593687rollingHashWithSalt salt = Fold step initial extract
688+
594689 where
690+
595691 k = 2891336453 :: Int64
692+
596693 initial = return salt
597− step cksum a = return $ cksum * k + fromIntegral (fromEnum a)
694+
695+ step cksum a = return $ Partial $ cksum * k + fromIntegral (fromEnum a)
696+
598697 extract = return
599698
600699-- | A default salt used in the implementation of 'rollingHash'.
@@ -631,7 +730,7 @@ rollingHashFirstN n = ltake n rollingHash
631730--
632731{-# INLINE sconcat #-}
633732sconcat :: (Monad m, Semigroup a) => a -> Fold m a a
634−sconcat i = Fold (\x a -> return $ x <> a) (return i) return
733+sconcat i = Fold (\x a -> return $ Partial $ x <> a) (return i) return
635734
636735-- | Fold an input stream consisting of monoidal elements using 'mappend'
637736-- and 'mempty'.
@@ -640,11 +739,12 @@ sconcat i = Fold (\x a -> return $ x <> a) (return i) return
640739--
641740-- @since 0.7.0
642741{-# INLINE mconcat #-}
643−mconcat :: (Monad m, Monoid a
644−#if __GLASGOW_HASKELL__ < 804
742+mconcat ::
743+ ( Monad m
744+#if !MIN_VERSION_base(4,11,0)
645745 , Semigroup a
646746#endif
647− ) => Fold m a a
747+ , Monoid a) => Fold m a a
648748mconcat = sconcat mempty
649749
650750-- |
@@ -676,12 +776,16 @@ foldMap f = lmap f mconcat
676776{-# INLINABLE foldMapM #-}
677777foldMapM :: (Monad m, Monoid b) => (a -> m b) -> Fold m a b
678778foldMapM act = Fold step begin done
779+
679780 where
781+
680782 done = return
783+
681784 begin = return mempty
785+
682786 step m a = do
683787 m' <- act a
684− return $! mappend m m'
788+ return $! Partial $! mappend m m'
685789
686790------------------------------------------------------------------------------
687791-- To Containers
@@ -690,14 +794,15 @@ foldMapM act = Fold step begin done
690794-- | Folds the input stream to a list.
691795--
692796-- /Warning!/ working on large lists accumulated as buffers in memory could be
693−-- very inefficient, consider using "Streamly.Data.Array.Storable.Foreign" instead.
797+-- very inefficient, consider using "Streamly.Data.Array.Storable.Foreign"
798+-- instead.
694799--
695800-- @since 0.7.0
696801
697802-- id . (x1 :) . (x2 :) . (x3 :) . ... . (xn :) $ []
698803{-# INLINABLE toList #-}
699804toList :: Monad m => Fold m a [a]
700−toList = Fold (\f x -> return $ f . (x :))
805+toList = Fold (\f x -> return $ Partial $ f . (x :))
701806 (return id)
702807 (return . ($ []))
703808
@@ -713,9 +818,9 @@ drainN n = ltake n drain
713818
714819-- | A fold that drains elements of its input as long as the predicate succeeds,
715820-- running the effects and discarding the results.
716−{-# INLINABLE drainWhile #-}
717−drainWhile :: Monad m => (a -> Bool) -> Fold m a ()
718−drainWhile p = ltakeWhile p drain
821+{-# INLINABLE drainSepBy #-}
822+drainSepBy :: Monad m => (a -> Bool) -> Fold m a ()
823+drainSepBy p = takeSepBy p drain
719824
720825------------------------------------------------------------------------------
721826-- To Elements
@@ -726,18 +831,15 @@ drainWhile p = ltakeWhile p drain
726831-- @since 0.7.0
727832{-# INLINABLE genericIndex #-}
728833genericIndex :: (Integral i, Monad m) => i -> Fold m a (Maybe a)
729−genericIndex i = Fold step (return $ Left' 0) done
730− where
731− step x a = return $
732− case x of
733− Left' j -> if i == j
734− then Right' a
735− else Left' (j + 1)
736− _ -> x
737− done x = return $
738− case x of
739− Left' _ -> Nothing
740− Right' a -> Just a
834+genericIndex i = Fold step (return 0) (const (return Nothing))
835+
836+ where
837+
838+ step j a =
839+ return
840+ $ if i == j
841+ then Done $ Just a
842+ else Partial (j + 1)
741843
742844-- | Lookup the element at the given index.
743845--
@@ -765,14 +867,15 @@ head = _Fold1 const
765867-- @since 0.7.0
766868{-# INLINABLE find #-}
767869find :: Monad m => (a -> Bool) -> Fold m a (Maybe a)
768−find predicate = Fold step (return Nothing') (return . toMaybe)
769− where
770− step x a = return $
771− case x of
772− Nothing' -> if predicate a
773− then Just' a
774− else Nothing'
775− _ -> x
870+find predicate = Fold step (return ()) (const (return Nothing))
871+
872+ where
873+
874+ step () a =
875+ return
876+ $ if predicate a
877+ then Done (Just a)
878+ else Partial ()
776879
777880-- | In a stream of (key-value) pairs @(a, b)@, return the value @b@ of the
778881-- first pair where the key equals the given value @a@.
@@ -782,35 +885,30 @@ find predicate = Fold step (return Nothing') (return . toMaybe)
782885-- @since 0.7.0
783886{-# INLINABLE lookup #-}
784887lookup :: (Eq a, Monad m) => a -> Fold m (a,b) (Maybe b)
785−lookup a0 = Fold step (return Nothing') (return . toMaybe)
786− where
787− step x (a,b) = return $
788− case x of
789− Nothing' -> if a == a0
790− then Just' b
791− else Nothing'
792− _ -> x
888+lookup a0 = Fold step (return ()) (const (return Nothing))
793889
794−-- | Convert strict 'Either'' to lazy 'Maybe'
795−{-# INLINABLE hush #-}
796−hush :: Either' a b -> Maybe b
797−hush (Left' _) = Nothing
798−hush (Right' b) = Just b
890+ where
891+
892+ step () (a, b) =
893+ return
894+ $ if a == a0
895+ then Done $ Just b
896+ else Partial ()
799897
800898-- | Returns the first index that satisfies the given predicate.
801899--
802900-- @since 0.7.0
803901{-# INLINABLE findIndex #-}
804902findIndex :: Monad m => (a -> Bool) -> Fold m a (Maybe Int)
805−findIndex predicate = Fold step (return $ Left' 0) (return . hush)
806− where
807− step x a = return $
808− case x of
809− Left' i ->
810− if predicate a
811− then Right' i
812− else Left' (i + 1)
813− _ -> x
903+findIndex predicate = Fold step (return 0) (const (return Nothing))
904+
905+ where
906+
907+ step i a =
908+ return
909+ $ if predicate a
910+ then Done $ Just i
911+ else Partial (i + 1)
814912
815913-- | Returns the first index where a given value is found in the stream.
816914--
@@ -832,7 +930,7 @@ elemIndex a = findIndex (a ==)
832930-- @since 0.7.0
833931{-# INLINABLE null #-}
834932null :: Monad m => Fold m a Bool
835−null = Fold (\_ _ -> return False) (return True) return
933+null = Fold (\() _ -> return $ Done False) (return ()) (\() -> return True)
836934
837935-- |
838936-- > any p = lmap p or
@@ -841,9 +939,19 @@ null = Fold (\_ _ -> return False) (return True) return
841939-- | Returns 'True' if any of the elements of a stream satisfies a predicate.
842940--
843941-- @since 0.7.0
844−{-# INLINABLE any #-}
942+{-# INLINE any #-}
845943any :: Monad m => (a -> Bool) -> Fold m a Bool
846−any predicate = Fold (\x a -> return $ x || predicate a) (return False) return
944+any predicate = Fold step initial return
945+
946+ where
947+
948+ initial = return False
949+
950+ step _ a =
951+ return
952+ $ if predicate a
953+ then Done True
954+ else Partial False
847955
848956-- | Return 'True' if the given element is present in the stream.
849957--
@@ -863,7 +971,17 @@ elem a = any (a ==)
863971-- @since 0.7.0
864972{-# INLINABLE all #-}
865973all :: Monad m => (a -> Bool) -> Fold m a Bool
866−all predicate = Fold (\x a -> return $ x && predicate a) (return True) return
974+all predicate = Fold step initial return
975+
976+ where
977+
978+ initial = return True
979+
980+ step _ a =
981+ return
982+ $ if predicate a
983+ then Partial True
984+ else Done False
867985
868986-- | Returns 'True' if the given element is not present in the stream.
869987--
@@ -880,9 +998,9 @@ notElem a = all (a /=)
880998-- > and = fmap getAll . FL.foldMap All
881999--
8821000-- @since 0.7.0
883−{-# INLINABLE and #-}
1001+{-# INLINE and #-}
8841002and :: Monad m => Fold m Bool Bool
885−and = Fold (\x a -> return $ x && a) (return True) return
1003+and = all (== True)
8861004
8871005-- | Returns 'True' if any element is 'True', 'False' otherwise
8881006--
@@ -890,9 +1008,9 @@ and = Fold (\x a -> return $ x && a) (return True) return
8901008-- > or = fmap getAny . FL.foldMap Any
8911009--
8921010-- @since 0.7.0
893−{-# INLINABLE or #-}
1011+{-# INLINE or #-}
8941012or :: Monad m => Fold m Bool Bool
895−or = Fold (\x a -> return $ x || a) (return False) return
1013+or = any (== True)
8961014
8971015------------------------------------------------------------------------------
8981016-- Grouping/Splitting
@@ -905,7 +1023,7 @@ or = Fold (\x a -> return $ x || a) (return False) return
9051023------------------------------------------------------------------------------
9061024-- Binary APIs
9071025------------------------------------------------------------------------------
908−--
1026+
9091027-- XXX These would just be applicative compositions of terminating folds.
9101028
9111029-- | @splitAt n f1 f2@ composes folds @f1@ and @f2@ such that first @n@
@@ -944,17 +1062,7 @@ splitAt
9441062 -> Fold m a b
9451063 -> Fold m a c
9461064 -> Fold m a (b, c)
947−splitAt n (Fold stepL initialL extractL) (Fold stepR initialR extractR) =
948− Fold step initial extract
949− where
950− initial = Tuple3' n <$> initialL <*> initialR
951−
952− step (Tuple3' i xL xR) input =
953− if i > 0
954− then stepL xL input >>= (\a -> return (Tuple3' (i - 1) a xR))
955− else stepR xR input >>= (return . Tuple3' i xL)
956−
957− extract (Tuple3' _ a b) = (,) <$> extractL a <*> extractR b
1065+splitAt n fld1 = splitWith (,) (ltake n fld1)
9581066
9591067------------------------------------------------------------------------------
9601068-- Element Aware APIs
@@ -964,12 +1072,71 @@ splitAt n (Fold stepL initialL extractL) (Fold stepR initialR extractR) =
9641072-- Binary APIs
9651073------------------------------------------------------------------------------
9661074
1075+-- | Stops the fold at an infixed separator element, dropping the separator.
1076+--
1077+-- @
1078+-- "--.--" => "--" "--"
1079+-- "--." => "--" ""
1080+-- ".--" => "" "--"
1081+-- @
1082+--
1083+-- * Stops - when the predicate succeeds.
1084+--
1085+-- /Internal/
1086+{-# INLINE sliceSepBy #-}
1087+sliceSepBy :: Monad m => (a -> Bool) -> Fold m a b -> Fold m a b
1088+sliceSepBy predicate (Fold fstep finitial fextract) = Fold step initial fextract
1089+
1090+ where
1091+
1092+ initial = finitial
1093+
1094+ step s a =
1095+ if not (predicate a)
1096+ then fstep s a
1097+ else Done <$> fextract s
1098+
1099+-- | Like 'sliceSepBy' but does not drop the seperator element.
1100+--
1101+-- @
1102+-- "--.--" => "--." "--"
1103+-- "--." => "--." ""
1104+-- ".--" => "." "--"
1105+-- @
1106+--
1107+-- * Stops - when the predicate succeeds.
1108+--
1109+-- /Internal/
1110+{-# INLINE sliceSepWith #-}
1111+sliceSepWith :: Monad m => (a -> Bool) -> Fold m a b -> Fold m a b
1112+sliceSepWith predicate (Fold fstep finitial fextract) =
1113+ Fold step initial fextract
1114+
1115+ where
1116+
1117+ initial = finitial
1118+
1119+ step s a =
1120+ if not (predicate a)
1121+ then fstep s a
1122+ else do
1123+ res <- fstep s a
1124+ case res of
1125+ Partial sres -> Done <$> fextract sres
1126+ Done bres -> return $ Done bres
1127+
1128+data SpanByState a bl fl fr
1129+ = SpanByLeft0 !fl !fr
1130+ | SpanByLeft !a !fl !fr
1131+ | SpanByRight !bl !fr
1132+
9671133-- | Break the input stream into two groups, the first group takes the input as
9681134-- long as the predicate applied to the first element of the stream and next
9691135-- input element holds 'True', the second group takes the rest of the input.
9701136--
9711137-- /Internal/
9721138--
1139+{-# INLINE spanBy #-}
9731140spanBy
9741141 :: Monad m
9751142 => (a -> a -> Bool)
@@ -980,29 +1147,52 @@ spanBy cmp (Fold stepL initialL extractL) (Fold stepR initialR extractR) =
9801147 Fold step initial extract
9811148
9821149 where
983− initial = Tuple3' <$> initialL <*> initialR <*> return (Tuple' Nothing True)
9841150
985− step (Tuple3' a b (Tuple' (Just frst) isFirstG)) input =
986− if cmp frst input && isFirstG
987− then stepL a input
988− >>= (\a' -> return (Tuple3' a' b (Tuple' (Just frst) isFirstG)))
989− else stepR b input
990− >>= (\a' -> return (Tuple3' a a' (Tuple' Nothing False)))
991−
992− step (Tuple3' a b (Tuple' Nothing isFirstG)) input =
993− if isFirstG
994− then stepL a input
995− >>= (\a' -> return (Tuple3' a' b (Tuple' (Just input) isFirstG)))
996− else stepR b input
997− >>= (\a' -> return (Tuple3' a a' (Tuple' Nothing False)))
998−
999− extract (Tuple3' a b _) = (,) <$> extractL a <*> extractR b
1151+ initial = SpanByLeft0 <$> initialL <*> initialR
1152+
1153+ step (SpanByLeft0 fl fr) input = do
1154+ sfl <- stepL fl input
1155+ return
1156+ $ Partial
1157+ $ case sfl of
1158+ Partial fl1 -> SpanByLeft input fl1 fr
1159+ Done bl -> SpanByRight bl fr
1160+ step (SpanByLeft frst fl fr) input =
1161+ if cmp frst input
1162+ then do
1163+ sfl <- stepL fl input
1164+ return
1165+ $ Partial
1166+ $ case sfl of
1167+ Partial fl1 -> SpanByLeft frst fl1 fr
1168+ Done bl -> SpanByRight bl fr
1169+ else do
1170+ bl <- extractL fl
1171+ sfr <- stepR fr input
1172+ return
1173+ $ case sfr of
1174+ Partial fr1 -> Partial $ SpanByRight bl fr1
1175+ Done br -> Done (bl, br)
1176+ step (SpanByRight bl fr) input = do
1177+ sfr <- stepR fr input
1178+ return
1179+ $ case sfr of
1180+ Partial fr1 -> Partial $ SpanByRight bl fr1
1181+ Done br -> Done (bl, br)
1182+
1183+ extract (SpanByLeft0 fl fr) = (,) <$> extractL fl <*> extractR fr
1184+ extract (SpanByLeft _ fl fr) = (,) <$> extractL fl <*> extractR fr
1185+ extract (SpanByRight bl fr) = (bl, ) <$> extractR fr
1186+
1187+data SpanState bl fl fr
1188+ = SpanLeft !fl !fr
1189+ | SpanRight !bl !fr
10001190
10011191-- | @span p f1 f2@ composes folds @f1@ and @f2@ such that @f1@ consumes the
10021192-- input as long as the predicate @p@ is 'True'. @f2@ consumes the rest of the
10031193-- input.
10041194--
1005−-- > let span_ p xs = S.fold (S.span p FL.toList FL.toList) $ S.fromList xs
1195+-- > let span_ p xs = S.fold (FL.span p FL.toList FL.toList) $ S.fromList xs
10061196--
10071197-- >>> span_ (< 1) [1,2,3]
10081198-- > ([],[1,2,3])
@@ -1029,14 +1219,33 @@ span p (Fold stepL initialL extractL) (Fold stepR initialR extractR) =
10291219
10301220 where
10311221
1032− initial = Tuple3' <$> initialL <*> initialR <*> return True
1033−
1034− step (Tuple3' a b isFirstG) input =
1035− if isFirstG && p input
1036− then stepL a input >>= (\a' -> return (Tuple3' a' b True))
1037− else stepR b input >>= (\a' -> return (Tuple3' a a' False))
1222+ initial = SpanLeft <$> initialL <*> initialR
10381223
1039− extract (Tuple3' a b _) = (,) <$> extractL a <*> extractR b
1224+ step (SpanLeft fl fr) input =
1225+ if p input
1226+ then do
1227+ sfl <- stepL fl input
1228+ return
1229+ $ Partial
1230+ $ case sfl of
1231+ Partial fl1 -> SpanLeft fl1 fr
1232+ Done bl -> SpanRight bl fr
1233+ else do
1234+ bl <- extractL fl
1235+ sfr <- stepR fr input
1236+ return
1237+ $ case sfr of
1238+ Partial fr1 -> Partial $ SpanRight bl fr1
1239+ Done br -> Done (bl, br)
1240+ step (SpanRight bl fr) input = do
1241+ sfr <- stepR fr input
1242+ return
1243+ $ case sfr of
1244+ Partial fr1 -> Partial $ SpanRight bl fr1
1245+ Done br -> Done (bl, br)
1246+
1247+ extract (SpanLeft fl fr) = (,) <$> extractL fl <*> extractR fr
1248+ extract (SpanRight bl fr) = (bl, ) <$> extractR fr
10401249
10411250-- |
10421251-- > break p = span (not . p)
@@ -1082,18 +1291,43 @@ spanByRolling
10821291spanByRolling cmp (Fold stepL initialL extractL) (Fold stepR initialR extractR) =
10831292 Fold step initial extract
10841293
1085− where
1086− initial = Tuple3' <$> initialL <*> initialR <*> return Nothing
1087−
1088− step (Tuple3' a b (Just frst)) input =
1089− if cmp input frst
1090− then stepL a input >>= (\a' -> return (Tuple3' a' b (Just input)))
1091− else stepR b input >>= (\b' -> return (Tuple3' a b' (Just input)))
1092−
1093− step (Tuple3' a b Nothing) input =
1094− stepL a input >>= (\a' -> return (Tuple3' a' b (Just input)))
1294+ where
10951295
1096− extract (Tuple3' a b _) = (,) <$> extractL a <*> extractR b
1296+ initial = SpanByLeft0 <$> initialL <*> initialR
1297+
1298+ step (SpanByLeft0 fl fr) input = do
1299+ sfl <- stepL fl input
1300+ return
1301+ $ Partial
1302+ $ case sfl of
1303+ Partial fl1 -> SpanByLeft input fl1 fr
1304+ Done bl -> SpanByRight bl fr
1305+ step (SpanByLeft frst fl fr) input =
1306+ if cmp frst input
1307+ then do
1308+ sfl <- stepL fl input
1309+ return
1310+ $ Partial
1311+ $ case sfl of
1312+ Partial fl1 -> SpanByLeft input fl1 fr
1313+ Done bl -> SpanByRight bl fr
1314+ else do
1315+ bl <- extractL fl
1316+ sfr <- stepR fr input
1317+ return
1318+ $ case sfr of
1319+ Partial fr1 -> Partial $ SpanByRight bl fr1
1320+ Done br -> Done (bl, br)
1321+ step (SpanByRight bl fr) input = do
1322+ sfr <- stepR fr input
1323+ return
1324+ $ case sfr of
1325+ Partial fr1 -> Partial $ SpanByRight bl fr1
1326+ Done br -> Done (bl, br)
1327+
1328+ extract (SpanByLeft0 fl fr) = (,) <$> extractL fl <*> extractR fr
1329+ extract (SpanByLeft _ fl fr) = (,) <$> extractL fl <*> extractR fr
1330+ extract (SpanByRight bl fr) = (bl, ) <$> extractR fr
10971331
10981332------------------------------------------------------------------------------
10991333-- Binary splitting on a separator
@@ -1132,25 +1366,19 @@ breakOn pat f m = undefined
11321366-- @since 0.7.0
11331367{-# INLINE tee #-}
11341368tee :: Monad m => Fold m a b -> Fold m a c -> Fold m a (b,c)
1135−tee f1 f2 = (,) <$> f1 <*> f2
1369+tee = teeWith (,)
11361370
11371371{-# INLINE foldNil #-}
11381372foldNil :: Monad m => Fold m a [b]
11391373foldNil = Fold step begin done where
11401374 begin = return []
1141− step _ _ = return []
1375+ step _ _ = return $ Partial []
11421376 done = return
11431377
1378+-- XXX How is the performance?
11441379{-# INLINE foldCons #-}
11451380foldCons :: Monad m => Fold m a b -> Fold m a [b] -> Fold m a [b]
1146−foldCons (Fold stepL beginL doneL) (Fold stepR beginR doneR) =
1147− Fold step begin done
1148−
1149− where
1150−
1151− begin = Tuple' <$> beginL <*> beginR
1152− step (Tuple' xL xR) a = Tuple' <$> stepL xL a <*> stepR xR a
1153− done (Tuple' xL xR) = (:) <$> doneL xL <*> doneR xR
1381+foldCons = teeWith (:)
11541382
11551383-- XXX use "List" instead of "[]"?, use Array for output to scale it to a large
11561384-- number of consumers? For polymorphic case a vector could be helpful. For
@@ -1178,27 +1406,53 @@ foldCons (Fold stepL beginL doneL) (Fold stepR beginR doneR) =
11781406distribute :: Monad m => [Fold m a b] -> Fold m a [b]
11791407distribute = foldr foldCons foldNil
11801408
1409+-- XXX This is 2x times faster wiithout the terminating condition.
11811410-- | Like 'distribute' but for folds that return (), this can be more efficient
1182−-- than 'distribute' as it does not need to maintain state.
1411+-- than 'distribute' as it does not need to maintain state. This fold terminates
1412+-- when all the folds terminate.
11831413--
11841414{-# INLINE distribute_ #-}
11851415distribute_ :: Monad m => [Fold m a ()] -> Fold m a ()
11861416distribute_ fs = Fold step initial extract
1417+
11871418 where
1188− initial = Prelude.mapM (\(Fold s i e) ->
1189− i >>= \r -> return (Fold s (return r) e)) fs
1190− step ss a = do
1191− Prelude.mapM_ (\(Fold s i _) -> i >>= \r -> void (s r a)) ss
1192− return ss
1193− extract =
1194− Prelude.mapM_ (\(Fold _ i e) -> i >>= \r -> e r)
1419+
1420+ initial = Prelude.mapM initialize fs
1421+
1422+ -- XXX This itself can be a recursive function but we might face inlining
1423+ -- problems.
1424+ step [] _ = return $ Done ()
1425+ step ss a = do
1426+ !ss1 <- go ss a
1427+ return
1428+ $ if Prelude.null ss1
1429+ then Done ()
1430+ else Partial ss1
1431+
1432+ go [] _ = return []
1433+ go ((Fold s i d):ss) a = do
1434+ i1 <- i
1435+ res <- s i1 a
1436+ case res of
1437+ Partial i2 -> do
1438+ rst <- go ss a
1439+ let fld = Fold s (return i2) d
1440+ return $ fld : rst
1441+ Done () -> go ss a
1442+
1443+ extract ss = go1 ss
1444+
1445+ go1 [] = return ()
1446+ go1 ((Fold _ i d):ss) = (i >>= d) >> go1 ss
11951447
11961448------------------------------------------------------------------------------
11971449-- Partitioning
11981450------------------------------------------------------------------------------
1199−--
1451+
12001452-- | Partition the input over two folds using an 'Either' partitioning
1201−-- predicate.
1453+-- predicate. This fold terminates when both the folds terminate.
1454+--
1455+-- See also: 'partitionByFstM' and 'partitionByMinM'.
12021456--
12031457-- @
12041458--
@@ -1241,23 +1495,85 @@ distribute_ fs = Fold step initial extract
12411495partitionByM :: Monad m
12421496 => (a -> m (Either b c)) -> Fold m b x -> Fold m c y -> Fold m a (x, y)
12431497partitionByM f (Fold stepL beginL doneL) (Fold stepR beginR doneR) =
1244−
12451498 Fold step begin done
12461499
12471500 where
12481501
1249− begin = Tuple' <$> beginL <*> beginR
1250− step (Tuple' xL xR) a = do
1502+ begin = do
1503+ sL <- beginL
1504+ sR <- beginR
1505+ return $ RunBoth sL sR
1506+
1507+ step (RunBoth sL sR) a = do
1508+ r <- f a
1509+ case r of
1510+ Left b -> do
1511+ res <- stepL sL b
1512+ return
1513+ $ Partial
1514+ $ case res of
1515+ Partial sres -> RunBoth sres sR
1516+ Done bres -> RunRight bres sR
1517+ Right c -> do
1518+ res <- stepR sR c
1519+ return
1520+ $ Partial
1521+ $ case res of
1522+ Partial sres -> RunBoth sL sres
1523+ Done bres -> RunLeft sL bres
1524+ step (RunLeft sL bR) a = do
1525+ r <- f a
1526+ case r of
1527+ Left b -> do
1528+ res <- stepL sL b
1529+ return
1530+ $ case res of
1531+ Partial sres -> Partial $ RunLeft sres bR
1532+ Done bres -> Done (bres, bR)
1533+ Right _ -> return $ Partial $ RunLeft sL bR
1534+ step (RunRight bL sR) a = do
12511535 r <- f a
12521536 case r of
1253− Left b -> Tuple' <$> stepL xL b <*> return xR
1254− Right c -> Tuple' xL <$> stepR xR c
1255− done (Tuple' xL xR) = (,) <$> doneL xL <*> doneR xR
1537+ Left _ -> return $ Partial $ RunRight bL sR
1538+ Right c -> do
1539+ res <- stepR sR c
1540+ return
1541+ $ case res of
1542+ Partial sres -> Partial $ RunRight bL sres
1543+ Done bres -> Done (bL, bres)
1544+
1545+ done (RunBoth sL sR) = do
1546+ bL <- doneL sL
1547+ bR <- doneR sR
1548+ return (bL, bR)
1549+ done (RunLeft sL bR) = do
1550+ bL <- doneL sL
1551+ return (bL, bR)
1552+ done (RunRight bL sR) = do
1553+ bR <- doneR sR
1554+ return (bL, bR)
1555+
1556+-- | Similar to 'partitionByM' but terminates when the first fold terminates.
1557+--
1558+-- /Unimplemented/
1559+--
1560+{-# INLINE partitionByFstM #-}
1561+partitionByFstM :: -- Monad m =>
1562+ (a -> m (Either b c)) -> Fold m b x -> Fold m c y -> Fold m a (x, y)
1563+partitionByFstM = undefined
1564+
1565+-- | Similar to 'partitionByM' but terminates when any fold terminates.
1566+--
1567+-- /Unimplemented/
1568+--
1569+{-# INLINE partitionByMinM #-}
1570+partitionByMinM :: -- Monad m =>
1571+ (a -> m (Either b c)) -> Fold m b x -> Fold m c y -> Fold m a (x, y)
1572+partitionByMinM = undefined
12561573
12571574-- Note: we could use (a -> Bool) instead of (a -> Either b c), but the latter
12581575-- makes the signature clearer as to which case belongs to which fold.
12591576-- XXX need to check the performance in both cases.
1260−
12611577-- | Same as 'partitionByM' but with a pure partition function.
12621578--
12631579-- Count even and odd numbers in a stream:
@@ -1331,31 +1647,36 @@ demuxWith f kv = Fold step initial extract
13311647
13321648 where
13331649
1334− initial = return kv
1335−-- alterF is available only since containers version 0.5.8.2
1336−#if MIN_VERSION_containers(0,5,8)
1337− step mp a = case f a of
1338− (k, a') -> Map.alterF twiddle k mp
1339− -- XXX should we raise an exception in Nothing case?
1340− -- Ideally we should enforce that it is a total map over k so that look
1341− -- up never fails
1342− -- XXX we could use a monadic update function for a single lookup and
1343− -- update in the map.
1344− where
1345− twiddle Nothing = pure Nothing
1346− twiddle (Just (Fold step' acc extract')) = do
1347− !r <- acc >>= \x -> step' x a'
1348− pure . Just $ Fold step' (return r) extract'
1349−#else
1350− step mp a =
1351− let (k, a') = f a
1352− in case Map.lookup k mp of
1353− Nothing -> return mp
1354− Just (Fold step' acc extract') -> do
1355− !r <- acc >>= \x -> step' x a'
1356− return $ Map.insert k (Fold step' (return r) extract') mp
1357−#endif
1358− extract = Prelude.mapM (\(Fold _ acc e) -> acc >>= e)
1650+ initial = Tuple' Map.empty <$> Prelude.mapM initialize kv
1651+
1652+ step (Tuple' bmp mp) a =
1653+ if Map.size mp > 0
1654+ then do
1655+ let (k, a') = f a
1656+ case Map.lookup k mp of
1657+ Nothing -> return $ Partial $ Tuple' bmp mp
1658+ Just (Fold stp ini dn) -> do
1659+ !st <- ini
1660+ !res <- stp st a'
1661+ return
1662+ $ case res of
1663+ Partial sres ->
1664+ let fld = Fold stp (return sres) dn
1665+ mp1 = Map.insert k fld mp
1666+ in Partial $ Tuple' bmp mp1
1667+ Done bres ->
1668+ let mp1 = Map.delete k mp
1669+ bmp1 = Map.insert k bres bmp
1670+ in if Map.size mp1 == 0
1671+ then Done bmp1
1672+ else Partial $ Tuple' bmp1 mp1
1673+ -- XXX This state is reached even for an empty Map. That will change
1674+ -- once we have Step in the initial value
1675+ else error "Illegal state"
1676+
1677+ extract (Tuple' bmp mp) = do
1678+ mpe <- Prelude.mapM (\(Fold _ i d) -> i >>= d) mp
1679+ return $ bmp `Map.union` mpe
13591680
13601681-- | Fold a stream of key value pairs using a map of specific folds for each
13611682-- key into a map from keys to the results of fold outputs of the corresponding
@@ -1374,6 +1695,9 @@ demux :: (Monad m, Ord k)
13741695 => Map k (Fold m a b) -> Fold m (k, a) (Map k b)
13751696demux = demuxWith id
13761697
1698+-- data DemuxState m s = DemuxP !m !s | DemuxingOnlyMap !m
1699+data DemuxState s m1 m2 = DemuxingWithDefault !s !m1 !m2 | DemuxingOnlyMap !m2
1700+
13771701{-# INLINE demuxWithDefault_ #-}
13781702demuxWithDefault_ :: (Monad m, Ord k)
13791703 => (a -> (k, a')) -> Map k (Fold m a' b) -> Fold m (k, a') b -> Fold m a ()
@@ -1382,23 +1706,82 @@ demuxWithDefault_ f kv (Fold dstep dinitial dextract) =
13821706
13831707 where
13841708
1385− initFold (Fold s i e) = i >>= \r -> return (Fold s (return r) e)
13861709 initial = do
1387− mp <- Prelude.mapM initFold kv
1388− Tuple' mp <$> dinitial
1710+ ini <- dinitial
1711+ mp <- Prelude.mapM initialize kv
1712+ return $ DemuxingWithDefault ini Set.empty mp
13891713
1390− step (Tuple' mp dacc) a
1391− | (k, a') <- f a
1392− = case Map.lookup k mp of
1393− Nothing -> do
1394− acc <- dstep dacc (k, a')
1395− return (Tuple' mp acc)
1396− Just (Fold step' acc _) -> do
1397− _ <- acc >>= \x -> step' x a'
1398− return (Tuple' mp dacc)
1399− extract (Tuple' mp dacc) = do
1714+ -- XXX Some code can probably be abstracted
1715+ step (DemuxingWithDefault dacc bst mp) a =
1716+ if Map.size mp > 0
1717+ then do
1718+ let (k, a1) = f a
1719+ case Map.lookup k mp of
1720+ Nothing ->
1721+ if k `Set.member` bst
1722+ then return $ Partial $ DemuxingWithDefault dacc bst mp
1723+ else do
1724+ res <- dstep dacc (k, a1)
1725+ return
1726+ $ Partial
1727+ $ case res of
1728+ Partial sres ->
1729+ DemuxingWithDefault sres bst mp
1730+ Done _ -> DemuxingOnlyMap mp
1731+ Just (Fold stp ini dn) -> do
1732+ !st <- ini
1733+ !res <- stp st a1
1734+ return
1735+ $ case res of
1736+ Partial sres ->
1737+ let fld = Fold stp (return sres) dn
1738+ mp1 = Map.insert k fld mp
1739+ in Partial $ DemuxingWithDefault dacc bst mp1
1740+ Done _ ->
1741+ -- XXX We can skip the check here
1742+ if Map.size mp == 1
1743+ then Done ()
1744+ else Partial
1745+ $ DemuxingWithDefault
1746+ dacc
1747+ (Set.insert k bst)
1748+ (Map.delete k mp)
1749+ -- XXX This state is reached even for an empty Map. That will change
1750+ -- once we have Step in the initial value
1751+ else error "Illegal state"
1752+ -- XXX Reduce code duplication?
1753+ step (DemuxingOnlyMap mp) a =
1754+ if Map.size mp > 0
1755+ then do
1756+ let (k, a1) = f a
1757+ case Map.lookup k mp of
1758+ Nothing -> return $ Partial $ DemuxingOnlyMap mp
1759+ Just (Fold stp ini dn) -> do
1760+ !st <- ini
1761+ !res <- stp st a1
1762+ return
1763+ $ case res of
1764+ Partial sres ->
1765+ Partial
1766+ $ DemuxingOnlyMap
1767+ $ Map.insert
1768+ k
1769+ (Fold stp (return sres) dn)
1770+ mp
1771+ Done _ ->
1772+ if Map.size mp == 1
1773+ then Done ()
1774+ else Partial
1775+ $ DemuxingOnlyMap $ Map.delete k mp
1776+ -- XXX This state is reached even for an empty Map. That will change
1777+ -- once we have Step in the initial value
1778+ else error "Illegal state"
1779+
1780+ extract (DemuxingWithDefault dacc _ mp) = do
14001781 void $ dextract dacc
1401− Prelude.mapM_ (\(Fold _ acc e) -> acc >>= e) mp
1782+ Prelude.mapM_ (\(Fold _ i d) -> i >>= d) mp
1783+ extract (DemuxingOnlyMap mp) = do
1784+ Prelude.mapM_ (\(Fold _ i d) -> i >>= d) mp
14021785
14031786-- | Split the input stream based on a key field and fold each split using a
14041787-- specific fold without collecting the results. Useful for cases like protocol
@@ -1426,20 +1809,28 @@ demuxWith_ f kv = Fold step initial extract
14261809
14271810 where
14281811
1429− initial =
1430− Prelude.mapM (\(Fold s i e) ->
1431− i >>= \r -> return (Fold s (return r) e)) kv
1432− step mp a
1433− -- XXX should we raise an exception in Nothing case?
1434− -- Ideally we should enforce that it is a total map over k so that look
1435− -- up never fails
1436− | (k, a') <- f a
1437− = case Map.lookup k mp of
1438− Nothing -> return mp
1439− Just (Fold step' acc _) -> do
1440− _ <- acc >>= \x -> step' x a'
1441− return mp
1442− extract = Prelude.mapM_ (\(Fold _ acc e) -> acc >>= e)
1812+ initial = Tuple' (Map.size kv) <$> Prelude.mapM initialize kv
1813+
1814+ step (Tuple' n mp) a
1815+ | (k, a') <- f a =
1816+ case Map.lookup k mp of
1817+ Nothing -> return $ Partial $ Tuple' n mp
1818+ Just (Fold stp ini dn) -> do
1819+ !st <- ini
1820+ !res <- stp st a'
1821+ let n1 = n - 1
1822+ return
1823+ $ case res of
1824+ Partial sres ->
1825+ let fld = Fold stp (return sres) dn
1826+ in Partial $ Tuple' n $ Map.insert k fld mp
1827+ Done _ ->
1828+ if n1 == 0
1829+ then Done ()
1830+ else Partial $ Tuple' n1 $ Map.delete k mp
1831+
1832+ extract (Tuple' _ mp) = do
1833+ Prelude.mapM_ (\(Fold _ i d) -> i >>= d) mp
14431834
14441835-- | Given a stream of key value pairs and a map from keys to folds, fold the
14451836-- values for each key using the corresponding folds, discarding the outputs.
@@ -1485,17 +1876,39 @@ classifyWith f (Fold step initial extract) = Fold step' initial' extract'
14851876 where
14861877
14871878 initial' = return Map.empty
1879+
14881880 step' kv a =
1489− let k = f a
1490− in case Map.lookup k kv of
1881+ case Map.lookup k kv of
14911882 Nothing -> do
14921883 x <- initial
14931884 r <- step x a
1494− return $ Map.insert k r kv
1885+ return
1886+ $ Partial
1887+ $ flip (Map.insert k) kv
1888+ $ case r of
1889+ Partial sr -> Left' sr
1890+ Done br -> Right' br
14951891 Just x -> do
1496− r <- step x a
1497− return $ Map.insert k r kv
1498− extract' = Prelude.mapM extract
1892+ case x of
1893+ Left' s -> do
1894+ r <- step s a
1895+ return
1896+ $ Partial
1897+ $ flip (Map.insert k) kv
1898+ $ case r of
1899+ Partial sr -> Left' sr
1900+ Done br -> Right' br
1901+ Right' _ -> return $ Partial kv
1902+
1903+ where
1904+
1905+ k = f a
1906+
1907+ extract' =
1908+ Prelude.mapM
1909+ (\case
1910+ Left' s -> extract s
1911+ Right' b -> return b)
14991912
15001913-- | Given an input stream of key value pairs and a fold for values, fold all
15011914-- the values belonging to each key. Useful for map/reduce, bucketizing the
@@ -1520,9 +1933,11 @@ classify fld = classifyWith fst (lmap snd fld)
15201933------------------------------------------------------------------------------
15211934-- Unzipping
15221935------------------------------------------------------------------------------
1523−--
1936+
15241937-- | Like 'unzipWith' but with a monadic splitter function.
15251938--
1939+-- -- @unzipWithM f fld1 fld2 = lmapM f (unzip fld1 fld2)@
1940+--
15261941-- @since 0.7.0
15271942{-# INLINE unzipWithM #-}
15281943unzipWithM :: Monad m
@@ -1532,15 +1947,68 @@ unzipWithM f (Fold stepL beginL doneL) (Fold stepR beginR doneR) =
15321947
15331948 where
15341949
1535− step (Tuple' xL xR) a = do
1536− (b,c) <- f a
1537− Tuple' <$> stepL xL b <*> stepR xR c
1538− begin = Tuple' <$> beginL <*> beginR
1539− done (Tuple' xL xR) = (,) <$> doneL xL <*> doneR xR
1950+ begin = RunBoth <$> beginL <*> beginR
1951+
1952+ step (RunBoth sL sR) a = do
1953+ (b, c) <- f a
1954+ resL <- stepL sL b
1955+ resR <- stepR sR c
1956+ case resL of
1957+ Partial sresL ->
1958+ return
1959+ $ Partial
1960+ $ case resR of
1961+ Partial sresR -> RunBoth sresL sresR
1962+ Done bresR -> RunLeft sresL bresR
1963+ Done bresL ->
1964+ return
1965+ $ case resR of
1966+ Partial sresR -> Partial $ RunRight bresL sresR
1967+ Done bresR -> Done (bresL, bresR)
1968+ step (RunLeft sL bR) a = do
1969+ (b, _) <- f a
1970+ resL <- stepL sL b
1971+ return
1972+ $ case resL of
1973+ Partial sresL -> Partial $ RunLeft sresL bR
1974+ Done bresL -> Done (bresL, bR)
1975+ step (RunRight bL sR) a = do
1976+ (_, c) <- f a
1977+ resR <- stepR sR c
1978+ return
1979+ $ case resR of
1980+ Partial sresR -> Partial $ RunRight bL sresR
1981+ Done bresR -> Done (bL, bresR)
1982+
1983+ done (RunBoth sL sR) = (,) <$> doneL sL <*> doneR sR
1984+ done (RunLeft sL bR) = (,) <$> doneL sL <*> return bR
1985+ done (RunRight bL sR) = (,) bL <$> doneR sR
1986+
1987+-- | Similar to 'unzipWithM' but terminates when the first fold terminates.
1988+--
1989+-- /Unimplemented/
1990+--
1991+{-# INLINE unzipWithFstM #-}
1992+unzipWithFstM :: -- Monad m =>
1993+ (a -> m (b,c)) -> Fold m b x -> Fold m c y -> Fold m a (x,y)
1994+unzipWithFstM = undefined
1995+
1996+-- | Similar to 'unzipWithM' but terminates when any fold terminates.
1997+--
1998+-- /Unimplemented/
1999+--
2000+{-# INLINE unzipWithMinM #-}
2001+unzipWithMinM :: -- Monad m =>
2002+ (a -> m (b,c)) -> Fold m b x -> Fold m c y -> Fold m a (x,y)
2003+unzipWithMinM = undefined
15402004
15412005-- | Split elements in the input stream into two parts using a pure splitter
15422006-- function, direct each part to a different fold and zip the results.
15432007--
2008+-- @unzipWith f fld1 fld2 = lmap f (unzip fld1 fld2)@
2009+--
2010+-- This fold terminates when both the input folds terminate.
2011+--
15442012-- @since 0.7.0
15452013{-# INLINE unzipWith #-}
15462014unzipWith :: Monad m
@@ -1569,35 +2037,9 @@ unzip = unzipWith id
15692037-- Nesting
15702038------------------------------------------------------------------------------
15712039
1572−{-
1573−-- All the stream flattening transformations can also be applied to a fold
1574−-- input stream.
1575−
1576−-- | This can be used to apply all the stream generation operations on folds.
1577−lconcatMap ::(IsStream t, Monad m) => (a -> t m b)
1578− -> Fold m b c
1579− -> Fold m a c
1580−lconcatMap s f1 f2 = undefined
1581−-}
1582−
1583−-- All the grouping transformation that we apply to a stream can also be
1584−-- applied to a fold input stream. groupBy et al can be written as terminating
1585−-- folds and then we can apply foldChunks to use those repeatedly on a stream.
1586−
1587−-- | Apply a terminating fold repeatedly to the input of another fold.
1588−--
1589−-- Compare with: Streamly.Prelude.concatMap, Streamly.Prelude.foldChunks
1590−--
1591−-- /Unimplemented/
1592−--
1593−{-# INLINABLE foldChunks #-}
1594−foldChunks ::
1595− -- Monad m =>
1596− Fold m a b -> Fold m b c -> Fold m a c
1597−foldChunks = undefined
15982040
15992041{-
1600−-- XXX this would be an application of foldChunks using a terminating fold.
2042+-- XXX this would be an application of "many" using a terminating fold.
16012043--
16022044-- | Group the input stream into groups of elements between @low@ and @high@.
16032045-- Collection starts in chunks of @low@ and then keeps doubling until we reach
@@ -1625,39 +2067,43 @@ lchunksInRange low high (Fold step1 initial1 extract1)
16252067{-# INLINE toParallelSVar #-}
16262068toParallelSVar :: MonadIO m => SVar t m a -> Maybe WorkerInfo -> Fold m a ()
16272069toParallelSVar svar winfo = Fold step initial extract
2070+
16282071 where
16292072
16302073 initial = return ()
16312074
1632− step () x = liftIO $ do
1633− -- XXX we can have a separate fold for unlimited buffer case to avoid a
1634− -- branch in the step here.
1635− decrementBufferLimit svar
1636− void $ send svar (ChildYield x)
2075+ -- XXX we can have a separate fold for unlimited buffer case to avoid a
2076+ -- branch in the step here.
2077+ step () x =
2078+ liftIO $ do
2079+ decrementBufferLimit svar
2080+ void $ send svar (ChildYield x)
2081+ return $ FL.Partial ()
16372082
1638− extract () = liftIO $
1639− sendStop svar winfo
2083+ extract () = liftIO $ sendStop svar winfo
16402084
16412085{-# INLINE toParallelSVarLimited #-}
16422086toParallelSVarLimited :: MonadIO m
16432087 => SVar t m a -> Maybe WorkerInfo -> Fold m a ()
16442088toParallelSVarLimited svar winfo = Fold step initial extract
2089+
16452090 where
16462091
16472092 initial = return True
16482093
1649− step True x = liftIO $ do
1650− yieldLimitOk <- decrementYieldLimit svar
1651− if yieldLimitOk
1652− then do
1653− decrementBufferLimit svar
1654− void $ send svar (ChildYield x)
1655− return True
1656− else do
1657− cleanupSVarFromWorker svar
1658− sendStop svar winfo
1659− return False
1660− step False _ = return False
2094+ step True x =
2095+ liftIO $ do
2096+ yieldLimitOk <- decrementYieldLimit svar
2097+ if yieldLimitOk
2098+ then do
2099+ decrementBufferLimit svar
2100+ void $ send svar (ChildYield x)
2101+ return $ FL.Partial True
2102+ else do
2103+ cleanupSVarFromWorker svar
2104+ sendStop svar winfo
2105+ return $ FL.Done ()
2106+ step False _ = return $ FL.Done ()
16612107
16622108 extract True = liftIO $ sendStop svar winfo
16632109 extract False = return ()