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