git.iko.soy/streamly
on a322c07 help sign in
streamly/src/Streamly/Internal/Data/Stream/StreamD/Transform.hs
1227 lines 41 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
111 -- * Rolling map
112 -- | Map using the previous element.
113 , rollingMap
114 , rollingMapM
116 -- * Maybe Streams
117 , mapMaybe
118 , mapMaybeM
119 )
120where
122#include "inline.hs"
124import Control.Concurrent (killThread, threadDelay)
125import Control.Exception (AsyncException)
126import Control.Monad (void, when)
127import Control.Monad.Catch (MonadCatch, throwM)
128import Control.Monad.IO.Class (MonadIO(..))
129import Control.Monad.Trans.Class (MonadTrans(lift))
130import Data.IORef (newIORef, mkWeakIORef)
131import Data.Maybe (fromJust, isJust)
132import Foreign.Storable (Storable(..))
133import GHC.Types (SPEC(..))
134import qualified Control.Monad.Catch as MC
136import Streamly.Internal.Data.Fold.Type (Fold(..))
137import Streamly.Internal.Data.Pipe.Type (Pipe(..), PipeState(..))
138import Streamly.Internal.Data.Time.Clock (Clock(Monotonic), getTime)
139import Streamly.Internal.Data.Time.Units
140 (TimeUnit64, toRelTime64, diffAbsTime64)
141import Streamly.Internal.System.IO (defaultChunkSize)
143import qualified Streamly.Internal.Data.Array.Foreign.Type as A
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
147import qualified Streamly.Internal.Data.Stream.StreamK as K
149import Prelude hiding
150 ( drop, dropWhile, filter, map, mapM, reverse
151 , scanl, scanl1, sequence, take, takeWhile)
153import Streamly.Internal.Data.Stream.StreamD.Type
154import Streamly.Internal.Data.SVar
156------------------------------------------------------------------------------
157-- Piping
158------------------------------------------------------------------------------
160{-# INLINE_NORMAL transform #-}
161transform :: Monad m => Pipe m a b -> Stream m a -> Stream m b
162transform (Pipe pstep1 pstep2 pstate) (Stream step state) =
163 Stream step' (Consume pstate, state)
165 where
167 {-# INLINE_LATE step' #-}
169 step' gst (Consume pst, st) = pst `seq` do
170 r <- step (adaptState gst) st
171 case r of
172 Yield x s -> do
173 res <- pstep1 pst x
174 case res of
175 Pipe.Yield b pst' -> return $ Yield b (pst', s)
176 Pipe.Continue pst' -> return $ Skip (pst', s)
177 Skip s -> return $ Skip (Consume pst, s)
178 Stop -> return Stop
180 step' _ (Produce pst, st) = pst `seq` do
181 res <- pstep2 pst
182 case res of
183 Pipe.Yield b pst' -> return $ Yield b (pst', st)
184 Pipe.Continue pst' -> return $ Skip (pst', st)
186------------------------------------------------------------------------------
187-- Transformation Folds
188------------------------------------------------------------------------------
190{-# INLINE_NORMAL foldlT #-}
191foldlT :: (Monad m, Monad (s m), MonadTrans s)
192 => (s m b -> a -> s m b) -> s m b -> Stream m a -> s m b
193foldlT fstep begin (Stream step state) = go SPEC begin state
194 where
195 go !_ acc st = do
196 r <- lift $ step defState st
197 case r of
198 Yield x s -> go SPEC (fstep acc x) s
199 Skip s -> go SPEC acc s
200 Stop -> acc
202-- Note, this is going to have horrible performance, because of the nature of
203-- the stream type (i.e. direct stream vs CPS). Its only for reference, it is
204-- likely be practically unusable.
205{-# INLINE_NORMAL foldlS #-}
206foldlS :: Monad m
207 => (Stream m b -> a -> Stream m b) -> Stream m b -> Stream m a -> Stream m b
208foldlS fstep begin (Stream step state) = Stream step' (Left (state, begin))
209 where
210 step' gst (Left (st, acc)) = do
211 r <- step (adaptState gst) st
212 return $ case r of
213 Yield x s -> Skip (Left (s, fstep acc x))
214 Skip s -> Skip (Left (s, acc))
215 Stop -> Skip (Right acc)
217 step' gst (Right (Stream stp stt)) = do
218 r <- stp (adaptState gst) stt
219 return $ case r of
220 Yield x s -> Yield x (Right (Stream stp s))
221 Skip s -> Skip (Right (Stream stp s))
222 Stop -> Stop
224------------------------------------------------------------------------------
225-- Transformation by Mapping
226------------------------------------------------------------------------------
228{-# INLINE_NORMAL sequence #-}
229sequence :: Monad m => Stream m (m a) -> Stream m a
230sequence (Stream step state) = Stream step' state
231 where
232 {-# INLINE_LATE step' #-}
233 step' gst st = do
234 r <- step (adaptState gst) st
235 case r of
236 Yield x s -> x >>= \a -> return (Yield a s)
237 Skip s -> return $ Skip s
238 Stop -> return Stop
240------------------------------------------------------------------------------
241-- Mapping side effects
242------------------------------------------------------------------------------
244data TapState fs st a
245 = TapInit | Tapping !fs st | TapDone st
247-- XXX Multiple yield points
248{-# INLINE tap #-}
249tap :: Monad m => Fold m a b -> Stream m a -> Stream m a
250tap (Fold fstep initial extract) (Stream step state) = Stream step' TapInit
252 where
254 step' _ TapInit = do
255 res <- initial
256 return
257 $ Skip
258 $ case res of
259 FL.Partial s -> Tapping s state
260 FL.Done _ -> TapDone state
261 step' gst (Tapping acc st) = do
262 r <- step gst st
263 case r of
264 Yield x s -> do
265 res <- fstep acc x
266 return
267 $ Yield x
268 $ case res of
269 FL.Partial fs -> Tapping fs s
270 FL.Done _ -> TapDone s
271 Skip s -> return $ Skip (Tapping acc s)
272 Stop -> do
273 void $ extract acc
274 return Stop
275 step' gst (TapDone st) = do
276 r <- step gst st
277 return
278 $ case r of
279 Yield x s -> Yield x (TapDone s)
280 Skip s -> Skip (TapDone s)
281 Stop -> Stop
283data TapOffState fs s a
284 = TapOffInit
285 | TapOffTapping !fs s Int
286 | TapOffDone s
288-- XXX Multiple yield points
289{-# INLINE_NORMAL tapOffsetEvery #-}
290tapOffsetEvery :: Monad m
291 => Int -> Int -> Fold m a b -> Stream m a -> Stream m a
292tapOffsetEvery offset n (Fold fstep initial extract) (Stream step state) =
293 Stream step' TapOffInit
295 where
297 {-# INLINE_LATE step' #-}
298 step' _ TapOffInit = do
299 res <- initial
300 return
301 $ Skip
302 $ case res of
303 FL.Partial s -> TapOffTapping s state (offset `mod` n)
304 FL.Done _ -> TapOffDone state
305 step' gst (TapOffTapping acc st count) = do
306 r <- step gst st
307 case r of
308 Yield x s -> do
309 next <-
310 if count <= 0
311 then do
312 res <- fstep acc x
313 return
314 $ case res of
315 FL.Partial sres ->
316 TapOffTapping sres s (n - 1)
317 FL.Done _ -> TapOffDone s
318 else return $ TapOffTapping acc s (count - 1)
319 return $ Yield x next
320 Skip s -> return $ Skip (TapOffTapping acc s count)
321 Stop -> do
322 void $ extract acc
323 return Stop
324 step' gst (TapOffDone st) = do
325 r <- step gst st
326 return
327 $ case r of
328 Yield x s -> Yield x (TapOffDone s)
329 Skip s -> Skip (TapOffDone s)
330 Stop -> Stop
332{-# INLINE_NORMAL pollCounts #-}
333pollCounts
334 :: MonadAsync m
335 => (a -> Bool)
336 -> (Stream m Int -> Stream m Int)
337 -> Fold m Int b
338 -> Stream m a
339 -> Stream m a
340pollCounts predicate transf fld (Stream step state) = Stream step' Nothing
341 where
343 {-# INLINE_LATE step' #-}
344 step' _ Nothing = do
345 -- As long as we are using an "Int" for counts lockfree reads from
346 -- Var should work correctly on both 32-bit and 64-bit machines.
347 -- However, an Int on a 32-bit machine may overflow quickly.
348 countVar <- liftIO $ Prim.newIORef (0 :: Int)
349 tid <- forkManaged
350 $ void $ fold fld
351 $ transf $ Prim.toStreamD countVar
352 return $ Skip (Just (countVar, tid, state))
354 step' gst (Just (countVar, tid, st)) = do
355 r <- step gst st
356 case r of
357 Yield x s -> do
358 when (predicate x) $ liftIO $ Prim.modifyIORef' countVar (+ 1)
359 return $ Yield x (Just (countVar, tid, s))
360 Skip s -> return $ Skip (Just (countVar, tid, s))
361 Stop -> do
362 liftIO $ killThread tid
363 return Stop
365{-# INLINE_NORMAL tapRate #-}
366tapRate ::
367 (MonadAsync m, MonadCatch m)
368 => Double
369 -> (Int -> m b)
370 -> Stream m a
371 -> Stream m a
372tapRate samplingRate action (Stream step state) = Stream step' Nothing
373 where
374 {-# NOINLINE loop #-}
375 loop countVar prev = do
376 i <-
377 MC.catch
378 (do liftIO $ threadDelay (round $ samplingRate * 1000000)
379 i <- liftIO $ Prim.readIORef countVar
380 let !diff = i - prev
381 void $ action diff
382 return i)
383 (\(e :: AsyncException) -> do
384 i <- liftIO $ Prim.readIORef countVar
385 let !diff = i - prev
386 void $ action diff
387 throwM (MC.toException e))
388 loop countVar i
390 {-# INLINE_LATE step' #-}
391 step' _ Nothing = do
392 countVar <- liftIO $ Prim.newIORef 0
393 tid <- fork $ loop countVar 0
394 ref <- liftIO $ newIORef ()
395 _ <- liftIO $ mkWeakIORef ref (killThread tid)
396 return $ Skip (Just (countVar, tid, state, ref))
398 step' gst (Just (countVar, tid, st, ref)) = do
399 r <- step gst st
400 case r of
401 Yield x s -> do
402 liftIO $ Prim.modifyIORef' countVar (+ 1)
403 return $ Yield x (Just (countVar, tid, s, ref))
404 Skip s -> return $ Skip (Just (countVar, tid, s, ref))
405 Stop -> do
406 liftIO $ killThread tid
407 return Stop
409------------------------------------------------------------------------------
410-- Scanning with a Fold
411------------------------------------------------------------------------------
413data ScanState s f = ScanInit s | ScanDo s !f | ScanDone
415{-# INLINE_NORMAL postscanOnce #-}
416postscanOnce :: Monad m => FL.Fold m a b -> Stream m a -> Stream m b
417postscanOnce (FL.Fold fstep initial extract) (Stream sstep state) =
418 Stream step (ScanInit state)
420 where
422 {-# INLINE_LATE step #-}
423 step _ (ScanInit st) = do
424 res <- initial
425 return
426 $ case res of
427 FL.Partial fs -> Skip $ ScanDo st fs
428 FL.Done b -> Yield b ScanDone
429 step gst (ScanDo st fs) = do
430 res <- sstep (adaptState gst) st
431 case res of
432 Yield x s -> do
433 r <- fstep fs x
434 case r of
435 FL.Partial fs1 -> do
436 !b <- extract fs1
437 return $ Yield b $ ScanDo s fs1
438 FL.Done b -> return $ Yield b ScanDone
439 Skip s -> return $ Skip $ ScanDo s fs
440 Stop -> return Stop
441 step _ ScanDone = return Stop
443{-# INLINE scanOnce #-}
444scanOnce :: Monad m
445 => FL.Fold m a b -> Stream m a -> Stream m b
446scanOnce (FL.Fold fstep initial extract) (Stream sstep state) =
447 Stream step (ScanInit state)
449 where
451 {-# INLINE_LATE step #-}
452 step _ (ScanInit st) = do
453 res <- initial
454 case res of
455 FL.Partial fs -> do
456 !b <- extract fs
457 return $ Yield b $ ScanDo st fs
458 FL.Done b -> return $ Yield b ScanDone
459 step gst (ScanDo st fs) = do
460 res <- sstep (adaptState gst) st
461 case res of
462 Yield x s -> do
463 r <- fstep fs x
464 case r of
465 FL.Partial fs1 -> do
466 !b <- extract fs1
467 return $ Yield b $ ScanDo s fs1
468 FL.Done b -> return $ Yield b ScanDone
469 Skip s -> return $ Skip $ ScanDo s fs
470 Stop -> return Stop
471 step _ ScanDone = return Stop
473------------------------------------------------------------------------------
474-- Scanning - Prescans
475------------------------------------------------------------------------------
477-- Adapted from the vector package.
478--
479-- XXX Is a prescan useful, discarding the last step does not sound useful? I
480-- am not sure about the utility of this function, so this is implemented but
481-- not exposed. We can expose it if someone provides good reasons why this is
482-- useful.
483--
484-- XXX We have to execute the stream one step ahead to know that we are at the
485-- last step. The vector implementation of prescan executes the last fold step
486-- but does not yield the result. This means we have executed the effect but
487-- discarded value. This does not sound right. In this implementation we are
488-- not executing the last fold step.
489{-# INLINE_NORMAL prescanlM' #-}
490prescanlM' :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
491prescanlM' f mz (Stream step state) = Stream step' (state, mz)
492 where
493 {-# INLINE_LATE step' #-}
494 step' gst (st, prev) = do
495 r <- step (adaptState gst) st
496 case r of
497 Yield x s -> do
498 acc <- prev
499 return $ Yield acc (s, f acc x)
500 Skip s -> return $ Skip (s, prev)
501 Stop -> return Stop
503{-# INLINE prescanl' #-}
504prescanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
505prescanl' f z = prescanlM' (\a b -> return (f a b)) (return z)
507------------------------------------------------------------------------------
508-- Monolithic postscans (postscan followed by a map)
509------------------------------------------------------------------------------
511-- The performance of a modular postscan followed by a map seems to be
512-- equivalent to this monolithic scan followed by map therefore we may not need
513-- this implementation. We just have it for performance comparison and in case
514-- modular version does not perform well in some situation.
515--
516{-# INLINE_NORMAL postscanlMx' #-}
517postscanlMx' :: Monad m
518 => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> Stream m b
519postscanlMx' fstep begin done (Stream step state) = do
520 Stream step' (state, begin)
521 where
522 {-# INLINE_LATE step' #-}
523 step' gst (st, acc) = do
524 r <- step (adaptState gst) st
525 case r of
526 Yield x s -> do
527 old <- acc
528 y <- fstep old x
529 v <- done y
530 v `seq` y `seq` return (Yield v (s, return y))
531 Skip s -> return $ Skip (s, acc)
532 Stop -> return Stop
534{-# INLINE_NORMAL postscanlx' #-}
535postscanlx' :: Monad m
536 => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> Stream m b
537postscanlx' fstep begin done s =
538 postscanlMx' (\b a -> return (fstep b a)) (return begin) (return . done) s
540-- XXX do we need consM strict to evaluate the begin value?
541{-# INLINE scanlMx' #-}
542scanlMx' :: Monad m
543 => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> Stream m b
544scanlMx' fstep begin done s =
545 (begin >>= \x -> x `seq` done x) `consM` postscanlMx' fstep begin done s
547{-# INLINE scanlx' #-}
548scanlx' :: Monad m
549 => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> Stream m b
550scanlx' fstep begin done s =
551 scanlMx' (\b a -> return (fstep b a)) (return begin) (return . done) s
553------------------------------------------------------------------------------
554-- postscans
555------------------------------------------------------------------------------
557-- Adapted from the vector package.
558{-# INLINE_NORMAL postscanlM' #-}
559postscanlM' :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
560postscanlM' fstep begin (Stream step state) =
561 Stream step' Nothing
562 where
563 {-# INLINE_LATE step' #-}
564 step' _ Nothing = do
565 !x <- begin
566 return $ Skip (Just (state, x))
568 step' gst (Just (st, acc)) = do
569 r <- step (adaptState gst) st
570 case r of
571 Yield x s -> do
572 !y <- fstep acc x
573 return $ Yield y (Just (s, y))
574 Skip s -> return $ Skip (Just (s, acc))
575 Stop -> return Stop
577{-# INLINE_NORMAL postscanl' #-}
578postscanl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
579postscanl' f seed = postscanlM' (\a b -> return (f a b)) (return seed)
581-- We can possibly have the "done" function as a Maybe to provide an option to
582-- emit or not emit the accumulator when the stream stops.
583--
584-- TBD: use a single Yield point
585--
586{-# INLINE_NORMAL postscanlMAfter' #-}
587postscanlMAfter' :: Monad m
588 => (b -> a -> m b) -> m b -> (b -> m b) -> Stream m a -> Stream m b
589postscanlMAfter' fstep initial done (Stream step1 state1) = do
590 Stream step (Just (state1, initial))
592 where
594 {-# INLINE_LATE step #-}
595 step gst (Just (st, acc)) = do
596 r <- step1 (adaptState gst) st
597 case r of
598 Yield x s -> do
599 old <- acc
600 y <- fstep old x
601 y `seq` return (Yield y (Just (s, return y)))
602 Skip s -> return $ Skip $ Just (s, acc)
603 Stop -> acc >>= done >>= \res -> return (Yield res Nothing)
604 step _ Nothing = return Stop
606{-# INLINE_NORMAL postscanlM #-}
607postscanlM :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
608postscanlM fstep begin (Stream step state) = Stream step' Nothing
609 where
610 {-# INLINE_LATE step' #-}
611 step' _ Nothing = do
612 r <- begin
613 return $ Skip (Just (state, r))
615 step' gst (Just (st, acc)) = do
616 r <- step (adaptState gst) st
617 case r of
618 Yield x s -> do
619 y <- fstep acc x
620 return (Yield y (Just (s, y)))
621 Skip s -> return $ Skip (Just (s, acc))
622 Stop -> return Stop
624{-# INLINE_NORMAL postscanl #-}
625postscanl :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
626postscanl f seed = postscanlM (\a b -> return (f a b)) (return seed)
628{-# INLINE_NORMAL scanlM' #-}
629scanlM' :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
630scanlM' fstep begin (Stream step state) = Stream step' Nothing
631 where
632 {-# INLINE_LATE step' #-}
633 step' _ Nothing = do
634 !x <- begin
635 return $ Yield x (Just (state, x))
636 step' gst (Just (st, acc)) = do
637 r <- step (adaptState gst) st
638 case r of
639 Yield x s -> do
640 !y <- fstep acc x
641 return $ Yield y (Just (s, y))
642 Skip s -> return $ Skip (Just (s, acc))
643 Stop -> return Stop
645{-# INLINE scanlMAfter' #-}
646scanlMAfter' :: Monad m
647 => (b -> a -> m b) -> m b -> (b -> m b) -> Stream m a -> Stream m b
648scanlMAfter' fstep initial done s =
649 (initial >>= \x -> x `seq` return x) `consM`
650 postscanlMAfter' fstep initial done s
652{-# INLINE scanl' #-}
653scanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
654scanl' f seed = scanlM' (\a b -> return (f a b)) (return seed)
656{-# INLINE_NORMAL scanlM #-}
657scanlM :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
658scanlM fstep begin (Stream step state) = Stream step' Nothing
659 where
660 {-# INLINE_LATE step' #-}
661 step' _ Nothing = do
662 x <- begin
663 return $ Yield x (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 scanl #-}
674scanl :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
675scanl f seed = scanlM (\a b -> return (f a b)) (return seed)
677-- Adapted from the vector package
678{-# INLINE_NORMAL scanl1M #-}
679scanl1M :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
680scanl1M fstep (Stream step state) = Stream step' (state, Nothing)
681 where
682 {-# INLINE_LATE step' #-}
683 step' gst (st, Nothing) = do
684 r <- step gst st
685 case r of
686 Yield x s -> return $ Yield x (s, Just x)
687 Skip s -> return $ Skip (s, Nothing)
688 Stop -> return Stop
690 step' gst (st, Just acc) = do
691 r <- step gst st
692 case r of
693 Yield y s -> do
694 z <- fstep acc y
695 return $ Yield z (s, Just z)
696 Skip s -> return $ Skip (s, Just acc)
697 Stop -> return Stop
699{-# INLINE scanl1 #-}
700scanl1 :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
701scanl1 f = scanl1M (\x y -> return (f x y))
703-- Adapted from the vector package
704{-# INLINE_NORMAL scanl1M' #-}
705scanl1M' :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
706scanl1M' fstep (Stream step state) = Stream step' (state, Nothing)
707 where
708 {-# INLINE_LATE step' #-}
709 step' gst (st, Nothing) = do
710 r <- step gst st
711 case r of
712 Yield x s -> x `seq` return $ Yield x (s, Just x)
713 Skip s -> return $ Skip (s, Nothing)
714 Stop -> return Stop
716 step' gst (st, Just acc) = acc `seq` do
717 r <- step gst st
718 case r of
719 Yield y s -> do
720 z <- fstep acc y
721 z `seq` return $ Yield z (s, Just z)
722 Skip s -> return $ Skip (s, Just acc)
723 Stop -> return Stop
725{-# INLINE scanl1' #-}
726scanl1' :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
727scanl1' f = scanl1M' (\x y -> return (f x y))
729-------------------------------------------------------------------------------
730-- Filtering
731-------------------------------------------------------------------------------
733-- Adapted from the vector package
734{-# INLINE_NORMAL filterM #-}
735filterM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
736filterM f (Stream step state) = Stream step' state
737 where
738 {-# INLINE_LATE step' #-}
739 step' gst st = do
740 r <- step gst st
741 case r of
742 Yield x s -> do
743 b <- f x
744 return $ if b
745 then Yield x s
746 else Skip s
747 Skip s -> return $ Skip s
748 Stop -> return Stop
750{-# INLINE filter #-}
751filter :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
752filter f = filterM (return . f)
754-- Adapted from the vector package
755{-# INLINE_NORMAL uniq #-}
756uniq :: (Eq a, Monad m) => Stream m a -> Stream m a
757uniq (Stream step state) = Stream step' (Nothing, state)
758 where
759 {-# INLINE_LATE step' #-}
760 step' gst (Nothing, st) = do
761 r <- step gst st
762 case r of
763 Yield x s -> return $ Yield x (Just x, s)
764 Skip s -> return $ Skip (Nothing, s)
765 Stop -> return Stop
766 step' gst (Just x, st) = do
767 r <- step gst st
768 case r of
769 Yield y s | x == y -> return $ Skip (Just x, s)
770 | otherwise -> return $ Yield y (Just y, s)
771 Skip s -> return $ Skip (Just x, s)
772 Stop -> return Stop
774{-# INLINE_NORMAL deleteBy #-}
775deleteBy :: Monad m => (a -> a -> Bool) -> a -> Stream m a -> Stream m a
776deleteBy eq x (Stream step state) = Stream step' (state, False)
777 where
778 {-# INLINE_LATE step' #-}
779 step' gst (st, False) = do
780 r <- step gst st
781 case r of
782 Yield y s -> return $
783 if eq x y then Skip (s, True) else Yield y (s, False)
784 Skip s -> return $ Skip (s, False)
785 Stop -> return Stop
787 step' gst (st, True) = do
788 r <- step gst st
789 case r of
790 Yield y s -> return $ Yield y (s, True)
791 Skip s -> return $ Skip (s, True)
792 Stop -> return Stop
794------------------------------------------------------------------------------
795-- Trimming
796------------------------------------------------------------------------------
798-- XXX using getTime in the loop can be pretty expensive especially for
799-- computations where iterations are lightweight. We have the following
800-- options:
801--
802-- 1) Run a timeout thread updating a flag asynchronously and check that
803-- flag here, that way we can have a cheap termination check.
804--
805-- 2) Use COARSE clock to get time with lower resolution but more efficiently.
806--
807-- 3) Use rdtscp/rdtsc to get time directly from the processor, compute the
808-- termination value of rdtsc in the beginning and then in each iteration just
809-- get rdtsc and check if we should terminate.
810--
811data TakeByTime st s
812 = TakeByTimeInit st
813 | TakeByTimeCheck st s
814 | TakeByTimeYield st s
816{-# INLINE_NORMAL takeByTime #-}
817takeByTime :: (MonadIO m, TimeUnit64 t) => t -> Stream m a -> Stream m a
818takeByTime duration (Stream step1 state1) = Stream step (TakeByTimeInit state1)
819 where
821 lim = toRelTime64 duration
823 {-# INLINE_LATE step #-}
824 step _ (TakeByTimeInit _) | lim == 0 = return Stop
825 step _ (TakeByTimeInit st) = do
826 t0 <- liftIO $ getTime Monotonic
827 return $ Skip (TakeByTimeYield st t0)
828 step _ (TakeByTimeCheck st t0) = do
829 t <- liftIO $ getTime Monotonic
830 return $
831 if diffAbsTime64 t t0 > lim
832 then Stop
833 else Skip (TakeByTimeYield st t0)
834 step gst (TakeByTimeYield st t0) = do
835 r <- step1 gst st
836 return $ case r of
837 Yield x s -> Yield x (TakeByTimeCheck s t0)
838 Skip s -> Skip (TakeByTimeCheck s t0)
839 Stop -> Stop
841data DropByTime st s x
842 = DropByTimeInit st
843 | DropByTimeGen st s
844 | DropByTimeCheck st s x
845 | DropByTimeYield st
847{-# INLINE_NORMAL dropByTime #-}
848dropByTime :: (MonadIO m, TimeUnit64 t) => t -> Stream m a -> Stream m a
849dropByTime duration (Stream step1 state1) = Stream step (DropByTimeInit state1)
850 where
852 lim = toRelTime64 duration
854 {-# INLINE_LATE step #-}
855 step _ (DropByTimeInit st) = do
856 t0 <- liftIO $ getTime Monotonic
857 return $ Skip (DropByTimeGen st t0)
858 step gst (DropByTimeGen st t0) = do
859 r <- step1 gst st
860 return $ case r of
861 Yield x s -> Skip (DropByTimeCheck s t0 x)
862 Skip s -> Skip (DropByTimeGen s t0)
863 Stop -> Stop
864 step _ (DropByTimeCheck st t0 x) = do
865 t <- liftIO $ getTime Monotonic
866 if diffAbsTime64 t t0 <= lim
867 then return $ Skip $ DropByTimeGen st t0
868 else return $ Yield x $ DropByTimeYield st
869 step gst (DropByTimeYield st) = do
870 r <- step1 gst st
871 return $ case r of
872 Yield x s -> Yield x (DropByTimeYield s)
873 Skip s -> Skip (DropByTimeYield s)
874 Stop -> Stop
876-- Adapted from the vector package
877{-# INLINE_NORMAL drop #-}
878drop :: Monad m => Int -> Stream m a -> Stream m a
879drop n (Stream step state) = Stream step' (state, Just n)
880 where
881 {-# INLINE_LATE step' #-}
882 step' gst (st, Just i)
883 | i > 0 = do
884 r <- step gst st
885 return $
886 case r of
887 Yield _ s -> Skip (s, Just (i - 1))
888 Skip s -> Skip (s, Just i)
889 Stop -> Stop
890 | otherwise = return $ Skip (st, Nothing)
892 step' gst (st, Nothing) = do
893 r <- step gst st
894 return $
895 case r of
896 Yield x s -> Yield x (s, Nothing)
897 Skip s -> Skip (s, Nothing)
898 Stop -> Stop
900-- Adapted from the vector package
901data DropWhileState s a
902 = DropWhileDrop s
903 | DropWhileYield a s
904 | DropWhileNext s
906{-# INLINE_NORMAL dropWhileM #-}
907dropWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
908dropWhileM f (Stream step state) = Stream step' (DropWhileDrop state)
909 where
910 {-# INLINE_LATE step' #-}
911 step' gst (DropWhileDrop st) = do
912 r <- step gst st
913 case r of
914 Yield x s -> do
915 b <- f x
916 if b
917 then return $ Skip (DropWhileDrop s)
918 else return $ Skip (DropWhileYield x s)
919 Skip s -> return $ Skip (DropWhileDrop s)
920 Stop -> return Stop
922 step' gst (DropWhileNext st) = do
923 r <- step gst st
924 case r of
925 Yield x s -> return $ Skip (DropWhileYield x s)
926 Skip s -> return $ Skip (DropWhileNext s)
927 Stop -> return Stop
929 step' _ (DropWhileYield x st) = return $ Yield x (DropWhileNext st)
931{-# INLINE dropWhile #-}
932dropWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
933dropWhile f = dropWhileM (return . f)
935------------------------------------------------------------------------------
936-- Inserting Elements
937------------------------------------------------------------------------------
939{-# INLINE_NORMAL insertBy #-}
940insertBy :: Monad m => (a -> a -> Ordering) -> a -> Stream m a -> Stream m a
941insertBy cmp a (Stream step state) = Stream step' (state, False, Nothing)
942 where
943 {-# INLINE_LATE step' #-}
944 step' gst (st, False, _) = do
945 r <- step gst st
946 case r of
947 Yield x s -> case cmp a x of
948 GT -> return $ Yield x (s, False, Nothing)
949 _ -> return $ Yield a (s, True, Just x)
950 Skip s -> return $ Skip (s, False, Nothing)
951 Stop -> return $ Yield a (st, True, Nothing)
953 step' _ (_, True, Nothing) = return Stop
955 step' gst (st, True, Just prev) = do
956 r <- step gst st
957 case r of
958 Yield x s -> return $ Yield prev (s, True, Just x)
959 Skip s -> return $ Skip (s, True, Just prev)
960 Stop -> return $ Yield prev (st, True, Nothing)
962data LoopState x s = FirstYield s
963 | InterspersingYield s
964 | YieldAndCarry x s
966{-# INLINE_NORMAL intersperseM #-}
967intersperseM :: Monad m => m a -> Stream m a -> Stream m a
968intersperseM m (Stream step state) = Stream step' (FirstYield state)
969 where
970 {-# INLINE_LATE step' #-}
971 step' gst (FirstYield st) = do
972 r <- step gst st
973 return $
974 case r of
975 Yield x s -> Skip (YieldAndCarry x s)
976 Skip s -> Skip (FirstYield s)
977 Stop -> Stop
979 step' gst (InterspersingYield st) = do
980 r <- step gst st
981 case r of
982 Yield x s -> do
983 a <- m
984 return $ Yield a (YieldAndCarry x s)
985 Skip s -> return $ Skip $ InterspersingYield s
986 Stop -> return Stop
988 step' _ (YieldAndCarry x st) = return $ Yield x (InterspersingYield st)
990{-# INLINE intersperse #-}
991intersperse :: Monad m => a -> Stream m a -> Stream m a
992intersperse a = intersperseM (return a)
994{-# INLINE_NORMAL intersperseM_ #-}
995intersperseM_ :: Monad m => m b -> Stream m a -> Stream m a
996intersperseM_ m (Stream step1 state1) = Stream step (Left (pure (), state1))
997 where
998 {-# INLINE_LATE step #-}
999 step gst (Left (eff, st)) = do
1000 r <- step1 gst st
1001 case r of
1002 Yield x s -> eff >> return (Yield x (Right s))
1003 Skip s -> return $ Skip (Left (eff, s))
1004 Stop -> return Stop
1006 step _ (Right st) = return $ Skip $ Left (void m, st)
1008data SuffixState s a
1009 = SuffixElem s
1010 | SuffixSuffix s
1011 | SuffixYield a (SuffixState s a)
1013{-# INLINE_NORMAL intersperseSuffix #-}
1014intersperseSuffix :: forall m a. Monad m => m a -> Stream m a -> Stream m a
1015intersperseSuffix action (Stream step state) = Stream step' (SuffixElem state)
1016 where
1017 {-# INLINE_LATE step' #-}
1018 step' gst (SuffixElem st) = do
1019 r <- step gst st
1020 return $ case r of
1021 Yield x s -> Skip (SuffixYield x (SuffixSuffix s))
1022 Skip s -> Skip (SuffixElem s)
1023 Stop -> Stop
1025 step' _ (SuffixSuffix st) = do
1026 action >>= \r -> return $ Skip (SuffixYield r (SuffixElem st))
1028 step' _ (SuffixYield x next) = return $ Yield x next
1030{-# INLINE_NORMAL intersperseSuffix_ #-}
1031intersperseSuffix_ :: Monad m => m b -> Stream m a -> Stream m a
1032intersperseSuffix_ m (Stream step1 state1) = Stream step (Left state1)
1033 where
1034 {-# INLINE_LATE step #-}
1035 step gst (Left st) = do
1036 r <- step1 gst st
1037 case r of
1038 Yield x s -> return $ Yield x (Right s)
1039 Skip s -> return $ Skip $ Left s
1040 Stop -> return Stop
1042 step _ (Right st) = m >> return (Skip (Left st))
1044data SuffixSpanState s a
1045 = SuffixSpanElem s Int
1046 | SuffixSpanSuffix s
1047 | SuffixSpanYield a (SuffixSpanState s a)
1048 | SuffixSpanLast
1049 | SuffixSpanStop
1051-- | intersperse after every n items
1052{-# INLINE_NORMAL intersperseSuffixBySpan #-}
1053intersperseSuffixBySpan :: forall m a. Monad m
1054 => Int -> m a -> Stream m a -> Stream m a
1055intersperseSuffixBySpan n action (Stream step state) =
1056 Stream step' (SuffixSpanElem state n)
1057 where
1058 {-# INLINE_LATE step' #-}
1059 step' gst (SuffixSpanElem st i) | i > 0 = do
1060 r <- step gst st
1061 return $ case r of
1062 Yield x s -> Skip (SuffixSpanYield x (SuffixSpanElem s (i - 1)))
1063 Skip s -> Skip (SuffixSpanElem s i)
1064 Stop -> if i == n then Stop else Skip SuffixSpanLast
1065 step' _ (SuffixSpanElem st _) = return $ Skip (SuffixSpanSuffix st)
1067 step' _ (SuffixSpanSuffix st) = do
1068 action >>= \r -> return $ Skip (SuffixSpanYield r (SuffixSpanElem st n))
1070 step' _ (SuffixSpanLast) = do
1071 action >>= \r -> return $ Skip (SuffixSpanYield r SuffixSpanStop)
1073 step' _ (SuffixSpanYield x next) = return $ Yield x next
1075 step' _ (SuffixSpanStop) = return Stop
1077------------------------------------------------------------------------------
1078-- Reordering
1079------------------------------------------------------------------------------
1081-- We can implement reverse as:
1083-- > reverse = foldlS (flip cons) nil
1085-- However, this implementation is unusable because of the horrible performance
1086-- of cons. So we just convert it to a list first and then stream from the
1087-- list.
1089-- XXX Maybe we can use an Array instead of a list here?
1090{-# INLINE_NORMAL reverse #-}
1091reverse :: Monad m => Stream m a -> Stream m a
1092reverse m = Stream step Nothing
1093 where
1094 {-# INLINE_LATE step #-}
1095 step _ Nothing = do
1096 xs <- foldl' (flip (:)) [] m
1097 return $ Skip (Just xs)
1098 step _ (Just (x:xs)) = return $ Yield x (Just xs)
1099 step _ (Just []) = return Stop
1101-- Much faster reverse for Storables
1102{-# INLINE_NORMAL reverse' #-}
1103reverse' :: forall m a. (MonadIO m, Storable a) => Stream m a -> Stream m a
1105-- This commented implementation copies the whole stream into one single array
1106-- and then streams from that array, this has exactly the same performance as
1107-- the chunked code that follows. Though this could be problematic due to
1108-- unbounded large allocations. However, if we use an idiomatic implementation
1109-- of arraysOf instead of the custom implementation then the chunked code
1110-- becomes worse by 6 times. Need to investigate if that can be improved.
1111import Foreign.ForeignPtr (touchForeignPtr)
1112import Foreign.ForeignPtr.Unsafe (unsafeForeignPtrToPtr)
1113import Foreign.Ptr (Ptr, plusPtr)
1114reverse' m = Stream step Nothing
1115 where
1116 {-# INLINE_LATE step #-}
1117 step _ Nothing = do
1118 arr <- A.fromStreamD m
1119 let p = A.aEnd arr `plusPtr` negate (sizeOf (undefined :: a))
1120 return $ Skip $ Just (A.aStart arr, p)
1122 step _ (Just (start, p)) | p < unsafeForeignPtrToPtr start = return Stop
1124 step _ (Just (start, p)) = do
1125 let !x = A.unsafeInlineIO $ do
1126 r <- peek p
1127 touchForeignPtr start
1128 return r
1129 next = p `plusPtr` negate (sizeOf (undefined :: a))
1130 return $ Yield x (Just (start, next))
1132reverse' =
1133 A.flattenArraysRev -- unfoldMany A.readRev
1134 . fromStreamK
1135 . K.reverse
1136 . toStreamK
1137 . A.arraysOf defaultChunkSize
1139------------------------------------------------------------------------------
1140-- Position Indexing
1141------------------------------------------------------------------------------
1143-- Adapted from the vector package
1144{-# INLINE_NORMAL indexed #-}
1145indexed :: Monad m => Stream m a -> Stream m (Int, a)
1146indexed (Stream step state) = Stream step' (state, 0)
1147 where
1148 {-# INLINE_LATE step' #-}
1149 step' gst (st, i) = i `seq` do
1150 r <- step (adaptState gst) st
1151 case r of
1152 Yield x s -> return $ Yield (i, x) (s, i+1)
1153 Skip s -> return $ Skip (s, i)
1154 Stop -> return Stop
1156-- Adapted from the vector package
1157{-# INLINE_NORMAL indexedR #-}
1158indexedR :: Monad m => Int -> Stream m a -> Stream m (Int, a)
1159indexedR m (Stream step state) = Stream step' (state, m)
1160 where
1161 {-# INLINE_LATE step' #-}
1162 step' gst (st, i) = i `seq` do
1163 r <- step (adaptState gst) st
1164 case r of
1165 Yield x s -> let i' = i - 1
1166 in return $ Yield (i, x) (s, i')
1167 Skip s -> return $ Skip (s, i)
1168 Stop -> return Stop
1170------------------------------------------------------------------------------
1171-- Searching
1172------------------------------------------------------------------------------
1174{-# INLINE_NORMAL findIndices #-}
1175findIndices :: Monad m => (a -> Bool) -> Stream m a -> Stream m Int
1176findIndices p (Stream step state) = Stream step' (state, 0)
1177 where
1178 {-# INLINE_LATE step' #-}
1179 step' gst (st, i) = i `seq` do
1180 r <- step (adaptState gst) st
1181 return $ case r of
1182 Yield x s -> if p x then Yield i (s, i+1) else Skip (s, i+1)
1183 Skip s -> Skip (s, i)
1184 Stop -> Stop
1186------------------------------------------------------------------------------
1187-- Rolling map
1188------------------------------------------------------------------------------
1190data RollingMapState s a = RollingMapInit s | RollingMapGo s a
1192{-# INLINE rollingMapM #-}
1193rollingMapM :: Monad m => (a -> a -> m b) -> Stream m a -> Stream m b
1194rollingMapM f (Stream step1 state1) = Stream step (RollingMapInit state1)
1195 where
1196 step gst (RollingMapInit st) = do
1197 r <- step1 (adaptState gst) st
1198 return $ case r of
1199 Yield x s -> Skip $ RollingMapGo s x
1200 Skip s -> Skip $ RollingMapInit s
1201 Stop -> Stop
1203 step gst (RollingMapGo s1 x1) = do
1204 r <- step1 (adaptState gst) s1
1205 case r of
1206 Yield x s -> do
1207 !res <- f x x1
1208 return $ Yield res $ RollingMapGo s x
1209 Skip s -> return $ Skip $ RollingMapGo s x1
1210 Stop -> return $ Stop
1212{-# INLINE rollingMap #-}
1213rollingMap :: Monad m => (a -> a -> b) -> Stream m a -> Stream m b
1214rollingMap f = rollingMapM (\x y -> return $ f x y)
1216------------------------------------------------------------------------------
1217-- Maybe Streams
1218------------------------------------------------------------------------------
1220-- XXX Will this always fuse properly?
1221{-# INLINE_NORMAL mapMaybe #-}
1222mapMaybe :: Monad m => (a -> Maybe b) -> Stream m a -> Stream m b
1223mapMaybe f = fmap fromJust . filter isJust . map f
1225{-# INLINE_NORMAL mapMaybeM #-}
1226mapMaybeM :: Monad m => (a -> m (Maybe b)) -> Stream m a -> Stream m b
1227mapMaybeM f = fmap fromJust . filter isJust . mapM f