git.iko.soy/streamly
on master help sign in
streamly/core/src/Streamly/Internal/Data/Stream/Transform.hs
1971 lines 64 kB history raw
1{-# LANGUAGE CPP #-}
2-- |
3-- Module : Streamly.Internal.Data.Stream.Transform
4-- Copyright : (c) 2018 Composewell Technologies
5-- (c) Roman Leshchinskiy 2008-2010
6-- License : BSD-3-Clause
7-- Maintainer : streamly@composewell.com
8-- Stability : experimental
9-- Portability : GHC
10--
11-- "Streamly.Internal.Data.Pipe" might ultimately replace this module.
13-- A few functions in this module have been adapted from the vector package
14-- (c) Roman Leshchinskiy. See the notes in specific combinators.
16module Streamly.Internal.Data.Stream.Transform
17 (
18 -- * Piping
19 -- | Pass through a 'Pipe'.
20 transform
22 -- * Mapping
23 -- | Stateless one-to-one maps.
24 , sequence
26 -- * Mapping Effects
27 , tap
28 , tapOffsetEvery
29 , trace
30 , trace_
32 -- * Folding
33 , foldrS
34 , foldlS
36 -- * Scanning By 'Fold'
37 , postscan
38 , scan
39 , scanMany
41 -- * Splitting
42 , splitOn
44 -- * Scanning
45 -- | Left scans. Stateful, mostly one-to-one maps.
46 , scanlM'
47 , scanlMAfter'
48 , scanl'
49 , scanlM
50 , scanl
51 , scanl1M'
52 , scanl1'
53 , scanl1M
54 , scanl1
56 , prescanl'
57 , prescanlM'
59 , postscanl
60 , postscanlM
61 , postscanl'
62 , postscanlM'
63 , postscanlMAfter'
65 , postscanlx'
66 , postscanlMx'
67 , scanlMx'
68 , scanlx'
70 -- * Filtering
71 -- | Produce a subset of the stream.
72 , with
73 , scanMaybe
74 , filter
75 , filterM
76 , deleteBy -- deleteOnceBy?
77 , uniqBy
78 , uniq
79 , prune
80 , repeated
82 -- * Trimming
83 -- | Produce a subset of the stream trimmed at ends.
84 , initNonEmpty
85 , tailNonEmpty
86 , takeWhileLast
87 , takeWhileAround
88 , drop
89 , dropWhile
90 , dropWhileM
91 , dropLast
92 , dropWhileLast
93 , dropWhileAround
95 -- * Inserting Elements
96 -- | Produce a superset of the stream.
97 , insertBy
98 , intersperse
99 , intersperseM
100 , intersperseMWith
101 , intersperseMSuffix
102 , intersperseMSuffixWith
104 -- * Inserting Side Effects
105 , intersperseM_
106 , intersperseMSuffix_
107 , intersperseMPrefix_
109 , delay
110 , delayPre
111 , delayPost
113 -- * Reordering
114 -- | Produce strictly the same set but reordered.
115 , reverse
116 , reverseUnbox
117 , reassembleBy
119 -- * Position Indexing
120 , indexed
121 , indexedR
123 -- * Time Indexing
124 , timestampWith
125 , timestamped
126 , timeIndexWith
127 , timeIndexed
129 -- * Searching
130 , findIndices
131 , elemIndices
133 -- * Rolling map
134 -- | Map using the previous element.
135 , rollingMap
136 , rollingMapM
137 , rollingMap2
139 -- * Maybe Streams
140 , mapMaybe
141 , mapMaybeM
142 , catMaybes
144 -- * Either Streams
145 , catLefts
146 , catRights
147 , catEithers
148 )
149where
151#include "inline.hs"
153import Control.Concurrent (threadDelay)
154import Control.Monad (void)
155import Control.Monad.IO.Class (MonadIO (liftIO))
156import Data.Either (fromLeft, isLeft, isRight, fromRight)
157import Data.Functor ((<&>))
158import Data.Maybe (fromJust, isJust)
159import Fusion.Plugin.Types (Fuse(..))
161import Streamly.Internal.Data.Fold.Type (Fold(..))
162import Streamly.Internal.Data.Pipe.Type (Pipe(..), PipeState(..))
163import Streamly.Internal.Data.SVar.Type (adaptState)
164import Streamly.Internal.Data.Time.Units (AbsTime, RelTime64)
165import Streamly.Internal.Data.Unbox (Unbox)
166import Streamly.Internal.System.IO (defaultChunkSize)
168-- import qualified Data.List as List
169import qualified Streamly.Internal.Data.Array.Type as A
170import qualified Streamly.Internal.Data.Fold as FL
171import qualified Streamly.Internal.Data.Pipe.Type as Pipe
172import qualified Streamly.Internal.Data.StreamK.Type as K
174import Prelude hiding
175 ( drop, dropWhile, filter, map, mapM, reverse
176 , scanl, scanl1, sequence, take, takeWhile, zipWith)
178import Streamly.Internal.Data.Stream.Generate
179 (absTimesWith, relTimesWith)
180import Streamly.Internal.Data.Stream.Type
182#include "DocTestDataStream.hs"
184------------------------------------------------------------------------------
185-- Piping
186------------------------------------------------------------------------------
188-- | Use a 'Pipe' to transform a stream.
189--
190-- /Pre-release/
191--
192{-# INLINE_NORMAL transform #-}
193transform :: Monad m => Pipe m a b -> Stream m a -> Stream m b
194transform (Pipe pstep1 pstep2 pstate) (Stream step state) =
195 Stream step' (Consume pstate, state)
197 where
199 {-# INLINE_LATE step' #-}
201 step' gst (Consume pst, st) = pst `seq` do
202 r <- step (adaptState gst) st
203 case r of
204 Yield x s -> do
205 res <- pstep1 pst x
206 case res of
207 Pipe.Yield b pst' -> return $ Yield b (pst', s)
208 Pipe.Continue pst' -> return $ Skip (pst', s)
209 Skip s -> return $ Skip (Consume pst, s)
210 Stop -> return Stop
212 step' _ (Produce pst, st) = pst `seq` do
213 res <- pstep2 pst
214 case res of
215 Pipe.Yield b pst' -> return $ Yield b (pst', st)
216 Pipe.Continue pst' -> return $ Skip (pst', st)
218------------------------------------------------------------------------------
219-- Transformation Folds
220------------------------------------------------------------------------------
222-- Note, this is going to have horrible performance, because of the nature of
223-- the stream type (i.e. direct stream vs CPS). Its only for reference, it is
224-- likely be practically unusable.
225{-# INLINE_NORMAL foldlS #-}
226foldlS :: Monad m
227 => (Stream m b -> a -> Stream m b) -> Stream m b -> Stream m a -> Stream m b
228foldlS fstep begin (Stream step state) = Stream step' (Left (state, begin))
229 where
230 step' gst (Left (st, acc)) = do
231 r <- step (adaptState gst) st
232 return $ case r of
233 Yield x s -> Skip (Left (s, fstep acc x))
234 Skip s -> Skip (Left (s, acc))
235 Stop -> Skip (Right acc)
237 step' gst (Right (Stream stp stt)) = do
238 r <- stp (adaptState gst) stt
239 return $ case r of
240 Yield x s -> Yield x (Right (Stream stp s))
241 Skip s -> Skip (Right (Stream stp s))
242 Stop -> Stop
244------------------------------------------------------------------------------
245-- Transformation by Mapping
246------------------------------------------------------------------------------
248-- |
249-- >>> sequence = Stream.mapM id
250--
251-- Replace the elements of a stream of monadic actions with the outputs of
252-- those actions.
253--
254-- >>> s = Stream.fromList [putStr "a", putStr "b", putStrLn "c"]
255-- >>> Stream.fold Fold.drain $ Stream.sequence s
256-- abc
257--
258{-# INLINE_NORMAL sequence #-}
259sequence :: Monad m => Stream m (m a) -> Stream m a
260sequence (Stream step state) = Stream step' state
261 where
262 {-# INLINE_LATE step' #-}
263 step' gst st = do
264 r <- step (adaptState gst) st
265 case r of
266 Yield x s -> x >>= \a -> return (Yield a s)
267 Skip s -> return $ Skip s
268 Stop -> return Stop
270------------------------------------------------------------------------------
271-- Mapping side effects
272------------------------------------------------------------------------------
274data TapState fs st a
275 = TapInit | Tapping !fs st | TapDone st
277-- XXX Multiple yield points
279-- | Tap the data flowing through a stream into a 'Fold'. For example, you may
280-- add a tap to log the contents flowing through the stream. The fold is used
281-- only for effects, its result is discarded.
282--
283-- @
284-- Fold m a b
285-- |
286-- -----stream m a ---------------stream m a-----
287--
288-- @
289--
290-- >>> s = Stream.enumerateFromTo 1 2
291-- >>> Stream.fold Fold.drain $ Stream.tap (Fold.drainMapM print) s
292-- 1
293-- 2
294--
295-- Compare with 'trace'.
296--
297{-# INLINE tap #-}
298tap :: Monad m => Fold m a b -> Stream m a -> Stream m a
299tap (Fold fstep initial _ final) (Stream step state) = Stream step' TapInit
301 where
303 step' _ TapInit = do
304 res <- initial
305 return
306 $ Skip
307 $ case res of
308 FL.Partial s -> Tapping s state
309 FL.Done _ -> TapDone state
310 step' gst (Tapping acc st) = do
311 r <- step gst st
312 case r of
313 Yield x s -> do
314 res <- fstep acc x
315 return
316 $ Yield x
317 $ case res of
318 FL.Partial fs -> Tapping fs s
319 FL.Done _ -> TapDone s
320 Skip s -> return $ Skip (Tapping acc s)
321 Stop -> do
322 void $ final acc
323 return Stop
324 step' gst (TapDone st) = do
325 r <- step gst st
326 return
327 $ case r of
328 Yield x s -> Yield x (TapDone s)
329 Skip s -> Skip (TapDone s)
330 Stop -> Stop
332data TapOffState fs s a
333 = TapOffInit
334 | TapOffTapping !fs s Int
335 | TapOffDone s
337-- XXX Multiple yield points
338{-# INLINE_NORMAL tapOffsetEvery #-}
339tapOffsetEvery :: Monad m
340 => Int -> Int -> Fold m a b -> Stream m a -> Stream m a
341tapOffsetEvery offset n (Fold fstep initial _ final) (Stream step state) =
342 Stream step' TapOffInit
344 where
346 {-# INLINE_LATE step' #-}
347 step' _ TapOffInit = do
348 res <- initial
349 return
350 $ Skip
351 $ case res of
352 FL.Partial s -> TapOffTapping s state (offset `mod` n)
353 FL.Done _ -> TapOffDone state
354 step' gst (TapOffTapping acc st count) = do
355 r <- step gst st
356 case r of
357 Yield x s -> do
358 next <-
359 if count <= 0
360 then do
361 res <- fstep acc x
362 return
363 $ case res of
364 FL.Partial sres ->
365 TapOffTapping sres s (n - 1)
366 FL.Done _ -> TapOffDone s
367 else return $ TapOffTapping acc s (count - 1)
368 return $ Yield x next
369 Skip s -> return $ Skip (TapOffTapping acc s count)
370 Stop -> do
371 void $ final acc
372 return Stop
373 step' gst (TapOffDone st) = do
374 r <- step gst st
375 return
376 $ case r of
377 Yield x s -> Yield x (TapOffDone s)
378 Skip s -> Skip (TapOffDone s)
379 Stop -> Stop
381-- | Apply a monadic function to each element flowing through the stream and
382-- discard the results.
383--
384-- >>> s = Stream.enumerateFromTo 1 2
385-- >>> Stream.fold Fold.drain $ Stream.trace print s
386-- 1
387-- 2
388--
389-- Compare with 'tap'.
390--
391{-# INLINE trace #-}
392trace :: Monad m => (a -> m b) -> Stream m a -> Stream m a
393trace f = mapM (\x -> void (f x) >> return x)
395-- | Perform a side effect before yielding each element of the stream and
396-- discard the results.
397--
398-- >>> s = Stream.enumerateFromTo 1 2
399-- >>> Stream.fold Fold.drain $ Stream.trace_ (print "got here") s
400-- "got here"
401-- "got here"
402--
403-- Same as 'intersperseMPrefix_' but always serial.
404--
405-- See also: 'trace'
406--
407-- /Pre-release/
408{-# INLINE trace_ #-}
409trace_ :: Monad m => m b -> Stream m a -> Stream m a
410trace_ eff = mapM (\x -> eff >> return x)
412------------------------------------------------------------------------------
413-- Scanning with a Fold
414------------------------------------------------------------------------------
416data ScanState s f = ScanInit s | ScanDo s !f | ScanDone
418-- | Postscan a stream using the given monadic fold.
419--
420-- The following example extracts the input stream up to a point where the
421-- running average of elements is no more than 10:
422--
423-- >>> import Data.Maybe (fromJust)
424-- >>> let avg = Fold.teeWith (/) Fold.sum (fmap fromIntegral Fold.length)
425-- >>> s = Stream.enumerateFromTo 1.0 100.0
426-- >>> :{
427-- Stream.fold Fold.toList
428-- $ fmap (fromJust . fst)
429-- $ Stream.takeWhile (\(_,x) -> x <= 10)
430-- $ Stream.postscan (Fold.tee Fold.latest avg) s
431-- :}
432-- [1.0,2.0,3.0,4.0,5.0,6.0,7.0,8.0,9.0,10.0,11.0,12.0,13.0,14.0,15.0,16.0,17.0,18.0,19.0]
433--
434{-# INLINE_NORMAL postscan #-}
435postscan :: Monad m => FL.Fold m a b -> Stream m a -> Stream m b
436postscan (FL.Fold fstep initial extract final) (Stream sstep state) =
437 Stream step (ScanInit state)
439 where
441 {-# INLINE_LATE step #-}
442 step _ (ScanInit st) = do
443 res <- initial
444 return
445 $ case res of
446 FL.Partial fs -> Skip $ ScanDo st fs
447 FL.Done b -> Yield b ScanDone
448 step gst (ScanDo st fs) = do
449 res <- sstep (adaptState gst) st
450 case res of
451 Yield x s -> do
452 r <- fstep fs x
453 case r of
454 FL.Partial fs1 -> do
455 !b <- extract fs1
456 return $ Yield b $ ScanDo s fs1
457 FL.Done b -> return $ Yield b ScanDone
458 Skip s -> return $ Skip $ ScanDo s fs
459 Stop -> final fs >> return Stop
460 step _ ScanDone = return Stop
462{-# INLINE scanWith #-}
463scanWith :: Monad m
464 => Bool -> Fold m a b -> Stream m a -> Stream m b
465scanWith restart (Fold fstep initial extract final) (Stream sstep state) =
466 Stream step (ScanInit state)
468 where
470 {-# INLINE runStep #-}
471 runStep st action = do
472 res <- action
473 case res of
474 FL.Partial fs -> do
475 !b <- extract fs
476 return $ Yield b $ ScanDo st fs
477 FL.Done b ->
478 let next = if restart then ScanInit st else ScanDone
479 in return $ Yield b next
481 {-# INLINE_LATE step #-}
482 step _ (ScanInit st) = runStep st initial
483 step gst (ScanDo st fs) = do
484 res <- sstep (adaptState gst) st
485 case res of
486 Yield x s -> runStep s (fstep fs x)
487 Skip s -> return $ Skip $ ScanDo s fs
488 Stop -> final fs >> return Stop
489 step _ ScanDone = return Stop
491-- XXX It may be useful to have a version of scan where we can keep the
492-- accumulator independent of the value emitted. So that we do not necessarily
493-- have to keep a value in the accumulator which we are not using. We can pass
494-- an extraction function that will take the accumulator and the current value
495-- of the element and emit the next value in the stream. That will also make it
496-- possible to modify the accumulator after using it. In fact, the step function
497-- can return new accumulator and the value to be emitted. The signature would
498-- be more like mapAccumL.
500-- | Strict left scan. Scan a stream using the given monadic fold.
501--
502-- >>> s = Stream.fromList [1..10]
503-- >>> Stream.fold Fold.toList $ Stream.takeWhile (< 10) $ Stream.scan Fold.sum s
504-- [0,1,3,6]
505--
506-- See also: 'usingStateT'
507--
509-- EXPLANATION:
510-- >>> scanl' step z = Stream.scan (Fold.foldl' step z)
511--
512-- Like 'map', 'scanl'' too is a one to one transformation,
513-- however it adds an extra element.
514--
515-- >>> s = Stream.fromList [1,2,3,4]
516-- >>> Stream.fold Fold.toList $ scanl' (+) 0 s
517-- [0,1,3,6,10]
518--
519-- >>> Stream.fold Fold.toList $ scanl' (flip (:)) [] s
520-- [[],[1],[2,1],[3,2,1],[4,3,2,1]]
521--
522-- The output of 'scanl'' is the initial value of the accumulator followed by
523-- all the intermediate steps and the final result of 'foldl''.
524--
525-- By streaming the accumulated state after each fold step, we can share the
526-- state across multiple stages of stream composition. Each stage can modify or
527-- extend the state, do some processing with it and emit it for the next stage,
528-- thus modularizing the stream processing. This can be useful in
529-- stateful or event-driven programming.
530--
531-- Consider the following monolithic example, computing the sum and the product
532-- of the elements in a stream in one go using a @foldl'@:
533--
534-- >>> foldl' step z = Stream.fold (Fold.foldl' step z)
535-- >>> foldl' (\(s, p) x -> (s + x, p * x)) (0,1) s
536-- (10,24)
537--
538-- Using @scanl'@ we can make it modular by computing the sum in the first
539-- stage and passing it down to the next stage for computing the product:
540--
541-- >>> :{
542-- foldl' (\(_, p) (s, x) -> (s, p * x)) (0,1)
543-- $ scanl' (\(s, _) x -> (s + x, x)) (0,1)
544-- $ Stream.fromList [1,2,3,4]
545-- :}
546-- (10,24)
547--
548-- IMPORTANT: 'scanl'' evaluates the accumulator to WHNF. To avoid building
549-- lazy expressions inside the accumulator, it is recommended that a strict
550-- data structure is used for accumulator.
551--
552{-# INLINE_NORMAL scan #-}
553scan :: Monad m
554 => FL.Fold m a b -> Stream m a -> Stream m b
555scan = scanWith False
557-- | Like 'scan' but restarts scanning afresh when the scanning fold
558-- terminates.
559--
560{-# INLINE_NORMAL scanMany #-}
561scanMany :: Monad m
562 => FL.Fold m a b -> Stream m a -> Stream m b
563scanMany = scanWith True
565------------------------------------------------------------------------------
566-- Scanning - Prescans
567------------------------------------------------------------------------------
569-- Adapted from the vector package.
570--
571-- XXX Is a prescan useful, discarding the last step does not sound useful? I
572-- am not sure about the utility of this function, so this is implemented but
573-- not exposed. We can expose it if someone provides good reasons why this is
574-- useful.
575--
576-- XXX We have to execute the stream one step ahead to know that we are at the
577-- last step. The vector implementation of prescan executes the last fold step
578-- but does not yield the result. This means we have executed the effect but
579-- discarded value. This does not sound right. In this implementation we are
580-- not executing the last fold step.
581{-# INLINE_NORMAL prescanlM' #-}
582prescanlM' :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
583prescanlM' f mz (Stream step state) = Stream step' (state, mz)
584 where
585 {-# INLINE_LATE step' #-}
586 step' gst (st, prev) = do
587 r <- step (adaptState gst) st
588 case r of
589 Yield x s -> do
590 acc <- prev
591 return $ Yield acc (s, f acc x)
592 Skip s -> return $ Skip (s, prev)
593 Stop -> return Stop
595{-# INLINE prescanl' #-}
596prescanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
597prescanl' f z = prescanlM' (\a b -> return (f a b)) (return z)
599------------------------------------------------------------------------------
600-- Monolithic postscans (postscan followed by a map)
601------------------------------------------------------------------------------
603-- The performance of a modular postscan followed by a map seems to be
604-- equivalent to this monolithic scan followed by map therefore we may not need
605-- this implementation. We just have it for performance comparison and in case
606-- modular version does not perform well in some situation.
607--
608{-# INLINE_NORMAL postscanlMx' #-}
609postscanlMx' :: Monad m
610 => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> Stream m b
611postscanlMx' fstep begin done (Stream step state) = do
612 Stream step' (state, begin)
613 where
614 {-# INLINE_LATE step' #-}
615 step' gst (st, acc) = do
616 r <- step (adaptState gst) st
617 case r of
618 Yield x s -> do
619 old <- acc
620 y <- fstep old x
621 v <- done y
622 v `seq` y `seq` return (Yield v (s, return y))
623 Skip s -> return $ Skip (s, acc)
624 Stop -> return Stop
626{-# INLINE_NORMAL postscanlx' #-}
627postscanlx' :: Monad m
628 => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> Stream m b
629postscanlx' fstep begin done =
630 postscanlMx' (\b a -> return (fstep b a)) (return begin) (return . done)
632-- XXX do we need consM strict to evaluate the begin value?
633{-# INLINE scanlMx' #-}
634scanlMx' :: Monad m
635 => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> Stream m b
636scanlMx' fstep begin done s =
637 (begin >>= \x -> x `seq` done x) `consM` postscanlMx' fstep begin done s
639{-# INLINE scanlx' #-}
640scanlx' :: Monad m
641 => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> Stream m b
642scanlx' fstep begin done =
643 scanlMx' (\b a -> return (fstep b a)) (return begin) (return . done)
645------------------------------------------------------------------------------
646-- postscans
647------------------------------------------------------------------------------
649-- Adapted from the vector package.
650{-# INLINE_NORMAL postscanlM' #-}
651postscanlM' :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
652postscanlM' fstep begin (Stream step state) =
653 Stream step' Nothing
654 where
655 {-# INLINE_LATE step' #-}
656 step' _ Nothing = do
657 !x <- begin
658 return $ Skip (Just (state, x))
660 step' gst (Just (st, acc)) = do
661 r <- step (adaptState gst) st
662 case r of
663 Yield x s -> do
664 !y <- fstep acc x
665 return $ Yield y (Just (s, y))
666 Skip s -> return $ Skip (Just (s, acc))
667 Stop -> return Stop
669{-# INLINE_NORMAL postscanl' #-}
670postscanl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
671postscanl' f seed = postscanlM' (\a b -> return (f a b)) (return seed)
673{-# ANN type PScanAfterState Fuse #-}
674data PScanAfterState m st acc =
675 PScanAfterStep st (m acc)
676 | PScanAfterYield acc (PScanAfterState m st acc)
677 | PScanAfterStop
679-- We can possibly have the "done" function as a Maybe to provide an option to
680-- emit or not emit the accumulator when the stream stops.
681--
682-- TBD: use a single Yield point
683--
684{-# INLINE_NORMAL postscanlMAfter' #-}
685postscanlMAfter' :: Monad m
686 => (b -> a -> m b) -> m b -> (b -> m b) -> Stream m a -> Stream m b
687postscanlMAfter' fstep initial done (Stream step1 state1) = do
688 Stream step (PScanAfterStep state1 initial)
690 where
692 {-# INLINE_LATE step #-}
693 step gst (PScanAfterStep st acc) = do
694 r <- step1 (adaptState gst) st
695 case r of
696 Yield x s -> do
697 !old <- acc
698 !y <- fstep old x
699 return (Skip $ PScanAfterYield y (PScanAfterStep s (return y)))
700 Skip s -> return $ Skip $ PScanAfterStep s acc
701 -- Strictness is important for fusion
702 Stop -> do
703 !v <- acc
704 !res <- done v
705 return (Skip $ PScanAfterYield res PScanAfterStop)
706 step _ (PScanAfterYield acc next) = return $ Yield acc next
707 step _ PScanAfterStop = return Stop
709{-# INLINE_NORMAL postscanlM #-}
710postscanlM :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
711postscanlM fstep begin (Stream step state) = Stream step' Nothing
712 where
713 {-# INLINE_LATE step' #-}
714 step' _ Nothing = do
715 r <- begin
716 return $ Skip (Just (state, r))
718 step' gst (Just (st, acc)) = do
719 r <- step (adaptState gst) st
720 case r of
721 Yield x s -> do
722 y <- fstep acc x
723 return (Yield y (Just (s, y)))
724 Skip s -> return $ Skip (Just (s, acc))
725 Stop -> return Stop
727{-# INLINE_NORMAL postscanl #-}
728postscanl :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
729postscanl f seed = postscanlM (\a b -> return (f a b)) (return seed)
731-- | Like 'scanl'' but with a monadic step function and a monadic seed.
732--
733{-# INLINE_NORMAL scanlM' #-}
734scanlM' :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
735scanlM' fstep begin (Stream step state) = Stream step' Nothing
736 where
737 {-# INLINE_LATE step' #-}
738 step' _ Nothing = do
739 !x <- begin
740 return $ Yield x (Just (state, x))
741 step' gst (Just (st, acc)) = do
742 r <- step (adaptState gst) st
743 case r of
744 Yield x s -> do
745 !y <- fstep acc x
746 return $ Yield y (Just (s, y))
747 Skip s -> return $ Skip (Just (s, acc))
748 Stop -> return Stop
750-- | @scanlMAfter' accumulate initial done stream@ is like 'scanlM'' except
751-- that it provides an additional @done@ function to be applied on the
752-- accumulator when the stream stops. The result of @done@ is also emitted in
753-- the stream.
754--
755-- This function can be used to allocate a resource in the beginning of the
756-- scan and release it when the stream ends or to flush the internal state of
757-- the scan at the end.
758--
759-- /Pre-release/
760--
761{-# INLINE scanlMAfter' #-}
762scanlMAfter' :: Monad m
763 => (b -> a -> m b) -> m b -> (b -> m b) -> Stream m a -> Stream m b
764scanlMAfter' fstep initial done s =
765 initial `consM` postscanlMAfter' fstep initial done s
767-- >>> scanl' f z xs = z `Stream.cons` postscanl' f z xs
769-- | Strict left scan. Like 'map', 'scanl'' too is a one to one transformation,
770-- however it adds an extra element.
771--
772-- >>> Stream.toList $ Stream.scanl' (+) 0 $ Stream.fromList [1,2,3,4]
773-- [0,1,3,6,10]
774--
775-- >>> Stream.toList $ Stream.scanl' (flip (:)) [] $ Stream.fromList [1,2,3,4]
776-- [[],[1],[2,1],[3,2,1],[4,3,2,1]]
777--
778-- The output of 'scanl'' is the initial value of the accumulator followed by
779-- all the intermediate steps and the final result of 'foldl''.
780--
781-- By streaming the accumulated state after each fold step, we can share the
782-- state across multiple stages of stream composition. Each stage can modify or
783-- extend the state, do some processing with it and emit it for the next stage,
784-- thus modularizing the stream processing. This can be useful in
785-- stateful or event-driven programming.
786--
787-- Consider the following monolithic example, computing the sum and the product
788-- of the elements in a stream in one go using a @foldl'@:
789--
790-- >>> Stream.fold (Fold.foldl' (\(s, p) x -> (s + x, p * x)) (0,1)) $ Stream.fromList [1,2,3,4]
791-- (10,24)
792--
793-- Using @scanl'@ we can make it modular by computing the sum in the first
794-- stage and passing it down to the next stage for computing the product:
795--
796-- >>> :{
797-- Stream.fold (Fold.foldl' (\(_, p) (s, x) -> (s, p * x)) (0,1))
798-- $ Stream.scanl' (\(s, _) x -> (s + x, x)) (0,1)
799-- $ Stream.fromList [1,2,3,4]
800-- :}
801-- (10,24)
802--
803-- IMPORTANT: 'scanl'' evaluates the accumulator to WHNF. To avoid building
804-- lazy expressions inside the accumulator, it is recommended that a strict
805-- data structure is used for accumulator.
806--
807-- >>> scanl' step z = Stream.scan (Fold.foldl' step z)
808-- >>> scanl' f z xs = Stream.scanlM' (\a b -> return (f a b)) (return z) xs
809--
810-- See also: 'usingStateT'
811--
812{-# INLINE scanl' #-}
813scanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
814scanl' f seed = scanlM' (\a b -> return (f a b)) (return seed)
816{-# INLINE_NORMAL scanlM #-}
817scanlM :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
818scanlM fstep begin (Stream step state) = Stream step' Nothing
819 where
820 {-# INLINE_LATE step' #-}
821 step' _ Nothing = do
822 x <- begin
823 return $ Yield x (Just (state, x))
824 step' gst (Just (st, acc)) = do
825 r <- step (adaptState gst) st
826 case r of
827 Yield x s -> do
828 y <- fstep acc x
829 return $ Yield y (Just (s, y))
830 Skip s -> return $ Skip (Just (s, acc))
831 Stop -> return Stop
833{-# INLINE scanl #-}
834scanl :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
835scanl f seed = scanlM (\a b -> return (f a b)) (return seed)
837-- Adapted from the vector package
838{-# INLINE_NORMAL scanl1M #-}
839scanl1M :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
840scanl1M fstep (Stream step state) = Stream step' (state, Nothing)
841 where
842 {-# INLINE_LATE step' #-}
843 step' gst (st, Nothing) = do
844 r <- step gst st
845 case r of
846 Yield x s -> return $ Yield x (s, Just x)
847 Skip s -> return $ Skip (s, Nothing)
848 Stop -> return Stop
850 step' gst (st, Just acc) = do
851 r <- step gst st
852 case r of
853 Yield y s -> do
854 z <- fstep acc y
855 return $ Yield z (s, Just z)
856 Skip s -> return $ Skip (s, Just acc)
857 Stop -> return Stop
859{-# INLINE scanl1 #-}
860scanl1 :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
861scanl1 f = scanl1M (\x y -> return (f x y))
863-- Adapted from the vector package
865-- | Like 'scanl1'' but with a monadic step function.
866--
867{-# INLINE_NORMAL scanl1M' #-}
868scanl1M' :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
869scanl1M' fstep (Stream step state) = Stream step' (state, Nothing)
870 where
871 {-# INLINE_LATE step' #-}
872 step' gst (st, Nothing) = do
873 r <- step gst st
874 case r of
875 Yield x s -> x `seq` return $ Yield x (s, Just x)
876 Skip s -> return $ Skip (s, Nothing)
877 Stop -> return Stop
879 step' gst (st, Just acc) = acc `seq` do
880 r <- step gst st
881 case r of
882 Yield y s -> do
883 z <- fstep acc y
884 z `seq` return $ Yield z (s, Just z)
885 Skip s -> return $ Skip (s, Just acc)
886 Stop -> return Stop
888-- | Like 'scanl'' but for a non-empty stream. The first element of the stream
889-- is used as the initial value of the accumulator. Does nothing if the stream
890-- is empty.
891--
892-- >>> Stream.toList $ Stream.scanl1' (+) $ Stream.fromList [1,2,3,4]
893-- [1,3,6,10]
894--
895{-# INLINE scanl1' #-}
896scanl1' :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
897scanl1' f = scanl1M' (\x y -> return (f x y))
899-------------------------------------------------------------------------------
900-- Filtering
901-------------------------------------------------------------------------------
903-- | Modify a @Stream m a -> Stream m a@ stream transformation that accepts a
904-- predicate @(a -> b)@ to accept @((s, a) -> b)@ instead, provided a
905-- transformation @Stream m a -> Stream m (s, a)@. Convenient to filter with
906-- index or time.
907--
908-- >>> filterWithIndex = Stream.with Stream.indexed Stream.filter
909--
910-- /Pre-release/
911{-# INLINE with #-}
912with :: Monad m =>
913 (Stream m a -> Stream m (s, a))
914 -> (((s, a) -> b) -> Stream m (s, a) -> Stream m (s, a))
915 -> (((s, a) -> b) -> Stream m a -> Stream m a)
916with f comb g = fmap snd . comb g . f
918-- Adapted from the vector package
920-- | Same as 'filter' but with a monadic predicate.
921--
922-- >>> f p x = p x >>= \r -> return $ if r then Just x else Nothing
923-- >>> filterM p = Stream.mapMaybeM (f p)
924--
925{-# INLINE_NORMAL filterM #-}
926filterM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
927filterM f (Stream step state) = Stream step' state
928 where
929 {-# INLINE_LATE step' #-}
930 step' gst st = do
931 r <- step gst st
932 case r of
933 Yield x s -> do
934 b <- f x
935 return $ if b
936 then Yield x s
937 else Skip s
938 Skip s -> return $ Skip s
939 Stop -> return Stop
941-- | Include only those elements that pass a predicate.
942--
943-- >>> filter p = Stream.filterM (return . p)
944-- >>> filter p = Stream.mapMaybe (\x -> if p x then Just x else Nothing)
945-- >>> filter p = Stream.scanMaybe (Fold.filtering p)
946--
947{-# INLINE filter #-}
948filter :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
949filter f = filterM (return . f)
950-- filter p = scanMaybe (FL.filtering p)
952-- | Drop repeated elements that are adjacent to each other using the supplied
953-- comparison function.
954--
955-- >>> uniq = Stream.uniqBy (==)
956--
957-- To strip duplicate path separators:
958--
959-- >>> input = Stream.fromList "//a//b"
960-- >>> f x y = x == '/' && y == '/'
961-- >>> Stream.fold Fold.toList $ Stream.uniqBy f input
962-- "/a/b"
963--
964-- Space: @O(1)@
965--
966-- /Pre-release/
967--
968{-# INLINE uniqBy #-}
969uniqBy :: Monad m =>
970 (a -> a -> Bool) -> Stream m a -> Stream m a
971-- uniqBy eq = scanMaybe (FL.uniqBy eq)
972uniqBy eq = catMaybes . rollingMap f
974 where
976 f pre curr =
977 case pre of
978 Nothing -> Just curr
979 Just x -> if x `eq` curr then Nothing else Just curr
981-- Adapted from the vector package
983-- | Drop repeated elements that are adjacent to each other.
984--
985-- >>> uniq = Stream.uniqBy (==)
986--
987{-# INLINE_NORMAL uniq #-}
988uniq :: (Eq a, Monad m) => Stream m a -> Stream m a
989-- uniq = scanMaybe FL.uniq
990uniq (Stream step state) = Stream step' (Nothing, state)
991 where
992 {-# INLINE_LATE step' #-}
993 step' gst (Nothing, st) = do
994 r <- step gst st
995 case r of
996 Yield x s -> return $ Yield x (Just x, s)
997 Skip s -> return $ Skip (Nothing, s)
998 Stop -> return Stop
999 step' gst (Just x, st) = do
1000 r <- step gst st
1001 case r of
1002 Yield y s | x == y -> return $ Skip (Just x, s)
1003 | otherwise -> return $ Yield y (Just y, s)
1004 Skip s -> return $ Skip (Just x, s)
1005 Stop -> return Stop
1007-- | Deletes the first occurrence of the element in the stream that satisfies
1008-- the given equality predicate.
1010-- >>> input = Stream.fromList [1,3,3,5]
1011-- >>> Stream.fold Fold.toList $ Stream.deleteBy (==) 3 input
1012-- [1,3,5]
1014{-# INLINE_NORMAL deleteBy #-}
1015deleteBy :: Monad m => (a -> a -> Bool) -> a -> Stream m a -> Stream m a
1016-- deleteBy cmp x = scanMaybe (FL.deleteBy cmp x)
1017deleteBy eq x (Stream step state) = Stream step' (state, False)
1018 where
1019 {-# INLINE_LATE step' #-}
1020 step' gst (st, False) = do
1021 r <- step gst st
1022 case r of
1023 Yield y s -> return $
1024 if eq x y then Skip (s, True) else Yield y (s, False)
1025 Skip s -> return $ Skip (s, False)
1026 Stop -> return Stop
1028 step' gst (st, True) = do
1029 r <- step gst st
1030 case r of
1031 Yield y s -> return $ Yield y (s, True)
1032 Skip s -> return $ Skip (s, True)
1033 Stop -> return Stop
1035-- | Strip all leading and trailing occurrences of an element passing a
1036-- predicate and make all other consecutive occurrences uniq.
1038-- >> prune p = Stream.dropWhileAround p $ Stream.uniqBy (x y -> p x && p y)
1040-- @
1041-- > Stream.prune isSpace (Stream.fromList " hello world! ")
1042-- "hello world!"
1044-- @
1046-- Space: @O(1)@
1048-- /Unimplemented/
1049{-# INLINE prune #-}
1050prune ::
1051 -- (Monad m, Eq a) =>
1052 (a -> Bool) -> Stream m a -> Stream m a
1053prune = error "Not implemented yet!"
1055-- Possible implementation:
1056-- @repeated =
1057-- Stream.catMaybes . Stream.parseMany (Parser.groupBy (==) Fold.repeated)@
1059-- 'Fold.repeated' should return 'Just' when repeated, and 'Nothing' for a
1060-- single element.
1062-- | Emit only repeated elements, once.
1064-- /Unimplemented/
1065repeated :: -- (Monad m, Eq a) =>
1066 Stream m a -> Stream m a
1067repeated = undefined
1069------------------------------------------------------------------------------
1070-- Trimming
1071------------------------------------------------------------------------------
1073-- | init for non-empty streams, fails for empty stream case.
1075{-# INLINE initNonEmpty #-}
1076initNonEmpty :: Monad m => Stream m a -> Stream m a
1077initNonEmpty (Stream step1 state1) = Stream step (Nothing, state1)
1079 where
1081 step gst (Nothing, s1) = do
1082 r <- step1 (adaptState gst) s1
1083 return $
1084 case r of
1085 Yield x s -> Skip (Just x, s)
1086 Skip s -> Skip (Nothing, s)
1087 Stop -> error "initNonEmpty: empty Stream"
1089 step gst (Just a, s1) = do
1090 r <- step1 (adaptState gst) s1
1091 return $
1092 case r of
1093 Yield x s -> Yield a (Just x, s)
1094 Skip s -> Skip (Just a, s)
1095 Stop -> Stop
1097-- | tail for non-empty streams, fails for empty stream case.
1099-- See also 'tail' for a non-partial version of this function..
1100{-# INLINE tailNonEmpty #-}
1101tailNonEmpty :: Monad m => Stream m a -> Stream m a
1102tailNonEmpty (Stream step1 state1) = Stream step (Nothing, state1)
1104 where
1106 step gst (Nothing, s1) = do
1107 r <- step1 (adaptState gst) s1
1108 return $
1109 case r of
1110 Yield x s -> Skip (Just x, s)
1111 Skip s -> Skip (Nothing, s)
1112 Stop -> error "tailNonEmpty: empty Stream"
1114 step gst (Just a, s1) = do
1115 r <- step1 (adaptState gst) s1
1116 return $
1117 case r of
1118 Yield x s -> Yield x (Just x, s)
1119 Skip s -> Skip (Just a, s)
1120 Stop -> Stop
1122-- | Take all consecutive elements at the end of the stream for which the
1123-- predicate is true.
1125-- O(n) space, where n is the number elements taken.
1127-- /Unimplemented/
1128{-# INLINE takeWhileLast #-}
1129takeWhileLast :: -- Monad m =>
1130 (a -> Bool) -> Stream m a -> Stream m a
1131takeWhileLast = undefined -- fromStreamD $ D.takeWhileLast n $ toStreamD m
1133-- | Like 'takeWhile' and 'takeWhileLast' combined.
1135-- O(n) space, where n is the number elements taken from the end.
1137-- /Unimplemented/
1138{-# INLINE takeWhileAround #-}
1139takeWhileAround :: -- Monad m =>
1140 (a -> Bool) -> Stream m a -> Stream m a
1141takeWhileAround = undefined -- fromStreamD $ D.takeWhileAround n $ toStreamD m
1143-- Adapted from the vector package
1145-- | Discard first 'n' elements from the stream and take the rest.
1147{-# INLINE_NORMAL drop #-}
1148drop :: Monad m => Int -> Stream m a -> Stream m a
1149drop n (Stream step state) = Stream step' (state, Just n)
1150 where
1151 {-# INLINE_LATE step' #-}
1152 step' gst (st, Just i)
1153 | i > 0 = do
1154 r <- step gst st
1155 return $
1156 case r of
1157 Yield _ s -> Skip (s, Just (i - 1))
1158 Skip s -> Skip (s, Just i)
1159 Stop -> Stop
1160 | otherwise = return $ Skip (st, Nothing)
1162 step' gst (st, Nothing) = do
1163 r <- step gst st
1164 return $
1165 case r of
1166 Yield x s -> Yield x (s, Nothing)
1167 Skip s -> Skip (s, Nothing)
1168 Stop -> Stop
1170-- Adapted from the vector package
1171data DropWhileState s a
1172 = DropWhileDrop s
1173 | DropWhileYield a s
1174 | DropWhileNext s
1176-- | Same as 'dropWhile' but with a monadic predicate.
1178{-# INLINE_NORMAL dropWhileM #-}
1179dropWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
1180-- dropWhileM p = scanMaybe (FL.droppingWhileM p)
1181dropWhileM f (Stream step state) = Stream step' (DropWhileDrop state)
1182 where
1183 {-# INLINE_LATE step' #-}
1184 step' gst (DropWhileDrop st) = do
1185 r <- step gst st
1186 case r of
1187 Yield x s -> do
1188 b <- f x
1189 if b
1190 then return $ Skip (DropWhileDrop s)
1191 else return $ Skip (DropWhileYield x s)
1192 Skip s -> return $ Skip (DropWhileDrop s)
1193 Stop -> return Stop
1195 step' gst (DropWhileNext st) = do
1196 r <- step gst st
1197 case r of
1198 Yield x s -> return $ Skip (DropWhileYield x s)
1199 Skip s -> return $ Skip (DropWhileNext s)
1200 Stop -> return Stop
1202 step' _ (DropWhileYield x st) = return $ Yield x (DropWhileNext st)
1204-- | Drop elements in the stream as long as the predicate succeeds and then
1205-- take the rest of the stream.
1207{-# INLINE dropWhile #-}
1208dropWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
1209-- dropWhile p = scanMaybe (FL.droppingWhile p)
1210dropWhile f = dropWhileM (return . f)
1212-- | Drop @n@ elements at the end of the stream.
1214-- O(n) space, where n is the number elements dropped.
1216-- /Unimplemented/
1217{-# INLINE dropLast #-}
1218dropLast :: -- Monad m =>
1219 Int -> Stream m a -> Stream m a
1220dropLast = undefined -- fromStreamD $ D.dropLast n $ toStreamD m
1222-- | Drop all consecutive elements at the end of the stream for which the
1223-- predicate is true.
1225-- O(n) space, where n is the number elements dropped.
1227-- /Unimplemented/
1228{-# INLINE dropWhileLast #-}
1229dropWhileLast :: -- Monad m =>
1230 (a -> Bool) -> Stream m a -> Stream m a
1231dropWhileLast = undefined -- fromStreamD $ D.dropWhileLast n $ toStreamD m
1233-- | Like 'dropWhile' and 'dropWhileLast' combined.
1235-- O(n) space, where n is the number elements dropped from the end.
1237-- /Unimplemented/
1238{-# INLINE dropWhileAround #-}
1239dropWhileAround :: -- Monad m =>
1240 (a -> Bool) -> Stream m a -> Stream m a
1241dropWhileAround = undefined -- fromStreamD $ D.dropWhileAround n $ toStreamD m
1243------------------------------------------------------------------------------
1244-- Inserting Elements
1245------------------------------------------------------------------------------
1247-- | @insertBy cmp elem stream@ inserts @elem@ before the first element in
1248-- @stream@ that is less than @elem@ when compared using @cmp@.
1250-- >>> insertBy cmp x = Stream.mergeBy cmp (Stream.fromPure x)
1252-- >>> input = Stream.fromList [1,3,5]
1253-- >>> Stream.fold Fold.toList $ Stream.insertBy compare 2 input
1254-- [1,2,3,5]
1256{-# INLINE_NORMAL insertBy #-}
1257insertBy :: Monad m => (a -> a -> Ordering) -> a -> Stream m a -> Stream m a
1258insertBy cmp a (Stream step state) = Stream step' (state, False, Nothing)
1259 where
1260 {-# INLINE_LATE step' #-}
1261 step' gst (st, False, _) = do
1262 r <- step gst st
1263 case r of
1264 Yield x s -> case cmp a x of
1265 GT -> return $ Yield x (s, False, Nothing)
1266 _ -> return $ Yield a (s, True, Just x)
1267 Skip s -> return $ Skip (s, False, Nothing)
1268 Stop -> return $ Yield a (st, True, Nothing)
1270 step' _ (_, True, Nothing) = return Stop
1272 step' gst (st, True, Just prev) = do
1273 r <- step gst st
1274 case r of
1275 Yield x s -> return $ Yield prev (s, True, Just x)
1276 Skip s -> return $ Skip (s, True, Just prev)
1277 Stop -> return $ Yield prev (st, True, Nothing)
1279data LoopState x s = FirstYield s
1280 | InterspersingYield s
1281 | YieldAndCarry x s
1283-- intersperseM = intersperseMWith 1
1285-- | Insert an effect and its output before consuming an element of a stream
1286-- except the first one.
1288-- >>> input = Stream.fromList "hello"
1289-- >>> Stream.fold Fold.toList $ Stream.trace putChar $ Stream.intersperseM (putChar '.' >> return ',') input
1290-- h.,e.,l.,l.,o"h,e,l,l,o"
1292-- Be careful about the order of effects. In the above example we used trace
1293-- after the intersperse, if we use it before the intersperse the output would
1294-- be he.l.l.o."h,e,l,l,o".
1296-- >>> Stream.fold Fold.toList $ Stream.intersperseM (putChar '.' >> return ',') $ Stream.trace putChar input
1297-- he.l.l.o."h,e,l,l,o"
1299{-# INLINE_NORMAL intersperseM #-}
1300intersperseM :: Monad m => m a -> Stream m a -> Stream m a
1301intersperseM m (Stream step state) = Stream step' (FirstYield state)
1302 where
1303 {-# INLINE_LATE step' #-}
1304 step' gst (FirstYield st) = do
1305 r <- step gst st
1306 return $
1307 case r of
1308 Yield x s -> Skip (YieldAndCarry x s)
1309 Skip s -> Skip (FirstYield s)
1310 Stop -> Stop
1312 step' gst (InterspersingYield st) = do
1313 r <- step gst st
1314 case r of
1315 Yield x s -> do
1316 a <- m
1317 return $ Yield a (YieldAndCarry x s)
1318 Skip s -> return $ Skip $ InterspersingYield s
1319 Stop -> return Stop
1321 step' _ (YieldAndCarry x st) = return $ Yield x (InterspersingYield st)
1323-- | Insert a pure value between successive elements of a stream.
1325-- >>> input = Stream.fromList "hello"
1326-- >>> Stream.fold Fold.toList $ Stream.intersperse ',' input
1327-- "h,e,l,l,o"
1329{-# INLINE intersperse #-}
1330intersperse :: Monad m => a -> Stream m a -> Stream m a
1331intersperse a = intersperseM (return a)
1333-- | Insert a side effect before consuming an element of a stream except the
1334-- first one.
1336-- >>> input = Stream.fromList "hello"
1337-- >>> Stream.fold Fold.drain $ Stream.trace putChar $ Stream.intersperseM_ (putChar '.') input
1338-- h.e.l.l.o
1340-- /Pre-release/
1341{-# INLINE_NORMAL intersperseM_ #-}
1342intersperseM_ :: Monad m => m b -> Stream m a -> Stream m a
1343intersperseM_ m (Stream step1 state1) = Stream step (Left (pure (), state1))
1344 where
1345 {-# INLINE_LATE step #-}
1346 step gst (Left (eff, st)) = do
1347 r <- step1 gst st
1348 case r of
1349 Yield x s -> eff >> return (Yield x (Right s))
1350 Skip s -> return $ Skip (Left (eff, s))
1351 Stop -> return Stop
1353 step _ (Right st) = return $ Skip $ Left (void m, st)
1355-- | Intersperse a monadic action into the input stream after every @n@
1356-- elements.
1358-- >> input = Stream.fromList "hello"
1359-- >> Stream.fold Fold.toList $ Stream.intersperseMWith 2 (return ',') input
1360-- "he,ll,o"
1362-- /Unimplemented/
1363{-# INLINE intersperseMWith #-}
1364intersperseMWith :: -- Monad m =>
1365 Int -> m a -> Stream m a -> Stream m a
1366intersperseMWith _n _f _xs = undefined
1368data SuffixState s a
1369 = SuffixElem s
1370 | SuffixSuffix s
1371 | SuffixYield a (SuffixState s a)
1373-- | Insert an effect and its output after consuming an element of a stream.
1375-- >>> input = Stream.fromList "hello"
1376-- >>> Stream.fold Fold.toList $ Stream.trace putChar $ Stream.intersperseMSuffix (putChar '.' >> return ',') input
1377-- h.,e.,l.,l.,o.,"h,e,l,l,o,"
1379-- /Pre-release/
1380{-# INLINE_NORMAL intersperseMSuffix #-}
1381intersperseMSuffix :: forall m a. Monad m => m a -> Stream m a -> Stream m a
1382intersperseMSuffix action (Stream step state) = Stream step' (SuffixElem state)
1383 where
1384 {-# INLINE_LATE step' #-}
1385 step' gst (SuffixElem st) = do
1386 r <- step gst st
1387 return $ case r of
1388 Yield x s -> Skip (SuffixYield x (SuffixSuffix s))
1389 Skip s -> Skip (SuffixElem s)
1390 Stop -> Stop
1392 step' _ (SuffixSuffix st) = do
1393 action >>= \r -> return $ Skip (SuffixYield r (SuffixElem st))
1395 step' _ (SuffixYield x next) = return $ Yield x next
1397-- | Insert a side effect after consuming an element of a stream.
1399-- >>> input = Stream.fromList "hello"
1400-- >>> Stream.fold Fold.toList $ Stream.intersperseMSuffix_ (threadDelay 1000000) input
1401-- "hello"
1403-- /Pre-release/
1405{-# INLINE_NORMAL intersperseMSuffix_ #-}
1406intersperseMSuffix_ :: Monad m => m b -> Stream m a -> Stream m a
1407intersperseMSuffix_ m (Stream step1 state1) = Stream step (Left state1)
1408 where
1409 {-# INLINE_LATE step #-}
1410 step gst (Left st) = do
1411 r <- step1 gst st
1412 case r of
1413 Yield x s -> return $ Yield x (Right s)
1414 Skip s -> return $ Skip $ Left s
1415 Stop -> return Stop
1417 step _ (Right st) = m >> return (Skip (Left st))
1419data SuffixSpanState s a
1420 = SuffixSpanElem s Int
1421 | SuffixSpanSuffix s
1422 | SuffixSpanYield a (SuffixSpanState s a)
1423 | SuffixSpanLast
1424 | SuffixSpanStop
1426-- | Like 'intersperseMSuffix' but intersperses an effectful action into the
1427-- input stream after every @n@ elements and after the last element.
1429-- >>> input = Stream.fromList "hello"
1430-- >>> Stream.fold Fold.toList $ Stream.intersperseMSuffixWith 2 (return ',') input
1431-- "he,ll,o,"
1433-- /Pre-release/
1435{-# INLINE_NORMAL intersperseMSuffixWith #-}
1436intersperseMSuffixWith :: forall m a. Monad m
1437 => Int -> m a -> Stream m a -> Stream m a
1438intersperseMSuffixWith n action (Stream step state) =
1439 Stream step' (SuffixSpanElem state n)
1440 where
1441 {-# INLINE_LATE step' #-}
1442 step' gst (SuffixSpanElem st i) | i > 0 = do
1443 r <- step gst st
1444 return $ case r of
1445 Yield x s -> Skip (SuffixSpanYield x (SuffixSpanElem s (i - 1)))
1446 Skip s -> Skip (SuffixSpanElem s i)
1447 Stop -> if i == n then Stop else Skip SuffixSpanLast
1448 step' _ (SuffixSpanElem st _) = return $ Skip (SuffixSpanSuffix st)
1450 step' _ (SuffixSpanSuffix st) = do
1451 action >>= \r -> return $ Skip (SuffixSpanYield r (SuffixSpanElem st n))
1453 step' _ SuffixSpanLast = do
1454 action >>= \r -> return $ Skip (SuffixSpanYield r SuffixSpanStop)
1456 step' _ (SuffixSpanYield x next) = return $ Yield x next
1458 step' _ SuffixSpanStop = return Stop
1460-- | Insert a side effect before consuming an element of a stream.
1462-- Definition:
1464-- >>> intersperseMPrefix_ m = Stream.mapM (\x -> void m >> return x)
1466-- >>> input = Stream.fromList "hello"
1467-- >>> Stream.fold Fold.toList $ Stream.trace putChar $ Stream.intersperseMPrefix_ (putChar '.' >> return ',') input
1468-- .h.e.l.l.o"hello"
1470-- Same as 'trace_'.
1472-- /Pre-release/
1474{-# INLINE intersperseMPrefix_ #-}
1475intersperseMPrefix_ :: Monad m => m b -> Stream m a -> Stream m a
1476intersperseMPrefix_ m = mapM (\x -> void m >> return x)
1478------------------------------------------------------------------------------
1479-- Inserting Time
1480------------------------------------------------------------------------------
1482-- XXX This should be in Prelude, should we export this as a helper function?
1484-- | Block the current thread for specified number of seconds.
1485{-# INLINE sleep #-}
1486sleep :: MonadIO m => Double -> m ()
1487sleep n = liftIO $ threadDelay $ round $ n * 1000000
1489-- | Introduce a delay of specified seconds between elements of the stream.
1491-- Definition:
1493-- >>> sleep n = liftIO $ threadDelay $ round $ n * 1000000
1494-- >>> delay = Stream.intersperseM_ . sleep
1496-- Example:
1498-- >>> input = Stream.enumerateFromTo 1 3
1499-- >>> Stream.fold (Fold.drainMapM print) $ Stream.delay 1 input
1500-- 1
1501-- 2
1502-- 3
1504{-# INLINE delay #-}
1505delay :: MonadIO m => Double -> Stream m a -> Stream m a
1506delay = intersperseM_ . sleep
1508-- | Introduce a delay of specified seconds after consuming an element of a
1509-- stream.
1511-- Definition:
1513-- >>> sleep n = liftIO $ threadDelay $ round $ n * 1000000
1514-- >>> delayPost = Stream.intersperseMSuffix_ . sleep
1516-- Example:
1518-- >>> input = Stream.enumerateFromTo 1 3
1519-- >>> Stream.fold (Fold.drainMapM print) $ Stream.delayPost 1 input
1520-- 1
1521-- 2
1522-- 3
1524-- /Pre-release/
1526{-# INLINE delayPost #-}
1527delayPost :: MonadIO m => Double -> Stream m a -> Stream m a
1528delayPost n = intersperseMSuffix_ $ liftIO $ threadDelay $ round $ n * 1000000
1530-- | Introduce a delay of specified seconds before consuming an element of a
1531-- stream.
1533-- Definition:
1535-- >>> sleep n = liftIO $ threadDelay $ round $ n * 1000000
1536-- >>> delayPre = Stream.intersperseMPrefix_. sleep
1538-- Example:
1540-- >>> input = Stream.enumerateFromTo 1 3
1541-- >>> Stream.fold (Fold.drainMapM print) $ Stream.delayPre 1 input
1542-- 1
1543-- 2
1544-- 3
1546-- /Pre-release/
1548{-# INLINE delayPre #-}
1549delayPre :: MonadIO m => Double -> Stream m a -> Stream m a
1550delayPre = intersperseMPrefix_. sleep
1552------------------------------------------------------------------------------
1553-- Reordering
1554------------------------------------------------------------------------------
1556-- | Returns the elements of the stream in reverse order. The stream must be
1557-- finite. Note that this necessarily buffers the entire stream in memory.
1559-- Definition:
1561-- >>> reverse m = Stream.concatEffect $ Stream.fold Fold.toListRev m >>= return . Stream.fromList
1563{-# INLINE_NORMAL reverse #-}
1564reverse :: Monad m => Stream m a -> Stream m a
1565reverse m = concatEffect $ fold FL.toListRev m <&> fromList
1567reverse m = Stream step Nothing
1568 where
1569 {-# INLINE_LATE step #-}
1570 step _ Nothing = do
1571 xs <- foldl' (flip (:)) [] m
1572 return $ Skip (Just xs)
1573 step _ (Just (x:xs)) = return $ Yield x (Just xs)
1574 step _ (Just []) = return Stop
1577-- | Like 'reverse' but several times faster, requires an 'Unbox' instance.
1579-- /O(n) space/
1581-- /Pre-release/
1582{-# INLINE reverseUnbox #-}
1583reverseUnbox :: (MonadIO m, Unbox a) => Stream m a -> Stream m a
1584reverseUnbox =
1585 A.concatRev -- unfoldMany A.readerRev
1586 . fromStreamK
1587 . K.reverse
1588 . toStreamK
1589 . A.chunksOf defaultChunkSize
1591-- | Buffer until the next element in sequence arrives. The function argument
1592-- determines the difference in sequence numbers. This could be useful in
1593-- implementing sequenced streams, for example, TCP reassembly.
1595-- /Unimplemented/
1597{-# INLINE reassembleBy #-}
1598reassembleBy
1599 :: -- Monad m =>
1600 Fold m a b
1601 -> (a -> a -> Int)
1602 -> Stream m a
1603 -> Stream m b
1604reassembleBy = undefined
1606------------------------------------------------------------------------------
1607-- Position Indexing
1608------------------------------------------------------------------------------
1610-- Adapted from the vector package
1612-- |
1613-- >>> f = Fold.foldl' (\(i, _) x -> (i + 1, x)) (-1,undefined)
1614-- >>> indexed = Stream.postscan f
1615-- >>> indexed = Stream.zipWith (,) (Stream.enumerateFrom 0)
1616-- >>> indexedR n = fmap (\(i, a) -> (n - i, a)) . indexed
1618-- Pair each element in a stream with its index, starting from index 0.
1620-- >>> Stream.fold Fold.toList $ Stream.indexed $ Stream.fromList "hello"
1621-- [(0,'h'),(1,'e'),(2,'l'),(3,'l'),(4,'o')]
1623{-# INLINE_NORMAL indexed #-}
1624indexed :: Monad m => Stream m a -> Stream m (Int, a)
1625-- indexed = scanMaybe FL.indexing
1626indexed (Stream step state) = Stream step' (state, 0)
1627 where
1628 {-# INLINE_LATE step' #-}
1629 step' gst (st, i) = i `seq` do
1630 r <- step (adaptState gst) st
1631 case r of
1632 Yield x s -> return $ Yield (i, x) (s, i+1)
1633 Skip s -> return $ Skip (s, i)
1634 Stop -> return Stop
1636-- Adapted from the vector package
1638-- |
1639-- >>> f n = Fold.foldl' (\(i, _) x -> (i - 1, x)) (n + 1,undefined)
1640-- >>> indexedR n = Stream.postscan (f n)
1642-- >>> s n = Stream.enumerateFromThen n (n - 1)
1643-- >>> indexedR n = Stream.zipWith (,) (s n)
1645-- Pair each element in a stream with its index, starting from the
1646-- given index @n@ and counting down.
1648-- >>> Stream.fold Fold.toList $ Stream.indexedR 10 $ Stream.fromList "hello"
1649-- [(10,'h'),(9,'e'),(8,'l'),(7,'l'),(6,'o')]
1651{-# INLINE_NORMAL indexedR #-}
1652indexedR :: Monad m => Int -> Stream m a -> Stream m (Int, a)
1653-- indexedR n = scanMaybe (FL.indexingRev n)
1654indexedR m (Stream step state) = Stream step' (state, m)
1655 where
1656 {-# INLINE_LATE step' #-}
1657 step' gst (st, i) = i `seq` do
1658 r <- step (adaptState gst) st
1659 case r of
1660 Yield x s -> let i' = i - 1
1661 in return $ Yield (i, x) (s, i')
1662 Skip s -> return $ Skip (s, i)
1663 Stop -> return Stop
1665-------------------------------------------------------------------------------
1666-- Time Indexing
1667-------------------------------------------------------------------------------
1669-- Note: The timestamp stream must be the second stream in the zip so that the
1670-- timestamp is generated after generating the stream element and not before.
1671-- If we do not do that then the following example will generate the same
1672-- timestamp for first two elements:
1674-- Stream.fold Fold.toList $ Stream.timestamped $ Stream.delay $ Stream.enumerateFromTo 1 3
1676-- | Pair each element in a stream with an absolute timestamp, using a clock of
1677-- specified granularity. The timestamp is generated just before the element
1678-- is consumed.
1680-- >>> Stream.fold Fold.toList $ Stream.timestampWith 0.01 $ Stream.delay 1 $ Stream.enumerateFromTo 1 3
1681-- [(AbsTime (TimeSpec {sec = ..., nsec = ...}),1),(AbsTime (TimeSpec {sec = ..., nsec = ...}),2),(AbsTime (TimeSpec {sec = ..., nsec = ...}),3)]
1683-- /Pre-release/
1685{-# INLINE timestampWith #-}
1686timestampWith :: (MonadIO m)
1687 => Double -> Stream m a -> Stream m (AbsTime, a)
1688timestampWith g stream = zipWith (flip (,)) stream (absTimesWith g)
1690-- TBD: check performance vs a custom implementation without using zipWith.
1692-- /Pre-release/
1694{-# INLINE timestamped #-}
1695timestamped :: (MonadIO m)
1696 => Stream m a -> Stream m (AbsTime, a)
1697timestamped = timestampWith 0.01
1699-- | Pair each element in a stream with relative times starting from 0, using a
1700-- clock with the specified granularity. The time is measured just before the
1701-- element is consumed.
1703-- >>> Stream.fold Fold.toList $ Stream.timeIndexWith 0.01 $ Stream.delay 1 $ Stream.enumerateFromTo 1 3
1704-- [(RelTime64 (NanoSecond64 ...),1),(RelTime64 (NanoSecond64 ...),2),(RelTime64 (NanoSecond64 ...),3)]
1706-- /Pre-release/
1708{-# INLINE timeIndexWith #-}
1709timeIndexWith :: (MonadIO m)
1710 => Double -> Stream m a -> Stream m (RelTime64, a)
1711timeIndexWith g stream = zipWith (flip (,)) stream (relTimesWith g)
1713-- | Pair each element in a stream with relative times starting from 0, using a
1714-- 10 ms granularity clock. The time is measured just before the element is
1715-- consumed.
1717-- >>> Stream.fold Fold.toList $ Stream.timeIndexed $ Stream.delay 1 $ Stream.enumerateFromTo 1 3
1718-- [(RelTime64 (NanoSecond64 ...),1),(RelTime64 (NanoSecond64 ...),2),(RelTime64 (NanoSecond64 ...),3)]
1720-- /Pre-release/
1722{-# INLINE timeIndexed #-}
1723timeIndexed :: (MonadIO m)
1724 => Stream m a -> Stream m (RelTime64, a)
1725timeIndexed = timeIndexWith 0.01
1727------------------------------------------------------------------------------
1728-- Searching
1729------------------------------------------------------------------------------
1731-- | Find all the indices where the element in the stream satisfies the given
1732-- predicate.
1734-- >>> findIndices p = Stream.scanMaybe (Fold.findIndices p)
1736{-# INLINE_NORMAL findIndices #-}
1737findIndices :: Monad m => (a -> Bool) -> Stream m a -> Stream m Int
1738findIndices p (Stream step state) = Stream step' (state, 0)
1739 where
1740 {-# INLINE_LATE step' #-}
1741 step' gst (st, i) = i `seq` do
1742 r <- step (adaptState gst) st
1743 return $ case r of
1744 Yield x s -> if p x then Yield i (s, i+1) else Skip (s, i+1)
1745 Skip s -> Skip (s, i)
1746 Stop -> Stop
1748-- | Find all the indices where the value of the element in the stream is equal
1749-- to the given value.
1751-- >>> elemIndices a = Stream.findIndices (== a)
1753{-# INLINE elemIndices #-}
1754elemIndices :: (Monad m, Eq a) => a -> Stream m a -> Stream m Int
1755elemIndices a = findIndices (== a)
1757------------------------------------------------------------------------------
1758-- Rolling map
1759------------------------------------------------------------------------------
1761data RollingMapState s a = RollingMapGo s a
1763-- | Like 'rollingMap' but with an effectful map function.
1765-- /Pre-release/
1767{-# INLINE rollingMapM #-}
1768rollingMapM :: Monad m => (Maybe a -> a -> m b) -> Stream m a -> Stream m b
1769-- rollingMapM f = scanMaybe (FL.slide2 $ Window.rollingMapM f)
1770rollingMapM f (Stream step1 state1) = Stream step (RollingMapGo state1 Nothing)
1772 where
1774 step gst (RollingMapGo s1 curr) = do
1775 r <- step1 (adaptState gst) s1
1776 case r of
1777 Yield x s -> do
1778 !res <- f curr x
1779 return $ Yield res $ RollingMapGo s (Just x)
1780 Skip s -> return $ Skip $ RollingMapGo s curr
1781 Stop -> return Stop
1783-- rollingMap is a special case of an incremental sliding fold. It can be
1784-- written as:
1786-- > fld f = slidingWindow 1 (Fold.foldl' (\_ (x,y) -> f y x)
1787-- > rollingMap f = Stream.postscan (fld f) undefined
1789-- | Apply a function on every two successive elements of a stream. The first
1790-- argument of the map function is the previous element and the second argument
1791-- is the current element. When the current element is the first element, the
1792-- previous element is 'Nothing'.
1794-- /Pre-release/
1796{-# INLINE rollingMap #-}
1797rollingMap :: Monad m => (Maybe a -> a -> b) -> Stream m a -> Stream m b
1798-- rollingMap f = scanMaybe (FL.slide2 $ Window.rollingMap f)
1799rollingMap f = rollingMapM (\x y -> return $ f x y)
1801-- | Like 'rollingMap' but requires at least two elements in the stream,
1802-- returns an empty stream otherwise.
1804-- This is the stream equivalent of the list idiom @zipWith f xs (tail xs)@.
1806-- /Pre-release/
1808{-# INLINE rollingMap2 #-}
1809rollingMap2 :: Monad m => (a -> a -> b) -> Stream m a -> Stream m b
1810rollingMap2 f = catMaybes . rollingMap g
1812 where
1814 g Nothing _ = Nothing
1815 g (Just x) y = Just (f x y)
1817------------------------------------------------------------------------------
1818-- Maybe Streams
1819------------------------------------------------------------------------------
1821-- XXX Will this always fuse properly?
1823-- | Map a 'Maybe' returning function to a stream, filter out the 'Nothing'
1824-- elements, and return a stream of values extracted from 'Just'.
1826-- Equivalent to:
1828-- >>> mapMaybe f = Stream.catMaybes . fmap f
1830{-# INLINE_NORMAL mapMaybe #-}
1831mapMaybe :: Monad m => (a -> Maybe b) -> Stream m a -> Stream m b
1832mapMaybe f = fmap fromJust . filter isJust . map f
1834-- | Like 'mapMaybe' but maps a monadic function.
1836-- Equivalent to:
1838-- >>> mapMaybeM f = Stream.catMaybes . Stream.mapM f
1840-- >>> mapM f = Stream.mapMaybeM (\x -> Just <$> f x)
1842{-# INLINE_NORMAL mapMaybeM #-}
1843mapMaybeM :: Monad m => (a -> m (Maybe b)) -> Stream m a -> Stream m b
1844mapMaybeM f = fmap fromJust . filter isJust . mapM f
1846-- | In a stream of 'Maybe's, discard 'Nothing's and unwrap 'Just's.
1848-- >>> catMaybes = Stream.mapMaybe id
1849-- >>> catMaybes = fmap fromJust . Stream.filter isJust
1851-- /Pre-release/
1853{-# INLINE catMaybes #-}
1854catMaybes :: Monad m => Stream m (Maybe a) -> Stream m a
1855-- catMaybes = fmap fromJust . filter isJust
1856catMaybes (Stream step state) = Stream step1 state
1858 where
1860 {-# INLINE_LATE step1 #-}
1861 step1 gst st = do
1862 r <- step (adaptState gst) st
1863 case r of
1864 Yield x s -> do
1865 return
1866 $ case x of
1867 Just a -> Yield a s
1868 Nothing -> Skip s
1869 Skip s -> return $ Skip s
1870 Stop -> return Stop
1872-- | Use a filtering fold on a stream.
1874-- >>> scanMaybe f = Stream.catMaybes . Stream.postscan f
1876{-# INLINE scanMaybe #-}
1877scanMaybe :: Monad m => Fold m a (Maybe b) -> Stream m a -> Stream m b
1878scanMaybe f = catMaybes . postscan f
1880------------------------------------------------------------------------------
1881-- Either streams
1882------------------------------------------------------------------------------
1884-- | Discard 'Right's and unwrap 'Left's in an 'Either' stream.
1886-- >>> catLefts = fmap (fromLeft undefined) . Stream.filter isLeft
1888-- /Pre-release/
1890{-# INLINE catLefts #-}
1891catLefts :: Monad m => Stream m (Either a b) -> Stream m a
1892catLefts = fmap (fromLeft undefined) . filter isLeft
1894-- | Discard 'Left's and unwrap 'Right's in an 'Either' stream.
1896-- >>> catRights = fmap (fromRight undefined) . Stream.filter isRight
1898-- /Pre-release/
1900{-# INLINE catRights #-}
1901catRights :: Monad m => Stream m (Either a b) -> Stream m b
1902catRights = fmap (fromRight undefined) . filter isRight
1904-- | Remove the either wrapper and flatten both lefts and as well as rights in
1905-- the output stream.
1907-- >>> catEithers = fmap (either id id)
1909-- /Pre-release/
1911{-# INLINE catEithers #-}
1912catEithers :: Monad m => Stream m (Either a a) -> Stream m a
1913catEithers = fmap (either id id)
1915------------------------------------------------------------------------------
1916-- Splitting
1917------------------------------------------------------------------------------
1919-- | Split on an infixed separator element, dropping the separator. The
1920-- supplied 'Fold' is applied on the split segments. Splits the stream on
1921-- separator elements determined by the supplied predicate, separator is
1922-- considered as infixed between two segments:
1924-- >>> splitOn' p xs = Stream.fold Fold.toList $ Stream.splitOn p Fold.toList (Stream.fromList xs)
1925-- >>> splitOn' (== '.') "a.b"
1926-- ["a","b"]
1928-- An empty stream is folded to the default value of the fold:
1930-- >>> splitOn' (== '.') ""
1931-- [""]
1933-- If one or both sides of the separator are missing then the empty segment on
1934-- that side is folded to the default output of the fold:
1936-- >>> splitOn' (== '.') "."
1937-- ["",""]
1939-- >>> splitOn' (== '.') ".a"
1940-- ["","a"]
1942-- >>> splitOn' (== '.') "a."
1943-- ["a",""]
1945-- >>> splitOn' (== '.') "a..b"
1946-- ["a","","b"]
1948-- splitOn is an inverse of intercalating single element:
1950-- > Stream.intercalate (Stream.fromPure '.') Unfold.fromList . Stream.splitOn (== '.') Fold.toList === id
1952-- Assuming the input stream does not contain the separator:
1954-- > Stream.splitOn (== '.') Fold.toList . Stream.intercalate (Stream.fromPure '.') Unfold.fromList === id
1956{-# INLINE splitOn #-}
1957splitOn :: Monad m => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1958splitOn predicate f =
1959 -- We can express the infix splitting in terms of optional suffix split
1960 -- fold. After applying a suffix split fold repeatedly if the last segment
1961 -- ends with a suffix then we need to return the default output of the fold
1962 -- after that to make it an infix split.
1963 --
1964 -- Alternately, we can also express it using an optional prefix split fold.
1965 -- If the first segment starts with a prefix then we need to emit the
1966 -- default output of the fold before that to make it an infix split, and
1967 -- then apply prefix split fold repeatedly.
1968 --
1969 -- Since a suffix split fold can be easily expressed using a
1970 -- non-backtracking fold, we use that.
1971 foldManyPost (FL.takeEndBy_ predicate f)