on 79df417 help sign in
1945 lines 63 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 , map
25 , mapM
26 , sequence
28 -- * Mapping Effects
29 , tap
30 , tapOffsetEvery
31 , trace
32 , trace_
34 -- * Folding
35 , foldrS
36 , foldlS
38 -- * Scanning By 'Fold'
39 , postscan
40 , scan
41 , scanMany
43 -- * Splitting
44 , splitOn
46 -- * Scanning
47 -- | Left scans. Stateful, mostly one-to-one maps.
48 , scanlM'
49 , scanlMAfter'
50 , scanl'
51 , scanlM
52 , scanl
53 , scanl1M'
54 , scanl1'
55 , scanl1M
56 , scanl1
58 , prescanl'
59 , prescanlM'
61 , postscanl
62 , postscanlM
63 , postscanl'
64 , postscanlM'
65 , postscanlMAfter'
67 , postscanlx'
68 , postscanlMx'
69 , scanlMx'
70 , scanlx'
72 -- * Filtering
73 -- | Produce a subset of the stream.
74 , with
75 , scanMaybe
76 , filter
77 , filterM
78 , deleteBy
79 , uniqBy
80 , uniq
81 , prune
82 , repeated
84 -- * Trimming
85 -- | Produce a subset of the stream trimmed at ends.
86 , take
87 , takeWhile
88 , takeWhileM
89 , takeWhileLast
90 , takeWhileAround
91 , drop
92 , dropWhile
93 , dropWhileM
94 , dropLast
95 , dropWhileLast
96 , dropWhileAround
98 -- * Inserting Elements
99 -- | Produce a superset of the stream.
100 , insertBy
101 , intersperse
102 , intersperseM
103 , intersperseMWith
104 , intersperseMSuffix
105 , intersperseMSuffixWith
107 -- * Inserting Side Effects
108 , intersperseM_
109 , intersperseMSuffix_
110 , intersperseMPrefix_
112 , delay
113 , delayPre
114 , delayPost
116 -- * Reordering
117 -- | Produce strictly the same set but reordered.
118 , reverse
119 , reverseUnbox
120 , reassembleBy
122 -- * Position Indexing
123 , indexed
124 , indexedR
126 -- * Time Indexing
127 , timestampWith
128 , timestamped
129 , timeIndexWith
130 , timeIndexed
132 -- * Searching
133 , findIndices
134 , elemIndices
135 , slicesBy
137 -- * Rolling map
138 -- | Map using the previous element.
139 , rollingMap
140 , rollingMapM
141 , rollingMap2
143 -- * Maybe Streams
144 , mapMaybe
145 , mapMaybeM
146 , catMaybes
148 -- * Either Streams
149 , catLefts
150 , catRights
151 , catEithers
152 )
153where
155#include "inline.hs"
157import Control.Concurrent (threadDelay)
158import Control.Monad (void)
159import Control.Monad.IO.Class (MonadIO (liftIO))
160import Data.Either (fromLeft, isLeft, isRight, fromRight)
161import Data.Functor ((<&>))
162import Data.Maybe (fromJust, isJust)
163import Fusion.Plugin.Types (Fuse(..))
165import Streamly.Internal.Data.Fold.Type (Fold(..))
166import Streamly.Internal.Data.Pipe.Type (Pipe(..), PipeState(..))
167import Streamly.Internal.Data.SVar.Type (adaptState)
168import Streamly.Internal.Data.Time.Units (AbsTime, RelTime64)
169import Streamly.Internal.Data.Unbox (Unbox)
170import Streamly.Internal.System.IO (defaultChunkSize)
172-- import qualified Data.List as List
173import qualified Streamly.Internal.Data.Array.Type as A
174import qualified Streamly.Internal.Data.Fold as FL
175import qualified Streamly.Internal.Data.Pipe.Type as Pipe
176import qualified Streamly.Internal.Data.StreamK.Type as K
178import Prelude hiding
179 ( drop, dropWhile, filter, map, mapM, reverse
180 , scanl, scanl1, sequence, take, takeWhile, zipWith)
182import Streamly.Internal.Data.Stream.Generate
183 (absTimesWith, relTimesWith)
184import Streamly.Internal.Data.Stream.Type
186#include "DocTestDataStream.hs"
188------------------------------------------------------------------------------
189-- Piping
190------------------------------------------------------------------------------
192-- | Use a 'Pipe' to transform a stream.
193--
194-- /Pre-release/
195--
196{-# INLINE_NORMAL transform #-}
197transform :: Monad m => Pipe m a b -> Stream m a -> Stream m b
198transform (Pipe pstep1 pstep2 pstate) (Stream step state) =
199 Stream step' (Consume pstate, state)
201 where
203 {-# INLINE_LATE step' #-}
205 step' gst (Consume pst, st) = pst `seq` do
206 r <- step (adaptState gst) st
207 case r of
208 Yield x s -> do
209 res <- pstep1 pst x
210 case res of
211 Pipe.Yield b pst' -> return $ Yield b (pst', s)
212 Pipe.Continue pst' -> return $ Skip (pst', s)
213 Skip s -> return $ Skip (Consume pst, s)
214 Stop -> return Stop
216 step' _ (Produce pst, st) = pst `seq` do
217 res <- pstep2 pst
218 case res of
219 Pipe.Yield b pst' -> return $ Yield b (pst', st)
220 Pipe.Continue pst' -> return $ Skip (pst', st)
222------------------------------------------------------------------------------
223-- Transformation Folds
224------------------------------------------------------------------------------
226-- Note, this is going to have horrible performance, because of the nature of
227-- the stream type (i.e. direct stream vs CPS). Its only for reference, it is
228-- likely be practically unusable.
229{-# INLINE_NORMAL foldlS #-}
230foldlS :: Monad m
231 => (Stream m b -> a -> Stream m b) -> Stream m b -> Stream m a -> Stream m b
232foldlS fstep begin (Stream step state) = Stream step' (Left (state, begin))
233 where
234 step' gst (Left (st, acc)) = do
235 r <- step (adaptState gst) st
236 return $ case r of
237 Yield x s -> Skip (Left (s, fstep acc x))
238 Skip s -> Skip (Left (s, acc))
239 Stop -> Skip (Right acc)
241 step' gst (Right (Stream stp stt)) = do
242 r <- stp (adaptState gst) stt
243 return $ case r of
244 Yield x s -> Yield x (Right (Stream stp s))
245 Skip s -> Skip (Right (Stream stp s))
246 Stop -> Stop
248------------------------------------------------------------------------------
249-- Transformation by Mapping
250------------------------------------------------------------------------------
252-- |
253-- >>> sequence = Stream.mapM id
254--
255-- Replace the elements of a stream of monadic actions with the outputs of
256-- those actions.
257--
258-- >>> s = Stream.fromList [putStr "a", putStr "b", putStrLn "c"]
259-- >>> Stream.fold Fold.drain $ Stream.sequence s
260-- abc
261--
262{-# INLINE_NORMAL sequence #-}
263sequence :: Monad m => Stream m (m a) -> Stream m a
264sequence (Stream step state) = Stream step' state
265 where
266 {-# INLINE_LATE step' #-}
267 step' gst st = do
268 r <- step (adaptState gst) st
269 case r of
270 Yield x s -> x >>= \a -> return (Yield a s)
271 Skip s -> return $ Skip s
272 Stop -> return Stop
274------------------------------------------------------------------------------
275-- Mapping side effects
276------------------------------------------------------------------------------
278data TapState fs st a
279 = TapInit | Tapping !fs st | TapDone st
281-- XXX Multiple yield points
283-- | Tap the data flowing through a stream into a 'Fold'. For example, you may
284-- add a tap to log the contents flowing through the stream. The fold is used
285-- only for effects, its result is discarded.
286--
287-- @
288-- Fold m a b
289-- |
290-- -----stream m a ---------------stream m a-----
291--
292-- @
293--
294-- >>> s = Stream.enumerateFromTo 1 2
295-- >>> Stream.fold Fold.drain $ Stream.tap (Fold.drainMapM print) s
296-- 1
297-- 2
298--
299-- Compare with 'trace'.
300--
301{-# INLINE tap #-}
302tap :: Monad m => Fold m a b -> Stream m a -> Stream m a
303tap (Fold fstep initial extract) (Stream step state) = Stream step' TapInit
305 where
307 step' _ TapInit = do
308 res <- initial
309 return
310 $ Skip
311 $ case res of
312 FL.Partial s -> Tapping s state
313 FL.Done _ -> TapDone state
314 step' gst (Tapping acc st) = do
315 r <- step gst st
316 case r of
317 Yield x s -> do
318 res <- fstep acc x
319 return
320 $ Yield x
321 $ case res of
322 FL.Partial fs -> Tapping fs s
323 FL.Done _ -> TapDone s
324 Skip s -> return $ Skip (Tapping acc s)
325 Stop -> do
326 void $ extract acc
327 return Stop
328 step' gst (TapDone st) = do
329 r <- step gst st
330 return
331 $ case r of
332 Yield x s -> Yield x (TapDone s)
333 Skip s -> Skip (TapDone s)
334 Stop -> Stop
336data TapOffState fs s a
337 = TapOffInit
338 | TapOffTapping !fs s Int
339 | TapOffDone s
341-- XXX Multiple yield points
342{-# INLINE_NORMAL tapOffsetEvery #-}
343tapOffsetEvery :: Monad m
344 => Int -> Int -> Fold m a b -> Stream m a -> Stream m a
345tapOffsetEvery offset n (Fold fstep initial extract) (Stream step state) =
346 Stream step' TapOffInit
348 where
350 {-# INLINE_LATE step' #-}
351 step' _ TapOffInit = do
352 res <- initial
353 return
354 $ Skip
355 $ case res of
356 FL.Partial s -> TapOffTapping s state (offset `mod` n)
357 FL.Done _ -> TapOffDone state
358 step' gst (TapOffTapping acc st count) = do
359 r <- step gst st
360 case r of
361 Yield x s -> do
362 next <-
363 if count <= 0
364 then do
365 res <- fstep acc x
366 return
367 $ case res of
368 FL.Partial sres ->
369 TapOffTapping sres s (n - 1)
370 FL.Done _ -> TapOffDone s
371 else return $ TapOffTapping acc s (count - 1)
372 return $ Yield x next
373 Skip s -> return $ Skip (TapOffTapping acc s count)
374 Stop -> do
375 void $ extract acc
376 return Stop
377 step' gst (TapOffDone st) = do
378 r <- step gst st
379 return
380 $ case r of
381 Yield x s -> Yield x (TapOffDone s)
382 Skip s -> Skip (TapOffDone s)
383 Stop -> Stop
385-- | Apply a monadic function to each element flowing through the stream and
386-- discard the results.
387--
388-- >>> s = Stream.enumerateFromTo 1 2
389-- >>> Stream.fold Fold.drain $ Stream.trace print s
390-- 1
391-- 2
392--
393-- Compare with 'tap'.
394--
395{-# INLINE trace #-}
396trace :: Monad m => (a -> m b) -> Stream m a -> Stream m a
397trace f = mapM (\x -> void (f x) >> return x)
399-- | Perform a side effect before yielding each element of the stream and
400-- discard the results.
401--
402-- >>> s = Stream.enumerateFromTo 1 2
403-- >>> Stream.fold Fold.drain $ Stream.trace_ (print "got here") s
404-- "got here"
405-- "got here"
406--
407-- Same as 'intersperseMPrefix_' but always serial.
408--
409-- See also: 'trace'
410--
411-- /Pre-release/
412{-# INLINE trace_ #-}
413trace_ :: Monad m => m b -> Stream m a -> Stream m a
414trace_ eff = mapM (\x -> eff >> return x)
416------------------------------------------------------------------------------
417-- Scanning with a Fold
418------------------------------------------------------------------------------
420data ScanState s f = ScanInit s | ScanDo s !f | ScanDone
422-- | Postscan a stream using the given monadic fold.
423--
424-- The following example extracts the input stream up to a point where the
425-- running average of elements is no more than 10:
426--
427-- >>> import Data.Maybe (fromJust)
428-- >>> let avg = Fold.teeWith (/) Fold.sum (fmap fromIntegral Fold.length)
429-- >>> s = Stream.enumerateFromTo 1.0 100.0
430-- >>> :{
431-- Stream.fold Fold.toList
432-- $ fmap (fromJust . fst)
433-- $ Stream.takeWhile (\(_,x) -> x <= 10)
434-- $ Stream.postscan (Fold.tee Fold.latest avg) s
435-- :}
436-- [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]
437--
438{-# INLINE_NORMAL postscan #-}
439postscan :: Monad m => FL.Fold m a b -> Stream m a -> Stream m b
440postscan (FL.Fold fstep initial extract) (Stream sstep state) =
441 Stream step (ScanInit state)
443 where
445 {-# INLINE_LATE step #-}
446 step _ (ScanInit st) = do
447 res <- initial
448 return
449 $ case res of
450 FL.Partial fs -> Skip $ ScanDo st fs
451 FL.Done b -> Yield b ScanDone
452 step gst (ScanDo st fs) = do
453 res <- sstep (adaptState gst) st
454 case res of
455 Yield x s -> do
456 r <- fstep fs x
457 case r of
458 FL.Partial fs1 -> do
459 !b <- extract fs1
460 return $ Yield b $ ScanDo s fs1
461 FL.Done b -> return $ Yield b ScanDone
462 Skip s -> return $ Skip $ ScanDo s fs
463 Stop -> return Stop
464 step _ ScanDone = return Stop
466{-# INLINE scanWith #-}
467scanWith :: Monad m
468 => Bool -> Fold m a b -> Stream m a -> Stream m b
469scanWith restart (Fold fstep initial extract) (Stream sstep state) =
470 Stream step (ScanInit state)
472 where
474 {-# INLINE runStep #-}
475 runStep st action = do
476 res <- action
477 case res of
478 FL.Partial fs -> do
479 !b <- extract fs
480 return $ Yield b $ ScanDo st fs
481 FL.Done b ->
482 let next = if restart then ScanInit st else ScanDone
483 in return $ Yield b next
485 {-# INLINE_LATE step #-}
486 step _ (ScanInit st) = runStep st initial
487 step gst (ScanDo st fs) = do
488 res <- sstep (adaptState gst) st
489 case res of
490 Yield x s -> runStep s (fstep fs x)
491 Skip s -> return $ Skip $ ScanDo s fs
492 Stop -> return Stop
493 step _ ScanDone = return Stop
495-- XXX It may be useful to have a version of scan where we can keep the
496-- accumulator independent of the value emitted. So that we do not necessarily
497-- have to keep a value in the accumulator which we are not using. We can pass
498-- an extraction function that will take the accumulator and the current value
499-- of the element and emit the next value in the stream. That will also make it
500-- possible to modify the accumulator after using it. In fact, the step function
501-- can return new accumulator and the value to be emitted. The signature would
502-- be more like mapAccumL.
504-- | Strict left scan. Scan a stream using the given monadic fold.
505--
506-- >>> s = Stream.fromList [1..10]
507-- >>> Stream.fold Fold.toList $ Stream.takeWhile (< 10) $ Stream.scan Fold.sum s
508-- [0,1,3,6]
509--
510-- See also: 'usingStateT'
511--
513-- EXPLANATION:
514-- >>> scanl' step z = Stream.scan (Fold.foldl' step z)
515--
516-- Like 'map', 'scanl'' too is a one to one transformation,
517-- however it adds an extra element.
518--
519-- >>> s = Stream.fromList [1,2,3,4]
520-- >>> Stream.fold Fold.toList $ scanl' (+) 0 s
521-- [0,1,3,6,10]
522--
523-- >>> Stream.fold Fold.toList $ scanl' (flip (:)) [] s
524-- [[],[1],[2,1],[3,2,1],[4,3,2,1]]
525--
526-- The output of 'scanl'' is the initial value of the accumulator followed by
527-- all the intermediate steps and the final result of 'foldl''.
528--
529-- By streaming the accumulated state after each fold step, we can share the
530-- state across multiple stages of stream composition. Each stage can modify or
531-- extend the state, do some processing with it and emit it for the next stage,
532-- thus modularizing the stream processing. This can be useful in
533-- stateful or event-driven programming.
534--
535-- Consider the following monolithic example, computing the sum and the product
536-- of the elements in a stream in one go using a @foldl'@:
537--
538-- >>> foldl' step z = Stream.fold (Fold.foldl' step z)
539-- >>> foldl' (\(s, p) x -> (s + x, p * x)) (0,1) s
540-- (10,24)
541--
542-- Using @scanl'@ we can make it modular by computing the sum in the first
543-- stage and passing it down to the next stage for computing the product:
544--
545-- >>> :{
546-- foldl' (\(_, p) (s, x) -> (s, p * x)) (0,1)
547-- $ scanl' (\(s, _) x -> (s + x, x)) (0,1)
548-- $ Stream.fromList [1,2,3,4]
549-- :}
550-- (10,24)
551--
552-- IMPORTANT: 'scanl'' evaluates the accumulator to WHNF. To avoid building
553-- lazy expressions inside the accumulator, it is recommended that a strict
554-- data structure is used for accumulator.
555--
556{-# INLINE_NORMAL scan #-}
557scan :: Monad m
558 => FL.Fold m a b -> Stream m a -> Stream m b
559scan = scanWith False
561-- | Like 'scan' but restarts scanning afresh when the scanning fold
562-- terminates.
563--
564{-# INLINE_NORMAL scanMany #-}
565scanMany :: Monad m
566 => FL.Fold m a b -> Stream m a -> Stream m b
567scanMany = scanWith True
569------------------------------------------------------------------------------
570-- Scanning - Prescans
571------------------------------------------------------------------------------
573-- Adapted from the vector package.
574--
575-- XXX Is a prescan useful, discarding the last step does not sound useful? I
576-- am not sure about the utility of this function, so this is implemented but
577-- not exposed. We can expose it if someone provides good reasons why this is
578-- useful.
579--
580-- XXX We have to execute the stream one step ahead to know that we are at the
581-- last step. The vector implementation of prescan executes the last fold step
582-- but does not yield the result. This means we have executed the effect but
583-- discarded value. This does not sound right. In this implementation we are
584-- not executing the last fold step.
585{-# INLINE_NORMAL prescanlM' #-}
586prescanlM' :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
587prescanlM' f mz (Stream step state) = Stream step' (state, mz)
588 where
589 {-# INLINE_LATE step' #-}
590 step' gst (st, prev) = do
591 r <- step (adaptState gst) st
592 case r of
593 Yield x s -> do
594 acc <- prev
595 return $ Yield acc (s, f acc x)
596 Skip s -> return $ Skip (s, prev)
597 Stop -> return Stop
599{-# INLINE prescanl' #-}
600prescanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
601prescanl' f z = prescanlM' (\a b -> return (f a b)) (return z)
603------------------------------------------------------------------------------
604-- Monolithic postscans (postscan followed by a map)
605------------------------------------------------------------------------------
607-- The performance of a modular postscan followed by a map seems to be
608-- equivalent to this monolithic scan followed by map therefore we may not need
609-- this implementation. We just have it for performance comparison and in case
610-- modular version does not perform well in some situation.
611--
612{-# INLINE_NORMAL postscanlMx' #-}
613postscanlMx' :: Monad m
614 => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> Stream m b
615postscanlMx' fstep begin done (Stream step state) = do
616 Stream step' (state, begin)
617 where
618 {-# INLINE_LATE step' #-}
619 step' gst (st, acc) = do
620 r <- step (adaptState gst) st
621 case r of
622 Yield x s -> do
623 old <- acc
624 y <- fstep old x
625 v <- done y
626 v `seq` y `seq` return (Yield v (s, return y))
627 Skip s -> return $ Skip (s, acc)
628 Stop -> return Stop
630{-# INLINE_NORMAL postscanlx' #-}
631postscanlx' :: Monad m
632 => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> Stream m b
633postscanlx' fstep begin done =
634 postscanlMx' (\b a -> return (fstep b a)) (return begin) (return . done)
636-- XXX do we need consM strict to evaluate the begin value?
637{-# INLINE scanlMx' #-}
638scanlMx' :: Monad m
639 => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> Stream m b
640scanlMx' fstep begin done s =
641 (begin >>= \x -> x `seq` done x) `consM` postscanlMx' fstep begin done s
643{-# INLINE scanlx' #-}
644scanlx' :: Monad m
645 => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> Stream m b
646scanlx' fstep begin done =
647 scanlMx' (\b a -> return (fstep b a)) (return begin) (return . done)
649------------------------------------------------------------------------------
650-- postscans
651------------------------------------------------------------------------------
653-- Adapted from the vector package.
654{-# INLINE_NORMAL postscanlM' #-}
655postscanlM' :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
656postscanlM' fstep begin (Stream step state) =
657 Stream step' Nothing
658 where
659 {-# INLINE_LATE step' #-}
660 step' _ Nothing = do
661 !x <- begin
662 return $ Skip (Just (state, x))
664 step' gst (Just (st, acc)) = do
665 r <- step (adaptState gst) st
666 case r of
667 Yield x s -> do
668 !y <- fstep acc x
669 return $ Yield y (Just (s, y))
670 Skip s -> return $ Skip (Just (s, acc))
671 Stop -> return Stop
673{-# INLINE_NORMAL postscanl' #-}
674postscanl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
675postscanl' f seed = postscanlM' (\a b -> return (f a b)) (return seed)
677{-# ANN type PScanAfterState Fuse #-}
678data PScanAfterState m st acc =
679 PScanAfterStep st (m acc)
680 | PScanAfterYield acc (PScanAfterState m st acc)
681 | PScanAfterStop
683-- We can possibly have the "done" function as a Maybe to provide an option to
684-- emit or not emit the accumulator when the stream stops.
685--
686-- TBD: use a single Yield point
687--
688{-# INLINE_NORMAL postscanlMAfter' #-}
689postscanlMAfter' :: Monad m
690 => (b -> a -> m b) -> m b -> (b -> m b) -> Stream m a -> Stream m b
691postscanlMAfter' fstep initial done (Stream step1 state1) = do
692 Stream step (PScanAfterStep state1 initial)
694 where
696 {-# INLINE_LATE step #-}
697 step gst (PScanAfterStep st acc) = do
698 r <- step1 (adaptState gst) st
699 case r of
700 Yield x s -> do
701 !old <- acc
702 !y <- fstep old x
703 return (Skip $ PScanAfterYield y (PScanAfterStep s (return y)))
704 Skip s -> return $ Skip $ PScanAfterStep s acc
705 -- Strictness is important for fusion
706 Stop -> do
707 !v <- acc
708 !res <- done v
709 return (Skip $ PScanAfterYield res PScanAfterStop)
710 step _ (PScanAfterYield acc next) = return $ Yield acc next
711 step _ PScanAfterStop = return Stop
713{-# INLINE_NORMAL postscanlM #-}
714postscanlM :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
715postscanlM fstep begin (Stream step state) = Stream step' Nothing
716 where
717 {-# INLINE_LATE step' #-}
718 step' _ Nothing = do
719 r <- begin
720 return $ Skip (Just (state, r))
722 step' gst (Just (st, acc)) = do
723 r <- step (adaptState gst) st
724 case r of
725 Yield x s -> do
726 y <- fstep acc x
727 return (Yield y (Just (s, y)))
728 Skip s -> return $ Skip (Just (s, acc))
729 Stop -> return Stop
731{-# INLINE_NORMAL postscanl #-}
732postscanl :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
733postscanl f seed = postscanlM (\a b -> return (f a b)) (return seed)
735-- | Like 'scanl'' but with a monadic step function and a monadic seed.
736--
737{-# INLINE_NORMAL scanlM' #-}
738scanlM' :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
739scanlM' fstep begin (Stream step state) = Stream step' Nothing
740 where
741 {-# INLINE_LATE step' #-}
742 step' _ Nothing = do
743 !x <- begin
744 return $ Yield x (Just (state, x))
745 step' gst (Just (st, acc)) = do
746 r <- step (adaptState gst) st
747 case r of
748 Yield x s -> do
749 !y <- fstep acc x
750 return $ Yield y (Just (s, y))
751 Skip s -> return $ Skip (Just (s, acc))
752 Stop -> return Stop
754-- | @scanlMAfter' accumulate initial done stream@ is like 'scanlM'' except
755-- that it provides an additional @done@ function to be applied on the
756-- accumulator when the stream stops. The result of @done@ is also emitted in
757-- the stream.
758--
759-- This function can be used to allocate a resource in the beginning of the
760-- scan and release it when the stream ends or to flush the internal state of
761-- the scan at the end.
762--
763-- /Pre-release/
764--
765{-# INLINE scanlMAfter' #-}
766scanlMAfter' :: Monad m
767 => (b -> a -> m b) -> m b -> (b -> m b) -> Stream m a -> Stream m b
768scanlMAfter' fstep initial done s =
769 initial `consM` postscanlMAfter' fstep initial done s
771-- >>> scanl' f z xs = z `Stream.cons` postscanl' f z xs
773-- | Strict left scan. Like 'map', 'scanl'' too is a one to one transformation,
774-- however it adds an extra element.
775--
776-- >>> Stream.toList $ Stream.scanl' (+) 0 $ Stream.fromList [1,2,3,4]
777-- [0,1,3,6,10]
778--
779-- >>> Stream.toList $ Stream.scanl' (flip (:)) [] $ Stream.fromList [1,2,3,4]
780-- [[],[1],[2,1],[3,2,1],[4,3,2,1]]
781--
782-- The output of 'scanl'' is the initial value of the accumulator followed by
783-- all the intermediate steps and the final result of 'foldl''.
784--
785-- By streaming the accumulated state after each fold step, we can share the
786-- state across multiple stages of stream composition. Each stage can modify or
787-- extend the state, do some processing with it and emit it for the next stage,
788-- thus modularizing the stream processing. This can be useful in
789-- stateful or event-driven programming.
790--
791-- Consider the following monolithic example, computing the sum and the product
792-- of the elements in a stream in one go using a @foldl'@:
793--
794-- >>> Stream.fold (Fold.foldl' (\(s, p) x -> (s + x, p * x)) (0,1)) $ Stream.fromList [1,2,3,4]
795-- (10,24)
796--
797-- Using @scanl'@ we can make it modular by computing the sum in the first
798-- stage and passing it down to the next stage for computing the product:
799--
800-- >>> :{
801-- Stream.fold (Fold.foldl' (\(_, p) (s, x) -> (s, p * x)) (0,1))
802-- $ Stream.scanl' (\(s, _) x -> (s + x, x)) (0,1)
803-- $ Stream.fromList [1,2,3,4]
804-- :}
805-- (10,24)
806--
807-- IMPORTANT: 'scanl'' evaluates the accumulator to WHNF. To avoid building
808-- lazy expressions inside the accumulator, it is recommended that a strict
809-- data structure is used for accumulator.
810--
811-- >>> scanl' step z = Stream.scan (Fold.foldl' step z)
812-- >>> scanl' f z xs = Stream.scanlM' (\a b -> return (f a b)) (return z) xs
813--
814-- See also: 'usingStateT'
815--
816{-# INLINE scanl' #-}
817scanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
818scanl' f seed = scanlM' (\a b -> return (f a b)) (return seed)
820{-# INLINE_NORMAL scanlM #-}
821scanlM :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
822scanlM fstep begin (Stream step state) = Stream step' Nothing
823 where
824 {-# INLINE_LATE step' #-}
825 step' _ Nothing = do
826 x <- begin
827 return $ Yield x (Just (state, x))
828 step' gst (Just (st, acc)) = do
829 r <- step (adaptState gst) st
830 case r of
831 Yield x s -> do
832 y <- fstep acc x
833 return $ Yield y (Just (s, y))
834 Skip s -> return $ Skip (Just (s, acc))
835 Stop -> return Stop
837{-# INLINE scanl #-}
838scanl :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
839scanl f seed = scanlM (\a b -> return (f a b)) (return seed)
841-- Adapted from the vector package
842{-# INLINE_NORMAL scanl1M #-}
843scanl1M :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
844scanl1M fstep (Stream step state) = Stream step' (state, Nothing)
845 where
846 {-# INLINE_LATE step' #-}
847 step' gst (st, Nothing) = do
848 r <- step gst st
849 case r of
850 Yield x s -> return $ Yield x (s, Just x)
851 Skip s -> return $ Skip (s, Nothing)
852 Stop -> return Stop
854 step' gst (st, Just acc) = do
855 r <- step gst st
856 case r of
857 Yield y s -> do
858 z <- fstep acc y
859 return $ Yield z (s, Just z)
860 Skip s -> return $ Skip (s, Just acc)
861 Stop -> return Stop
863{-# INLINE scanl1 #-}
864scanl1 :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
865scanl1 f = scanl1M (\x y -> return (f x y))
867-- Adapted from the vector package
869-- | Like 'scanl1'' but with a monadic step function.
870--
871{-# INLINE_NORMAL scanl1M' #-}
872scanl1M' :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
873scanl1M' fstep (Stream step state) = Stream step' (state, Nothing)
874 where
875 {-# INLINE_LATE step' #-}
876 step' gst (st, Nothing) = do
877 r <- step gst st
878 case r of
879 Yield x s -> x `seq` return $ Yield x (s, Just x)
880 Skip s -> return $ Skip (s, Nothing)
881 Stop -> return Stop
883 step' gst (st, Just acc) = acc `seq` do
884 r <- step gst st
885 case r of
886 Yield y s -> do
887 z <- fstep acc y
888 z `seq` return $ Yield z (s, Just z)
889 Skip s -> return $ Skip (s, Just acc)
890 Stop -> return Stop
892-- | Like 'scanl'' but for a non-empty stream. The first element of the stream
893-- is used as the initial value of the accumulator. Does nothing if the stream
894-- is empty.
895--
896-- >>> Stream.toList $ Stream.scanl1' (+) $ Stream.fromList [1,2,3,4]
897-- [1,3,6,10]
898--
899{-# INLINE scanl1' #-}
900scanl1' :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
901scanl1' f = scanl1M' (\x y -> return (f x y))
903-------------------------------------------------------------------------------
904-- Filtering
905-------------------------------------------------------------------------------
907-- | Modify a @Stream m a -> Stream m a@ stream transformation that accepts a
908-- predicate @(a -> b)@ to accept @((s, a) -> b)@ instead, provided a
909-- transformation @Stream m a -> Stream m (s, a)@. Convenient to filter with
910-- index or time.
911--
912-- >>> filterWithIndex = Stream.with Stream.indexed Stream.filter
913--
914-- /Pre-release/
915{-# INLINE with #-}
916with :: Monad m =>
917 (Stream m a -> Stream m (s, a))
918 -> (((s, a) -> b) -> Stream m (s, a) -> Stream m (s, a))
919 -> (((s, a) -> b) -> Stream m a -> Stream m a)
920with f comb g = fmap snd . comb g . f
922-- Adapted from the vector package
924-- | Same as 'filter' but with a monadic predicate.
925--
926-- >>> f p x = p x >>= \r -> return $ if r then Just x else Nothing
927-- >>> filterM p = Stream.mapMaybeM (f p)
928--
929{-# INLINE_NORMAL filterM #-}
930filterM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
931filterM f (Stream step state) = Stream step' state
932 where
933 {-# INLINE_LATE step' #-}
934 step' gst st = do
935 r <- step gst st
936 case r of
937 Yield x s -> do
938 b <- f x
939 return $ if b
940 then Yield x s
941 else Skip s
942 Skip s -> return $ Skip s
943 Stop -> return Stop
945-- | Include only those elements that pass a predicate.
946--
947-- >>> filter p = Stream.filterM (return . p)
948-- >>> filter p = Stream.mapMaybe (\x -> if p x then Just x else Nothing)
949-- >>> filter p = Stream.scanMaybe (Fold.filtering p)
950--
951{-# INLINE filter #-}
952filter :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
953filter f = filterM (return . f)
954-- filter p = scanMaybe (FL.filtering p)
956-- | Drop repeated elements that are adjacent to each other using the supplied
957-- comparison function.
958--
959-- >>> uniq = Stream.uniqBy (==)
960--
961-- To strip duplicate path separators:
962--
963-- >>> input = Stream.fromList "//a//b"
964-- >>> f x y = x == '/' && y == '/'
965-- >>> Stream.fold Fold.toList $ Stream.uniqBy f input
966-- "/a/b"
967--
968-- Space: @O(1)@
969--
970-- /Pre-release/
971--
972{-# INLINE uniqBy #-}
973uniqBy :: Monad m =>
974 (a -> a -> Bool) -> Stream m a -> Stream m a
975-- uniqBy eq = scanMaybe (FL.uniqBy eq)
976uniqBy eq = catMaybes . rollingMap f
978 where
980 f pre curr =
981 case pre of
982 Nothing -> Just curr
983 Just x -> if x `eq` curr then Nothing else Just curr
985-- Adapted from the vector package
987-- | Drop repeated elements that are adjacent to each other.
988--
989-- >>> uniq = Stream.uniqBy (==)
990--
991{-# INLINE_NORMAL uniq #-}
992uniq :: (Eq a, Monad m) => Stream m a -> Stream m a
993-- uniq = scanMaybe FL.uniq
994uniq (Stream step state) = Stream step' (Nothing, state)
995 where
996 {-# INLINE_LATE step' #-}
997 step' gst (Nothing, st) = do
998 r <- step gst st
999 case r of
1000 Yield x s -> return $ Yield x (Just x, s)
1001 Skip s -> return $ Skip (Nothing, s)
1002 Stop -> return Stop
1003 step' gst (Just x, st) = do
1004 r <- step gst st
1005 case r of
1006 Yield y s | x == y -> return $ Skip (Just x, s)
1007 | otherwise -> return $ Yield y (Just y, s)
1008 Skip s -> return $ Skip (Just x, s)
1009 Stop -> return Stop
1011-- | Deletes the first occurrence of the element in the stream that satisfies
1012-- the given equality predicate.
1014-- >>> input = Stream.fromList [1,3,3,5]
1015-- >>> Stream.fold Fold.toList $ Stream.deleteBy (==) 3 input
1016-- [1,3,5]
1018{-# INLINE_NORMAL deleteBy #-}
1019deleteBy :: Monad m => (a -> a -> Bool) -> a -> Stream m a -> Stream m a
1020-- deleteBy cmp x = scanMaybe (FL.deleteBy cmp x)
1021deleteBy eq x (Stream step state) = Stream step' (state, False)
1022 where
1023 {-# INLINE_LATE step' #-}
1024 step' gst (st, False) = do
1025 r <- step gst st
1026 case r of
1027 Yield y s -> return $
1028 if eq x y then Skip (s, True) else Yield y (s, False)
1029 Skip s -> return $ Skip (s, False)
1030 Stop -> return Stop
1032 step' gst (st, True) = do
1033 r <- step gst st
1034 case r of
1035 Yield y s -> return $ Yield y (s, True)
1036 Skip s -> return $ Skip (s, True)
1037 Stop -> return Stop
1039-- | Strip all leading and trailing occurrences of an element passing a
1040-- predicate and make all other consecutive occurrences uniq.
1042-- >> prune p = Stream.dropWhileAround p $ Stream.uniqBy (x y -> p x && p y)
1044-- @
1045-- > Stream.prune isSpace (Stream.fromList " hello world! ")
1046-- "hello world!"
1048-- @
1050-- Space: @O(1)@
1052-- /Unimplemented/
1053{-# INLINE prune #-}
1054prune ::
1055 -- (Monad m, Eq a) =>
1056 (a -> Bool) -> Stream m a -> Stream m a
1057prune = error "Not implemented yet!"
1059-- Possible implementation:
1060-- @repeated =
1061-- Stream.catMaybes . Stream.parseMany (Parser.groupBy (==) Fold.repeated)@
1063-- 'Fold.repeated' should return 'Just' when repeated, and 'Nothing' for a
1064-- single element.
1066-- | Emit only repeated elements, once.
1068-- /Unimplemented/
1069repeated :: -- (Monad m, Eq a) =>
1070 Stream m a -> Stream m a
1071repeated = undefined
1073------------------------------------------------------------------------------
1074-- Trimming
1075------------------------------------------------------------------------------
1077-- | Take all consecutive elements at the end of the stream for which the
1078-- predicate is true.
1080-- O(n) space, where n is the number elements taken.
1082-- /Unimplemented/
1083{-# INLINE takeWhileLast #-}
1084takeWhileLast :: -- Monad m =>
1085 (a -> Bool) -> Stream m a -> Stream m a
1086takeWhileLast = undefined -- fromStreamD $ D.takeWhileLast n $ toStreamD m
1088-- | Like 'takeWhile' and 'takeWhileLast' combined.
1090-- O(n) space, where n is the number elements taken from the end.
1092-- /Unimplemented/
1093{-# INLINE takeWhileAround #-}
1094takeWhileAround :: -- Monad m =>
1095 (a -> Bool) -> Stream m a -> Stream m a
1096takeWhileAround = undefined -- fromStreamD $ D.takeWhileAround n $ toStreamD m
1098-- Adapted from the vector package
1100-- | Discard first 'n' elements from the stream and take the rest.
1102{-# INLINE_NORMAL drop #-}
1103drop :: Monad m => Int -> Stream m a -> Stream m a
1104drop n (Stream step state) = Stream step' (state, Just n)
1105 where
1106 {-# INLINE_LATE step' #-}
1107 step' gst (st, Just i)
1108 | i > 0 = do
1109 r <- step gst st
1110 return $
1111 case r of
1112 Yield _ s -> Skip (s, Just (i - 1))
1113 Skip s -> Skip (s, Just i)
1114 Stop -> Stop
1115 | otherwise = return $ Skip (st, Nothing)
1117 step' gst (st, Nothing) = do
1118 r <- step gst st
1119 return $
1120 case r of
1121 Yield x s -> Yield x (s, Nothing)
1122 Skip s -> Skip (s, Nothing)
1123 Stop -> Stop
1125-- Adapted from the vector package
1126data DropWhileState s a
1127 = DropWhileDrop s
1128 | DropWhileYield a s
1129 | DropWhileNext s
1131-- | Same as 'dropWhile' but with a monadic predicate.
1133{-# INLINE_NORMAL dropWhileM #-}
1134dropWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
1135-- dropWhileM p = scanMaybe (FL.droppingWhileM p)
1136dropWhileM f (Stream step state) = Stream step' (DropWhileDrop state)
1137 where
1138 {-# INLINE_LATE step' #-}
1139 step' gst (DropWhileDrop st) = do
1140 r <- step gst st
1141 case r of
1142 Yield x s -> do
1143 b <- f x
1144 if b
1145 then return $ Skip (DropWhileDrop s)
1146 else return $ Skip (DropWhileYield x s)
1147 Skip s -> return $ Skip (DropWhileDrop s)
1148 Stop -> return Stop
1150 step' gst (DropWhileNext st) = do
1151 r <- step gst st
1152 case r of
1153 Yield x s -> return $ Skip (DropWhileYield x s)
1154 Skip s -> return $ Skip (DropWhileNext s)
1155 Stop -> return Stop
1157 step' _ (DropWhileYield x st) = return $ Yield x (DropWhileNext st)
1159-- | Drop elements in the stream as long as the predicate succeeds and then
1160-- take the rest of the stream.
1162{-# INLINE dropWhile #-}
1163dropWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
1164-- dropWhile p = scanMaybe (FL.droppingWhile p)
1165dropWhile f = dropWhileM (return . f)
1167-- | Drop @n@ elements at the end of the stream.
1169-- O(n) space, where n is the number elements dropped.
1171-- /Unimplemented/
1172{-# INLINE dropLast #-}
1173dropLast :: -- Monad m =>
1174 Int -> Stream m a -> Stream m a
1175dropLast = undefined -- fromStreamD $ D.dropLast n $ toStreamD m
1177-- | Drop all consecutive elements at the end of the stream for which the
1178-- predicate is true.
1180-- O(n) space, where n is the number elements dropped.
1182-- /Unimplemented/
1183{-# INLINE dropWhileLast #-}
1184dropWhileLast :: -- Monad m =>
1185 (a -> Bool) -> Stream m a -> Stream m a
1186dropWhileLast = undefined -- fromStreamD $ D.dropWhileLast n $ toStreamD m
1188-- | Like 'dropWhile' and 'dropWhileLast' combined.
1190-- O(n) space, where n is the number elements dropped from the end.
1192-- /Unimplemented/
1193{-# INLINE dropWhileAround #-}
1194dropWhileAround :: -- Monad m =>
1195 (a -> Bool) -> Stream m a -> Stream m a
1196dropWhileAround = undefined -- fromStreamD $ D.dropWhileAround n $ toStreamD m
1198------------------------------------------------------------------------------
1199-- Inserting Elements
1200------------------------------------------------------------------------------
1202-- | @insertBy cmp elem stream@ inserts @elem@ before the first element in
1203-- @stream@ that is less than @elem@ when compared using @cmp@.
1205-- >>> insertBy cmp x = Stream.mergeBy cmp (Stream.fromPure x)
1207-- >>> input = Stream.fromList [1,3,5]
1208-- >>> Stream.fold Fold.toList $ Stream.insertBy compare 2 input
1209-- [1,2,3,5]
1211{-# INLINE_NORMAL insertBy #-}
1212insertBy :: Monad m => (a -> a -> Ordering) -> a -> Stream m a -> Stream m a
1213insertBy cmp a (Stream step state) = Stream step' (state, False, Nothing)
1214 where
1215 {-# INLINE_LATE step' #-}
1216 step' gst (st, False, _) = do
1217 r <- step gst st
1218 case r of
1219 Yield x s -> case cmp a x of
1220 GT -> return $ Yield x (s, False, Nothing)
1221 _ -> return $ Yield a (s, True, Just x)
1222 Skip s -> return $ Skip (s, False, Nothing)
1223 Stop -> return $ Yield a (st, True, Nothing)
1225 step' _ (_, True, Nothing) = return Stop
1227 step' gst (st, True, Just prev) = do
1228 r <- step gst st
1229 case r of
1230 Yield x s -> return $ Yield prev (s, True, Just x)
1231 Skip s -> return $ Skip (s, True, Just prev)
1232 Stop -> return $ Yield prev (st, True, Nothing)
1234data LoopState x s = FirstYield s
1235 | InterspersingYield s
1236 | YieldAndCarry x s
1238-- intersperseM = intersperseMWith 1
1240-- | Insert an effect and its output before consuming an element of a stream
1241-- except the first one.
1243-- >>> input = Stream.fromList "hello"
1244-- >>> Stream.fold Fold.toList $ Stream.trace putChar $ Stream.intersperseM (putChar '.' >> return ',') input
1245-- h.,e.,l.,l.,o"h,e,l,l,o"
1247-- Be careful about the order of effects. In the above example we used trace
1248-- after the intersperse, if we use it before the intersperse the output would
1249-- be he.l.l.o."h,e,l,l,o".
1251-- >>> Stream.fold Fold.toList $ Stream.intersperseM (putChar '.' >> return ',') $ Stream.trace putChar input
1252-- he.l.l.o."h,e,l,l,o"
1254{-# INLINE_NORMAL intersperseM #-}
1255intersperseM :: Monad m => m a -> Stream m a -> Stream m a
1256intersperseM m (Stream step state) = Stream step' (FirstYield state)
1257 where
1258 {-# INLINE_LATE step' #-}
1259 step' gst (FirstYield st) = do
1260 r <- step gst st
1261 return $
1262 case r of
1263 Yield x s -> Skip (YieldAndCarry x s)
1264 Skip s -> Skip (FirstYield s)
1265 Stop -> Stop
1267 step' gst (InterspersingYield st) = do
1268 r <- step gst st
1269 case r of
1270 Yield x s -> do
1271 a <- m
1272 return $ Yield a (YieldAndCarry x s)
1273 Skip s -> return $ Skip $ InterspersingYield s
1274 Stop -> return Stop
1276 step' _ (YieldAndCarry x st) = return $ Yield x (InterspersingYield st)
1278-- | Insert a pure value between successive elements of a stream.
1280-- >>> input = Stream.fromList "hello"
1281-- >>> Stream.fold Fold.toList $ Stream.intersperse ',' input
1282-- "h,e,l,l,o"
1284{-# INLINE intersperse #-}
1285intersperse :: Monad m => a -> Stream m a -> Stream m a
1286intersperse a = intersperseM (return a)
1288-- | Insert a side effect before consuming an element of a stream except the
1289-- first one.
1291-- >>> input = Stream.fromList "hello"
1292-- >>> Stream.fold Fold.drain $ Stream.trace putChar $ Stream.intersperseM_ (putChar '.') input
1293-- h.e.l.l.o
1295-- /Pre-release/
1296{-# INLINE_NORMAL intersperseM_ #-}
1297intersperseM_ :: Monad m => m b -> Stream m a -> Stream m a
1298intersperseM_ m (Stream step1 state1) = Stream step (Left (pure (), state1))
1299 where
1300 {-# INLINE_LATE step #-}
1301 step gst (Left (eff, st)) = do
1302 r <- step1 gst st
1303 case r of
1304 Yield x s -> eff >> return (Yield x (Right s))
1305 Skip s -> return $ Skip (Left (eff, s))
1306 Stop -> return Stop
1308 step _ (Right st) = return $ Skip $ Left (void m, st)
1310-- | Intersperse a monadic action into the input stream after every @n@
1311-- elements.
1313-- >> input = Stream.fromList "hello"
1314-- >> Stream.fold Fold.toList $ Stream.intersperseMWith 2 (return ',') input
1315-- "he,ll,o"
1317-- /Unimplemented/
1318{-# INLINE intersperseMWith #-}
1319intersperseMWith :: -- Monad m =>
1320 Int -> m a -> Stream m a -> Stream m a
1321intersperseMWith _n _f _xs = undefined
1323data SuffixState s a
1324 = SuffixElem s
1325 | SuffixSuffix s
1326 | SuffixYield a (SuffixState s a)
1328-- | Insert an effect and its output after consuming an element of a stream.
1330-- >>> input = Stream.fromList "hello"
1331-- >>> Stream.fold Fold.toList $ Stream.trace putChar $ Stream.intersperseMSuffix (putChar '.' >> return ',') input
1332-- h.,e.,l.,l.,o.,"h,e,l,l,o,"
1334-- /Pre-release/
1335{-# INLINE_NORMAL intersperseMSuffix #-}
1336intersperseMSuffix :: forall m a. Monad m => m a -> Stream m a -> Stream m a
1337intersperseMSuffix action (Stream step state) = Stream step' (SuffixElem state)
1338 where
1339 {-# INLINE_LATE step' #-}
1340 step' gst (SuffixElem st) = do
1341 r <- step gst st
1342 return $ case r of
1343 Yield x s -> Skip (SuffixYield x (SuffixSuffix s))
1344 Skip s -> Skip (SuffixElem s)
1345 Stop -> Stop
1347 step' _ (SuffixSuffix st) = do
1348 action >>= \r -> return $ Skip (SuffixYield r (SuffixElem st))
1350 step' _ (SuffixYield x next) = return $ Yield x next
1352-- | Insert a side effect after consuming an element of a stream.
1354-- >>> input = Stream.fromList "hello"
1355-- >>> Stream.fold Fold.toList $ Stream.intersperseMSuffix_ (threadDelay 1000000) input
1356-- "hello"
1358-- /Pre-release/
1360{-# INLINE_NORMAL intersperseMSuffix_ #-}
1361intersperseMSuffix_ :: Monad m => m b -> Stream m a -> Stream m a
1362intersperseMSuffix_ m (Stream step1 state1) = Stream step (Left state1)
1363 where
1364 {-# INLINE_LATE step #-}
1365 step gst (Left st) = do
1366 r <- step1 gst st
1367 case r of
1368 Yield x s -> return $ Yield x (Right s)
1369 Skip s -> return $ Skip $ Left s
1370 Stop -> return Stop
1372 step _ (Right st) = m >> return (Skip (Left st))
1374data SuffixSpanState s a
1375 = SuffixSpanElem s Int
1376 | SuffixSpanSuffix s
1377 | SuffixSpanYield a (SuffixSpanState s a)
1378 | SuffixSpanLast
1379 | SuffixSpanStop
1381-- | Like 'intersperseMSuffix' but intersperses an effectful action into the
1382-- input stream after every @n@ elements and after the last element.
1384-- >>> input = Stream.fromList "hello"
1385-- >>> Stream.fold Fold.toList $ Stream.intersperseMSuffixWith 2 (return ',') input
1386-- "he,ll,o,"
1388-- /Pre-release/
1390{-# INLINE_NORMAL intersperseMSuffixWith #-}
1391intersperseMSuffixWith :: forall m a. Monad m
1392 => Int -> m a -> Stream m a -> Stream m a
1393intersperseMSuffixWith n action (Stream step state) =
1394 Stream step' (SuffixSpanElem state n)
1395 where
1396 {-# INLINE_LATE step' #-}
1397 step' gst (SuffixSpanElem st i) | i > 0 = do
1398 r <- step gst st
1399 return $ case r of
1400 Yield x s -> Skip (SuffixSpanYield x (SuffixSpanElem s (i - 1)))
1401 Skip s -> Skip (SuffixSpanElem s i)
1402 Stop -> if i == n then Stop else Skip SuffixSpanLast
1403 step' _ (SuffixSpanElem st _) = return $ Skip (SuffixSpanSuffix st)
1405 step' _ (SuffixSpanSuffix st) = do
1406 action >>= \r -> return $ Skip (SuffixSpanYield r (SuffixSpanElem st n))
1408 step' _ SuffixSpanLast = do
1409 action >>= \r -> return $ Skip (SuffixSpanYield r SuffixSpanStop)
1411 step' _ (SuffixSpanYield x next) = return $ Yield x next
1413 step' _ SuffixSpanStop = return Stop
1415-- | Insert a side effect before consuming an element of a stream.
1417-- Definition:
1419-- >>> intersperseMPrefix_ m = Stream.mapM (\x -> void m >> return x)
1421-- >>> input = Stream.fromList "hello"
1422-- >>> Stream.fold Fold.toList $ Stream.trace putChar $ Stream.intersperseMPrefix_ (putChar '.' >> return ',') input
1423-- .h.e.l.l.o"hello"
1425-- Same as 'trace_'.
1427-- /Pre-release/
1429{-# INLINE intersperseMPrefix_ #-}
1430intersperseMPrefix_ :: Monad m => m b -> Stream m a -> Stream m a
1431intersperseMPrefix_ m = mapM (\x -> void m >> return x)
1433------------------------------------------------------------------------------
1434-- Inserting Time
1435------------------------------------------------------------------------------
1437-- XXX This should be in Prelude, should we export this as a helper function?
1439-- | Block the current thread for specified number of seconds.
1440{-# INLINE sleep #-}
1441sleep :: MonadIO m => Double -> m ()
1442sleep n = liftIO $ threadDelay $ round $ n * 1000000
1444-- | Introduce a delay of specified seconds between elements of the stream.
1446-- Definition:
1448-- >>> sleep n = liftIO $ threadDelay $ round $ n * 1000000
1449-- >>> delay = Stream.intersperseM_ . sleep
1451-- Example:
1453-- >>> input = Stream.enumerateFromTo 1 3
1454-- >>> Stream.fold (Fold.drainMapM print) $ Stream.delay 1 input
1455-- 1
1456-- 2
1457-- 3
1459{-# INLINE delay #-}
1460delay :: MonadIO m => Double -> Stream m a -> Stream m a
1461delay = intersperseM_ . sleep
1463-- | Introduce a delay of specified seconds after consuming an element of a
1464-- stream.
1466-- Definition:
1468-- >>> sleep n = liftIO $ threadDelay $ round $ n * 1000000
1469-- >>> delayPost = Stream.intersperseMSuffix_ . sleep
1471-- Example:
1473-- >>> input = Stream.enumerateFromTo 1 3
1474-- >>> Stream.fold (Fold.drainMapM print) $ Stream.delayPost 1 input
1475-- 1
1476-- 2
1477-- 3
1479-- /Pre-release/
1481{-# INLINE delayPost #-}
1482delayPost :: MonadIO m => Double -> Stream m a -> Stream m a
1483delayPost n = intersperseMSuffix_ $ liftIO $ threadDelay $ round $ n * 1000000
1485-- | Introduce a delay of specified seconds before consuming an element of a
1486-- stream.
1488-- Definition:
1490-- >>> sleep n = liftIO $ threadDelay $ round $ n * 1000000
1491-- >>> delayPre = Stream.intersperseMPrefix_. sleep
1493-- Example:
1495-- >>> input = Stream.enumerateFromTo 1 3
1496-- >>> Stream.fold (Fold.drainMapM print) $ Stream.delayPre 1 input
1497-- 1
1498-- 2
1499-- 3
1501-- /Pre-release/
1503{-# INLINE delayPre #-}
1504delayPre :: MonadIO m => Double -> Stream m a -> Stream m a
1505delayPre = intersperseMPrefix_. sleep
1507------------------------------------------------------------------------------
1508-- Reordering
1509------------------------------------------------------------------------------
1511-- | Returns the elements of the stream in reverse order. The stream must be
1512-- finite. Note that this necessarily buffers the entire stream in memory.
1514-- Definition:
1516-- >>> reverse m = Stream.concatEffect $ Stream.fold Fold.toListRev m >>= return . Stream.fromList
1518{-# INLINE_NORMAL reverse #-}
1519reverse :: Monad m => Stream m a -> Stream m a
1520reverse m = concatEffect $ fold FL.toListRev m <&> fromList
1522reverse m = Stream step Nothing
1523 where
1524 {-# INLINE_LATE step #-}
1525 step _ Nothing = do
1526 xs <- foldl' (flip (:)) [] m
1527 return $ Skip (Just xs)
1528 step _ (Just (x:xs)) = return $ Yield x (Just xs)
1529 step _ (Just []) = return Stop
1532-- | Like 'reverse' but several times faster, requires an 'Unbox' instance.
1534-- /O(n) space/
1536-- /Pre-release/
1537{-# INLINE reverseUnbox #-}
1538reverseUnbox :: (MonadIO m, Unbox a) => Stream m a -> Stream m a
1539reverseUnbox =
1540 A.flattenArraysRev -- unfoldMany A.readRev
1541 . fromStreamK
1542 . K.reverse
1543 . toStreamK
1544 . A.chunksOf defaultChunkSize
1546-- | Buffer until the next element in sequence arrives. The function argument
1547-- determines the difference in sequence numbers. This could be useful in
1548-- implementing sequenced streams, for example, TCP reassembly.
1550-- /Unimplemented/
1552{-# INLINE reassembleBy #-}
1553reassembleBy
1554 :: -- Monad m =>
1555 Fold m a b
1556 -> (a -> a -> Int)
1557 -> Stream m a
1558 -> Stream m b
1559reassembleBy = undefined
1561------------------------------------------------------------------------------
1562-- Position Indexing
1563------------------------------------------------------------------------------
1565-- Adapted from the vector package
1567-- |
1568-- >>> f = Fold.foldl' (\(i, _) x -> (i + 1, x)) (-1,undefined)
1569-- >>> indexed = Stream.postscan f
1570-- >>> indexed = Stream.zipWith (,) (Stream.enumerateFrom 0)
1571-- >>> indexedR n = fmap (\(i, a) -> (n - i, a)) . indexed
1573-- Pair each element in a stream with its index, starting from index 0.
1575-- >>> Stream.fold Fold.toList $ Stream.indexed $ Stream.fromList "hello"
1576-- [(0,'h'),(1,'e'),(2,'l'),(3,'l'),(4,'o')]
1578{-# INLINE_NORMAL indexed #-}
1579indexed :: Monad m => Stream m a -> Stream m (Int, a)
1580-- indexed = scanMaybe FL.indexing
1581indexed (Stream step state) = Stream step' (state, 0)
1582 where
1583 {-# INLINE_LATE step' #-}
1584 step' gst (st, i) = i `seq` do
1585 r <- step (adaptState gst) st
1586 case r of
1587 Yield x s -> return $ Yield (i, x) (s, i+1)
1588 Skip s -> return $ Skip (s, i)
1589 Stop -> return Stop
1591-- Adapted from the vector package
1593-- |
1594-- >>> f n = Fold.foldl' (\(i, _) x -> (i - 1, x)) (n + 1,undefined)
1595-- >>> indexedR n = Stream.postscan (f n)
1597-- >>> s n = Stream.enumerateFromThen n (n - 1)
1598-- >>> indexedR n = Stream.zipWith (,) (s n)
1600-- Pair each element in a stream with its index, starting from the
1601-- given index @n@ and counting down.
1603-- >>> Stream.fold Fold.toList $ Stream.indexedR 10 $ Stream.fromList "hello"
1604-- [(10,'h'),(9,'e'),(8,'l'),(7,'l'),(6,'o')]
1606{-# INLINE_NORMAL indexedR #-}
1607indexedR :: Monad m => Int -> Stream m a -> Stream m (Int, a)
1608-- indexedR n = scanMaybe (FL.indexingRev n)
1609indexedR m (Stream step state) = Stream step' (state, m)
1610 where
1611 {-# INLINE_LATE step' #-}
1612 step' gst (st, i) = i `seq` do
1613 r <- step (adaptState gst) st
1614 case r of
1615 Yield x s -> let i' = i - 1
1616 in return $ Yield (i, x) (s, i')
1617 Skip s -> return $ Skip (s, i)
1618 Stop -> return Stop
1620-------------------------------------------------------------------------------
1621-- Time Indexing
1622-------------------------------------------------------------------------------
1624-- Note: The timestamp stream must be the second stream in the zip so that the
1625-- timestamp is generated after generating the stream element and not before.
1626-- If we do not do that then the following example will generate the same
1627-- timestamp for first two elements:
1629-- Stream.fold Fold.toList $ Stream.timestamped $ Stream.delay $ Stream.enumerateFromTo 1 3
1631-- | Pair each element in a stream with an absolute timestamp, using a clock of
1632-- specified granularity. The timestamp is generated just before the element
1633-- is consumed.
1635-- >>> Stream.fold Fold.toList $ Stream.timestampWith 0.01 $ Stream.delay 1 $ Stream.enumerateFromTo 1 3
1636-- [(AbsTime (TimeSpec {sec = ..., nsec = ...}),1),(AbsTime (TimeSpec {sec = ..., nsec = ...}),2),(AbsTime (TimeSpec {sec = ..., nsec = ...}),3)]
1638-- /Pre-release/
1640{-# INLINE timestampWith #-}
1641timestampWith :: (MonadIO m)
1642 => Double -> Stream m a -> Stream m (AbsTime, a)
1643timestampWith g stream = zipWith (flip (,)) stream (absTimesWith g)
1645-- TBD: check performance vs a custom implementation without using zipWith.
1647-- /Pre-release/
1649{-# INLINE timestamped #-}
1650timestamped :: (MonadIO m)
1651 => Stream m a -> Stream m (AbsTime, a)
1652timestamped = timestampWith 0.01
1654-- | Pair each element in a stream with relative times starting from 0, using a
1655-- clock with the specified granularity. The time is measured just before the
1656-- element is consumed.
1658-- >>> Stream.fold Fold.toList $ Stream.timeIndexWith 0.01 $ Stream.delay 1 $ Stream.enumerateFromTo 1 3
1659-- [(RelTime64 (NanoSecond64 ...),1),(RelTime64 (NanoSecond64 ...),2),(RelTime64 (NanoSecond64 ...),3)]
1661-- /Pre-release/
1663{-# INLINE timeIndexWith #-}
1664timeIndexWith :: (MonadIO m)
1665 => Double -> Stream m a -> Stream m (RelTime64, a)
1666timeIndexWith g stream = zipWith (flip (,)) stream (relTimesWith g)
1668-- | Pair each element in a stream with relative times starting from 0, using a
1669-- 10 ms granularity clock. The time is measured just before the element is
1670-- consumed.
1672-- >>> Stream.fold Fold.toList $ Stream.timeIndexed $ Stream.delay 1 $ Stream.enumerateFromTo 1 3
1673-- [(RelTime64 (NanoSecond64 ...),1),(RelTime64 (NanoSecond64 ...),2),(RelTime64 (NanoSecond64 ...),3)]
1675-- /Pre-release/
1677{-# INLINE timeIndexed #-}
1678timeIndexed :: (MonadIO m)
1679 => Stream m a -> Stream m (RelTime64, a)
1680timeIndexed = timeIndexWith 0.01
1682------------------------------------------------------------------------------
1683-- Searching
1684------------------------------------------------------------------------------
1686-- | Find all the indices where the element in the stream satisfies the given
1687-- predicate.
1689-- >>> findIndices p = Stream.scanMaybe (Fold.findIndices p)
1691{-# INLINE_NORMAL findIndices #-}
1692findIndices :: Monad m => (a -> Bool) -> Stream m a -> Stream m Int
1693findIndices p (Stream step state) = Stream step' (state, 0)
1694 where
1695 {-# INLINE_LATE step' #-}
1696 step' gst (st, i) = i `seq` do
1697 r <- step (adaptState gst) st
1698 return $ case r of
1699 Yield x s -> if p x then Yield i (s, i+1) else Skip (s, i+1)
1700 Skip s -> Skip (s, i)
1701 Stop -> Stop
1703-- | Find all the indices where the value of the element in the stream is equal
1704-- to the given value.
1706-- >>> elemIndices a = Stream.findIndices (== a)
1708{-# INLINE elemIndices #-}
1709elemIndices :: (Monad m, Eq a) => a -> Stream m a -> Stream m Int
1710elemIndices a = findIndices (== a)
1712{-# INLINE_NORMAL slicesBy #-}
1713slicesBy :: Monad m => (a -> Bool) -> Stream m a -> Stream m (Int, Int)
1714slicesBy p (Stream step1 state1) = Stream step (Just (state1, 0, 0))
1716 where
1718 {-# INLINE_LATE step #-}
1719 step gst (Just (st, i, len)) = i `seq` len `seq` do
1720 r <- step1 (adaptState gst) st
1721 return
1722 $ case r of
1723 Yield x s ->
1724 if p x
1725 then Yield (i, len + 1) (Just (s, i + len + 1, 0))
1726 else Skip (Just (s, i, len + 1))
1727 Skip s -> Skip (Just (s, i, len))
1728 Stop -> if len == 0 then Stop else Yield (i, len) Nothing
1729 step _ Nothing = return Stop
1731------------------------------------------------------------------------------
1732-- Rolling map
1733------------------------------------------------------------------------------
1735data RollingMapState s a = RollingMapGo s a
1737-- | Like 'rollingMap' but with an effectful map function.
1739-- /Pre-release/
1741{-# INLINE rollingMapM #-}
1742rollingMapM :: Monad m => (Maybe a -> a -> m b) -> Stream m a -> Stream m b
1743-- rollingMapM f = scanMaybe (FL.slide2 $ Window.rollingMapM f)
1744rollingMapM f (Stream step1 state1) = Stream step (RollingMapGo state1 Nothing)
1746 where
1748 step gst (RollingMapGo s1 curr) = do
1749 r <- step1 (adaptState gst) s1
1750 case r of
1751 Yield x s -> do
1752 !res <- f curr x
1753 return $ Yield res $ RollingMapGo s (Just x)
1754 Skip s -> return $ Skip $ RollingMapGo s curr
1755 Stop -> return Stop
1757-- rollingMap is a special case of an incremental sliding fold. It can be
1758-- written as:
1760-- > fld f = slidingWindow 1 (Fold.foldl' (\_ (x,y) -> f y x)
1761-- > rollingMap f = Stream.postscan (fld f) undefined
1763-- | Apply a function on every two successive elements of a stream. The first
1764-- argument of the map function is the previous element and the second argument
1765-- is the current element. When the current element is the first element, the
1766-- previous element is 'Nothing'.
1768-- /Pre-release/
1770{-# INLINE rollingMap #-}
1771rollingMap :: Monad m => (Maybe a -> a -> b) -> Stream m a -> Stream m b
1772-- rollingMap f = scanMaybe (FL.slide2 $ Window.rollingMap f)
1773rollingMap f = rollingMapM (\x y -> return $ f x y)
1775-- | Like 'rollingMap' but requires at least two elements in the stream,
1776-- returns an empty stream otherwise.
1778-- This is the stream equivalent of the list idiom @zipWith f xs (tail xs)@.
1780-- /Pre-release/
1782{-# INLINE rollingMap2 #-}
1783rollingMap2 :: Monad m => (a -> a -> b) -> Stream m a -> Stream m b
1784rollingMap2 f = catMaybes . rollingMap g
1786 where
1788 g Nothing _ = Nothing
1789 g (Just x) y = Just (f x y)
1791------------------------------------------------------------------------------
1792-- Maybe Streams
1793------------------------------------------------------------------------------
1795-- XXX Will this always fuse properly?
1797-- | Map a 'Maybe' returning function to a stream, filter out the 'Nothing'
1798-- elements, and return a stream of values extracted from 'Just'.
1800-- Equivalent to:
1802-- >>> mapMaybe f = Stream.catMaybes . fmap f
1804{-# INLINE_NORMAL mapMaybe #-}
1805mapMaybe :: Monad m => (a -> Maybe b) -> Stream m a -> Stream m b
1806mapMaybe f = fmap fromJust . filter isJust . map f
1808-- | Like 'mapMaybe' but maps a monadic function.
1810-- Equivalent to:
1812-- >>> mapMaybeM f = Stream.catMaybes . Stream.mapM f
1814-- >>> mapM f = Stream.mapMaybeM (\x -> Just <$> f x)
1816{-# INLINE_NORMAL mapMaybeM #-}
1817mapMaybeM :: Monad m => (a -> m (Maybe b)) -> Stream m a -> Stream m b
1818mapMaybeM f = fmap fromJust . filter isJust . mapM f
1820-- | In a stream of 'Maybe's, discard 'Nothing's and unwrap 'Just's.
1822-- >>> catMaybes = Stream.mapMaybe id
1823-- >>> catMaybes = fmap fromJust . Stream.filter isJust
1825-- /Pre-release/
1827{-# INLINE catMaybes #-}
1828catMaybes :: Monad m => Stream m (Maybe a) -> Stream m a
1829-- catMaybes = fmap fromJust . filter isJust
1830catMaybes (Stream step state) = Stream step1 state
1832 where
1834 {-# INLINE_LATE step1 #-}
1835 step1 gst st = do
1836 r <- step (adaptState gst) st
1837 case r of
1838 Yield x s -> do
1839 return
1840 $ case x of
1841 Just a -> Yield a s
1842 Nothing -> Skip s
1843 Skip s -> return $ Skip s
1844 Stop -> return Stop
1846-- | Use a filtering fold on a stream.
1848-- >>> scanMaybe f = Stream.catMaybes . Stream.postscan f
1850{-# INLINE scanMaybe #-}
1851scanMaybe :: Monad m => Fold m a (Maybe b) -> Stream m a -> Stream m b
1852scanMaybe f = catMaybes . postscan f
1854------------------------------------------------------------------------------
1855-- Either streams
1856------------------------------------------------------------------------------
1858-- | Discard 'Right's and unwrap 'Left's in an 'Either' stream.
1860-- >>> catLefts = fmap (fromLeft undefined) . Stream.filter isLeft
1862-- /Pre-release/
1864{-# INLINE catLefts #-}
1865catLefts :: Monad m => Stream m (Either a b) -> Stream m a
1866catLefts = fmap (fromLeft undefined) . filter isLeft
1868-- | Discard 'Left's and unwrap 'Right's in an 'Either' stream.
1870-- >>> catRights = fmap (fromRight undefined) . Stream.filter isRight
1872-- /Pre-release/
1874{-# INLINE catRights #-}
1875catRights :: Monad m => Stream m (Either a b) -> Stream m b
1876catRights = fmap (fromRight undefined) . filter isRight
1878-- | Remove the either wrapper and flatten both lefts and as well as rights in
1879-- the output stream.
1881-- >>> catEithers = fmap (either id id)
1883-- /Pre-release/
1885{-# INLINE catEithers #-}
1886catEithers :: Monad m => Stream m (Either a a) -> Stream m a
1887catEithers = fmap (either id id)
1889------------------------------------------------------------------------------
1890-- Splitting
1891------------------------------------------------------------------------------
1893-- | Split on an infixed separator element, dropping the separator. The
1894-- supplied 'Fold' is applied on the split segments. Splits the stream on
1895-- separator elements determined by the supplied predicate, separator is
1896-- considered as infixed between two segments:
1898-- >>> splitOn' p xs = Stream.fold Fold.toList $ Stream.splitOn p Fold.toList (Stream.fromList xs)
1899-- >>> splitOn' (== '.') "a.b"
1900-- ["a","b"]
1902-- An empty stream is folded to the default value of the fold:
1904-- >>> splitOn' (== '.') ""
1905-- [""]
1907-- If one or both sides of the separator are missing then the empty segment on
1908-- that side is folded to the default output of the fold:
1910-- >>> splitOn' (== '.') "."
1911-- ["",""]
1913-- >>> splitOn' (== '.') ".a"
1914-- ["","a"]
1916-- >>> splitOn' (== '.') "a."
1917-- ["a",""]
1919-- >>> splitOn' (== '.') "a..b"
1920-- ["a","","b"]
1922-- splitOn is an inverse of intercalating single element:
1924-- > Stream.intercalate (Stream.fromPure '.') Unfold.fromList . Stream.splitOn (== '.') Fold.toList === id
1926-- Assuming the input stream does not contain the separator:
1928-- > Stream.splitOn (== '.') Fold.toList . Stream.intercalate (Stream.fromPure '.') Unfold.fromList === id
1930{-# INLINE splitOn #-}
1931splitOn :: Monad m => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1932splitOn predicate f =
1933 -- We can express the infix splitting in terms of optional suffix split
1934 -- fold. After applying a suffix split fold repeatedly if the last segment
1935 -- ends with a suffix then we need to return the default output of the fold
1936 -- after that to make it an infix split.
1937 --
1938 -- Alternately, we can also express it using an optional prefix split fold.
1939 -- If the first segment starts with a prefix then we need to emit the
1940 -- default output of the fold before that to make it an infix split, and
1941 -- then apply prefix split fold repeatedly.
1942 --
1943 -- Since a suffix split fold can be easily expressed using a
1944 -- non-backtracking fold, we use that.
1945 foldManyPost (FL.takeEndBy_ predicate f)