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