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