@@ -301,25 +301,6 @@ module Streamly.Internal.Data.Stream.StreamD
301301 )
302302where
303303
304−import Control.Concurrent (myThreadId)
305−import Control.Monad (void)
306−import Control.Monad.IO.Class (MonadIO(..))
307−import Control.Monad.Trans.Reader (ReaderT)
308−import Control.Monad.Trans.State.Strict (StateT)
309−import Control.Monad.Trans.Class (MonadTrans(lift))
310−import Data.Functor.Identity (Identity(..))
311−import Data.Maybe (fromJust, isJust)
312−import Streamly.Internal.Data.Time.Units
313− (TimeUnit64, toRelTime64, diffAbsTime64)
314−import Streamly.Internal.Data.Pipe.Types (Pipe(..), PipeState(..))
315−import Streamly.Internal.Data.Time.Clock (Clock(Monotonic), getTime)
316−
317−import qualified Streamly.Internal.Data.Pipe.Types as Pipe
318−import qualified Streamly.Internal.Data.Fold as FL
319−import qualified Streamly.Internal.Data.Stream.StreamK as K
320−import qualified Control.Monad.Trans.Reader as Reader
321−import qualified Control.Monad.Trans.State.Strict as State
322−
323304import Prelude hiding
324305 ( map, mapM, mapM_, repeat, foldr, last, take, filter
325306 , takeWhile, drop, dropWhile, all, any, maximum, minimum, elem
@@ -327,1079 +308,10 @@ import Prelude hiding
327308 , (!!), scanl, scanl1, concatMap, replicate, enumFromTo, concat
328309 , reverse, iterate, splitAt)
329310import Streamly.Internal.Data.Stream.StreamD.Type
330−import Streamly.Internal.Data.Stream.StreamD.Common (takeWhileM, takeWhile)
311+import Streamly.Internal.Data.Stream.StreamD.Common (uncons)
331312import Streamly.Internal.Data.Stream.StreamD.Generate
332313import Streamly.Internal.Data.Stream.StreamD.Eliminate
333314import Streamly.Internal.Data.Stream.StreamD.Exceptions
334315import Streamly.Internal.Data.Stream.StreamD.Flatten
335316import Streamly.Internal.Data.Stream.StreamD.SplitGroup
336−import Streamly.Internal.Data.SVar
337−
338−-------------------------------------------------------------------------------
339−-- Hoisting the inner monad
340−-------------------------------------------------------------------------------
341−
342−{-# INLINE_NORMAL hoist #-}
343−hoist :: Monad n => (forall x. m x -> n x) -> Stream m a -> Stream n a
344−hoist f (Stream step state) = (Stream step' state)
345− where
346− {-# INLINE_LATE step' #-}
347− step' gst st = do
348− r <- f $ step (adaptState gst) st
349− return $ case r of
350− Yield x s -> Yield x s
351− Skip s -> Skip s
352− Stop -> Stop
353−
354−{-# INLINE generally #-}
355−generally :: Monad m => Stream Identity a -> Stream m a
356−generally = hoist (return . runIdentity)
357−
358−{-# INLINE_NORMAL liftInner #-}
359−liftInner :: (Monad m, MonadTrans t, Monad (t m))
360− => Stream m a -> Stream (t m) a
361−liftInner (Stream step state) = Stream step' state
362− where
363− {-# INLINE_LATE step' #-}
364− step' gst st = do
365− r <- lift $ step (adaptState gst) st
366− return $ case r of
367− Yield x s -> Yield x s
368− Skip s -> Skip s
369− Stop -> Stop
370−
371−{-# INLINE_NORMAL runReaderT #-}
372−runReaderT :: Monad m => m s -> Stream (ReaderT s m) a -> Stream m a
373−runReaderT env (Stream step state) = Stream step' (state, env)
374− where
375− {-# INLINE_LATE step' #-}
376− step' gst (st, action) = do
377− sv <- action
378− r <- Reader.runReaderT (step (adaptState gst) st) sv
379− return $ case r of
380− Yield x s -> Yield x (s, return sv)
381− Skip s -> Skip (s, return sv)
382− Stop -> Stop
383−
384−{-# INLINE_NORMAL evalStateT #-}
385−evalStateT :: Monad m => m s -> Stream (StateT s m) a -> Stream m a
386−evalStateT initial (Stream step state) = Stream step' (state, initial)
387− where
388− {-# INLINE_LATE step' #-}
389− step' gst (st, action) = do
390− sv <- action
391− (r, sv') <- State.runStateT (step (adaptState gst) st) sv
392− return $ case r of
393− Yield x s -> Yield x (s, return sv')
394− Skip s -> Skip (s, return sv')
395− Stop -> Stop
396−
397−{-# INLINE_NORMAL runStateT #-}
398−runStateT :: Monad m => m s -> Stream (StateT s m) a -> Stream m (s, a)
399−runStateT initial (Stream step state) = Stream step' (state, initial)
400− where
401− {-# INLINE_LATE step' #-}
402− step' gst (st, action) = do
403− sv <- action
404− (r, sv') <- State.runStateT (step (adaptState gst) st) sv
405− return $ case r of
406− Yield x s -> Yield (sv', x) (s, return sv')
407− Skip s -> Skip (s, return sv')
408− Stop -> Stop
409−
410−------------------------------------------------------------------------------
411−-- Substreams
412−------------------------------------------------------------------------------
413−
414−{-# INLINE_NORMAL isPrefixOf #-}
415−isPrefixOf :: (Eq a, Monad m) => Stream m a -> Stream m a -> m Bool
416−isPrefixOf (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
417− where
418− go (sa, sb, Nothing) = do
419− r <- stepa defState sa
420− case r of
421− Yield x sa' -> go (sa', sb, Just x)
422− Skip sa' -> go (sa', sb, Nothing)
423− Stop -> return True
424−
425− go (sa, sb, Just x) = do
426− r <- stepb defState sb
427− case r of
428− Yield y sb' ->
429− if x == y
430− then go (sa, sb', Nothing)
431− else return False
432− Skip sb' -> go (sa, sb', Just x)
433− Stop -> return False
434−
435−{-# INLINE_NORMAL isSubsequenceOf #-}
436−isSubsequenceOf :: (Eq a, Monad m) => Stream m a -> Stream m a -> m Bool
437−isSubsequenceOf (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
438− where
439− go (sa, sb, Nothing) = do
440− r <- stepa defState sa
441− case r of
442− Yield x sa' -> go (sa', sb, Just x)
443− Skip sa' -> go (sa', sb, Nothing)
444− Stop -> return True
445−
446− go (sa, sb, Just x) = do
447− r <- stepb defState sb
448− case r of
449− Yield y sb' ->
450− if x == y
451− then go (sa, sb', Nothing)
452− else go (sa, sb', Just x)
453− Skip sb' -> go (sa, sb', Just x)
454− Stop -> return False
455−
456−{-# INLINE_NORMAL stripPrefix #-}
457−stripPrefix
458− :: (Eq a, Monad m)
459− => Stream m a -> Stream m a -> m (Maybe (Stream m a))
460−stripPrefix (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
461− where
462− go (sa, sb, Nothing) = do
463− r <- stepa defState sa
464− case r of
465− Yield x sa' -> go (sa', sb, Just x)
466− Skip sa' -> go (sa', sb, Nothing)
467− Stop -> return $ Just (Stream stepb sb)
468−
469− go (sa, sb, Just x) = do
470− r <- stepb defState sb
471− case r of
472− Yield y sb' ->
473− if x == y
474− then go (sa, sb', Nothing)
475− else return Nothing
476− Skip sb' -> go (sa, sb', Just x)
477− Stop -> return Nothing
478−
479−------------------------------------------------------------------------------
480−-- Map and Fold
481−------------------------------------------------------------------------------
482−
483−-- | Execute a monadic action for each element of the 'Stream'
484−{-# INLINE_NORMAL mapM_ #-}
485−mapM_ :: Monad m => (a -> m b) -> Stream m a -> m ()
486−mapM_ m = drain . mapM m
487−
488−-------------------------------------------------------------------------------
489−-- General transformation
490−-------------------------------------------------------------------------------
491−
492−{-# INLINE_NORMAL transform #-}
493−transform :: Monad m => Pipe m a b -> Stream m a -> Stream m b
494−transform (Pipe pstep1 pstep2 pstate) (Stream step state) =
495− Stream step' (Consume pstate, state)
496−
497− where
498−
499− {-# INLINE_LATE step' #-}
500−
501− step' gst (Consume pst, st) = pst `seq` do
502− r <- step (adaptState gst) st
503− case r of
504− Yield x s -> do
505− res <- pstep1 pst x
506− case res of
507− Pipe.Yield b pst' -> return $ Yield b (pst', s)
508− Pipe.Continue pst' -> return $ Skip (pst', s)
509− Skip s -> return $ Skip (Consume pst, s)
510− Stop -> return Stop
511−
512− step' _ (Produce pst, st) = pst `seq` do
513− res <- pstep2 pst
514− case res of
515− Pipe.Yield b pst' -> return $ Yield b (pst', st)
516− Pipe.Continue pst' -> return $ Skip (pst', st)
517−
518−------------------------------------------------------------------------------
519−-- Transformation by Folding (Scans)
520−------------------------------------------------------------------------------
521−
522−------------------------------------------------------------------------------
523−-- Prescans
524−------------------------------------------------------------------------------
525−
526−-- XXX Is a prescan useful, discarding the last step does not sound useful? I
527−-- am not sure about the utility of this function, so this is implemented but
528−-- not exposed. We can expose it if someone provides good reasons why this is
529−-- useful.
530−--
531−-- XXX We have to execute the stream one step ahead to know that we are at the
532−-- last step. The vector implementation of prescan executes the last fold step
533−-- but does not yield the result. This means we have executed the effect but
534−-- discarded value. This does not sound right. In this implementation we are
535−-- not executing the last fold step.
536−{-# INLINE_NORMAL prescanlM' #-}
537−prescanlM' :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
538−prescanlM' f mz (Stream step state) = Stream step' (state, mz)
539− where
540− {-# INLINE_LATE step' #-}
541− step' gst (st, prev) = do
542− r <- step (adaptState gst) st
543− case r of
544− Yield x s -> do
545− acc <- prev
546− return $ Yield acc (s, f acc x)
547− Skip s -> return $ Skip (s, prev)
548− Stop -> return Stop
549−
550−{-# INLINE prescanl' #-}
551−prescanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
552−prescanl' f z = prescanlM' (\a b -> return (f a b)) (return z)
553−
554−------------------------------------------------------------------------------
555−-- Monolithic postscans (postscan followed by a map)
556−------------------------------------------------------------------------------
557−
558−-- The performance of a modular postscan followed by a map seems to be
559−-- equivalent to this monolithic scan followed by map therefore we may not need
560−-- this implementation. We just have it for performance comparison and in case
561−-- modular version does not perform well in some situation.
562−--
563−{-# INLINE_NORMAL postscanlMx' #-}
564−postscanlMx' :: Monad m
565− => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> Stream m b
566−postscanlMx' fstep begin done (Stream step state) = do
567− Stream step' (state, begin)
568− where
569− {-# INLINE_LATE step' #-}
570− step' gst (st, acc) = do
571− r <- step (adaptState gst) st
572− case r of
573− Yield x s -> do
574− old <- acc
575− y <- fstep old x
576− v <- done y
577− v `seq` y `seq` return (Yield v (s, return y))
578− Skip s -> return $ Skip (s, acc)
579− Stop -> return Stop
580−
581−{-# INLINE_NORMAL postscanlx' #-}
582−postscanlx' :: Monad m
583− => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> Stream m b
584−postscanlx' fstep begin done s =
585− postscanlMx' (\b a -> return (fstep b a)) (return begin) (return . done) s
586−
587−-- XXX do we need consM strict to evaluate the begin value?
588−{-# INLINE scanlMx' #-}
589−scanlMx' :: Monad m
590− => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> Stream m b
591−scanlMx' fstep begin done s =
592− (begin >>= \x -> x `seq` done x) `consM` postscanlMx' fstep begin done s
593−
594−
595−data PostScanState s f = PostScan s !f
596−
597−{-# INLINE_NORMAL postscanOnce #-}
598−postscanOnce :: Monad m => FL.Fold m a b -> Stream m a -> Stream m b
599−postscanOnce (FL.Fold fstep begin done) (Stream step state) =
600− Stream step' (PostScan state begin)
601−
602− where
603−
604− {-# INLINE_LATE step' #-}
605− step' gst (PostScan st acc) = do
606− r <- step (adaptState gst) st
607− case r of
608− Yield x s -> do
609− old <- acc
610− res <- fstep old x
611− case res of
612− FL.Partial fs -> do
613− !v <- done fs
614− return $ Yield v $ PostScan s (return fs)
615− FL.Done _ -> return Stop
616− Skip s -> return $ Skip $ PostScan s acc
617− Stop -> return Stop
618−
619−{-# INLINE scanOnce #-}
620−scanOnce :: Monad m
621− => FL.Fold m a b -> Stream m a -> Stream m b
622−scanOnce fld@(FL.Fold _ begin done) s =
623− (begin >>= \x -> x `seq` done x) `consM` postscanOnce fld s
624−
625−{-# INLINE scanlx' #-}
626−scanlx' :: Monad m
627− => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> Stream m b
628−scanlx' fstep begin done s =
629− scanlMx' (\b a -> return (fstep b a)) (return begin) (return . done) s
630−
631−------------------------------------------------------------------------------
632−-- postscans
633−------------------------------------------------------------------------------
634−
635−{-# INLINE_NORMAL postscanlM' #-}
636−postscanlM' :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
637−postscanlM' fstep begin (Stream step state) =
638− Stream step' Nothing
639− where
640− {-# INLINE_LATE step' #-}
641− step' _ Nothing = do
642− !x <- begin
643− return $ Skip (Just (state, x))
644−
645− step' gst (Just (st, acc)) = do
646− r <- step (adaptState gst) st
647− case r of
648− Yield x s -> do
649− !y <- fstep acc x
650− return $ Yield y (Just (s, y))
651− Skip s -> return $ Skip (Just (s, acc))
652− Stop -> return Stop
653−
654−{-# INLINE_NORMAL postscanl' #-}
655−postscanl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
656−postscanl' f seed = postscanlM' (\a b -> return (f a b)) (return seed)
657−
658−-- We can possibly have the "done" function as a Maybe to provide an option to
659−-- emit or not emit the accumulator when the stream stops.
660−--
661−-- TBD: use a single Yield point
662−--
663−{-# INLINE_NORMAL postscanlMAfter' #-}
664−postscanlMAfter' :: Monad m
665− => (b -> a -> m b) -> m b -> (b -> m b) -> Stream m a -> Stream m b
666−postscanlMAfter' fstep initial done (Stream step1 state1) = do
667− Stream step (Just (state1, initial))
668−
669− where
670−
671− {-# INLINE_LATE step #-}
672− step gst (Just (st, acc)) = do
673− r <- step1 (adaptState gst) st
674− case r of
675− Yield x s -> do
676− old <- acc
677− y <- fstep old x
678− y `seq` return (Yield y (Just (s, return y)))
679− Skip s -> return $ Skip $ Just (s, acc)
680− Stop -> acc >>= done >>= \res -> return (Yield res Nothing)
681− step _ Nothing = return Stop
682−
683−{-# INLINE_NORMAL postscanlM #-}
684−postscanlM :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
685−postscanlM fstep begin (Stream step state) = Stream step' Nothing
686− where
687− {-# INLINE_LATE step' #-}
688− step' _ Nothing = do
689− r <- begin
690− return $ Skip (Just (state, r))
691−
692− step' gst (Just (st, acc)) = do
693− r <- step (adaptState gst) st
694− case r of
695− Yield x s -> do
696− y <- fstep acc x
697− return (Yield y (Just (s, y)))
698− Skip s -> return $ Skip (Just (s, acc))
699− Stop -> return Stop
700−
701−{-# INLINE_NORMAL postscanl #-}
702−postscanl :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
703−postscanl f seed = postscanlM (\a b -> return (f a b)) (return seed)
704−
705−{-# INLINE_NORMAL scanlM' #-}
706−scanlM' :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
707−scanlM' fstep begin (Stream step state) = Stream step' Nothing
708− where
709− {-# INLINE_LATE step' #-}
710− step' _ Nothing = do
711− !x <- begin
712− return $ Yield x (Just (state, x))
713− step' gst (Just (st, acc)) = do
714− r <- step (adaptState gst) st
715− case r of
716− Yield x s -> do
717− !y <- fstep acc x
718− return $ Yield y (Just (s, y))
719− Skip s -> return $ Skip (Just (s, acc))
720− Stop -> return Stop
721−
722−{-# INLINE scanlMAfter' #-}
723−scanlMAfter' :: Monad m
724− => (b -> a -> m b) -> m b -> (b -> m b) -> Stream m a -> Stream m b
725−scanlMAfter' fstep initial done s =
726− (initial >>= \x -> x `seq` return x) `consM`
727− postscanlMAfter' fstep initial done s
728−
729−{-# INLINE scanl' #-}
730−scanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
731−scanl' f seed = scanlM' (\a b -> return (f a b)) (return seed)
732−
733−{-# INLINE_NORMAL scanlM #-}
734−scanlM :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
735−scanlM fstep begin (Stream step state) = Stream step' Nothing
736− where
737− {-# INLINE_LATE step' #-}
738− step' _ Nothing = do
739− x <- begin
740− return $ Yield x (Just (state, x))
741− step' gst (Just (st, acc)) = do
742− r <- step (adaptState gst) st
743− case r of
744− Yield x s -> do
745− y <- fstep acc x
746− return $ Yield y (Just (s, y))
747− Skip s -> return $ Skip (Just (s, acc))
748− Stop -> return $ Stop
749−
750−{-# INLINE scanl #-}
751−scanl :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
752−scanl f seed = scanlM (\a b -> return (f a b)) (return seed)
753−
754−{-# INLINE_NORMAL scanl1M #-}
755−scanl1M :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
756−scanl1M fstep (Stream step state) = Stream step' (state, Nothing)
757− where
758− {-# INLINE_LATE step' #-}
759− step' gst (st, Nothing) = do
760− r <- step gst st
761− case r of
762− Yield x s -> return $ Yield x (s, Just x)
763− Skip s -> return $ Skip (s, Nothing)
764− Stop -> return Stop
765−
766− step' gst (st, Just acc) = do
767− r <- step gst st
768− case r of
769− Yield y s -> do
770− z <- fstep acc y
771− return $ Yield z (s, Just z)
772− Skip s -> return $ Skip (s, Just acc)
773− Stop -> return Stop
774−
775−{-# INLINE scanl1 #-}
776−scanl1 :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
777−scanl1 f = scanl1M (\x y -> return (f x y))
778−
779−{-# INLINE_NORMAL scanl1M' #-}
780−scanl1M' :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
781−scanl1M' fstep (Stream step state) = Stream step' (state, Nothing)
782− where
783− {-# INLINE_LATE step' #-}
784− step' gst (st, Nothing) = do
785− r <- step gst st
786− case r of
787− Yield x s -> x `seq` return $ Yield x (s, Just x)
788− Skip s -> return $ Skip (s, Nothing)
789− Stop -> return Stop
790−
791− step' gst (st, Just acc) = acc `seq` do
792− r <- step gst st
793− case r of
794− Yield y s -> do
795− z <- fstep acc y
796− z `seq` return $ Yield z (s, Just z)
797− Skip s -> return $ Skip (s, Just acc)
798− Stop -> return Stop
799−
800−{-# INLINE scanl1' #-}
801−scanl1' :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
802−scanl1' f = scanl1M' (\x y -> return (f x y))
803−
804−------------------------------------------------------------------------------
805−-- Stateful map/scan
806−------------------------------------------------------------------------------
807−
808−data RollingMapState s a = RollingMapInit s | RollingMapGo s a
809−
810−{-# INLINE rollingMapM #-}
811−rollingMapM :: Monad m => (a -> a -> m b) -> Stream m a -> Stream m b
812−rollingMapM f (Stream step1 state1) = Stream step (RollingMapInit state1)
813− where
814− step gst (RollingMapInit st) = do
815− r <- step1 (adaptState gst) st
816− return $ case r of
817− Yield x s -> Skip $ RollingMapGo s x
818− Skip s -> Skip $ RollingMapInit s
819− Stop -> Stop
820−
821− step gst (RollingMapGo s1 x1) = do
822− r <- step1 (adaptState gst) s1
823− case r of
824− Yield x s -> do
825− !res <- f x x1
826− return $ Yield res $ RollingMapGo s x
827− Skip s -> return $ Skip $ RollingMapGo s x1
828− Stop -> return $ Stop
829−
830−{-# INLINE rollingMap #-}
831−rollingMap :: Monad m => (a -> a -> b) -> Stream m a -> Stream m b
832−rollingMap f = rollingMapM (\x y -> return $ f x y)
833−
834−-------------------------------------------------------------------------------
835−-- Filtering
836−-------------------------------------------------------------------------------
837−
838−{-# INLINE_NORMAL drop #-}
839−drop :: Monad m => Int -> Stream m a -> Stream m a
840−drop n (Stream step state) = Stream step' (state, Just n)
841− where
842− {-# INLINE_LATE step' #-}
843− step' gst (st, Just i)
844− | i > 0 = do
845− r <- step gst st
846− return $
847− case r of
848− Yield _ s -> Skip (s, Just (i - 1))
849− Skip s -> Skip (s, Just i)
850− Stop -> Stop
851− | otherwise = return $ Skip (st, Nothing)
852−
853− step' gst (st, Nothing) = do
854− r <- step gst st
855− return $
856− case r of
857− Yield x s -> Yield x (s, Nothing)
858− Skip s -> Skip (s, Nothing)
859− Stop -> Stop
860−
861−data DropWhileState s a
862− = DropWhileDrop s
863− | DropWhileYield a s
864− | DropWhileNext s
865−
866−{-# INLINE_NORMAL dropWhileM #-}
867−dropWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
868−dropWhileM f (Stream step state) = Stream step' (DropWhileDrop state)
869− where
870− {-# INLINE_LATE step' #-}
871− step' gst (DropWhileDrop st) = do
872− r <- step gst st
873− case r of
874− Yield x s -> do
875− b <- f x
876− if b
877− then return $ Skip (DropWhileDrop s)
878− else return $ Skip (DropWhileYield x s)
879− Skip s -> return $ Skip (DropWhileDrop s)
880− Stop -> return Stop
881−
882− step' gst (DropWhileNext st) = do
883− r <- step gst st
884− case r of
885− Yield x s -> return $ Skip (DropWhileYield x s)
886− Skip s -> return $ Skip (DropWhileNext s)
887− Stop -> return Stop
888−
889− step' _ (DropWhileYield x st) = return $ Yield x (DropWhileNext st)
890−
891−{-# INLINE dropWhile #-}
892−dropWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
893−dropWhile f = dropWhileM (return . f)
894−
895−{-# INLINE_NORMAL filterM #-}
896−filterM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
897−filterM f (Stream step state) = Stream step' state
898− where
899− {-# INLINE_LATE step' #-}
900− step' gst st = do
901− r <- step gst st
902− case r of
903− Yield x s -> do
904− b <- f x
905− return $ if b
906− then Yield x s
907− else Skip s
908− Skip s -> return $ Skip s
909− Stop -> return Stop
910−
911−{-# INLINE filter #-}
912−filter :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
913−filter f = filterM (return . f)
914−
915−{-# INLINE_NORMAL uniq #-}
916−uniq :: (Eq a, Monad m) => Stream m a -> Stream m a
917−uniq (Stream step state) = Stream step' (Nothing, state)
918− where
919− {-# INLINE_LATE step' #-}
920− step' gst (Nothing, st) = do
921− r <- step gst st
922− case r of
923− Yield x s -> return $ Yield x (Just x, s)
924− Skip s -> return $ Skip (Nothing, s)
925− Stop -> return Stop
926− step' gst (Just x, st) = do
927− r <- step gst st
928− case r of
929− Yield y s | x == y -> return $ Skip (Just x, s)
930− | otherwise -> return $ Yield y (Just y, s)
931− Skip s -> return $ Skip (Just x, s)
932− Stop -> return Stop
933−
934−------------------------------------------------------------------------------
935−-- Transformation by Mapping
936−------------------------------------------------------------------------------
937−
938−{-# INLINE_NORMAL sequence #-}
939−sequence :: Monad m => Stream m (m a) -> Stream m a
940−sequence (Stream step state) = Stream step' state
941− where
942− {-# INLINE_LATE step' #-}
943− step' gst st = do
944− r <- step (adaptState gst) st
945− case r of
946− Yield x s -> x >>= \a -> return (Yield a s)
947− Skip s -> return $ Skip s
948− Stop -> return Stop
949−
950−------------------------------------------------------------------------------
951−-- Inserting
952−------------------------------------------------------------------------------
953−
954−data LoopState x s = FirstYield s
955− | InterspersingYield s
956− | YieldAndCarry x s
957−
958−{-# INLINE_NORMAL intersperseM #-}
959−intersperseM :: Monad m => m a -> Stream m a -> Stream m a
960−intersperseM m (Stream step state) = Stream step' (FirstYield state)
961− where
962− {-# INLINE_LATE step' #-}
963− step' gst (FirstYield st) = do
964− r <- step gst st
965− return $
966− case r of
967− Yield x s -> Skip (YieldAndCarry x s)
968− Skip s -> Skip (FirstYield s)
969− Stop -> Stop
970−
971− step' gst (InterspersingYield st) = do
972− r <- step gst st
973− case r of
974− Yield x s -> do
975− a <- m
976− return $ Yield a (YieldAndCarry x s)
977− Skip s -> return $ Skip $ InterspersingYield s
978− Stop -> return Stop
979−
980− step' _ (YieldAndCarry x st) = return $ Yield x (InterspersingYield st)
981−
982−{-# INLINE_NORMAL intersperseM_ #-}
983−intersperseM_ :: Monad m => m b -> Stream m a -> Stream m a
984−intersperseM_ m (Stream step1 state1) = Stream step (Left (pure (), state1))
985− where
986− {-# INLINE_LATE step #-}
987− step gst (Left (eff, st)) = do
988− r <- step1 gst st
989− case r of
990− Yield x s -> eff >> return (Yield x (Right s))
991− Skip s -> return $ Skip (Left (eff, s))
992− Stop -> return Stop
993−
994− step _ (Right st) = return $ Skip $ Left (void m, st)
995−
996−data SuffixState s a
997− = SuffixElem s
998− | SuffixSuffix s
999− | SuffixYield a (SuffixState s a)
1000−
1001−{-# INLINE_NORMAL intersperseSuffix #-}
1002−intersperseSuffix :: forall m a. Monad m => m a -> Stream m a -> Stream m a
1003−intersperseSuffix action (Stream step state) = Stream step' (SuffixElem state)
1004− where
1005− {-# INLINE_LATE step' #-}
1006− step' gst (SuffixElem st) = do
1007− r <- step gst st
1008− return $ case r of
1009− Yield x s -> Skip (SuffixYield x (SuffixSuffix s))
1010− Skip s -> Skip (SuffixElem s)
1011− Stop -> Stop
1012−
1013− step' _ (SuffixSuffix st) = do
1014− action >>= \r -> return $ Skip (SuffixYield r (SuffixElem st))
1015−
1016− step' _ (SuffixYield x next) = return $ Yield x next
1017−
1018−{-# INLINE_NORMAL intersperseSuffix_ #-}
1019−intersperseSuffix_ :: Monad m => m b -> Stream m a -> Stream m a
1020−intersperseSuffix_ m (Stream step1 state1) = Stream step (Left state1)
1021− where
1022− {-# INLINE_LATE step #-}
1023− step gst (Left st) = do
1024− r <- step1 gst st
1025− case r of
1026− Yield x s -> return $ Yield x (Right s)
1027− Skip s -> return $ Skip $ Left s
1028− Stop -> return Stop
1029−
1030− step _ (Right st) = m >> return (Skip (Left st))
1031−
1032−data SuffixSpanState s a
1033− = SuffixSpanElem s Int
1034− | SuffixSpanSuffix s
1035− | SuffixSpanYield a (SuffixSpanState s a)
1036− | SuffixSpanLast
1037− | SuffixSpanStop
1038−
1039−-- | intersperse after every n items
1040−{-# INLINE_NORMAL intersperseSuffixBySpan #-}
1041−intersperseSuffixBySpan :: forall m a. Monad m
1042− => Int -> m a -> Stream m a -> Stream m a
1043−intersperseSuffixBySpan n action (Stream step state) =
1044− Stream step' (SuffixSpanElem state n)
1045− where
1046− {-# INLINE_LATE step' #-}
1047− step' gst (SuffixSpanElem st i) | i > 0 = do
1048− r <- step gst st
1049− return $ case r of
1050− Yield x s -> Skip (SuffixSpanYield x (SuffixSpanElem s (i - 1)))
1051− Skip s -> Skip (SuffixSpanElem s i)
1052− Stop -> if i == n then Stop else Skip SuffixSpanLast
1053− step' _ (SuffixSpanElem st _) = return $ Skip (SuffixSpanSuffix st)
1054−
1055− step' _ (SuffixSpanSuffix st) = do
1056− action >>= \r -> return $ Skip (SuffixSpanYield r (SuffixSpanElem st n))
1057−
1058− step' _ (SuffixSpanLast) = do
1059− action >>= \r -> return $ Skip (SuffixSpanYield r SuffixSpanStop)
1060−
1061− step' _ (SuffixSpanYield x next) = return $ Yield x next
1062−
1063− step' _ (SuffixSpanStop) = return Stop
1064−
1065−{-# INLINE intersperse #-}
1066−intersperse :: Monad m => a -> Stream m a -> Stream m a
1067−intersperse a = intersperseM (return a)
1068−
1069−{-# INLINE_NORMAL insertBy #-}
1070−insertBy :: Monad m => (a -> a -> Ordering) -> a -> Stream m a -> Stream m a
1071−insertBy cmp a (Stream step state) = Stream step' (state, False, Nothing)
1072− where
1073− {-# INLINE_LATE step' #-}
1074− step' gst (st, False, _) = do
1075− r <- step gst st
1076− case r of
1077− Yield x s -> case cmp a x of
1078− GT -> return $ Yield x (s, False, Nothing)
1079− _ -> return $ Yield a (s, True, Just x)
1080− Skip s -> return $ Skip (s, False, Nothing)
1081− Stop -> return $ Yield a (st, True, Nothing)
1082−
1083− step' _ (_, True, Nothing) = return Stop
1084−
1085− step' gst (st, True, Just prev) = do
1086− r <- step gst st
1087− case r of
1088− Yield x s -> return $ Yield prev (s, True, Just x)
1089− Skip s -> return $ Skip (s, True, Just prev)
1090− Stop -> return $ Yield prev (st, True, Nothing)
1091−
1092−------------------------------------------------------------------------------
1093−-- Deleting
1094−------------------------------------------------------------------------------
1095−
1096−{-# INLINE_NORMAL deleteBy #-}
1097−deleteBy :: Monad m => (a -> a -> Bool) -> a -> Stream m a -> Stream m a
1098−deleteBy eq x (Stream step state) = Stream step' (state, False)
1099− where
1100− {-# INLINE_LATE step' #-}
1101− step' gst (st, False) = do
1102− r <- step gst st
1103− case r of
1104− Yield y s -> return $
1105− if eq x y then Skip (s, True) else Yield y (s, False)
1106− Skip s -> return $ Skip (s, False)
1107− Stop -> return Stop
1108−
1109− step' gst (st, True) = do
1110− r <- step gst st
1111− case r of
1112− Yield y s -> return $ Yield y (s, True)
1113− Skip s -> return $ Skip (s, True)
1114− Stop -> return Stop
1115−
1116−------------------------------------------------------------------------------
1117−-- Transformation by Map and Filter
1118−------------------------------------------------------------------------------
1119−
1120−-- XXX Will this always fuse properly?
1121−{-# INLINE_NORMAL mapMaybe #-}
1122−mapMaybe :: Monad m => (a -> Maybe b) -> Stream m a -> Stream m b
1123−mapMaybe f = fmap fromJust . filter isJust . map f
1124−
1125−{-# INLINE_NORMAL mapMaybeM #-}
1126−mapMaybeM :: Monad m => (a -> m (Maybe b)) -> Stream m a -> Stream m b
1127−mapMaybeM f = fmap fromJust . filter isJust . mapM f
1128−
1129−------------------------------------------------------------------------------
1130−-- Zipping
1131−------------------------------------------------------------------------------
1132−
1133−{-# INLINE_NORMAL indexed #-}
1134−indexed :: Monad m => Stream m a -> Stream m (Int, a)
1135−indexed (Stream step state) = Stream step' (state, 0)
1136− where
1137− {-# INLINE_LATE step' #-}
1138− step' gst (st, i) = i `seq` do
1139− r <- step (adaptState gst) st
1140− case r of
1141− Yield x s -> return $ Yield (i, x) (s, i+1)
1142− Skip s -> return $ Skip (s, i)
1143− Stop -> return Stop
1144−
1145−{-# INLINE_NORMAL indexedR #-}
1146−indexedR :: Monad m => Int -> Stream m a -> Stream m (Int, a)
1147−indexedR m (Stream step state) = Stream step' (state, m)
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 -> let i' = i - 1
1154− in return $ Yield (i, x) (s, i')
1155− Skip s -> return $ Skip (s, i)
1156− Stop -> return Stop
1157−
1158−{-# INLINE_NORMAL zipWithM #-}
1159−zipWithM :: Monad m
1160− => (a -> b -> m c) -> Stream m a -> Stream m b -> Stream m c
1161−zipWithM f (Stream stepa ta) (Stream stepb tb) = Stream step (ta, tb, Nothing)
1162− where
1163− {-# INLINE_LATE step #-}
1164− step gst (sa, sb, Nothing) = do
1165− r <- stepa (adaptState gst) sa
1166− return $
1167− case r of
1168− Yield x sa' -> Skip (sa', sb, Just x)
1169− Skip sa' -> Skip (sa', sb, Nothing)
1170− Stop -> Stop
1171−
1172− step gst (sa, sb, Just x) = do
1173− r <- stepb (adaptState gst) sb
1174− case r of
1175− Yield y sb' -> do
1176− z <- f x y
1177− return $ Yield z (sa, sb', Nothing)
1178− Skip sb' -> return $ Skip (sa, sb', Just x)
1179− Stop -> return Stop
1180−
1181−#if __GLASGOW_HASKELL__ >= 801
1182−{-# RULES "zipWithM xs xs"
1183− forall f xs. zipWithM @Identity f xs xs = mapM (\x -> f x x) xs #-}
1184−#endif
1185−
1186−{-# INLINE zipWith #-}
1187−zipWith :: Monad m => (a -> b -> c) -> Stream m a -> Stream m b -> Stream m c
1188−zipWith f = zipWithM (\a b -> return (f a b))
1189−
1190−------------------------------------------------------------------------------
1191−-- Merging
1192−------------------------------------------------------------------------------
1193−
1194−{-# INLINE_NORMAL mergeByM #-}
1195−mergeByM
1196− :: (Monad m)
1197− => (a -> a -> m Ordering) -> Stream m a -> Stream m a -> Stream m a
1198−mergeByM cmp (Stream stepa ta) (Stream stepb tb) =
1199− Stream step (Just ta, Just tb, Nothing, Nothing)
1200− where
1201− {-# INLINE_LATE step #-}
1202−
1203− -- one of the values is missing, and the corresponding stream is running
1204− step gst (Just sa, sb, Nothing, b) = do
1205− r <- stepa gst sa
1206− return $ case r of
1207− Yield a sa' -> Skip (Just sa', sb, Just a, b)
1208− Skip sa' -> Skip (Just sa', sb, Nothing, b)
1209− Stop -> Skip (Nothing, sb, Nothing, b)
1210−
1211− step gst (sa, Just sb, a, Nothing) = do
1212− r <- stepb gst sb
1213− return $ case r of
1214− Yield b sb' -> Skip (sa, Just sb', a, Just b)
1215− Skip sb' -> Skip (sa, Just sb', a, Nothing)
1216− Stop -> Skip (sa, Nothing, a, Nothing)
1217−
1218− -- both the values are available
1219− step _ (sa, sb, Just a, Just b) = do
1220− res <- cmp a b
1221− return $ case res of
1222− GT -> Yield b (sa, sb, Just a, Nothing)
1223− _ -> Yield a (sa, sb, Nothing, Just b)
1224−
1225− -- one of the values is missing, corresponding stream is done
1226− step _ (Nothing, sb, Nothing, Just b) =
1227− return $ Yield b (Nothing, sb, Nothing, Nothing)
1228−
1229− step _ (sa, Nothing, Just a, Nothing) =
1230− return $ Yield a (sa, Nothing, Nothing, Nothing)
1231−
1232− step _ (Nothing, Nothing, Nothing, Nothing) = return Stop
1233−
1234−{-# INLINE mergeBy #-}
1235−mergeBy
1236− :: (Monad m)
1237− => (a -> a -> Ordering) -> Stream m a -> Stream m a -> Stream m a
1238−mergeBy cmp = mergeByM (\a b -> return $ cmp a b)
1239−
1240−------------------------------------------------------------------------------
1241−-- Transformation comprehensions
1242−------------------------------------------------------------------------------
1243−
1244−{-# INLINE_NORMAL the #-}
1245−the :: (Eq a, Monad m) => Stream m a -> m (Maybe a)
1246−the (Stream step state) = go state
1247− where
1248− go st = do
1249− r <- step defState st
1250− case r of
1251− Yield x s -> go' x s
1252− Skip s -> go s
1253− Stop -> return Nothing
1254− go' n st = do
1255− r <- step defState st
1256− case r of
1257− Yield x s | x == n -> go' n s
1258− | otherwise -> return Nothing
1259− Skip s -> go' n s
1260− Stop -> return (Just n)
1261−
1262−-------------------------------------------------------------------------------
1263−-- Concurrent application and fold
1264−-------------------------------------------------------------------------------
1265−
1266−{-# INLINE_NORMAL mkParallelD #-}
1267−mkParallelD :: MonadAsync m => Stream m a -> Stream m a
1268−mkParallelD m = Stream step Nothing
1269− where
1270−
1271− step gst Nothing = do
1272− sv <- newParallelVar StopNone gst
1273− toSVarParallel gst sv m
1274− -- XXX use unfold instead?
1275− return $ Skip $ Just $ fromSVar sv
1276−
1277− step gst (Just (UnStream step1 st)) = do
1278− r <- step1 gst st
1279− return $ case r of
1280− Yield a s -> Yield a (Just $ Stream step1 s)
1281− Skip s -> Skip (Just $ Stream step1 s)
1282− Stop -> Stop
1283−
1284−-- Compare with mkAsync. mkAsync uses an Async style SVar whereas this uses a
1285−-- parallel style SVar for evaluation. Currently, parallel style cannot use
1286−-- rate control whereas Async style can use rate control. In async style SVar
1287−-- the worker thread terminates when the buffer is full whereas in Parallel
1288−-- style it blocks.
1289−--
1290−-- | Make the stream producer and consumer run concurrently by introducing a
1291−-- buffer between them. The producer thread evaluates the input stream until
1292−-- the buffer fills, it blocks if the buffer is full until there is space in
1293−-- the buffer. The consumer consumes the stream lazily from the buffer.
1294−--
1295−-- /Internal/
1296−--
1297−{-# INLINE_NORMAL mkParallel #-}
1298−mkParallel :: (K.IsStream t, MonadAsync m) => t m a -> t m a
1299−mkParallel = fromStreamD . mkParallelD . toStreamD
1300−
1301−-- Note: we can use another API with two callbacks stop and yield if we want
1302−-- the callback to be able to indicate end of stream.
1303−--
1304−-- | Generates a callback and a stream pair. The callback returned is used to
1305−-- queue values to the stream. The stream is infinite, there is no way for the
1306−-- callback to indicate that it is done now.
1307−--
1308−-- /Internal/
1309−--
1310−{-# INLINE_NORMAL newCallbackStream #-}
1311−newCallbackStream :: (K.IsStream t, MonadAsync m) => m ((a -> m ()), t m a)
1312−newCallbackStream = do
1313− sv <- newParallelVar StopNone defState
1314−
1315− -- XXX Add our own thread-id to the SVar as we can not know the callback's
1316− -- thread-id and the callback is not run in a managed worker. We need to
1317− -- handle this better.
1318− liftIO myThreadId >>= modifyThread sv
1319−
1320− let callback a = liftIO $ void $ send sv (ChildYield a)
1321− -- XXX we can return an SVar and then the consumer can unfold from the
1322− -- SVar?
1323− return (callback, fromStreamD (fromSVar sv))
1324−
1325−------------------------------------------------------------------------------
1326−-- Time related
1327−------------------------------------------------------------------------------
1328−
1329−-- XXX using getTime in the loop can be pretty expensive especially for
1330−-- computations where iterations are lightweight. We have the following
1331−-- options:
1332−--
1333−-- 1) Run a timeout thread updating a flag asynchronously and check that
1334−-- flag here, that way we can have a cheap termination check.
1335−--
1336−-- 2) Use COARSE clock to get time with lower resolution but more efficiently.
1337−--
1338−-- 3) Use rdtscp/rdtsc to get time directly from the processor, compute the
1339−-- termination value of rdtsc in the beginning and then in each iteration just
1340−-- get rdtsc and check if we should terminate.
1341−--
1342−data TakeByTime st s
1343− = TakeByTimeInit st
1344− | TakeByTimeCheck st s
1345− | TakeByTimeYield st s
1346−
1347−{-# INLINE_NORMAL takeByTime #-}
1348−takeByTime :: (MonadIO m, TimeUnit64 t) => t -> Stream m a -> Stream m a
1349−takeByTime duration (Stream step1 state1) = Stream step (TakeByTimeInit state1)
1350− where
1351−
1352− lim = toRelTime64 duration
1353−
1354− {-# INLINE_LATE step #-}
1355− step _ (TakeByTimeInit _) | lim == 0 = return Stop
1356− step _ (TakeByTimeInit st) = do
1357− t0 <- liftIO $ getTime Monotonic
1358− return $ Skip (TakeByTimeYield st t0)
1359− step _ (TakeByTimeCheck st t0) = do
1360− t <- liftIO $ getTime Monotonic
1361− return $
1362− if diffAbsTime64 t t0 > lim
1363− then Stop
1364− else Skip (TakeByTimeYield st t0)
1365− step gst (TakeByTimeYield st t0) = do
1366− r <- step1 gst st
1367− return $ case r of
1368− Yield x s -> Yield x (TakeByTimeCheck s t0)
1369− Skip s -> Skip (TakeByTimeCheck s t0)
1370− Stop -> Stop
1371−
1372−data DropByTime st s x
1373− = DropByTimeInit st
1374− | DropByTimeGen st s
1375− | DropByTimeCheck st s x
1376− | DropByTimeYield st
1377−
1378−{-# INLINE_NORMAL dropByTime #-}
1379−dropByTime :: (MonadIO m, TimeUnit64 t) => t -> Stream m a -> Stream m a
1380−dropByTime duration (Stream step1 state1) = Stream step (DropByTimeInit state1)
1381− where
1382−
1383− lim = toRelTime64 duration
1384−
1385− {-# INLINE_LATE step #-}
1386− step _ (DropByTimeInit st) = do
1387− t0 <- liftIO $ getTime Monotonic
1388− return $ Skip (DropByTimeGen st t0)
1389− step gst (DropByTimeGen st t0) = do
1390− r <- step1 gst st
1391− return $ case r of
1392− Yield x s -> Skip (DropByTimeCheck s t0 x)
1393− Skip s -> Skip (DropByTimeGen s t0)
1394− Stop -> Stop
1395− step _ (DropByTimeCheck st t0 x) = do
1396− t <- liftIO $ getTime Monotonic
1397− if diffAbsTime64 t t0 <= lim
1398− then return $ Skip $ DropByTimeGen st t0
1399− else return $ Yield x $ DropByTimeYield st
1400− step gst (DropByTimeYield st) = do
1401− r <- step1 gst st
1402− return $ case r of
1403− Yield x s -> Yield x (DropByTimeYield s)
1404− Skip s -> Skip (DropByTimeYield s)
1405− Stop -> Stop
317+import Streamly.Internal.Data.Stream.StreamD.Transform