git.iko.soy/streamly
on cfcf26d help sign in
streamly/src/Streamly/Internal/Data/Stream/StreamD/Transform.hs
1238 lines 42 kB history raw
1-- |
2-- Module : Streamly.Internal.Data.Stream.StreamD.Transform
3-- Copyright : (c) 2018 Composewell Technologies
4-- (c) Roman Leshchinskiy 2008-2010
5-- License : BSD-3-Clause
6-- Maintainer : streamly@composewell.com
7-- Stability : experimental
8-- Portability : GHC
9--
10-- "Streamly.Internal.Data.Pipe" might ultimately replace this module.
12-- A few functions in this module have been adapted from the vector package
13-- (c) Roman Leshchinskiy. See the notes in specific combinators.
15module Streamly.Internal.Data.Stream.StreamD.Transform
16 (
17 -- * Piping
18 -- | Pass through a 'Pipe'.
19 transform
21 -- * Mapping
22 -- | Stateless one-to-one maps.
23 , map
24 , mapM
25 , sequence
27 -- * Mapping Effects
28 , tap
29 , tapOffsetEvery
30 , tapRate
31 , pollCounts
33 -- * Folding
34 , foldrS
35 , foldrT
36 , foldlS
37 , foldlT
39 -- * Scanning By 'Fold'
40 , postscanOnce -- XXX rename to postscan
41 , scanOnce -- XXX rename to scan
43 -- * Scanning
44 -- | Left scans. Stateful, mostly one-to-one maps.
45 , scanlM'
46 , scanlMAfter'
47 , scanl'
48 , scanlM
49 , scanl
50 , scanl1M'
51 , scanl1'
52 , scanl1M
53 , scanl1
55 , prescanl'
56 , prescanlM'
58 , postscanl
59 , postscanlM
60 , postscanl'
61 , postscanlM'
62 , postscanlMAfter'
64 , postscanlx'
65 , postscanlMx'
66 , scanlMx'
67 , scanlx'
69 -- * Filtering
70 -- | Produce a subset of the stream.
71 , filter
72 , filterM
73 , deleteBy
74 , uniq
76 -- * Trimming
77 -- | Produce a subset of the stream trimmed at ends.
78 , take
79 , takeByTime
80 , takeWhile
81 , takeWhileM
82 , drop
83 , dropByTime
84 , dropWhile
85 , dropWhileM
87 -- * Inserting Elements
88 -- | Produce a superset of the stream.
89 , insertBy
90 , intersperse
91 , intersperseM
92 , intersperseSuffix
93 , intersperseSuffixBySpan
95 -- * Inserting Side Effects
96 , intersperseM_
97 , intersperseSuffix_
99 -- * Reordering
100 -- | Produce strictly the same set but reordered.
101 , reverse
102 -- , reverse'
104 -- * Position Indexing
105 , indexed
106 , indexedR
108 -- * Searching
109 , findIndices
110 , slicesBy
112 -- * Rolling map
113 -- | Map using the previous element.
114 , rollingMap
115 , rollingMapM
117 -- * Maybe Streams
118 , mapMaybe
119 , mapMaybeM
120 )
121where
123#include "inline.hs"
125import Control.Concurrent (killThread, threadDelay)
126import Control.Exception (AsyncException)
127import Control.Monad (void, when)
128import Control.Monad.Catch (MonadCatch, throwM)
129import Control.Monad.IO.Class (MonadIO(..))
130import Control.Monad.Trans.Class (MonadTrans(lift))
131import Data.IORef (newIORef, mkWeakIORef)
132import Data.Maybe (fromJust, isJust)
133import GHC.Types (SPEC(..))
134import qualified Control.Monad.Catch as MC
136import Streamly.Internal.Control.Concurrent (MonadAsync, fork, forkManaged)
137import Streamly.Internal.Data.Fold.Type (Fold(..))
138import Streamly.Internal.Data.Pipe.Type (Pipe(..), PipeState(..))
139import Streamly.Internal.Data.SVar.Type (defState, adaptState)
140import Streamly.Internal.Data.Time.Clock (Clock(Monotonic), getTime)
141import Streamly.Internal.Data.Time.Units
142 (TimeUnit64, toRelTime64, diffAbsTime64)
144import qualified Streamly.Internal.Data.Fold.Type as FL
145import qualified Streamly.Internal.Data.IORef.Prim as Prim
146import qualified Streamly.Internal.Data.Pipe.Type as Pipe
148import Prelude hiding
149 ( drop, dropWhile, filter, map, mapM, reverse
150 , scanl, scanl1, sequence, take, takeWhile)
152import Streamly.Internal.Data.Stream.StreamD.Type
154------------------------------------------------------------------------------
155-- Piping
156------------------------------------------------------------------------------
158{-# INLINE_NORMAL transform #-}
159transform :: Monad m => Pipe m a b -> Stream m a -> Stream m b
160transform (Pipe pstep1 pstep2 pstate) (Stream step state) =
161 Stream step' (Consume pstate, state)
163 where
165 {-# INLINE_LATE step' #-}
167 step' gst (Consume pst, st) = pst `seq` do
168 r <- step (adaptState gst) st
169 case r of
170 Yield x s -> do
171 res <- pstep1 pst x
172 case res of
173 Pipe.Yield b pst' -> return $ Yield b (pst', s)
174 Pipe.Continue pst' -> return $ Skip (pst', s)
175 Skip s -> return $ Skip (Consume pst, s)
176 Stop -> return Stop
178 step' _ (Produce pst, st) = pst `seq` do
179 res <- pstep2 pst
180 case res of
181 Pipe.Yield b pst' -> return $ Yield b (pst', st)
182 Pipe.Continue pst' -> return $ Skip (pst', st)
184------------------------------------------------------------------------------
185-- Transformation Folds
186------------------------------------------------------------------------------
188{-# INLINE_NORMAL foldlT #-}
189foldlT :: (Monad m, Monad (s m), MonadTrans s)
190 => (s m b -> a -> s m b) -> s m b -> Stream m a -> s m b
191foldlT fstep begin (Stream step state) = go SPEC begin state
192 where
193 go !_ acc st = do
194 r <- lift $ step defState st
195 case r of
196 Yield x s -> go SPEC (fstep acc x) s
197 Skip s -> go SPEC acc s
198 Stop -> acc
200-- Note, this is going to have horrible performance, because of the nature of
201-- the stream type (i.e. direct stream vs CPS). Its only for reference, it is
202-- likely be practically unusable.
203{-# INLINE_NORMAL foldlS #-}
204foldlS :: Monad m
205 => (Stream m b -> a -> Stream m b) -> Stream m b -> Stream m a -> Stream m b
206foldlS fstep begin (Stream step state) = Stream step' (Left (state, begin))
207 where
208 step' gst (Left (st, acc)) = do
209 r <- step (adaptState gst) st
210 return $ case r of
211 Yield x s -> Skip (Left (s, fstep acc x))
212 Skip s -> Skip (Left (s, acc))
213 Stop -> Skip (Right acc)
215 step' gst (Right (Stream stp stt)) = do
216 r <- stp (adaptState gst) stt
217 return $ case r of
218 Yield x s -> Yield x (Right (Stream stp s))
219 Skip s -> Skip (Right (Stream stp s))
220 Stop -> Stop
222------------------------------------------------------------------------------
223-- Transformation by Mapping
224------------------------------------------------------------------------------
226{-# INLINE_NORMAL sequence #-}
227sequence :: Monad m => Stream m (m a) -> Stream m a
228sequence (Stream step state) = Stream step' state
229 where
230 {-# INLINE_LATE step' #-}
231 step' gst st = do
232 r <- step (adaptState gst) st
233 case r of
234 Yield x s -> x >>= \a -> return (Yield a s)
235 Skip s -> return $ Skip s
236 Stop -> return Stop
238------------------------------------------------------------------------------
239-- Mapping side effects
240------------------------------------------------------------------------------
242data TapState fs st a
243 = TapInit | Tapping !fs st | TapDone st
245-- XXX Multiple yield points
246{-# INLINE tap #-}
247tap :: Monad m => Fold m a b -> Stream m a -> Stream m a
248tap (Fold fstep initial extract) (Stream step state) = Stream step' TapInit
250 where
252 step' _ TapInit = do
253 res <- initial
254 return
255 $ Skip
256 $ case res of
257 FL.Partial s -> Tapping s state
258 FL.Done _ -> TapDone state
259 step' gst (Tapping acc st) = do
260 r <- step gst st
261 case r of
262 Yield x s -> do
263 res <- fstep acc x
264 return
265 $ Yield x
266 $ case res of
267 FL.Partial fs -> Tapping fs s
268 FL.Done _ -> TapDone s
269 Skip s -> return $ Skip (Tapping acc s)
270 Stop -> do
271 void $ extract acc
272 return Stop
273 step' gst (TapDone st) = do
274 r <- step gst st
275 return
276 $ case r of
277 Yield x s -> Yield x (TapDone s)
278 Skip s -> Skip (TapDone s)
279 Stop -> Stop
281data TapOffState fs s a
282 = TapOffInit
283 | TapOffTapping !fs s Int
284 | TapOffDone s
286-- XXX Multiple yield points
287{-# INLINE_NORMAL tapOffsetEvery #-}
288tapOffsetEvery :: Monad m
289 => Int -> Int -> Fold m a b -> Stream m a -> Stream m a
290tapOffsetEvery offset n (Fold fstep initial extract) (Stream step state) =
291 Stream step' TapOffInit
293 where
295 {-# INLINE_LATE step' #-}
296 step' _ TapOffInit = do
297 res <- initial
298 return
299 $ Skip
300 $ case res of
301 FL.Partial s -> TapOffTapping s state (offset `mod` n)
302 FL.Done _ -> TapOffDone state
303 step' gst (TapOffTapping acc st count) = do
304 r <- step gst st
305 case r of
306 Yield x s -> do
307 next <-
308 if count <= 0
309 then do
310 res <- fstep acc x
311 return
312 $ case res of
313 FL.Partial sres ->
314 TapOffTapping sres s (n - 1)
315 FL.Done _ -> TapOffDone s
316 else return $ TapOffTapping acc s (count - 1)
317 return $ Yield x next
318 Skip s -> return $ Skip (TapOffTapping acc s count)
319 Stop -> do
320 void $ extract acc
321 return Stop
322 step' gst (TapOffDone st) = do
323 r <- step gst st
324 return
325 $ case r of
326 Yield x s -> Yield x (TapOffDone s)
327 Skip s -> Skip (TapOffDone s)
328 Stop -> Stop
330{-# INLINE_NORMAL pollCounts #-}
331pollCounts
332 :: MonadAsync m
333 => (a -> Bool)
334 -> (Stream m Int -> Stream m Int)
335 -> Fold m Int b
336 -> Stream m a
337 -> Stream m a
338pollCounts predicate transf fld (Stream step state) = Stream step' Nothing
339 where
341 {-# INLINE_LATE step' #-}
342 step' _ Nothing = do
343 -- As long as we are using an "Int" for counts lockfree reads from
344 -- Var should work correctly on both 32-bit and 64-bit machines.
345 -- However, an Int on a 32-bit machine may overflow quickly.
346 countVar <- liftIO $ Prim.newIORef (0 :: Int)
347 tid <- forkManaged
348 $ void $ fold fld
349 $ transf $ Prim.toStreamD countVar
350 return $ Skip (Just (countVar, tid, state))
352 step' gst (Just (countVar, tid, st)) = do
353 r <- step gst st
354 case r of
355 Yield x s -> do
356 when (predicate x) $ liftIO $ Prim.modifyIORef' countVar (+ 1)
357 return $ Yield x (Just (countVar, tid, s))
358 Skip s -> return $ Skip (Just (countVar, tid, s))
359 Stop -> do
360 liftIO $ killThread tid
361 return Stop
363{-# INLINE_NORMAL tapRate #-}
364tapRate ::
365 (MonadAsync m, MonadCatch m)
366 => Double
367 -> (Int -> m b)
368 -> Stream m a
369 -> Stream m a
370tapRate samplingRate action (Stream step state) = Stream step' Nothing
371 where
372 {-# NOINLINE loop #-}
373 loop countVar prev = do
374 i <-
375 MC.catch
376 (do liftIO $ threadDelay (round $ samplingRate * 1000000)
377 i <- liftIO $ Prim.readIORef countVar
378 let !diff = i - prev
379 void $ action diff
380 return i)
381 (\(e :: AsyncException) -> do
382 i <- liftIO $ Prim.readIORef countVar
383 let !diff = i - prev
384 void $ action diff
385 throwM (MC.toException e))
386 loop countVar i
388 {-# INLINE_LATE step' #-}
389 step' _ Nothing = do
390 countVar <- liftIO $ Prim.newIORef 0
391 tid <- fork $ loop countVar 0
392 ref <- liftIO $ newIORef ()
393 _ <- liftIO $ mkWeakIORef ref (killThread tid)
394 return $ Skip (Just (countVar, tid, state, ref))
396 step' gst (Just (countVar, tid, st, ref)) = do
397 r <- step gst st
398 case r of
399 Yield x s -> do
400 liftIO $ Prim.modifyIORef' countVar (+ 1)
401 return $ Yield x (Just (countVar, tid, s, ref))
402 Skip s -> return $ Skip (Just (countVar, tid, s, ref))
403 Stop -> do
404 liftIO $ killThread tid
405 return Stop
407------------------------------------------------------------------------------
408-- Scanning with a Fold
409------------------------------------------------------------------------------
411data ScanState s f = ScanInit s | ScanDo s !f | ScanDone
413{-# INLINE_NORMAL postscanOnce #-}
414postscanOnce :: Monad m => FL.Fold m a b -> Stream m a -> Stream m b
415postscanOnce (FL.Fold fstep initial extract) (Stream sstep state) =
416 Stream step (ScanInit state)
418 where
420 {-# INLINE_LATE step #-}
421 step _ (ScanInit st) = do
422 res <- initial
423 return
424 $ case res of
425 FL.Partial fs -> Skip $ ScanDo st fs
426 FL.Done b -> Yield b ScanDone
427 step gst (ScanDo st fs) = do
428 res <- sstep (adaptState gst) st
429 case res of
430 Yield x s -> do
431 r <- fstep fs x
432 case r of
433 FL.Partial fs1 -> do
434 !b <- extract fs1
435 return $ Yield b $ ScanDo s fs1
436 FL.Done b -> return $ Yield b ScanDone
437 Skip s -> return $ Skip $ ScanDo s fs
438 Stop -> return Stop
439 step _ ScanDone = return Stop
441{-# INLINE scanOnce #-}
442scanOnce :: Monad m
443 => FL.Fold m a b -> Stream m a -> Stream m b
444scanOnce (FL.Fold fstep initial extract) (Stream sstep state) =
445 Stream step (ScanInit state)
447 where
449 {-# INLINE_LATE step #-}
450 step _ (ScanInit st) = do
451 res <- initial
452 case res of
453 FL.Partial fs -> do
454 !b <- extract fs
455 return $ Yield b $ ScanDo st fs
456 FL.Done b -> return $ Yield b ScanDone
457 step gst (ScanDo st fs) = do
458 res <- sstep (adaptState gst) st
459 case res of
460 Yield x s -> do
461 r <- fstep fs x
462 case r of
463 FL.Partial fs1 -> do
464 !b <- extract fs1
465 return $ Yield b $ ScanDo s fs1
466 FL.Done b -> return $ Yield b ScanDone
467 Skip s -> return $ Skip $ ScanDo s fs
468 Stop -> return Stop
469 step _ ScanDone = return Stop
471------------------------------------------------------------------------------
472-- Scanning - Prescans
473------------------------------------------------------------------------------
475-- Adapted from the vector package.
476--
477-- XXX Is a prescan useful, discarding the last step does not sound useful? I
478-- am not sure about the utility of this function, so this is implemented but
479-- not exposed. We can expose it if someone provides good reasons why this is
480-- useful.
481--
482-- XXX We have to execute the stream one step ahead to know that we are at the
483-- last step. The vector implementation of prescan executes the last fold step
484-- but does not yield the result. This means we have executed the effect but
485-- discarded value. This does not sound right. In this implementation we are
486-- not executing the last fold step.
487{-# INLINE_NORMAL prescanlM' #-}
488prescanlM' :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
489prescanlM' f mz (Stream step state) = Stream step' (state, mz)
490 where
491 {-# INLINE_LATE step' #-}
492 step' gst (st, prev) = do
493 r <- step (adaptState gst) st
494 case r of
495 Yield x s -> do
496 acc <- prev
497 return $ Yield acc (s, f acc x)
498 Skip s -> return $ Skip (s, prev)
499 Stop -> return Stop
501{-# INLINE prescanl' #-}
502prescanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
503prescanl' f z = prescanlM' (\a b -> return (f a b)) (return z)
505------------------------------------------------------------------------------
506-- Monolithic postscans (postscan followed by a map)
507------------------------------------------------------------------------------
509-- The performance of a modular postscan followed by a map seems to be
510-- equivalent to this monolithic scan followed by map therefore we may not need
511-- this implementation. We just have it for performance comparison and in case
512-- modular version does not perform well in some situation.
513--
514{-# INLINE_NORMAL postscanlMx' #-}
515postscanlMx' :: Monad m
516 => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> Stream m b
517postscanlMx' fstep begin done (Stream step state) = do
518 Stream step' (state, begin)
519 where
520 {-# INLINE_LATE step' #-}
521 step' gst (st, acc) = do
522 r <- step (adaptState gst) st
523 case r of
524 Yield x s -> do
525 old <- acc
526 y <- fstep old x
527 v <- done y
528 v `seq` y `seq` return (Yield v (s, return y))
529 Skip s -> return $ Skip (s, acc)
530 Stop -> return Stop
532{-# INLINE_NORMAL postscanlx' #-}
533postscanlx' :: Monad m
534 => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> Stream m b
535postscanlx' fstep begin done =
536 postscanlMx' (\b a -> return (fstep b a)) (return begin) (return . done)
538-- XXX do we need consM strict to evaluate the begin value?
539{-# INLINE scanlMx' #-}
540scanlMx' :: Monad m
541 => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> Stream m b
542scanlMx' fstep begin done s =
543 (begin >>= \x -> x `seq` done x) `consM` postscanlMx' fstep begin done s
545{-# INLINE scanlx' #-}
546scanlx' :: Monad m
547 => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> Stream m b
548scanlx' fstep begin done =
549 scanlMx' (\b a -> return (fstep b a)) (return begin) (return . done)
551------------------------------------------------------------------------------
552-- postscans
553------------------------------------------------------------------------------
555-- Adapted from the vector package.
556{-# INLINE_NORMAL postscanlM' #-}
557postscanlM' :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
558postscanlM' fstep begin (Stream step state) =
559 Stream step' Nothing
560 where
561 {-# INLINE_LATE step' #-}
562 step' _ Nothing = do
563 !x <- begin
564 return $ Skip (Just (state, x))
566 step' gst (Just (st, acc)) = do
567 r <- step (adaptState gst) st
568 case r of
569 Yield x s -> do
570 !y <- fstep acc x
571 return $ Yield y (Just (s, y))
572 Skip s -> return $ Skip (Just (s, acc))
573 Stop -> return Stop
575{-# INLINE_NORMAL postscanl' #-}
576postscanl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
577postscanl' f seed = postscanlM' (\a b -> return (f a b)) (return seed)
579-- We can possibly have the "done" function as a Maybe to provide an option to
580-- emit or not emit the accumulator when the stream stops.
581--
582-- TBD: use a single Yield point
583--
584{-# INLINE_NORMAL postscanlMAfter' #-}
585postscanlMAfter' :: Monad m
586 => (b -> a -> m b) -> m b -> (b -> m b) -> Stream m a -> Stream m b
587postscanlMAfter' fstep initial done (Stream step1 state1) = do
588 Stream step (Just (state1, initial))
590 where
592 {-# INLINE_LATE step #-}
593 step gst (Just (st, acc)) = do
594 r <- step1 (adaptState gst) st
595 case r of
596 Yield x s -> do
597 old <- acc
598 y <- fstep old x
599 y `seq` return (Yield y (Just (s, return y)))
600 Skip s -> return $ Skip $ Just (s, acc)
601 Stop -> acc >>= done >>= \res -> return (Yield res Nothing)
602 step _ Nothing = return Stop
604{-# INLINE_NORMAL postscanlM #-}
605postscanlM :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
606postscanlM fstep begin (Stream step state) = Stream step' Nothing
607 where
608 {-# INLINE_LATE step' #-}
609 step' _ Nothing = do
610 r <- begin
611 return $ Skip (Just (state, r))
613 step' gst (Just (st, acc)) = do
614 r <- step (adaptState gst) st
615 case r of
616 Yield x s -> do
617 y <- fstep acc x
618 return (Yield y (Just (s, y)))
619 Skip s -> return $ Skip (Just (s, acc))
620 Stop -> return Stop
622{-# INLINE_NORMAL postscanl #-}
623postscanl :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
624postscanl f seed = postscanlM (\a b -> return (f a b)) (return seed)
626{-# INLINE_NORMAL scanlM' #-}
627scanlM' :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
628scanlM' fstep begin (Stream step state) = Stream step' Nothing
629 where
630 {-# INLINE_LATE step' #-}
631 step' _ Nothing = do
632 !x <- begin
633 return $ Yield x (Just (state, x))
634 step' gst (Just (st, acc)) = do
635 r <- step (adaptState gst) st
636 case r of
637 Yield x s -> do
638 !y <- fstep acc x
639 return $ Yield y (Just (s, y))
640 Skip s -> return $ Skip (Just (s, acc))
641 Stop -> return Stop
643{-# INLINE scanlMAfter' #-}
644scanlMAfter' :: Monad m
645 => (b -> a -> m b) -> m b -> (b -> m b) -> Stream m a -> Stream m b
646scanlMAfter' fstep initial done s =
647 (initial >>= \x -> x `seq` return x) `consM`
648 postscanlMAfter' fstep initial done s
650{-# INLINE scanl' #-}
651scanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
652scanl' f seed = scanlM' (\a b -> return (f a b)) (return seed)
654{-# INLINE_NORMAL scanlM #-}
655scanlM :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
656scanlM fstep begin (Stream step state) = Stream step' Nothing
657 where
658 {-# INLINE_LATE step' #-}
659 step' _ Nothing = do
660 x <- begin
661 return $ Yield x (Just (state, x))
662 step' gst (Just (st, acc)) = do
663 r <- step (adaptState gst) st
664 case r of
665 Yield x s -> do
666 y <- fstep acc x
667 return $ Yield y (Just (s, y))
668 Skip s -> return $ Skip (Just (s, acc))
669 Stop -> return Stop
671{-# INLINE scanl #-}
672scanl :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
673scanl f seed = scanlM (\a b -> return (f a b)) (return seed)
675-- Adapted from the vector package
676{-# INLINE_NORMAL scanl1M #-}
677scanl1M :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
678scanl1M fstep (Stream step state) = Stream step' (state, Nothing)
679 where
680 {-# INLINE_LATE step' #-}
681 step' gst (st, Nothing) = do
682 r <- step gst st
683 case r of
684 Yield x s -> return $ Yield x (s, Just x)
685 Skip s -> return $ Skip (s, Nothing)
686 Stop -> return Stop
688 step' gst (st, Just acc) = do
689 r <- step gst st
690 case r of
691 Yield y s -> do
692 z <- fstep acc y
693 return $ Yield z (s, Just z)
694 Skip s -> return $ Skip (s, Just acc)
695 Stop -> return Stop
697{-# INLINE scanl1 #-}
698scanl1 :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
699scanl1 f = scanl1M (\x y -> return (f x y))
701-- Adapted from the vector package
702{-# INLINE_NORMAL scanl1M' #-}
703scanl1M' :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
704scanl1M' fstep (Stream step state) = Stream step' (state, Nothing)
705 where
706 {-# INLINE_LATE step' #-}
707 step' gst (st, Nothing) = do
708 r <- step gst st
709 case r of
710 Yield x s -> x `seq` return $ Yield x (s, Just x)
711 Skip s -> return $ Skip (s, Nothing)
712 Stop -> return Stop
714 step' gst (st, Just acc) = acc `seq` do
715 r <- step gst st
716 case r of
717 Yield y s -> do
718 z <- fstep acc y
719 z `seq` return $ Yield z (s, Just z)
720 Skip s -> return $ Skip (s, Just acc)
721 Stop -> return Stop
723{-# INLINE scanl1' #-}
724scanl1' :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
725scanl1' f = scanl1M' (\x y -> return (f x y))
727-------------------------------------------------------------------------------
728-- Filtering
729-------------------------------------------------------------------------------
731-- Adapted from the vector package
732{-# INLINE_NORMAL filterM #-}
733filterM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
734filterM f (Stream step state) = Stream step' state
735 where
736 {-# INLINE_LATE step' #-}
737 step' gst st = do
738 r <- step gst st
739 case r of
740 Yield x s -> do
741 b <- f x
742 return $ if b
743 then Yield x s
744 else Skip s
745 Skip s -> return $ Skip s
746 Stop -> return Stop
748{-# INLINE filter #-}
749filter :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
750filter f = filterM (return . f)
752-- Adapted from the vector package
753{-# INLINE_NORMAL uniq #-}
754uniq :: (Eq a, Monad m) => Stream m a -> Stream m a
755uniq (Stream step state) = Stream step' (Nothing, state)
756 where
757 {-# INLINE_LATE step' #-}
758 step' gst (Nothing, st) = do
759 r <- step gst st
760 case r of
761 Yield x s -> return $ Yield x (Just x, s)
762 Skip s -> return $ Skip (Nothing, s)
763 Stop -> return Stop
764 step' gst (Just x, st) = do
765 r <- step gst st
766 case r of
767 Yield y s | x == y -> return $ Skip (Just x, s)
768 | otherwise -> return $ Yield y (Just y, s)
769 Skip s -> return $ Skip (Just x, s)
770 Stop -> return Stop
772{-# INLINE_NORMAL deleteBy #-}
773deleteBy :: Monad m => (a -> a -> Bool) -> a -> Stream m a -> Stream m a
774deleteBy eq x (Stream step state) = Stream step' (state, False)
775 where
776 {-# INLINE_LATE step' #-}
777 step' gst (st, False) = do
778 r <- step gst st
779 case r of
780 Yield y s -> return $
781 if eq x y then Skip (s, True) else Yield y (s, False)
782 Skip s -> return $ Skip (s, False)
783 Stop -> return Stop
785 step' gst (st, True) = do
786 r <- step gst st
787 case r of
788 Yield y s -> return $ Yield y (s, True)
789 Skip s -> return $ Skip (s, True)
790 Stop -> return Stop
792------------------------------------------------------------------------------
793-- Trimming
794------------------------------------------------------------------------------
796-- XXX using getTime in the loop can be pretty expensive especially for
797-- computations where iterations are lightweight. We have the following
798-- options:
799--
800-- 1) Run a timeout thread updating a flag asynchronously and check that
801-- flag here, that way we can have a cheap termination check.
802--
803-- 2) Use COARSE clock to get time with lower resolution but more efficiently.
804--
805-- 3) Use rdtscp/rdtsc to get time directly from the processor, compute the
806-- termination value of rdtsc in the beginning and then in each iteration just
807-- get rdtsc and check if we should terminate.
808--
809data TakeByTime st s
810 = TakeByTimeInit st
811 | TakeByTimeCheck st s
812 | TakeByTimeYield st s
814{-# INLINE_NORMAL takeByTime #-}
815takeByTime :: (MonadIO m, TimeUnit64 t) => t -> Stream m a -> Stream m a
816takeByTime duration (Stream step1 state1) = Stream step (TakeByTimeInit state1)
817 where
819 lim = toRelTime64 duration
821 {-# INLINE_LATE step #-}
822 step _ (TakeByTimeInit _) | lim == 0 = return Stop
823 step _ (TakeByTimeInit st) = do
824 t0 <- liftIO $ getTime Monotonic
825 return $ Skip (TakeByTimeYield st t0)
826 step _ (TakeByTimeCheck st t0) = do
827 t <- liftIO $ getTime Monotonic
828 return $
829 if diffAbsTime64 t t0 > lim
830 then Stop
831 else Skip (TakeByTimeYield st t0)
832 step gst (TakeByTimeYield st t0) = do
833 r <- step1 gst st
834 return $ case r of
835 Yield x s -> Yield x (TakeByTimeCheck s t0)
836 Skip s -> Skip (TakeByTimeCheck s t0)
837 Stop -> Stop
839data DropByTime st s x
840 = DropByTimeInit st
841 | DropByTimeGen st s
842 | DropByTimeCheck st s x
843 | DropByTimeYield st
845{-# INLINE_NORMAL dropByTime #-}
846dropByTime :: (MonadIO m, TimeUnit64 t) => t -> Stream m a -> Stream m a
847dropByTime duration (Stream step1 state1) = Stream step (DropByTimeInit state1)
848 where
850 lim = toRelTime64 duration
852 {-# INLINE_LATE step #-}
853 step _ (DropByTimeInit st) = do
854 t0 <- liftIO $ getTime Monotonic
855 return $ Skip (DropByTimeGen st t0)
856 step gst (DropByTimeGen st t0) = do
857 r <- step1 gst st
858 return $ case r of
859 Yield x s -> Skip (DropByTimeCheck s t0 x)
860 Skip s -> Skip (DropByTimeGen s t0)
861 Stop -> Stop
862 step _ (DropByTimeCheck st t0 x) = do
863 t <- liftIO $ getTime Monotonic
864 if diffAbsTime64 t t0 <= lim
865 then return $ Skip $ DropByTimeGen st t0
866 else return $ Yield x $ DropByTimeYield st
867 step gst (DropByTimeYield st) = do
868 r <- step1 gst st
869 return $ case r of
870 Yield x s -> Yield x (DropByTimeYield s)
871 Skip s -> Skip (DropByTimeYield s)
872 Stop -> Stop
874-- Adapted from the vector package
875{-# INLINE_NORMAL drop #-}
876drop :: Monad m => Int -> Stream m a -> Stream m a
877drop n (Stream step state) = Stream step' (state, Just n)
878 where
879 {-# INLINE_LATE step' #-}
880 step' gst (st, Just i)
881 | i > 0 = do
882 r <- step gst st
883 return $
884 case r of
885 Yield _ s -> Skip (s, Just (i - 1))
886 Skip s -> Skip (s, Just i)
887 Stop -> Stop
888 | otherwise = return $ Skip (st, Nothing)
890 step' gst (st, Nothing) = do
891 r <- step gst st
892 return $
893 case r of
894 Yield x s -> Yield x (s, Nothing)
895 Skip s -> Skip (s, Nothing)
896 Stop -> Stop
898-- Adapted from the vector package
899data DropWhileState s a
900 = DropWhileDrop s
901 | DropWhileYield a s
902 | DropWhileNext s
904{-# INLINE_NORMAL dropWhileM #-}
905dropWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
906dropWhileM f (Stream step state) = Stream step' (DropWhileDrop state)
907 where
908 {-# INLINE_LATE step' #-}
909 step' gst (DropWhileDrop st) = do
910 r <- step gst st
911 case r of
912 Yield x s -> do
913 b <- f x
914 if b
915 then return $ Skip (DropWhileDrop s)
916 else return $ Skip (DropWhileYield x s)
917 Skip s -> return $ Skip (DropWhileDrop s)
918 Stop -> return Stop
920 step' gst (DropWhileNext st) = do
921 r <- step gst st
922 case r of
923 Yield x s -> return $ Skip (DropWhileYield x s)
924 Skip s -> return $ Skip (DropWhileNext s)
925 Stop -> return Stop
927 step' _ (DropWhileYield x st) = return $ Yield x (DropWhileNext st)
929{-# INLINE dropWhile #-}
930dropWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
931dropWhile f = dropWhileM (return . f)
933------------------------------------------------------------------------------
934-- Inserting Elements
935------------------------------------------------------------------------------
937{-# INLINE_NORMAL insertBy #-}
938insertBy :: Monad m => (a -> a -> Ordering) -> a -> Stream m a -> Stream m a
939insertBy cmp a (Stream step state) = Stream step' (state, False, Nothing)
940 where
941 {-# INLINE_LATE step' #-}
942 step' gst (st, False, _) = do
943 r <- step gst st
944 case r of
945 Yield x s -> case cmp a x of
946 GT -> return $ Yield x (s, False, Nothing)
947 _ -> return $ Yield a (s, True, Just x)
948 Skip s -> return $ Skip (s, False, Nothing)
949 Stop -> return $ Yield a (st, True, Nothing)
951 step' _ (_, True, Nothing) = return Stop
953 step' gst (st, True, Just prev) = do
954 r <- step gst st
955 case r of
956 Yield x s -> return $ Yield prev (s, True, Just x)
957 Skip s -> return $ Skip (s, True, Just prev)
958 Stop -> return $ Yield prev (st, True, Nothing)
960data LoopState x s = FirstYield s
961 | InterspersingYield s
962 | YieldAndCarry x s
964{-# INLINE_NORMAL intersperseM #-}
965intersperseM :: Monad m => m a -> Stream m a -> Stream m a
966intersperseM m (Stream step state) = Stream step' (FirstYield state)
967 where
968 {-# INLINE_LATE step' #-}
969 step' gst (FirstYield st) = do
970 r <- step gst st
971 return $
972 case r of
973 Yield x s -> Skip (YieldAndCarry x s)
974 Skip s -> Skip (FirstYield s)
975 Stop -> Stop
977 step' gst (InterspersingYield st) = do
978 r <- step gst st
979 case r of
980 Yield x s -> do
981 a <- m
982 return $ Yield a (YieldAndCarry x s)
983 Skip s -> return $ Skip $ InterspersingYield s
984 Stop -> return Stop
986 step' _ (YieldAndCarry x st) = return $ Yield x (InterspersingYield st)
988{-# INLINE intersperse #-}
989intersperse :: Monad m => a -> Stream m a -> Stream m a
990intersperse a = intersperseM (return a)
992{-# INLINE_NORMAL intersperseM_ #-}
993intersperseM_ :: Monad m => m b -> Stream m a -> Stream m a
994intersperseM_ m (Stream step1 state1) = Stream step (Left (pure (), state1))
995 where
996 {-# INLINE_LATE step #-}
997 step gst (Left (eff, st)) = do
998 r <- step1 gst st
999 case r of
1000 Yield x s -> eff >> return (Yield x (Right s))
1001 Skip s -> return $ Skip (Left (eff, s))
1002 Stop -> return Stop
1004 step _ (Right st) = return $ Skip $ Left (void m, st)
1006data SuffixState s a
1007 = SuffixElem s
1008 | SuffixSuffix s
1009 | SuffixYield a (SuffixState s a)
1011{-# INLINE_NORMAL intersperseSuffix #-}
1012intersperseSuffix :: forall m a. Monad m => m a -> Stream m a -> Stream m a
1013intersperseSuffix action (Stream step state) = Stream step' (SuffixElem state)
1014 where
1015 {-# INLINE_LATE step' #-}
1016 step' gst (SuffixElem st) = do
1017 r <- step gst st
1018 return $ case r of
1019 Yield x s -> Skip (SuffixYield x (SuffixSuffix s))
1020 Skip s -> Skip (SuffixElem s)
1021 Stop -> Stop
1023 step' _ (SuffixSuffix st) = do
1024 action >>= \r -> return $ Skip (SuffixYield r (SuffixElem st))
1026 step' _ (SuffixYield x next) = return $ Yield x next
1028{-# INLINE_NORMAL intersperseSuffix_ #-}
1029intersperseSuffix_ :: Monad m => m b -> Stream m a -> Stream m a
1030intersperseSuffix_ m (Stream step1 state1) = Stream step (Left state1)
1031 where
1032 {-# INLINE_LATE step #-}
1033 step gst (Left st) = do
1034 r <- step1 gst st
1035 case r of
1036 Yield x s -> return $ Yield x (Right s)
1037 Skip s -> return $ Skip $ Left s
1038 Stop -> return Stop
1040 step _ (Right st) = m >> return (Skip (Left st))
1042data SuffixSpanState s a
1043 = SuffixSpanElem s Int
1044 | SuffixSpanSuffix s
1045 | SuffixSpanYield a (SuffixSpanState s a)
1046 | SuffixSpanLast
1047 | SuffixSpanStop
1049-- | intersperse after every n items
1050{-# INLINE_NORMAL intersperseSuffixBySpan #-}
1051intersperseSuffixBySpan :: forall m a. Monad m
1052 => Int -> m a -> Stream m a -> Stream m a
1053intersperseSuffixBySpan n action (Stream step state) =
1054 Stream step' (SuffixSpanElem state n)
1055 where
1056 {-# INLINE_LATE step' #-}
1057 step' gst (SuffixSpanElem st i) | i > 0 = do
1058 r <- step gst st
1059 return $ case r of
1060 Yield x s -> Skip (SuffixSpanYield x (SuffixSpanElem s (i - 1)))
1061 Skip s -> Skip (SuffixSpanElem s i)
1062 Stop -> if i == n then Stop else Skip SuffixSpanLast
1063 step' _ (SuffixSpanElem st _) = return $ Skip (SuffixSpanSuffix st)
1065 step' _ (SuffixSpanSuffix st) = do
1066 action >>= \r -> return $ Skip (SuffixSpanYield r (SuffixSpanElem st n))
1068 step' _ SuffixSpanLast = do
1069 action >>= \r -> return $ Skip (SuffixSpanYield r SuffixSpanStop)
1071 step' _ (SuffixSpanYield x next) = return $ Yield x next
1073 step' _ SuffixSpanStop = return Stop
1075------------------------------------------------------------------------------
1076-- Reordering
1077------------------------------------------------------------------------------
1079-- We can implement reverse as:
1081-- > reverse = foldlS (flip cons) nil
1083-- However, this implementation is unusable because of the horrible performance
1084-- of cons. So we just convert it to a list first and then stream from the
1085-- list.
1087-- XXX Maybe we can use an Array instead of a list here?
1088{-# INLINE_NORMAL reverse #-}
1089reverse :: Monad m => Stream m a -> Stream m a
1090reverse m = Stream step Nothing
1091 where
1092 {-# INLINE_LATE step #-}
1093 step _ Nothing = do
1094 xs <- foldl' (flip (:)) [] m
1095 return $ Skip (Just xs)
1096 step _ (Just (x:xs)) = return $ Yield x (Just xs)
1097 step _ (Just []) = return Stop
1099-- Much faster reverse for Storables
1101{-# INLINE_NORMAL reverse' #-}
1102reverse' :: forall m a. (MonadIO m, Storable a) => Stream m a -> Stream m a
1103-- This commented implementation copies the whole stream into one single array
1104-- and then streams from that array, this has exactly the same performance as
1105-- the chunked code in IsStream.Common.reverse' . Though this could be problematic due to
1106-- unbounded large allocations. However, if we use an idiomatic implementation
1107-- of arraysOf instead of the custom implementation then the chunked code
1108-- becomes worse by 6 times. Need to investigate if that can be improved.
1109import Foreign.ForeignPtr (touchForeignPtr)
1110import Foreign.ForeignPtr.Unsafe (unsafeForeignPtrToPtr)
1111import Foreign.Ptr (Ptr, plusPtr)
1112reverse' m = Stream step Nothing
1113 where
1114 {-# INLINE_LATE step #-}
1115 step _ Nothing = do
1116 arr <- A.fromStreamD m
1117 let p = A.aEnd arr `plusPtr` negate (sizeOf (undefined :: a))
1118 return $ Skip $ Just (A.aStart arr, p)
1120 step _ (Just (start, p)) | p < unsafeForeignPtrToPtr start = return Stop
1122 step _ (Just (start, p)) = do
1123 let !x = A.unsafeInlineIO $ do
1124 r <- peek p
1125 touchForeignPtr start
1126 return r
1127 next = p `plusPtr` negate (sizeOf (undefined :: a))
1128 return $ Yield x (Just (start, next))
1131------------------------------------------------------------------------------
1132-- Position Indexing
1133------------------------------------------------------------------------------
1135-- Adapted from the vector package
1136{-# INLINE_NORMAL indexed #-}
1137indexed :: Monad m => Stream m a -> Stream m (Int, a)
1138indexed (Stream step state) = Stream step' (state, 0)
1139 where
1140 {-# INLINE_LATE step' #-}
1141 step' gst (st, i) = i `seq` do
1142 r <- step (adaptState gst) st
1143 case r of
1144 Yield x s -> return $ Yield (i, x) (s, i+1)
1145 Skip s -> return $ Skip (s, i)
1146 Stop -> return Stop
1148-- Adapted from the vector package
1149{-# INLINE_NORMAL indexedR #-}
1150indexedR :: Monad m => Int -> Stream m a -> Stream m (Int, a)
1151indexedR m (Stream step state) = Stream step' (state, m)
1152 where
1153 {-# INLINE_LATE step' #-}
1154 step' gst (st, i) = i `seq` do
1155 r <- step (adaptState gst) st
1156 case r of
1157 Yield x s -> let i' = i - 1
1158 in return $ Yield (i, x) (s, i')
1159 Skip s -> return $ Skip (s, i)
1160 Stop -> return Stop
1162------------------------------------------------------------------------------
1163-- Searching
1164------------------------------------------------------------------------------
1166{-# INLINE_NORMAL findIndices #-}
1167findIndices :: Monad m => (a -> Bool) -> Stream m a -> Stream m Int
1168findIndices p (Stream step state) = Stream step' (state, 0)
1169 where
1170 {-# INLINE_LATE step' #-}
1171 step' gst (st, i) = i `seq` do
1172 r <- step (adaptState gst) st
1173 return $ case r of
1174 Yield x s -> if p x then Yield i (s, i+1) else Skip (s, i+1)
1175 Skip s -> Skip (s, i)
1176 Stop -> Stop
1178{-# INLINE_NORMAL slicesBy #-}
1179slicesBy :: Monad m => (a -> Bool) -> Stream m a -> Stream m (Int, Int)
1180slicesBy p (Stream step1 state1) = Stream step (Just (state1, 0, 0))
1182 where
1184 {-# INLINE_LATE step #-}
1185 step gst (Just (st, i, len)) = i `seq` len `seq` do
1186 r <- step1 (adaptState gst) st
1187 return
1188 $ case r of
1189 Yield x s ->
1190 if p x
1191 then Yield (i, len + 1) (Just (s, i + len + 1, 0))
1192 else Skip (Just (s, i, len + 1))
1193 Skip s -> Skip (Just (s, i, len))
1194 Stop -> if len == 0 then Stop else Yield (i, len) Nothing
1195 step _ Nothing = return Stop
1197------------------------------------------------------------------------------
1198-- Rolling map
1199------------------------------------------------------------------------------
1201data RollingMapState s a = RollingMapInit s | RollingMapGo s a
1203{-# INLINE rollingMapM #-}
1204rollingMapM :: Monad m => (a -> a -> m b) -> Stream m a -> Stream m b
1205rollingMapM f (Stream step1 state1) = Stream step (RollingMapInit state1)
1206 where
1207 step gst (RollingMapInit st) = do
1208 r <- step1 (adaptState gst) st
1209 return $ case r of
1210 Yield x s -> Skip $ RollingMapGo s x
1211 Skip s -> Skip $ RollingMapInit s
1212 Stop -> Stop
1214 step gst (RollingMapGo s1 x1) = do
1215 r <- step1 (adaptState gst) s1
1216 case r of
1217 Yield x s -> do
1218 !res <- f x x1
1219 return $ Yield res $ RollingMapGo s x
1220 Skip s -> return $ Skip $ RollingMapGo s x1
1221 Stop -> return Stop
1223{-# INLINE rollingMap #-}
1224rollingMap :: Monad m => (a -> a -> b) -> Stream m a -> Stream m b
1225rollingMap f = rollingMapM (\x y -> return $ f x y)
1227------------------------------------------------------------------------------
1228-- Maybe Streams
1229------------------------------------------------------------------------------
1231-- XXX Will this always fuse properly?
1232{-# INLINE_NORMAL mapMaybe #-}
1233mapMaybe :: Monad m => (a -> Maybe b) -> Stream m a -> Stream m b
1234mapMaybe f = fmap fromJust . filter isJust . map f
1236{-# INLINE_NORMAL mapMaybeM #-}
1237mapMaybeM :: Monad m => (a -> m (Maybe b)) -> Stream m a -> Stream m b
1238mapMaybeM f = fmap fromJust . filter isJust . mapM f