4-- Module : Streamly.Internal.Data.Stream.StreamD
5-- Copyright : (c) 2018 Composewell Technologies
6-- (c) Roman Leshchinskiy 2008-2010
7-- (c) The University of Glasgow, 2009
10-- Maintainer : streamly@composewell.com
11-- Stability : experimental
14-- Direct style re-implementation of CPS style stream in StreamK module. The
15-- symbol or suffix 'D' in this module denotes the "Direct" style. GHC is able
16-- to INLINE and fuse direct style better, providing better performance than
20-- import qualified Streamly.Internal.Data.Stream.StreamD as D
23-- Some of the functions in this file have been adapted from the vector
24-- library, https://hackage.haskell.org/package/vector.
26module Streamly.Internal.Data.Stream.StreamD
30 , Stream (Stream, UnStream)
46 -- ** Specialized Generation
47 -- | Generate a monadic stream from a seed.
60 , enumerateFromStepIntegral
61 , enumerateFromIntegral
62 , enumerateFromThenIntegral
63 , enumerateFromToIntegral
64 , enumerateFromThenToIntegral
66 , enumerateFromStepNum
69 , enumerateFromToFractional
70 , enumerateFromThenToFractional
76 -- | Transform an input structure into a stream.
77 -- | Direct style stream does not support @fromFoldable@.
111 -- ** Specialized Folds
138 -- ** Flattening nested streams
141 , ConcatMapUState (..)
143 , ConcatUnfoldInterleaveState (..)
144 , concatUnfoldInterleave
145 , concatUnfoldRoundrobin
148 , InterleaveState(..)
153 , roundRobin -- interleaveFair?/ParallelFair
186 -- | Transform a stream into another type.
203 -- ** By folding (scans)
254 , intersperseSuffixBySpan
278 -- * Transformation comprehensions
298 -- * Concurrent Application
307import Control.Concurrent (killThread, myThreadId, takeMVar, threadDelay)
308import Control.Exception
309 (assert, Exception, SomeException, AsyncException, fromException, mask_)
310import Control.Monad (void, when, forever)
311import Control.Monad.Catch (MonadCatch, MonadThrow, throwM)
312import Control.Monad.IO.Class (MonadIO(..))
313import Control.Monad.Trans.Reader (ReaderT)
314import Control.Monad.Trans.State.Strict (StateT)
315import Control.Monad.Trans.Class (MonadTrans(lift))
316import Control.Monad.Trans.Control (MonadBaseControl, liftBaseOp_, control)
317import Data.Bits (shiftR, shiftL, (.|.), (.&.))
318import Data.Functor.Identity (Identity(..))
319import Data.Int (Int64)
320import Data.IORef (newIORef, readIORef, mkWeakIORef, writeIORef, IORef)
321import Data.Maybe (fromJust, isJust, isNothing)
322import Data.Word (Word32)
323import Foreign.Ptr (Ptr)
324import Foreign.Storable (Storable(..))
325import GHC.Exts (SpecConstrAnnotation(..))
326import GHC.Types (SPEC(..))
327import System.Mem (performMajorGC)
328import Fusion.Plugin.Types (Fuse(..))
329import Streamly.Internal.Data.IORef.Prim (Prim)
330import Streamly.Internal.Data.Time.Units
331 (TimeUnit64, toRelTime64, diffAbsTime64, RelTime64)
332import Streamly.Internal.Data.Atomics (atomicModifyIORefCAS_)
333import Streamly.Internal.Data.Array.Storable.Foreign.Types (Array(..))
334import Streamly.Internal.Data.Fold.Types (Fold(..))
335import Streamly.Internal.Data.Parser (ParseError(..))
336import Streamly.Internal.Data.Pipe.Types (Pipe(..), PipeState(..))
337import Streamly.Internal.Data.Time.Clock (Clock(Monotonic), getTime)
338import Streamly.Internal.Data.Time.Units
339 (MicroSecond64(..), fromAbsTime, toAbsTime, AbsTime)
340import Streamly.Internal.Data.Unfold.Types (Unfold(..))
341import Streamly.Internal.Data.Tuple.Strict (Tuple3'(..))
342import Streamly.Internal.Data.Stream.SVar (fromConsumer, pushToFold)
344import qualified Streamly.Internal.Data.IORef.Prim as Prim
345import qualified Streamly.Internal.Data.Pipe.Types as Pipe
346import qualified Streamly.Internal.Data.Array.Storable.Foreign.Types as A
347import qualified Streamly.Internal.Data.Array.Storable.Foreign.Mut.Types as MA
348import qualified Streamly.Internal.Data.Fold as FL
349import qualified Streamly.Memory.Ring as RB
350import qualified Streamly.Internal.Data.Stream.StreamK as K
351import qualified Streamly.Internal.Data.Parser as PR
352import qualified Streamly.Internal.Data.Parser.ParserD as PRD
353import qualified Control.Monad.Catch as MC
354import qualified Control.Monad.Trans.Reader as Reader
355import qualified Control.Monad.Trans.State.Strict as State
356import qualified Prelude
359 ( map, mapM, mapM_, repeat, foldr, last, take, filter
360 , takeWhile, drop, dropWhile, all, any, maximum, minimum, elem
361 , notElem, null, head, tail, zipWith, lookup, foldr1, sequence
362 , (!!), scanl, scanl1, concatMap, replicate, enumFromTo, concat
363 , reverse, iterate, splitAt)
364import Streamly.Internal.Data.Stream.StreamD.Type
365import Streamly.Internal.Data.SVar
367------------------------------------------------------------------------------
369------------------------------------------------------------------------------
371-- | An empty 'Stream'.
372{-# INLINE_NORMAL nil #-}
373nil :: Monad m => Stream m a
374nil = Stream (\_ _ -> return Stop) ()
376-- | An empty 'Stream' with a side effect.
377{-# INLINE_NORMAL nilM #-}
378nilM :: Monad m => m b -> Stream m a
379nilM m = Stream (\_ _ -> m >> return Stop) ()
381{-# INLINE_NORMAL consM #-}
382consM :: Monad m => m a -> Stream m a -> Stream m a
383consM m (Stream step state) = Stream step1 Nothing
385 {-# INLINE_LATE step1 #-}
386 step1 _ Nothing = m >>= \x -> return $ Yield x (Just state)
387 step1 gst (Just st) = do
391 Yield a s -> Yield a (Just s)
392 Skip s -> Skip (Just s)
395-- XXX implement in terms of consM?
396-- cons x = consM (return x)
398-- | Can fuse but has O(n^2) complexity.
399{-# INLINE_NORMAL cons #-}
400cons :: Monad m => a -> Stream m a -> Stream m a
401cons x (Stream step state) = Stream step1 Nothing
403 {-# INLINE_LATE step1 #-}
404 step1 _ Nothing = return $ Yield x (Just state)
405 step1 gst (Just st) = do
409 Yield a s -> Yield a (Just s)
410 Skip s -> Skip (Just s)
413-------------------------------------------------------------------------------
415-------------------------------------------------------------------------------
417-- Does not fuse, has the same performance as the StreamK version.
418{-# INLINE_NORMAL uncons #-}
419uncons :: Monad m => Stream m a -> m (Maybe (a, Stream m a))
420uncons (UnStream step state) = go state
423 r <- step defState st
425 Yield x s -> return $ Just (x, Stream step s)
427 Stop -> return Nothing
429------------------------------------------------------------------------------
430-- Generation by unfold
431------------------------------------------------------------------------------
433{-# INLINE_NORMAL unfoldrM #-}
434unfoldrM :: Monad m => (s -> m (Maybe (a, s))) -> s -> Stream m a
435unfoldrM next state = Stream step state
437 {-# INLINE_LATE step #-}
441 Just (x, s) -> Yield x s
444{-# INLINE_LATE unfoldr #-}
445unfoldr :: Monad m => (s -> Maybe (a, s)) -> s -> Stream m a
446unfoldr f = unfoldrM (return . f)
448-- | Convert an 'Unfold' into a 'Stream' by supplying it a seed.
450{-# INLINE_NORMAL unfold #-}
451unfold :: Monad m => Unfold m a b -> a -> Stream m b
452unfold (Unfold ustep inject) seed = Stream step Nothing
454 {-# INLINE_LATE step #-}
455 step _ Nothing = inject seed >>= return . Skip . Just
456 step _ (Just st) = do
459 Yield x s -> Yield x (Just s)
460 Skip s -> Skip (Just s)
463------------------------------------------------------------------------------
464-- Specialized Generation
465------------------------------------------------------------------------------
467{-# INLINE_NORMAL repeatM #-}
468repeatM :: Monad m => m a -> Stream m a
469repeatM x = Stream (\_ _ -> x >>= \r -> return $ Yield r ()) ()
471{-# INLINE_NORMAL repeat #-}
472repeat :: Monad m => a -> Stream m a
473repeat x = Stream (\_ _ -> return $ Yield x ()) ()
475{-# INLINE_NORMAL iterateM #-}
476iterateM :: Monad m => (a -> m a) -> m a -> Stream m a
477iterateM step = Stream (\_ st -> st >>= \x -> return $ Yield x (step x))
479{-# INLINE_NORMAL iterate #-}
480iterate :: Monad m => (a -> a) -> a -> Stream m a
481iterate step st = iterateM (return . step) (return st)
483{-# INLINE_NORMAL replicateM #-}
484replicateM :: forall m a. Monad m => Int -> m a -> Stream m a
485replicateM n p = Stream step n
487 {-# INLINE_LATE step #-}
489 | i <= 0 = return Stop
492 return $ Yield x (i - 1)
494{-# INLINE_NORMAL replicate #-}
495replicate :: Monad m => Int -> a -> Stream m a
496replicate n x = replicateM n (return x)
498-- This would not work properly for floats, therefore we put an Integral
500-- | Can be used to enumerate unbounded integrals. This does not check for
501-- overflow or underflow for bounded integrals.
502{-# INLINE_NORMAL enumerateFromStepIntegral #-}
503enumerateFromStepIntegral :: (Integral a, Monad m) => a -> a -> Stream m a
504enumerateFromStepIntegral from stride =
505 from `seq` stride `seq` Stream step from
507 {-# INLINE_LATE step #-}
508 step _ !x = return $ Yield x $! (x + stride)
510-- We are assuming that "to" is constrained by the type to be within
512{-# INLINE enumerateFromToIntegral #-}
513enumerateFromToIntegral :: (Monad m, Integral a) => a -> a -> Stream m a
514enumerateFromToIntegral from to =
515 takeWhile (<= to) $ enumerateFromStepIntegral from 1
517{-# INLINE enumerateFromIntegral #-}
518enumerateFromIntegral :: (Monad m, Integral a, Bounded a) => a -> Stream m a
519enumerateFromIntegral from = enumerateFromToIntegral from maxBound
521data EnumState a = EnumInit | EnumYield a a a | EnumStop
523{-# INLINE_NORMAL enumerateFromThenToIntegralUp #-}
524enumerateFromThenToIntegralUp
525 :: (Monad m, Integral a)
526 => a -> a -> a -> Stream m a
527enumerateFromThenToIntegralUp from next to = Stream step EnumInit
529 {-# INLINE_LATE step #-}
535 else Yield from EnumStop
536 else -- from <= next <= to
537 let stride = next - from
538 in Skip $ EnumYield from stride (to - stride)
540 step _ (EnumYield x stride toMinus) =
543 then Yield x EnumStop
544 else Yield x $ EnumYield (x + stride) stride toMinus
546 step _ EnumStop = return Stop
548{-# INLINE_NORMAL enumerateFromThenToIntegralDn #-}
549enumerateFromThenToIntegralDn
550 :: (Monad m, Integral a)
551 => a -> a -> a -> Stream m a
552enumerateFromThenToIntegralDn from next to = Stream step EnumInit
554 {-# INLINE_LATE step #-}
556 return $ if to > next
559 else Yield from EnumStop
560 else -- from >= next >= to
561 let stride = next - from
562 in Skip $ EnumYield from stride (to - stride)
564 step _ (EnumYield x stride toMinus) =
567 then Yield x EnumStop
568 else Yield x $ EnumYield (x + stride) stride toMinus
570 step _ EnumStop = return Stop
572{-# INLINE_NORMAL enumerateFromThenToIntegral #-}
573enumerateFromThenToIntegral
574 :: (Monad m, Integral a)
575 => a -> a -> a -> Stream m a
576enumerateFromThenToIntegral from next to
577 | next >= from = enumerateFromThenToIntegralUp from next to
578 | otherwise = enumerateFromThenToIntegralDn from next to
580{-# INLINE_NORMAL enumerateFromThenIntegral #-}
581enumerateFromThenIntegral
582 :: (Monad m, Integral a, Bounded a)
583 => a -> a -> Stream m a
584enumerateFromThenIntegral from next =
586 then enumerateFromThenToIntegralUp from next maxBound
587 else enumerateFromThenToIntegralDn from next minBound
589-- For floating point numbers if the increment is less than the precision then
590-- it just gets lost. Therefore we cannot always increment it correctly by just
592-- 9007199254740992 + 1 + 1 :: Double => 9.007199254740992e15
593-- 9007199254740992 + 2 :: Double => 9.007199254740994e15
595-- Instead we accumulate the increment counter and compute the increment
596-- every time before adding it to the starting number.
598-- This works for Integrals as well as floating point numbers, but
599-- enumerateFromStepIntegral is faster for integrals.
600{-# INLINE_NORMAL enumerateFromStepNum #-}
601enumerateFromStepNum :: (Monad m, Num a) => a -> a -> Stream m a
602enumerateFromStepNum from stride = Stream step 0
604 {-# INLINE_LATE step #-}
605 step _ !i = return $ (Yield $! (from + i * stride)) $! (i + 1)
607{-# INLINE_NORMAL numFrom #-}
608numFrom :: (Monad m, Num a) => a -> Stream m a
609numFrom from = enumerateFromStepNum from 1
611{-# INLINE_NORMAL numFromThen #-}
612numFromThen :: (Monad m, Num a) => a -> a -> Stream m a
613numFromThen from next = enumerateFromStepNum from (next - from)
615-- We cannot write a general function for Num. The only way to write code
616-- portable between the two is to use a 'Real' constraint and convert between
617-- Fractional and Integral using fromRational which is horribly slow.
618{-# INLINE_NORMAL enumerateFromToFractional #-}
619enumerateFromToFractional
620 :: (Monad m, Fractional a, Ord a)
621 => a -> a -> Stream m a
622enumerateFromToFractional from to =
623 takeWhile (<= to + 1 / 2) $ enumerateFromStepNum from 1
625{-# INLINE_NORMAL enumerateFromThenToFractional #-}
626enumerateFromThenToFractional
627 :: (Monad m, Fractional a, Ord a)
628 => a -> a -> a -> Stream m a
629enumerateFromThenToFractional from next to =
630 takeWhile predicate $ numFromThen from next
632 mid = (next - from) / 2
633 predicate | next >= from = (<= to + mid)
634 | otherwise = (>= to + mid)
636-------------------------------------------------------------------------------
637-- Generation by Conversion
638-------------------------------------------------------------------------------
640{-# INLINE_NORMAL fromIndicesM #-}
641fromIndicesM :: Monad m => (Int -> m a) -> Stream m a
642fromIndicesM gen = Stream step 0
644 {-# INLINE_LATE step #-}
647 return $ Yield x (i + 1)
649{-# INLINE fromIndices #-}
650fromIndices :: Monad m => (Int -> a) -> Stream m a
651fromIndices gen = fromIndicesM (return . gen)
653{-# INLINE_NORMAL generateM #-}
654generateM :: Monad m => Int -> (Int -> m a) -> Stream m a
655generateM n gen = n `seq` Stream step 0
657 {-# INLINE_LATE step #-}
658 step _ i | i < n = do
660 return $ Yield x (i + 1)
661 | otherwise = return Stop
663{-# INLINE generate #-}
664generate :: Monad m => Int -> (Int -> a) -> Stream m a
665generate n gen = generateM n (return . gen)
667-- XXX we need the MonadAsync constraint because of a rewrite rule.
668-- | Convert a list of monadic actions to a 'Stream'
669{-# INLINE_LATE fromListM #-}
670fromListM :: MonadAsync m => [m a] -> Stream m a
671fromListM = Stream step
673 {-# INLINE_LATE step #-}
674 step _ (m:ms) = m >>= \x -> return $ Yield x ms
675 step _ [] = return Stop
677{-# INLINE toStreamD #-}
678toStreamD :: (K.IsStream t, Monad m) => t m a -> Stream m a
679toStreamD = fromStreamK . K.toStream
681{-# INLINE_NORMAL fromPrimIORef #-}
682fromPrimIORef :: (MonadIO m, Prim a) => Prim.IORef a -> Stream m a
683fromPrimIORef var = Stream step ()
685 {-# INLINE_LATE step #-}
686 step _ () = liftIO (Prim.readIORef var) >>= \x -> return $ Yield x ()
688-------------------------------------------------------------------------------
689-- Generation from SVar
690-------------------------------------------------------------------------------
692data FromSVarState t m a =
694 | FromSVarRead (SVar t m a)
695 | FromSVarLoop (SVar t m a) [ChildEvent a]
696 | FromSVarDone (SVar t m a)
698{-# INLINE_NORMAL fromSVar #-}
699fromSVar :: (MonadAsync m) => SVar t m a -> Stream m a
700fromSVar svar = Stream step FromSVarInit
703 {-# INLINE_LATE step #-}
704 step _ FromSVarInit = do
705 ref <- liftIO $ newIORef ()
706 _ <- liftIO $ mkWeakIORef ref hook
707 -- when this copy of svar gets garbage collected "ref" will get
708 -- garbage collected and our GC hook will be called.
709 let sv = svar{svarRef = Just ref}
710 return $ Skip (FromSVarRead sv)
714 {-# NOINLINE hook #-}
716 when (svarInspectMode svar) $ do
717 r <- liftIO $ readIORef (svarStopTime (svarStats svar))
719 printSVar svar "SVar Garbage Collected"
721 -- If there are any SVars referenced by this SVar a GC will prompt
722 -- them to be cleaned up quickly.
723 when (svarInspectMode svar) performMajorGC
725 step _ (FromSVarRead sv) = do
726 list <- readOutputQ sv
727 -- Reversing the output is important to guarantee that we process the
728 -- outputs in the same order as they were generated by the constituent
730 return $ Skip $ FromSVarLoop sv (Prelude.reverse list)
732 step _ (FromSVarLoop sv []) = do
733 done <- postProcess sv
734 return $ Skip $ if done
735 then (FromSVarDone sv)
736 else (FromSVarRead sv)
738 step _ (FromSVarLoop sv (ev : es)) = do
740 ChildYield a -> return $ Yield a (FromSVarLoop sv es)
741 ChildStop tid e -> do
745 stop <- shouldStop tid
748 liftIO (cleanupSVar sv)
749 return $ Skip (FromSVarDone sv)
750 else return $ Skip (FromSVarLoop sv es)
752 case fromException ex of
754 return $ Skip (FromSVarLoop sv es)
755 Nothing -> liftIO (cleanupSVar sv) >> throwM ex
759 case svarStopStyle sv of
760 StopNone -> return False
761 StopAny -> return True
763 sid <- liftIO $ readIORef (svarStopBy sv)
764 return $ if tid == sid then True else False
766 step _ (FromSVarDone sv) = do
767 when (svarInspectMode sv) $ do
768 t <- liftIO $ getTime Monotonic
769 liftIO $ writeIORef (svarStopTime (svarStats sv)) (Just t)
770 liftIO $ printSVar sv "SVar Done"
773-------------------------------------------------------------------------------
774-- Process events received by a fold consumer from a stream producer
775-------------------------------------------------------------------------------
777{-# INLINE_NORMAL fromProducer #-}
778fromProducer :: (MonadAsync m) => SVar t m a -> Stream m a
779fromProducer svar = Stream step (FromSVarRead svar)
782 {-# INLINE_LATE step #-}
783 step _ (FromSVarRead sv) = do
784 list <- readOutputQ sv
785 -- Reversing the output is important to guarantee that we process the
786 -- outputs in the same order as they were generated by the constituent
788 return $ Skip $ FromSVarLoop sv (Prelude.reverse list)
790 step _ (FromSVarLoop sv []) = return $ Skip $ FromSVarRead sv
791 step _ (FromSVarLoop sv (ev : es)) = do
793 ChildYield a -> return $ Yield a (FromSVarLoop sv es)
794 ChildStop tid e -> do
798 sendStopToProducer sv
799 return $ Skip (FromSVarDone sv)
800 Just _ -> error "Bug: fromProducer: received exception"
802 step _ (FromSVarDone sv) = do
803 when (svarInspectMode sv) $ do
804 t <- liftIO $ getTime Monotonic
805 liftIO $ writeIORef (svarStopTime (svarStats sv)) (Just t)
806 liftIO $ printSVar sv "SVar Done"
809 step _ FromSVarInit = undefined
811-------------------------------------------------------------------------------
812-- Hoisting the inner monad
813-------------------------------------------------------------------------------
815{-# INLINE_NORMAL hoist #-}
816hoist :: Monad n => (forall x. m x -> n x) -> Stream m a -> Stream n a
817hoist f (Stream step state) = (Stream step' state)
819 {-# INLINE_LATE step' #-}
821 r <- f $ step (adaptState gst) st
823 Yield x s -> Yield x s
827{-# INLINE generally #-}
828generally :: Monad m => Stream Identity a -> Stream m a
829generally = hoist (return . runIdentity)
831{-# INLINE_NORMAL liftInner #-}
832liftInner :: (Monad m, MonadTrans t, Monad (t m))
833 => Stream m a -> Stream (t m) a
834liftInner (Stream step state) = Stream step' state
836 {-# INLINE_LATE step' #-}
838 r <- lift $ step (adaptState gst) st
840 Yield x s -> Yield x s
844{-# INLINE_NORMAL runReaderT #-}
845runReaderT :: Monad m => m s -> Stream (ReaderT s m) a -> Stream m a
846runReaderT env (Stream step state) = Stream step' (state, env)
848 {-# INLINE_LATE step' #-}
849 step' gst (st, action) = do
851 r <- Reader.runReaderT (step (adaptState gst) st) sv
853 Yield x s -> Yield x (s, return sv)
854 Skip s -> Skip (s, return sv)
857{-# INLINE_NORMAL evalStateT #-}
858evalStateT :: Monad m => m s -> Stream (StateT s m) a -> Stream m a
859evalStateT initial (Stream step state) = Stream step' (state, initial)
861 {-# INLINE_LATE step' #-}
862 step' gst (st, action) = do
864 (r, sv') <- State.runStateT (step (adaptState gst) st) sv
866 Yield x s -> Yield x (s, return sv')
867 Skip s -> Skip (s, return sv')
870{-# INLINE_NORMAL runStateT #-}
871runStateT :: Monad m => m s -> Stream (StateT s m) a -> Stream m (s, a)
872runStateT initial (Stream step state) = Stream step' (state, initial)
874 {-# INLINE_LATE step' #-}
875 step' gst (st, action) = do
877 (r, sv') <- State.runStateT (step (adaptState gst) st) sv
879 Yield x s -> Yield (sv', x) (s, return sv')
880 Skip s -> Skip (s, return sv')
883------------------------------------------------------------------------------
884-- Elimination by Folds
885------------------------------------------------------------------------------
887------------------------------------------------------------------------------
889------------------------------------------------------------------------------
891{-# INLINE_NORMAL foldr1 #-}
892foldr1 :: Monad m => (a -> a -> a) -> Stream m a -> m (Maybe a)
896 Nothing -> return Nothing
897 Just (h, t) -> fmap Just (foldr f h t)
899------------------------------------------------------------------------------
901------------------------------------------------------------------------------
903{-# INLINE_NORMAL foldlT #-}
904foldlT :: (Monad m, Monad (s m), MonadTrans s)
905 => (s m b -> a -> s m b) -> s m b -> Stream m a -> s m b
906foldlT fstep begin (Stream step state) = go SPEC begin state
909 r <- lift $ step defState st
911 Yield x s -> go SPEC (fstep acc x) s
912 Skip s -> go SPEC acc s
915-- Note, this is going to have horrible performance, because of the nature of
916-- the stream type (i.e. direct stream vs CPS). Its only for reference, it is
917-- likely be practically unusable.
918{-# INLINE_NORMAL foldlS #-}
920 => (Stream m b -> a -> Stream m b) -> Stream m b -> Stream m a -> Stream m b
921foldlS fstep begin (Stream step state) = Stream step' (Left (state, begin))
923 step' gst (Left (st, acc)) = do
924 r <- step (adaptState gst) st
926 Yield x s -> Skip (Left (s, fstep acc x))
927 Skip s -> Skip (Left (s, acc))
928 Stop -> Skip (Right acc)
930 step' gst (Right (Stream stp stt)) = do
931 r <- stp (adaptState gst) stt
933 Yield x s -> Yield x (Right (Stream stp s))
934 Skip s -> Skip (Right (Stream stp s))
937------------------------------------------------------------------------------
939------------------------------------------------------------------------------
941-- Inlined definition. Without the inline "serially/parser/take" benchmark
942-- degrades and parseMany does not fuse. Even using "inline" at the callsite
944{-# INLINE splitAt #-}
945splitAt :: Int -> [a] -> ([a],[a])
948 | otherwise = splitAt' n ls
950 splitAt' :: Int -> [a] -> ([a], [a])
951 splitAt' _ [] = ([], [])
952 splitAt' 1 (x:xs) = ([x], xs)
953 splitAt' m (x:xs) = (x:xs', xs'')
955 (xs', xs'') = splitAt' (m - 1) xs
957-- GHC parser does not accept {-# ANN type [] NoSpecConstr #-}, so we need
959{-# ANN type List NoSpecConstr #-}
960newtype List a = List {getList :: [a]}
962-- | Run a 'Parse' over a stream.
963{-# INLINE_NORMAL parselMx' #-}
966 => (s -> a -> m (PR.Step s b))
971parselMx' pstep initial extract (Stream step state) = do
972 initial >>= go SPEC state (List [])
976 -- XXX currently we are using a dumb list based approach for backtracking
977 -- buffer. This can be replaced by a sliding/ring buffer using Data.Array.
978 -- That will allow us more efficient random back and forth movement.
980 go !_ st buf !pst = do
981 r <- step defState st
986 PR.Partial 0 pst1 -> go SPEC s (List []) pst1
987 PR.Partial n pst1 -> do
988 assert (n <= length (x:getList buf)) (return ())
989 let src0 = Prelude.take n (x:getList buf)
990 src = Prelude.reverse src0
991 gobuf SPEC s (List []) (List src) pst1
992 PR.Continue 0 pst1 -> go SPEC s (List (x:getList buf)) pst1
993 PR.Continue n pst1 -> do
994 assert (n <= length (x:getList buf)) (return ())
995 let (src0, buf1) = splitAt n (x:getList buf)
996 src = Prelude.reverse src0
997 gobuf SPEC s (List buf1) (List src) pst1
998 PR.Done _ b -> return b
999 PR.Error err -> throwM $ ParseError err
1000 Skip s -> go SPEC s buf pst
1003 gobuf !_ s buf (List []) !pst = go SPEC s buf pst
1004 gobuf !_ s buf (List (x:xs)) !pst = do
1007 PR.Partial 0 pst1 ->
1008 gobuf SPEC s (List []) (List xs) pst1
1009 PR.Partial n pst1 -> do
1010 assert (n <= length (x:getList buf)) (return ())
1011 let src0 = Prelude.take n (x:getList buf)
1012 src = Prelude.reverse src0 ++ xs
1013 gobuf SPEC s (List []) (List src) pst1
1014 PR.Continue 0 pst1 ->
1015 gobuf SPEC s (List (x:getList buf)) (List xs) pst1
1016 PR.Continue n pst1 -> do
1017 assert (n <= length (x:getList buf)) (return ())
1018 let (src0, buf1) = splitAt n (x:getList buf)
1019 src = Prelude.reverse src0 ++ xs
1020 gobuf SPEC s (List buf1) (List src) pst1
1021 PR.Done _ b -> return b
1022 PR.Error err -> throwM $ ParseError err
1024------------------------------------------------------------------------------
1026------------------------------------------------------------------------------
1028{-# ANN type ParseChunksState Fuse #-}
1029data ParseChunksState x inpBuf st pst =
1030 ParseChunksInit inpBuf st
1031 | ParseChunksInitLeftOver inpBuf
1032 | ParseChunksStream st inpBuf !pst
1033 | ParseChunksBuf inpBuf st inpBuf !pst
1034 | ParseChunksYield x (ParseChunksState x inpBuf st pst)
1036{-# INLINE_NORMAL parseMany #-}
1042parseMany (PRD.Parser pstep initial extract) (Stream step state) =
1043 Stream stepOuter (ParseChunksInit [] state)
1047 {-# INLINE_LATE stepOuter #-}
1048 -- Buffer is empty, get the first element from the stream, initialize the
1049 -- fold and then go to stream processing loop.
1050 stepOuter gst (ParseChunksInit [] st) = do
1051 r <- step (adaptState gst) st
1053 Yield x s -> initial >>= return . Skip . ParseChunksBuf [x] s []
1054 Skip s -> return $ Skip $ ParseChunksInit [] s
1057 -- Buffer is not empty, go to buffered processing loop
1058 stepOuter _ (ParseChunksInit src st) = do
1059 initial >>= return . Skip . ParseChunksBuf src st []
1061 -- XXX we just discard any leftover input at the end
1062 stepOuter _ (ParseChunksInitLeftOver _) = return Stop
1064 -- Buffer is empty, process elements from the stream
1065 stepOuter gst (ParseChunksStream st buf pst) = do
1066 r <- step (adaptState gst) st
1071 PR.Partial 0 pst1 ->
1072 return $ Skip $ ParseChunksStream s [] pst1
1073 PR.Partial n pst1 -> do
1074 assert (n <= length (x:buf)) (return ())
1075 let src0 = Prelude.take n (x:buf)
1076 src = Prelude.reverse src0
1077 return $ Skip $ ParseChunksBuf src s [] pst1
1078 PR.Continue 0 pst1 ->
1079 return $ Skip $ ParseChunksStream s (x:buf) pst1
1080 PR.Continue n pst1 -> do
1081 assert (n <= length (x:buf)) (return ())
1082 let (src0, buf1) = splitAt n (x:buf)
1083 src = Prelude.reverse src0
1084 return $ Skip $ ParseChunksBuf src s buf1 pst1
1087 ParseChunksYield b (ParseChunksInit [] s)
1089 assert (n <= length (x:buf)) (return ())
1090 let src = Prelude.reverse (Prelude.take n (x:buf))
1092 ParseChunksYield b (ParseChunksInit src s)
1093 PR.Error err -> throwM $ ParseError err
1094 Skip s -> return $ Skip $ ParseChunksStream s buf pst
1097 let src = Prelude.reverse buf
1099 ParseChunksYield b (ParseChunksInitLeftOver src)
1101 -- go back to stream processing mode
1102 stepOuter _ (ParseChunksBuf [] s buf pst) =
1103 return $ Skip $ ParseChunksStream s buf pst
1105 -- buffered processing loop
1106 stepOuter _ (ParseChunksBuf (x:xs) s buf pst) = do
1109 PR.Partial 0 pst1 ->
1110 return $ Skip $ ParseChunksBuf xs s [] pst1
1111 PR.Partial n pst1 -> do
1112 assert (n <= length (x:buf)) (return ())
1113 let src0 = Prelude.take n (x:buf)
1114 src = Prelude.reverse src0 ++ xs
1115 return $ Skip $ ParseChunksBuf src s [] pst1
1116 PR.Continue 0 pst1 ->
1117 return $ Skip $ ParseChunksBuf xs s (x:buf) pst1
1118 PR.Continue n pst1 -> do
1119 assert (n <= length (x:buf)) (return ())
1120 let (src0, buf1) = splitAt n (x:buf)
1121 src = Prelude.reverse src0 ++ xs
1122 return $ Skip $ ParseChunksBuf src s buf1 pst1
1124 return $ Skip $ ParseChunksYield b (ParseChunksInit xs s)
1126 assert (n <= length (x:buf)) (return ())
1127 let src = Prelude.reverse (Prelude.take n (x:buf)) ++ xs
1128 return $ Skip $ ParseChunksYield b (ParseChunksInit src s)
1129 PR.Error err -> throwM $ ParseError err
1131 stepOuter _ (ParseChunksYield a next) = return $ Yield a next
1133{-# ANN type ConcatParseState Fuse #-}
1134data ConcatParseState x inpBuf st p =
1135 ConcatParseInit inpBuf st p
1136 | ConcatParseInitLeftOver inpBuf
1137 | ConcatParseStream st inpBuf p
1138 | ConcatParseBuf inpBuf st inpBuf p
1139 | ConcatParseYield x (ConcatParseState x inpBuf st p)
1141{-# INLINE_NORMAL parseIterate #-}
1144 => (b -> PRD.Parser m a b)
1148parseIterate func seed (Stream step state) =
1149 Stream stepOuter (ConcatParseInit [] state (func seed))
1153 {-# INLINE_LATE stepOuter #-}
1154 -- Buffer is empty, go to stream processing loop
1155 stepOuter _ (ConcatParseInit [] st (PRD.Parser pstep initial extract)) = do
1156 initial >>= \r -> return $ Skip $ ConcatParseStream st []
1157 (PRD.Parser pstep (return r) extract)
1159 -- Buffer is not empty, go to buffered processing loop
1160 stepOuter _ (ConcatParseInit src st
1161 (PRD.Parser pstep initial extract)) = do
1162 initial >>= \r -> return $ Skip $ ConcatParseBuf src st []
1163 (PRD.Parser pstep (return r) extract)
1165 -- XXX we just discard any leftover input at the end
1166 stepOuter _ (ConcatParseInitLeftOver _) = return Stop
1168 -- Buffer is empty process elements from the stream
1169 stepOuter gst (ConcatParseStream st buf
1170 p@(PRD.Parser pstep initial extract)) = do
1171 r <- step (adaptState gst) st
1177 PR.Partial 0 pst1 ->
1178 return $ Skip $ ConcatParseStream s []
1179 (PRD.Parser pstep (return pst1) extract)
1180 PR.Partial n pst1 -> do
1181 assert (n <= length (x:buf)) (return ())
1182 let src0 = Prelude.take n (x:buf)
1183 src = Prelude.reverse src0
1184 return $ Skip $ ConcatParseBuf src s []
1185 (PRD.Parser pstep (return pst1) extract)
1186 -- PR.Continue 0 pst1 ->
1187 -- return $ Skip $ ConcatParseStream s (x:buf) pst1
1188 PR.Continue n pst1 -> do
1189 assert (n <= length (x:buf)) (return ())
1190 let (src0, buf1) = splitAt n (x:buf)
1191 src = Prelude.reverse src0
1192 return $ Skip $ ConcatParseBuf src s buf1
1193 (PRD.Parser pstep (return pst1) extract)
1194 -- XXX Specialize for Stop 0 common case?
1196 assert (n <= length (x:buf)) (return ())
1197 let src = Prelude.reverse (Prelude.take n (x:buf))
1199 ConcatParseYield b (ConcatParseInit src s (func b))
1200 PR.Error err -> throwM $ ParseError err
1201 Skip s -> return $ Skip $ ConcatParseStream s buf p
1205 let src = Prelude.reverse buf
1207 ConcatParseYield b (ConcatParseInitLeftOver src)
1209 -- go back to stream processing mode
1210 stepOuter _ (ConcatParseBuf [] s buf p) =
1211 return $ Skip $ ConcatParseStream s buf p
1213 -- buffered processing loop
1214 stepOuter _ (ConcatParseBuf (x:xs) s buf
1215 (PRD.Parser pstep initial extract)) = do
1219 PR.Partial 0 pst1 ->
1220 return $ Skip $ ConcatParseBuf xs s []
1221 (PRD.Parser pstep (return pst1) extract)
1222 PR.Partial n pst1 -> do
1223 assert (n <= length (x:buf)) (return ())
1224 let src0 = Prelude.take n (x:buf)
1225 src = Prelude.reverse src0 ++ xs
1226 return $ Skip $ ConcatParseBuf src s []
1227 (PRD.Parser pstep (return pst1) extract)
1228 -- PR.Continue 0 pst1 -> return $ Skip $ ConcatParseBuf xs s (x:buf) pst1
1229 PR.Continue n pst1 -> do
1230 assert (n <= length (x:buf)) (return ())
1231 let (src0, buf1) = splitAt n (x:buf)
1232 src = Prelude.reverse src0 ++ xs
1233 return $ Skip $ ConcatParseBuf src s buf1
1234 (PRD.Parser pstep (return pst1) extract)
1235 -- XXX Specialize for Stop 0 common case?
1237 assert (n <= length (x:buf)) (return ())
1238 let src = Prelude.reverse (Prelude.take n (x:buf)) ++ xs
1239 return $ Skip $ ConcatParseYield b
1240 (ConcatParseInit src s (func b))
1241 PR.Error err -> throwM $ ParseError err
1243 stepOuter _ (ConcatParseYield a next) = return $ Yield a next
1245------------------------------------------------------------------------------
1247------------------------------------------------------------------------------
1249-- | Run a streaming composition, discard the results.
1250{-# INLINE_LATE drain #-}
1251drain :: Monad m => Stream m a -> m ()
1252-- drain = foldrM (\_ xs -> xs) (return ())
1253drain (Stream step state) = go SPEC state
1256 r <- step defState st
1258 Yield _ s -> go SPEC s
1262{-# INLINE_NORMAL null #-}
1263null :: Monad m => Stream m a -> m Bool
1264null m = foldrM (\_ _ -> return False) (return True) m
1266{-# INLINE_NORMAL head #-}
1267head :: Monad m => Stream m a -> m (Maybe a)
1268head m = foldrM (\x _ -> return (Just x)) (return Nothing) m
1270{-# INLINE_NORMAL headElse #-}
1271headElse :: Monad m => a -> Stream m a -> m a
1272headElse a m = foldrM (\x _ -> return x) (return a) m
1274-- Does not fuse, has the same performance as the StreamK version.
1275{-# INLINE_NORMAL tail #-}
1276tail :: Monad m => Stream m a -> m (Maybe (Stream m a))
1277tail (UnStream step state) = go state
1280 r <- step defState st
1282 Yield _ s -> return (Just $ Stream step s)
1284 Stop -> return Nothing
1286-- XXX will it fuse? need custom impl?
1287{-# INLINE_NORMAL last #-}
1288last :: Monad m => Stream m a -> m (Maybe a)
1289last = foldl' (\_ y -> Just y) Nothing
1291{-# INLINE_NORMAL elem #-}
1292elem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
1293-- elem e m = foldrM (\x xs -> if x == e then return True else xs) (return False) m
1294elem e (Stream step state) = go state
1297 r <- step defState st
1300 | x == e -> return True
1303 Stop -> return False
1305{-# INLINE_NORMAL notElem #-}
1306notElem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
1307notElem e s = fmap not (elem e s)
1309{-# INLINE_NORMAL all #-}
1310all :: Monad m => (a -> Bool) -> Stream m a -> m Bool
1311-- all p m = foldrM (\x xs -> if p x then xs else return False) (return True) m
1312all p (Stream step state) = go state
1315 r <- step defState st
1319 | otherwise -> return False
1323{-# INLINE_NORMAL any #-}
1324any :: Monad m => (a -> Bool) -> Stream m a -> m Bool
1325-- any p m = foldrM (\x xs -> if p x then return True else xs) (return False) m
1326any p (Stream step state) = go state
1329 r <- step defState st
1332 | p x -> return True
1335 Stop -> return False
1337{-# INLINE_NORMAL maximum #-}
1338maximum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
1339maximum (Stream step state) = go Nothing state
1342 r <- step defState st
1344 Yield x s -> go (Just x) s
1345 Skip s -> go Nothing s
1346 Stop -> return Nothing
1347 go (Just acc) st = do
1348 r <- step defState st
1351 | acc <= x -> go (Just x) s
1352 | otherwise -> go (Just acc) s
1353 Skip s -> go (Just acc) s
1354 Stop -> return (Just acc)
1356{-# INLINE_NORMAL maximumBy #-}
1357maximumBy :: Monad m => (a -> a -> Ordering) -> Stream m a -> m (Maybe a)
1358maximumBy cmp (Stream step state) = go Nothing state
1361 r <- step defState st
1363 Yield x s -> go (Just x) s
1364 Skip s -> go Nothing s
1365 Stop -> return Nothing
1366 go (Just acc) st = do
1367 r <- step defState st
1369 Yield x s -> case cmp acc x of
1370 GT -> go (Just acc) s
1372 Skip s -> go (Just acc) s
1373 Stop -> return (Just acc)
1375{-# INLINE_NORMAL minimum #-}
1376minimum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
1377minimum (Stream step state) = go Nothing state
1380 r <- step defState st
1382 Yield x s -> go (Just x) s
1383 Skip s -> go Nothing s
1384 Stop -> return Nothing
1385 go (Just acc) st = do
1386 r <- step defState st
1389 | acc <= x -> go (Just acc) s
1390 | otherwise -> go (Just x) s
1391 Skip s -> go (Just acc) s
1392 Stop -> return (Just acc)
1394{-# INLINE_NORMAL minimumBy #-}
1395minimumBy :: Monad m => (a -> a -> Ordering) -> Stream m a -> m (Maybe a)
1396minimumBy cmp (Stream step state) = go Nothing state
1399 r <- step defState st
1401 Yield x s -> go (Just x) s
1402 Skip s -> go Nothing s
1403 Stop -> return Nothing
1404 go (Just acc) st = do
1405 r <- step defState st
1407 Yield x s -> case cmp acc x of
1409 _ -> go (Just acc) s
1410 Skip s -> go (Just acc) s
1411 Stop -> return (Just acc)
1413{-# INLINE_NORMAL (!!) #-}
1414(!!) :: (Monad m) => Stream m a -> Int -> m (Maybe a)
1415(Stream step state) !! i = go i state
1418 r <- step defState st
1420 Yield x s | n < 0 -> return Nothing
1421 | n == 0 -> return $ Just x
1422 | otherwise -> go (n - 1) s
1424 Stop -> return Nothing
1426{-# INLINE_NORMAL lookup #-}
1427lookup :: (Monad m, Eq a) => a -> Stream m (a, b) -> m (Maybe b)
1428lookup e m = foldrM (\(a, b) xs -> if e == a then return (Just b) else xs)
1431{-# INLINE_NORMAL findM #-}
1432findM :: Monad m => (a -> m Bool) -> Stream m a -> m (Maybe a)
1433findM p m = foldrM (\x xs -> p x >>= \r -> if r then return (Just x) else xs)
1437find :: Monad m => (a -> Bool) -> Stream m a -> m (Maybe a)
1438find p = findM (return . p)
1440{-# INLINE_NORMAL findIndices #-}
1441findIndices :: Monad m => (a -> Bool) -> Stream m a -> Stream m Int
1442findIndices p (Stream step state) = Stream step' (state, 0)
1444 {-# INLINE_LATE step' #-}
1445 step' gst (st, i) = i `seq` do
1446 r <- step (adaptState gst) st
1448 Yield x s -> if p x then Yield i (s, i+1) else Skip (s, i+1)
1449 Skip s -> Skip (s, i)
1452{-# INLINE toListRev #-}
1453toListRev :: Monad m => Stream m a -> m [a]
1454toListRev = foldl' (flip (:)) []
1456-- We can implement reverse as:
1458-- > reverse = foldlS (flip cons) nil
1460-- However, this implementation is unusable because of the horrible performance
1461-- of cons. So we just convert it to a list first and then stream from the
1464-- XXX Maybe we can use an Array instead of a list here?
1465{-# INLINE_NORMAL reverse #-}
1466reverse :: Monad m => Stream m a -> Stream m a
1467reverse m = Stream step Nothing
1469 {-# INLINE_LATE step #-}
1472 return $ Skip (Just xs)
1473 step _ (Just (x:xs)) = return $ Yield x (Just xs)
1474 step _ (Just []) = return Stop
1476-- Much faster reverse for Storables
1477{-# INLINE_NORMAL reverse' #-}
1478reverse' :: forall m a. (MonadIO m, Storable a) => Stream m a -> Stream m a
1480-- This commented implementation copies the whole stream into one single array
1481-- and then streams from that array, this is 3-4x faster than the chunked code
1482-- that follows. Though this could be problematic due to unbounded large
1483-- allocations. We need to figure out why the chunked code is slower and if we
1484-- can optimize the chunked code to work as fast as this one. It may be a
1486import Foreign.ForeignPtr (touchForeignPtr)
1487import Foreign.ForeignPtr.Unsafe (unsafeForeignPtrToPtr)
1488import Foreign.Ptr (Ptr, plusPtr)
1489reverse' m = Stream step Nothing
1491 {-# INLINE_LATE step #-}
1493 arr <- A.fromStreamD m
1494 let p = aEnd arr `plusPtr` negate (sizeOf (undefined :: a))
1495 return $ Skip $ Just (aStart arr, p)
1497 step _ (Just (start, p)) | p < unsafeForeignPtrToPtr start = return Stop
1499 step _ (Just (start, p)) = do
1500 let !x = A.unsafeInlineIO $ do
1502 touchForeignPtr start
1504 next = p `plusPtr` negate (sizeOf (undefined :: a))
1505 return $ Yield x (Just (start, next))
1512 $ A.fromStreamDArraysOf A.defaultChunkSize m
1515------------------------------------------------------------------------------
1516-- Grouping/Splitting
1517------------------------------------------------------------------------------
1519{-# INLINE_NORMAL splitSuffixBy' #-}
1520splitSuffixBy' :: Monad m
1521 => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1522splitSuffixBy' predicate f (Stream step state) =
1523 Stream (stepOuter f) (Just state)
1527 {-# INLINE_LATE stepOuter #-}
1528 stepOuter (Fold fstep initial done) gst (Just st) = do
1529 res <- step (adaptState gst) st
1535 then done acc' >>= \val -> return $ Yield val (Just s)
1538 Skip s -> return $ Skip $ Just s
1544 res <- step (adaptState gst) stt
1549 then done acc' >>= \val -> return $ Yield val (Just s)
1551 Skip s -> go SPEC s acc
1552 Stop -> done acc >>= \val -> return $ Yield val Nothing
1554 stepOuter _ _ Nothing = return Stop
1556{-# INLINE_NORMAL groupsBy #-}
1562groupsBy cmp f (Stream step state) = Stream (stepOuter f) (Just state, Nothing)
1566 {-# INLINE_LATE stepOuter #-}
1567 stepOuter (Fold fstep initial done) gst (Just st, Nothing) = do
1568 res <- step (adaptState gst) st
1575 Skip s -> return $ Skip $ (Just s, Nothing)
1580 go !_ prev stt !acc = do
1581 res <- step (adaptState gst) stt
1588 else done acc >>= \r -> return $ Yield r (Just s, Just x)
1589 Skip s -> go SPEC prev s acc
1590 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1592 stepOuter (Fold fstep initial done) gst (Just st, Just prev) = do
1594 acc' <- fstep acc prev
1599 -- XXX code duplicated from the previous equation
1601 res <- step (adaptState gst) stt
1608 else done acc >>= \r -> return $ Yield r (Just s, Just x)
1609 Skip s -> go SPEC s acc
1610 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1612 stepOuter _ _ (Nothing,_) = return Stop
1614{-# INLINE_NORMAL groupsRollingBy #-}
1615groupsRollingBy :: Monad m
1620groupsRollingBy cmp f (Stream step state) =
1621 Stream (stepOuter f) (Just state, Nothing)
1624 {-# INLINE_LATE stepOuter #-}
1625 stepOuter (Fold fstep initial done) gst (Just st, Nothing) = do
1626 res <- step (adaptState gst) st
1633 Skip s -> return $ Skip $ (Just s, Nothing)
1637 go !_ prev stt !acc = do
1638 res <- step (adaptState gst) stt
1646 done acc >>= \r -> return $ Yield r (Just s, Just x)
1647 Skip s -> go SPEC prev s acc
1648 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1650 stepOuter (Fold fstep initial done) gst (Just st, Just prev') = do
1652 acc' <- fstep acc prev'
1653 go SPEC prev' st acc'
1656 go !_ prevv stt !acc = do
1657 res <- step (adaptState gst) stt
1664 else done acc >>= \r -> return $ Yield r (Just s, Just x)
1665 Skip s -> go SPEC prevv s acc
1666 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1668 stepOuter _ _ (Nothing, _) = return Stop
1670{-# INLINE_NORMAL splitBy #-}
1671splitBy :: Monad m => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1672splitBy predicate f (Stream step state) = Stream (step' f) (Just state)
1676 {-# INLINE_LATE step' #-}
1677 step' (Fold fstep initial done) gst (Just st) = initial >>= go SPEC st
1682 res <- step (adaptState gst) stt
1686 then done acc >>= \r -> return $ Yield r (Just s)
1690 Skip s -> go SPEC s acc
1691 Stop -> done acc >>= \r -> return $ Yield r Nothing
1693 step' _ _ Nothing = return Stop
1695-- XXX requires -funfolding-use-threshold=150 in lines-unlines benchmark
1696{-# INLINE_NORMAL splitSuffixBy #-}
1697splitSuffixBy :: Monad m
1698 => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1699splitSuffixBy predicate f (Stream step state) = Stream (step' f) (Just state)
1703 {-# INLINE_LATE step' #-}
1704 step' (Fold fstep initial done) gst (Just st) = do
1705 res <- step (adaptState gst) st
1710 then done acc >>= \val -> return $ Yield val (Just s)
1715 Skip s -> return $ Skip $ Just s
1721 res <- step (adaptState gst) stt
1725 then done acc >>= \r -> return $ Yield r (Just s)
1729 Skip s -> go SPEC s acc
1730 Stop -> done acc >>= \r -> return $ Yield r Nothing
1732 step' _ _ Nothing = return Stop
1734{-# INLINE_NORMAL wordsBy #-}
1735wordsBy :: Monad m => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1736wordsBy predicate f (Stream step state) = Stream (stepOuter f) (Just state)
1740 {-# INLINE_LATE stepOuter #-}
1741 stepOuter (Fold fstep initial done) gst (Just st) = do
1742 res <- step (adaptState gst) st
1746 then return $ Skip (Just s)
1752 Skip s -> return $ Skip $ Just s
1758 res <- step (adaptState gst) stt
1762 then done acc >>= \r -> return $ Yield r (Just s)
1766 Skip s -> go SPEC s acc
1767 Stop -> done acc >>= \r -> return $ Yield r Nothing
1769 stepOuter _ _ Nothing = return Stop
1771-- String search algorithms:
1772-- http://www-igm.univ-mlv.fr/~lecroq/string/index.html
1775-- TODO can we unify the splitting operations using a splitting configuration
1776-- like in the split package.
1778data SplitStyle = Infix | Suffix | Prefix deriving (Eq, Show)
1780data SplitOptions = SplitOptions
1781 { style :: SplitStyle
1782 , withSep :: Bool -- ^ keep the separators in output
1783 -- , compact :: Bool -- ^ treat multiple consecutive separators as one
1784 -- , trimHead :: Bool -- ^ drop blank at head
1785 -- , trimTail :: Bool -- ^ drop blank at tail
1789data SplitOnState s a =
1794 | GO_KARP_RABIN s !(RB.Ring a) !(Ptr a)
1797{-# INLINE_NORMAL splitOn #-}
1799 :: forall m a b. (MonadIO m, Storable a, Enum a, Eq a)
1804splitOn patArr (Fold fstep initial done) (Stream step state) =
1805 Stream stepOuter GO_START
1809 patLen = A.length patArr
1810 maxIndex = patLen - 1
1811 elemBits = sizeOf (undefined :: a) * 8
1813 {-# INLINE_LATE stepOuter #-}
1814 stepOuter _ GO_START =
1816 then return $ Skip $ GO_EMPTY_PAT state
1819 r <- liftIO $ (A.unsafeIndexIO patArr 0)
1820 return $ Skip $ GO_SINGLE_PAT state r
1821 else if sizeOf (undefined :: a) * patLen
1822 <= sizeOf (undefined :: Word)
1823 then return $ Skip $ GO_SHORT_PAT state
1825 (rb, rhead) <- liftIO $ RB.new patLen
1826 return $ Skip $ GO_KARP_RABIN state rb rhead
1828 stepOuter gst (GO_SINGLE_PAT stt pat) = initial >>= go SPEC stt
1833 res <- step (adaptState gst) st
1839 return $ Yield r (GO_SINGLE_PAT s pat)
1840 else fstep acc x >>= go SPEC s
1841 Skip s -> go SPEC s acc
1842 Stop -> done acc >>= \r -> return $ Yield r GO_DONE
1844 stepOuter gst (GO_SHORT_PAT stt) = initial >>= go0 SPEC 0 (0 :: Word) stt
1849 mask = (1 `shiftL` (elemBits * patLen)) - 1
1851 addToWord wrd a = (wrd `shiftL` elemBits) .|. fromIntegral (fromEnum a)
1854 patWord = mask .&. A.foldl' addToWord 0 patArr
1856 go0 !_ !idx wrd st !acc = do
1857 res <- step (adaptState gst) st
1860 let wrd' = addToWord wrd x
1863 if wrd' .&. mask == patWord
1866 return $ Yield r (GO_SHORT_PAT s)
1867 else go1 SPEC wrd' s acc
1868 else go0 SPEC (idx + 1) wrd' s acc
1869 Skip s -> go0 SPEC idx wrd s acc
1872 then go2 wrd idx acc
1874 done acc' >>= \r -> return $ Yield r GO_DONE
1877 go1 !_ wrd st !acc = do
1878 res <- step (adaptState gst) st
1881 let wrd' = addToWord wrd x
1882 old = (mask .&. wrd) `shiftR` (elemBits * (patLen - 1))
1883 acc' <- fstep acc (toEnum $ fromIntegral old)
1884 if wrd' .&. mask == patWord
1885 then done acc' >>= \r -> return $ Yield r (GO_SHORT_PAT s)
1886 else go1 SPEC wrd' s acc'
1887 Skip s -> go1 SPEC wrd s acc
1889 acc' <- go2 wrd patLen acc
1890 done acc' >>= \r -> return $ Yield r GO_DONE
1892 go2 !wrd !n !acc | n > 0 = do
1893 let old = (mask .&. wrd) `shiftR` (elemBits * (n - 1))
1894 fstep acc (toEnum $ fromIntegral old) >>= go2 wrd (n - 1)
1895 go2 _ _ acc = return acc
1897 stepOuter gst (GO_KARP_RABIN stt rb rhead) = do
1898 initial >>= go0 SPEC 0 rhead stt
1902 k = 2891336453 :: Word32
1904 addCksum cksum a = cksum * k + fromIntegral (fromEnum a)
1905 deltaCksum cksum old new =
1906 addCksum cksum new - coeff * fromIntegral (fromEnum old)
1908 -- XXX shall we use a random starting hash or 1 instead of 0?
1909 patHash = A.foldl' addCksum 0 patArr
1912 go0 !_ !idx !rh st !acc = do
1913 res <- step (adaptState gst) st
1916 rh' <- liftIO $ RB.unsafeInsert rb rh x
1919 let fold = RB.unsafeFoldRing (RB.ringBound rb)
1920 let !ringHash = fold addCksum 0 rb
1921 if ringHash == patHash
1922 then go2 SPEC ringHash rh' s acc
1923 else go1 SPEC ringHash rh' s acc
1924 else go0 SPEC (idx + 1) rh' s acc
1925 Skip s -> go0 SPEC idx rh s acc
1927 !acc' <- if idx /= 0
1928 then RB.unsafeFoldRingM rh fstep acc rb
1930 done acc' >>= \r -> return $ Yield r GO_DONE
1932 -- XXX Theoretically this code can do 4 times faster if GHC generates
1933 -- optimal code. If we use just "(cksum' == patHash)" condition it goes
1934 -- 4x faster, as soon as we add the "RB.unsafeEqArray rb v" condition
1935 -- the generated code changes drastically and becomes 4x slower. Need
1936 -- to investigate what is going on with GHC.
1938 go1 !_ !cksum !rh st !acc = do
1939 res <- step (adaptState gst) st
1942 old <- liftIO $ peek rh
1943 let cksum' = deltaCksum cksum old x
1944 acc' <- fstep acc old
1946 if (cksum' == patHash)
1948 rh' <- liftIO (RB.unsafeInsert rb rh x)
1949 go2 SPEC cksum' rh' s acc'
1951 rh' <- liftIO (RB.unsafeInsert rb rh x)
1952 go1 SPEC cksum' rh' s acc'
1953 Skip s -> go1 SPEC cksum rh s acc
1955 acc' <- RB.unsafeFoldRingFullM rh fstep acc rb
1956 done acc' >>= \r -> return $ Yield r GO_DONE
1958 go2 !_ !cksum' !rh' s !acc' = do
1959 if RB.unsafeEqArray rb rh' patArr
1962 return $ Yield r (GO_KARP_RABIN s rb rhead)
1963 else go1 SPEC cksum' rh' s acc'
1965 stepOuter gst (GO_EMPTY_PAT st) = do
1966 res <- step (adaptState gst) st
1971 done acc' >>= \r -> return $ Yield r (GO_EMPTY_PAT s)
1972 Skip s -> return $ Skip (GO_EMPTY_PAT s)
1975 stepOuter _ GO_DONE = return Stop
1977{-# INLINE_NORMAL splitSuffixOn #-}
1979 :: forall m a b. (MonadIO m, Storable a, Enum a, Eq a)
1985splitSuffixOn withSep patArr (Fold fstep initial done)
1986 (Stream step state) =
1987 Stream stepOuter GO_START
1991 patLen = A.length patArr
1992 maxIndex = patLen - 1
1993 elemBits = sizeOf (undefined :: a) * 8
1995 {-# INLINE_LATE stepOuter #-}
1996 stepOuter _ GO_START =
1998 then return $ Skip $ GO_EMPTY_PAT state
2001 r <- liftIO $ (A.unsafeIndexIO patArr 0)
2002 return $ Skip $ GO_SINGLE_PAT state r
2003 else if sizeOf (undefined :: a) * patLen
2004 <= sizeOf (undefined :: Word)
2005 then return $ Skip $ GO_SHORT_PAT state
2007 (rb, rhead) <- liftIO $ RB.new patLen
2008 return $ Skip $ GO_KARP_RABIN state rb rhead
2010 stepOuter gst (GO_SINGLE_PAT stt pat) = do
2011 -- This first part is the only difference between splitOn and
2013 -- If the last element is a separator do not issue a blank segment.
2014 res <- step (adaptState gst) stt
2020 acc' <- if withSep then fstep acc x else return acc
2021 done acc' >>= \r -> return $ Yield r (GO_SINGLE_PAT s pat)
2022 else fstep acc x >>= go SPEC s
2023 Skip s -> return $ Skip $ (GO_SINGLE_PAT s pat)
2028 -- This is identical for splitOn and splitSuffixOn
2030 res <- step (adaptState gst) st
2035 acc' <- if withSep then fstep acc x else return acc
2037 return $ Yield r (GO_SINGLE_PAT s pat)
2038 else fstep acc x >>= go SPEC s
2039 Skip s -> go SPEC s acc
2040 Stop -> done acc >>= \r -> return $ Yield r GO_DONE
2042 stepOuter gst (GO_SHORT_PAT stt) = do
2044 -- Call "initial" only if the stream yields an element, otherwise we
2045 -- may call "initial" but never yield anything. initial may produce a
2046 -- side effect, therefore we will end up doing and discard a side
2051 res <- step (adaptState gst) stt
2055 let wrd' = addToWord wrd x
2056 acc' <- if withSep then fstep acc x else return acc
2059 if wrd' .&. mask == patWord
2060 then done acc' >>= \r -> return $ Yield r (GO_SHORT_PAT s)
2061 else go0 SPEC (idx + 1) wrd' s acc'
2062 else go0 SPEC (idx + 1) wrd' s acc'
2063 Skip s -> return $ Skip (GO_SHORT_PAT s)
2069 mask = (1 `shiftL` (elemBits * patLen)) - 1
2071 addToWord wrd a = (wrd `shiftL` elemBits) .|. fromIntegral (fromEnum a)
2074 patWord = mask .&. A.foldl' addToWord 0 patArr
2076 go0 !_ !idx wrd st !acc = do
2077 res <- step (adaptState gst) st
2080 let wrd' = addToWord wrd x
2081 acc' <- if withSep then fstep acc x else return acc
2084 if wrd' .&. mask == patWord
2087 return $ Yield r (GO_SHORT_PAT s)
2088 else go1 SPEC wrd' s acc'
2089 else go0 SPEC (idx + 1) wrd' s acc'
2090 Skip s -> go0 SPEC idx wrd s acc
2092 if (idx == maxIndex) && (wrd .&. mask == patWord)
2095 acc' <- if idx /= 0 && not withSep
2096 then go2 wrd idx acc
2098 done acc' >>= \r -> return $ Yield r GO_DONE
2101 go1 !_ wrd st !acc = do
2102 res <- step (adaptState gst) st
2105 let wrd' = addToWord wrd x
2106 old = (mask .&. wrd) `shiftR` (elemBits * (patLen - 1))
2109 else fstep acc (toEnum $ fromIntegral old)
2110 if wrd' .&. mask == patWord
2111 then done acc' >>= \r -> return $ Yield r (GO_SHORT_PAT s)
2112 else go1 SPEC wrd' s acc'
2113 Skip s -> go1 SPEC wrd s acc
2115 -- If the last sequence is a separator do not issue a blank
2117 if wrd .&. mask == patWord
2122 else go2 wrd patLen acc
2123 done acc' >>= \r -> return $ Yield r GO_DONE
2125 go2 !wrd !n !acc | n > 0 = do
2126 let old = (mask .&. wrd) `shiftR` (elemBits * (n - 1))
2127 fstep acc (toEnum $ fromIntegral old) >>= go2 wrd (n - 1)
2128 go2 _ _ acc = return acc
2130 stepOuter gst (GO_KARP_RABIN stt rb rhead) = do
2132 res <- step (adaptState gst) stt
2136 acc' <- if withSep then fstep acc x else return acc
2137 rh' <- liftIO (RB.unsafeInsert rb rhead x)
2140 let fold = RB.unsafeFoldRing (RB.ringBound rb)
2141 let !ringHash = fold addCksum 0 rb
2142 if ringHash == patHash
2143 then go2 SPEC ringHash rh' s acc'
2144 else go0 SPEC (idx + 1) rh' s acc'
2145 else go0 SPEC (idx + 1) rh' s acc'
2146 Skip s -> return $ Skip (GO_KARP_RABIN s rb rhead)
2151 k = 2891336453 :: Word32
2153 addCksum cksum a = cksum * k + fromIntegral (fromEnum a)
2154 deltaCksum cksum old new =
2155 addCksum cksum new - coeff * fromIntegral (fromEnum old)
2157 -- XXX shall we use a random starting hash or 1 instead of 0?
2158 patHash = A.foldl' addCksum 0 patArr
2161 go0 !_ !idx !rh st !acc = do
2162 res <- step (adaptState gst) st
2165 acc' <- if withSep then fstep acc x else return acc
2166 rh' <- liftIO (RB.unsafeInsert rb rh x)
2169 let fold = RB.unsafeFoldRing (RB.ringBound rb)
2170 let !ringHash = fold addCksum 0 rb
2171 if ringHash == patHash
2172 then go2 SPEC ringHash rh' s acc'
2173 else go1 SPEC ringHash rh' s acc'
2174 else go0 SPEC (idx + 1) rh' s acc'
2175 Skip s -> go0 SPEC idx rh s acc
2177 -- do not issue a blank segment when we end at pattern
2178 if (idx == maxIndex) && RB.unsafeEqArray rb rh patArr
2181 !acc' <- if idx /= 0 && not withSep
2182 then RB.unsafeFoldRingM rh fstep acc rb
2184 done acc' >>= \r -> return $ Yield r GO_DONE
2186 -- XXX Theoretically this code can do 4 times faster if GHC generates
2187 -- optimal code. If we use just "(cksum' == patHash)" condition it goes
2188 -- 4x faster, as soon as we add the "RB.unsafeEqArray rb v" condition
2189 -- the generated code changes drastically and becomes 4x slower. Need
2190 -- to investigate what is going on with GHC.
2192 go1 !_ !cksum !rh st !acc = do
2193 res <- step (adaptState gst) st
2196 old <- liftIO $ peek rh
2197 let cksum' = deltaCksum cksum old x
2202 if (cksum' == patHash)
2204 rh' <- liftIO (RB.unsafeInsert rb rh x)
2205 go2 SPEC cksum' rh' s acc'
2207 rh' <- liftIO (RB.unsafeInsert rb rh x)
2208 go1 SPEC cksum' rh' s acc'
2209 Skip s -> go1 SPEC cksum rh s acc
2211 if RB.unsafeEqArray rb rh patArr
2216 else RB.unsafeFoldRingFullM rh fstep acc rb
2217 done acc' >>= \r -> return $ Yield r GO_DONE
2219 go2 !_ !cksum' !rh' s !acc' = do
2220 if RB.unsafeEqArray rb rh' patArr
2223 return $ Yield r (GO_KARP_RABIN s rb rhead)
2224 else go1 SPEC cksum' rh' s acc'
2226 stepOuter gst (GO_EMPTY_PAT st) = do
2227 res <- step (adaptState gst) st
2232 done acc' >>= \r -> return $ Yield r (GO_EMPTY_PAT s)
2233 Skip s -> return $ Skip (GO_EMPTY_PAT s)
2236 stepOuter _ GO_DONE = return Stop
2238data SplitState s arr
2240 | SplitBuffering s arr
2241 | SplitSplitting s arr
2242 | SplitYielding arr (SplitState s arr)
2245-- XXX An alternative approach would be to use a partial fold (Fold m a b) to
2246-- split using a splitBy like combinator. The Fold would consume upto the
2247-- separator and return any leftover which can then be fed to the next fold.
2249-- We can revisit this once we have partial folds/parsers.
2251-- | Performs infix separator style splitting.
2252{-# INLINE_NORMAL splitInnerBy #-}
2255 => (f a -> m (f a, Maybe (f a))) -- splitter
2256 -> (f a -> f a -> m (f a)) -- joiner
2259splitInnerBy splitter joiner (Stream step1 state1) =
2260 (Stream step (SplitInitial state1))
2264 {-# INLINE_LATE step #-}
2265 step gst (SplitInitial st) = do
2269 (x1, mx2) <- splitter x
2270 return $ case mx2 of
2271 Nothing -> Skip (SplitBuffering s x1)
2272 Just x2 -> Skip (SplitYielding x1 (SplitSplitting s x2))
2273 Skip s -> return $ Skip (SplitInitial s)
2274 Stop -> return $ Stop
2276 step gst (SplitBuffering st buf) = do
2280 (x1, mx2) <- splitter x
2281 buf' <- joiner buf x1
2282 return $ case mx2 of
2283 Nothing -> Skip (SplitBuffering s buf')
2284 Just x2 -> Skip (SplitYielding buf' (SplitSplitting s x2))
2285 Skip s -> return $ Skip (SplitBuffering s buf)
2286 Stop -> return $ Skip (SplitYielding buf SplitFinishing)
2288 step _ (SplitSplitting st buf) = do
2289 (x1, mx2) <- splitter buf
2290 return $ case mx2 of
2291 Nothing -> Skip $ SplitBuffering st x1
2292 Just x2 -> Skip $ SplitYielding x1 (SplitSplitting st x2)
2294 step _ (SplitYielding x next) = return $ Yield x next
2295 step _ SplitFinishing = return $ Stop
2297-- | Performs infix separator style splitting.
2298{-# INLINE_NORMAL splitInnerBySuffix #-}
2300 :: (Monad m, Eq (f a), Monoid (f a))
2301 => (f a -> m (f a, Maybe (f a))) -- splitter
2302 -> (f a -> f a -> m (f a)) -- joiner
2305splitInnerBySuffix splitter joiner (Stream step1 state1) =
2306 (Stream step (SplitInitial state1))
2310 {-# INLINE_LATE step #-}
2311 step gst (SplitInitial st) = do
2315 (x1, mx2) <- splitter x
2316 return $ case mx2 of
2317 Nothing -> Skip (SplitBuffering s x1)
2318 Just x2 -> Skip (SplitYielding x1 (SplitSplitting s x2))
2319 Skip s -> return $ Skip (SplitInitial s)
2320 Stop -> return $ Stop
2322 step gst (SplitBuffering st buf) = do
2326 (x1, mx2) <- splitter x
2327 buf' <- joiner buf x1
2328 return $ case mx2 of
2329 Nothing -> Skip (SplitBuffering s buf')
2330 Just x2 -> Skip (SplitYielding buf' (SplitSplitting s x2))
2331 Skip s -> return $ Skip (SplitBuffering s buf)
2335 else Skip (SplitYielding buf SplitFinishing)
2337 step _ (SplitSplitting st buf) = do
2338 (x1, mx2) <- splitter buf
2339 return $ case mx2 of
2340 Nothing -> Skip $ SplitBuffering st x1
2341 Just x2 -> Skip $ SplitYielding x1 (SplitSplitting st x2)
2343 step _ (SplitYielding x next) = return $ Yield x next
2344 step _ SplitFinishing = return $ Stop
2346------------------------------------------------------------------------------
2348------------------------------------------------------------------------------
2350{-# INLINE_NORMAL isPrefixOf #-}
2351isPrefixOf :: (Eq a, Monad m) => Stream m a -> Stream m a -> m Bool
2352isPrefixOf (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
2354 go (sa, sb, Nothing) = do
2355 r <- stepa defState sa
2357 Yield x sa' -> go (sa', sb, Just x)
2358 Skip sa' -> go (sa', sb, Nothing)
2361 go (sa, sb, Just x) = do
2362 r <- stepb defState sb
2366 then go (sa, sb', Nothing)
2368 Skip sb' -> go (sa, sb', Just x)
2369 Stop -> return False
2371{-# INLINE_NORMAL isSubsequenceOf #-}
2372isSubsequenceOf :: (Eq a, Monad m) => Stream m a -> Stream m a -> m Bool
2373isSubsequenceOf (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
2375 go (sa, sb, Nothing) = do
2376 r <- stepa defState sa
2378 Yield x sa' -> go (sa', sb, Just x)
2379 Skip sa' -> go (sa', sb, Nothing)
2382 go (sa, sb, Just x) = do
2383 r <- stepb defState sb
2387 then go (sa, sb', Nothing)
2388 else go (sa, sb', Just x)
2389 Skip sb' -> go (sa, sb', Just x)
2390 Stop -> return False
2392{-# INLINE_NORMAL stripPrefix #-}
2395 => Stream m a -> Stream m a -> m (Maybe (Stream m a))
2396stripPrefix (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
2398 go (sa, sb, Nothing) = do
2399 r <- stepa defState sa
2401 Yield x sa' -> go (sa', sb, Just x)
2402 Skip sa' -> go (sa', sb, Nothing)
2403 Stop -> return $ Just (Stream stepb sb)
2405 go (sa, sb, Just x) = do
2406 r <- stepb defState sb
2410 then go (sa, sb', Nothing)
2412 Skip sb' -> go (sa, sb', Just x)
2413 Stop -> return Nothing
2415------------------------------------------------------------------------------
2417------------------------------------------------------------------------------
2419-- | Execute a monadic action for each element of the 'Stream'
2420{-# INLINE_NORMAL mapM_ #-}
2421mapM_ :: Monad m => (a -> m b) -> Stream m a -> m ()
2422mapM_ m = drain . mapM m
2424-------------------------------------------------------------------------------
2425-- Stream transformations using Unfolds
2426-------------------------------------------------------------------------------
2428-- Define a unique structure to use in inspection testing
2429data ConcatMapUState o i =
2431 | ConcatMapUInner o i
2433-- | @concatMapU unfold stream@ uses @unfold@ to map the input stream elements
2434-- to streams and then flattens the generated streams into a single output
2437-- This is like 'concatMap' but uses an unfold with an explicit state to
2438-- generate the stream instead of a 'Stream' type generator. This allows better
2439-- optimization via fusion. This can be many times more efficient than
2442{-# INLINE_NORMAL concatMapU #-}
2443concatMapU :: Monad m => Unfold m a b -> Stream m a -> Stream m b
2444concatMapU (Unfold istep inject) (Stream ostep ost) =
2445 Stream step (ConcatMapUOuter ost)
2447 {-# INLINE_LATE step #-}
2448 step gst (ConcatMapUOuter o) = do
2449 r <- ostep (adaptState gst) o
2453 i `seq` return (Skip (ConcatMapUInner o' i))
2454 Skip o' -> return $ Skip (ConcatMapUOuter o')
2455 Stop -> return $ Stop
2457 step _ (ConcatMapUInner o i) = do
2460 Yield x i' -> Yield x (ConcatMapUInner o i')
2461 Skip i' -> Skip (ConcatMapUInner o i')
2462 Stop -> Skip (ConcatMapUOuter o)
2464data ConcatUnfoldInterleaveState o i =
2465 ConcatUnfoldInterleaveOuter o [i]
2466 | ConcatUnfoldInterleaveInner o [i]
2467 | ConcatUnfoldInterleaveInnerL [i] [i]
2468 | ConcatUnfoldInterleaveInnerR [i] [i]
2470-- XXX use arrays to store state instead of lists.
2471-- XXX In general we can use different scheduling strategies e.g. how to
2472-- schedule the outer vs inner loop or assigning weights to different streams
2473-- or outer and inner loops.
2475-- After a yield, switch to the next stream. Do not switch streams on Skip.
2476-- Yield from outer stream switches to the inner stream.
2478-- There are two choices here, (1) exhaust the outer stream first and then
2479-- start yielding from the inner streams, this is much simpler to implement,
2480-- (2) yield at least one element from an inner stream before going back to
2481-- outer stream and opening the next stream from it.
2483-- Ideally, we need some scheduling bias to inner streams vs outer stream.
2484-- Maybe we can configure the behavior.
2486{-# INLINE_NORMAL concatUnfoldInterleave #-}
2487concatUnfoldInterleave :: Monad m => Unfold m a b -> Stream m a -> Stream m b
2488concatUnfoldInterleave (Unfold istep inject) (Stream ostep ost) =
2489 Stream step (ConcatUnfoldInterleaveOuter ost [])
2491 {-# INLINE_LATE step #-}
2492 step gst (ConcatUnfoldInterleaveOuter o ls) = do
2493 r <- ostep (adaptState gst) o
2497 i `seq` return (Skip (ConcatUnfoldInterleaveInner o' (i : ls)))
2498 Skip o' -> return $ Skip (ConcatUnfoldInterleaveOuter o' ls)
2499 Stop -> return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2501 step _ (ConcatUnfoldInterleaveInner _ []) = undefined
2502 step _ (ConcatUnfoldInterleaveInner o (st:ls)) = do
2505 Yield x s -> Yield x (ConcatUnfoldInterleaveOuter o (s:ls))
2506 Skip s -> Skip (ConcatUnfoldInterleaveInner o (s:ls))
2507 Stop -> Skip (ConcatUnfoldInterleaveOuter o ls)
2509 step _ (ConcatUnfoldInterleaveInnerL [] []) = return Stop
2510 step _ (ConcatUnfoldInterleaveInnerL [] rs) =
2511 return $ Skip (ConcatUnfoldInterleaveInnerR [] rs)
2513 step _ (ConcatUnfoldInterleaveInnerL (st:ls) rs) = do
2516 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerL ls (s:rs))
2517 Skip s -> Skip (ConcatUnfoldInterleaveInnerL (s:ls) rs)
2518 Stop -> Skip (ConcatUnfoldInterleaveInnerL ls rs)
2520 step _ (ConcatUnfoldInterleaveInnerR [] []) = return Stop
2521 step _ (ConcatUnfoldInterleaveInnerR ls []) =
2522 return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2524 step _ (ConcatUnfoldInterleaveInnerR ls (st:rs)) = do
2527 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerR (s:ls) rs)
2528 Skip s -> Skip (ConcatUnfoldInterleaveInnerR ls (s:rs))
2529 Stop -> Skip (ConcatUnfoldInterleaveInnerR ls rs)
2531-- XXX In general we can use different scheduling strategies e.g. how to
2532-- schedule the outer vs inner loop or assigning weights to different streams
2533-- or outer and inner loops.
2535-- This could be inefficient if the tasks are too small.
2537-- Compared to concatUnfoldInterleave this one switches streams on Skips.
2539{-# INLINE_NORMAL concatUnfoldRoundrobin #-}
2540concatUnfoldRoundrobin :: Monad m => Unfold m a b -> Stream m a -> Stream m b
2541concatUnfoldRoundrobin (Unfold istep inject) (Stream ostep ost) =
2542 Stream step (ConcatUnfoldInterleaveOuter ost [])
2544 {-# INLINE_LATE step #-}
2545 step gst (ConcatUnfoldInterleaveOuter o ls) = do
2546 r <- ostep (adaptState gst) o
2550 i `seq` return (Skip (ConcatUnfoldInterleaveInner o' (i : ls)))
2551 Skip o' -> return $ Skip (ConcatUnfoldInterleaveInner o' ls)
2552 Stop -> return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2554 step _ (ConcatUnfoldInterleaveInner o []) =
2555 return $ Skip (ConcatUnfoldInterleaveOuter o [])
2557 step _ (ConcatUnfoldInterleaveInner o (st:ls)) = do
2560 Yield x s -> Yield x (ConcatUnfoldInterleaveOuter o (s:ls))
2561 Skip s -> Skip (ConcatUnfoldInterleaveOuter o (s:ls))
2562 Stop -> Skip (ConcatUnfoldInterleaveOuter o ls)
2564 step _ (ConcatUnfoldInterleaveInnerL [] []) = return Stop
2565 step _ (ConcatUnfoldInterleaveInnerL [] rs) =
2566 return $ Skip (ConcatUnfoldInterleaveInnerR [] rs)
2568 step _ (ConcatUnfoldInterleaveInnerL (st:ls) rs) = do
2571 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerL ls (s:rs))
2572 Skip s -> Skip (ConcatUnfoldInterleaveInnerL ls (s:rs))
2573 Stop -> Skip (ConcatUnfoldInterleaveInnerL ls rs)
2575 step _ (ConcatUnfoldInterleaveInnerR [] []) = return Stop
2576 step _ (ConcatUnfoldInterleaveInnerR ls []) =
2577 return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2579 step _ (ConcatUnfoldInterleaveInnerR ls (st:rs)) = do
2582 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerR (s:ls) rs)
2583 Skip s -> Skip (ConcatUnfoldInterleaveInnerR (s:ls) rs)
2584 Stop -> Skip (ConcatUnfoldInterleaveInnerR ls rs)
2586data AppendState s1 s2 = AppendFirst s1 | AppendSecond s2
2588-- Note that this could be much faster compared to the CPS stream. However, as
2589-- the number of streams being composed increases this may become expensive.
2590-- Need to see where the breaking point is between the two.
2592{-# INLINE_NORMAL append #-}
2593append :: Monad m => Stream m a -> Stream m a -> Stream m a
2594append (Stream step1 state1) (Stream step2 state2) =
2595 Stream step (AppendFirst state1)
2599 {-# INLINE_LATE step #-}
2600 step gst (AppendFirst st) = do
2603 Yield a s -> Yield a (AppendFirst s)
2604 Skip s -> Skip (AppendFirst s)
2605 Stop -> Skip (AppendSecond state2)
2607 step gst (AppendSecond st) = do
2610 Yield a s -> Yield a (AppendSecond s)
2611 Skip s -> Skip (AppendSecond s)
2614data InterleaveState s1 s2 = InterleaveFirst s1 s2 | InterleaveSecond s1 s2
2615 | InterleaveSecondOnly s2 | InterleaveFirstOnly s1
2617{-# INLINE_NORMAL interleave #-}
2618interleave :: Monad m => Stream m a -> Stream m a -> Stream m a
2619interleave (Stream step1 state1) (Stream step2 state2) =
2620 Stream step (InterleaveFirst state1 state2)
2624 {-# INLINE_LATE step #-}
2625 step gst (InterleaveFirst st1 st2) = do
2628 Yield a s -> Yield a (InterleaveSecond s st2)
2629 Skip s -> Skip (InterleaveFirst s st2)
2630 Stop -> Skip (InterleaveSecondOnly st2)
2632 step gst (InterleaveSecond st1 st2) = do
2635 Yield a s -> Yield a (InterleaveFirst st1 s)
2636 Skip s -> Skip (InterleaveSecond st1 s)
2637 Stop -> Skip (InterleaveFirstOnly st1)
2639 step gst (InterleaveFirstOnly st1) = do
2642 Yield a s -> Yield a (InterleaveFirstOnly s)
2643 Skip s -> Skip (InterleaveFirstOnly s)
2646 step gst (InterleaveSecondOnly st2) = do
2649 Yield a s -> Yield a (InterleaveSecondOnly s)
2650 Skip s -> Skip (InterleaveSecondOnly s)
2653{-# INLINE_NORMAL interleaveMin #-}
2654interleaveMin :: Monad m => Stream m a -> Stream m a -> Stream m a
2655interleaveMin (Stream step1 state1) (Stream step2 state2) =
2656 Stream step (InterleaveFirst state1 state2)
2660 {-# INLINE_LATE step #-}
2661 step gst (InterleaveFirst st1 st2) = do
2664 Yield a s -> Yield a (InterleaveSecond s st2)
2665 Skip s -> Skip (InterleaveFirst s st2)
2668 step gst (InterleaveSecond st1 st2) = do
2671 Yield a s -> Yield a (InterleaveFirst st1 s)
2672 Skip s -> Skip (InterleaveSecond st1 s)
2675 step _ (InterleaveFirstOnly _) = undefined
2676 step _ (InterleaveSecondOnly _) = undefined
2678{-# INLINE_NORMAL interleaveSuffix #-}
2679interleaveSuffix :: Monad m => Stream m a -> Stream m a -> Stream m a
2680interleaveSuffix (Stream step1 state1) (Stream step2 state2) =
2681 Stream step (InterleaveFirst state1 state2)
2685 {-# INLINE_LATE step #-}
2686 step gst (InterleaveFirst st1 st2) = do
2689 Yield a s -> Yield a (InterleaveSecond s st2)
2690 Skip s -> Skip (InterleaveFirst s st2)
2693 step gst (InterleaveSecond st1 st2) = do
2696 Yield a s -> Yield a (InterleaveFirst st1 s)
2697 Skip s -> Skip (InterleaveSecond st1 s)
2698 Stop -> Skip (InterleaveFirstOnly st1)
2700 step gst (InterleaveFirstOnly st1) = do
2703 Yield a s -> Yield a (InterleaveFirstOnly s)
2704 Skip s -> Skip (InterleaveFirstOnly s)
2707 step _ (InterleaveSecondOnly _) = undefined
2709data InterleaveInfixState s1 s2 a
2710 = InterleaveInfixFirst s1 s2
2711 | InterleaveInfixSecondBuf s1 s2
2712 | InterleaveInfixSecondYield s1 s2 a
2713 | InterleaveInfixFirstYield s1 s2 a
2714 | InterleaveInfixFirstOnly s1
2716{-# INLINE_NORMAL interleaveInfix #-}
2717interleaveInfix :: Monad m => Stream m a -> Stream m a -> Stream m a
2718interleaveInfix (Stream step1 state1) (Stream step2 state2) =
2719 Stream step (InterleaveInfixFirst state1 state2)
2723 {-# INLINE_LATE step #-}
2724 step gst (InterleaveInfixFirst st1 st2) = do
2727 Yield a s -> Yield a (InterleaveInfixSecondBuf s st2)
2728 Skip s -> Skip (InterleaveInfixFirst s st2)
2731 step gst (InterleaveInfixSecondBuf st1 st2) = do
2734 Yield a s -> Skip (InterleaveInfixSecondYield st1 s a)
2735 Skip s -> Skip (InterleaveInfixSecondBuf st1 s)
2736 Stop -> Skip (InterleaveInfixFirstOnly st1)
2738 step gst (InterleaveInfixSecondYield st1 st2 x) = do
2741 Yield a s -> Yield x (InterleaveInfixFirstYield s st2 a)
2742 Skip s -> Skip (InterleaveInfixSecondYield s st2 x)
2745 step _ (InterleaveInfixFirstYield st1 st2 x) = do
2746 return $ Yield x (InterleaveInfixSecondBuf st1 st2)
2748 step gst (InterleaveInfixFirstOnly st1) = do
2751 Yield a s -> Yield a (InterleaveInfixFirstOnly s)
2752 Skip s -> Skip (InterleaveInfixFirstOnly s)
2755{-# INLINE_NORMAL roundRobin #-}
2756roundRobin :: Monad m => Stream m a -> Stream m a -> Stream m a
2757roundRobin (Stream step1 state1) (Stream step2 state2) =
2758 Stream step (InterleaveFirst state1 state2)
2762 {-# INLINE_LATE step #-}
2763 step gst (InterleaveFirst st1 st2) = do
2766 Yield a s -> Yield a (InterleaveSecond s st2)
2767 Skip s -> Skip (InterleaveSecond s st2)
2768 Stop -> Skip (InterleaveSecondOnly st2)
2770 step gst (InterleaveSecond st1 st2) = do
2773 Yield a s -> Yield a (InterleaveFirst st1 s)
2774 Skip s -> Skip (InterleaveFirst st1 s)
2775 Stop -> Skip (InterleaveFirstOnly st1)
2777 step gst (InterleaveSecondOnly st2) = do
2780 Yield a s -> Yield a (InterleaveSecondOnly s)
2781 Skip s -> Skip (InterleaveSecondOnly s)
2784 step gst (InterleaveFirstOnly st1) = do
2787 Yield a s -> Yield a (InterleaveFirstOnly s)
2788 Skip s -> Skip (InterleaveFirstOnly s)
2791data ICUState s1 s2 i1 i2 =
2796 | ICUFirstInner s1 s2 i1
2797 | ICUSecondInner s1 s2 i2
2798 | ICUFirstOnlyInner s1 i1
2799 | ICUSecondOnlyInner s2 i2
2801-- | Interleave streams (full streams, not the elements) unfolded from two
2802-- input streams and concat. Stop when the first stream stops. If the second
2803-- stream ends before the first one then first stream still keeps running alone
2804-- without any interleaving with the second stream.
2806-- [a1, a2, ... an] [b1, b2 ...]
2807-- => [streamA1, streamA2, ... streamAn] [streamB1, streamB2, ...]
2808-- => [streamA1, streamB1, streamA2...StreamAn, streamBn]
2809-- => [a11, a12, ...a1j, b11, b12, ...b1k, a21, a22, ...]
2811{-# INLINE_NORMAL gintercalateSuffix #-}
2814 => Unfold m a c -> Stream m a -> Unfold m b c -> Stream m b -> Stream m c
2816 (Unfold istep1 inject1) (Stream step1 state1)
2817 (Unfold istep2 inject2) (Stream step2 state2) =
2818 Stream step (ICUFirst state1 state2)
2822 {-# INLINE_LATE step #-}
2823 step gst (ICUFirst s1 s2) = do
2824 r <- step1 (adaptState gst) s1
2828 i `seq` return (Skip (ICUFirstInner s s2 i))
2829 Skip s -> return $ Skip (ICUFirst s s2)
2832 step gst (ICUFirstOnly s1) = do
2833 r <- step1 (adaptState gst) s1
2837 i `seq` return (Skip (ICUFirstOnlyInner s i))
2838 Skip s -> return $ Skip (ICUFirstOnly s)
2841 step _ (ICUFirstInner s1 s2 i1) = do
2844 Yield x i' -> Yield x (ICUFirstInner s1 s2 i')
2845 Skip i' -> Skip (ICUFirstInner s1 s2 i')
2846 Stop -> Skip (ICUSecond s1 s2)
2848 step _ (ICUFirstOnlyInner s1 i1) = do
2851 Yield x i' -> Yield x (ICUFirstOnlyInner s1 i')
2852 Skip i' -> Skip (ICUFirstOnlyInner s1 i')
2853 Stop -> Skip (ICUFirstOnly s1)
2855 step gst (ICUSecond s1 s2) = do
2856 r <- step2 (adaptState gst) s2
2860 i `seq` return (Skip (ICUSecondInner s1 s i))
2861 Skip s -> return $ Skip (ICUSecond s1 s)
2862 Stop -> return $ Skip (ICUFirstOnly s1)
2864 step _ (ICUSecondInner s1 s2 i2) = do
2867 Yield x i' -> Yield x (ICUSecondInner s1 s2 i')
2868 Skip i' -> Skip (ICUSecondInner s1 s2 i')
2869 Stop -> Skip (ICUFirst s1 s2)
2871 step _ (ICUSecondOnly _s2) = undefined
2872 step _ (ICUSecondOnlyInner _s2 _i2) = undefined
2874data InterposeSuffixState s1 i1 =
2875 InterposeSuffixFirst s1
2876 -- | InterposeSuffixFirstYield s1 i1
2877 | InterposeSuffixFirstInner s1 i1
2878 | InterposeSuffixSecond s1
2880-- Note that if an unfolded layer turns out to be nil we still emit the
2881-- separator effect. An alternate behavior could be to emit the separator
2882-- effect only if at least one element has been yielded by the unfolding.
2883-- However, that becomes a bit complicated, so we have chosen the former
2885{-# INLINE_NORMAL interposeSuffix #-}
2888 => m c -> Unfold m b c -> Stream m b -> Stream m c
2891 (Unfold istep1 inject1) (Stream step1 state1) =
2892 Stream step (InterposeSuffixFirst state1)
2896 {-# INLINE_LATE step #-}
2897 step gst (InterposeSuffixFirst s1) = do
2898 r <- step1 (adaptState gst) s1
2902 i `seq` return (Skip (InterposeSuffixFirstInner s i))
2903 -- i `seq` return (Skip (InterposeSuffixFirstYield s i))
2904 Skip s -> return $ Skip (InterposeSuffixFirst s)
2908 step _ (InterposeSuffixFirstYield s1 i1) = do
2911 Yield x i' -> Yield x (InterposeSuffixFirstInner s1 i')
2912 Skip i' -> Skip (InterposeSuffixFirstYield s1 i')
2913 Stop -> Skip (InterposeSuffixFirst s1)
2916 step _ (InterposeSuffixFirstInner s1 i1) = do
2919 Yield x i' -> Yield x (InterposeSuffixFirstInner s1 i')
2920 Skip i' -> Skip (InterposeSuffixFirstInner s1 i')
2921 Stop -> Skip (InterposeSuffixSecond s1)
2923 step _ (InterposeSuffixSecond s1) = do
2925 return $ Yield r (InterposeSuffixFirst s1)
2927data ICALState s1 s2 i1 i2 a =
2929 -- | ICALFirstYield s1 s2 i1
2930 | ICALFirstInner s1 s2 i1
2932 | ICALFirstOnlyInner s1 i1
2933 | ICALSecondInject s1 s2
2934 | ICALFirstInject s1 s2 i2
2935 -- | ICALFirstBuf s1 s2 i1 i2
2936 | ICALSecondInner s1 s2 i1 i2
2937 -- -- | ICALSecondInner s1 s2 i1 i2 a
2938 -- -- | ICALFirstResume s1 s2 i1 i2 a
2940-- | Interleave streams (full streams, not the elements) unfolded from two
2941-- input streams and concat. Stop when the first stream stops. If the second
2942-- stream ends before the first one then first stream still keeps running alone
2943-- without any interleaving with the second stream.
2945-- [a1, a2, ... an] [b1, b2 ...]
2946-- => [streamA1, streamA2, ... streamAn] [streamB1, streamB2, ...]
2947-- => [streamA1, streamB1, streamA2...StreamAn, streamBn]
2948-- => [a11, a12, ...a1j, b11, b12, ...b1k, a21, a22, ...]
2950{-# INLINE_NORMAL gintercalate #-}
2953 => Unfold m a c -> Stream m a -> Unfold m b c -> Stream m b -> Stream m c
2955 (Unfold istep1 inject1) (Stream step1 state1)
2956 (Unfold istep2 inject2) (Stream step2 state2) =
2957 Stream step (ICALFirst state1 state2)
2961 {-# INLINE_LATE step #-}
2962 step gst (ICALFirst s1 s2) = do
2963 r <- step1 (adaptState gst) s1
2967 i `seq` return (Skip (ICALFirstInner s s2 i))
2968 -- i `seq` return (Skip (ICALFirstYield s s2 i))
2969 Skip s -> return $ Skip (ICALFirst s s2)
2973 step _ (ICALFirstYield s1 s2 i1) = do
2976 Yield x i' -> Yield x (ICALFirstInner s1 s2 i')
2977 Skip i' -> Skip (ICALFirstYield s1 s2 i')
2978 Stop -> Skip (ICALFirst s1 s2)
2981 step _ (ICALFirstInner s1 s2 i1) = do
2984 Yield x i' -> Yield x (ICALFirstInner s1 s2 i')
2985 Skip i' -> Skip (ICALFirstInner s1 s2 i')
2986 Stop -> Skip (ICALSecondInject s1 s2)
2988 step gst (ICALFirstOnly s1) = do
2989 r <- step1 (adaptState gst) s1
2993 i `seq` return (Skip (ICALFirstOnlyInner s i))
2994 Skip s -> return $ Skip (ICALFirstOnly s)
2997 step _ (ICALFirstOnlyInner s1 i1) = do
3000 Yield x i' -> Yield x (ICALFirstOnlyInner s1 i')
3001 Skip i' -> Skip (ICALFirstOnlyInner s1 i')
3002 Stop -> Skip (ICALFirstOnly s1)
3004 -- We inject the second stream even before checking if the first stream
3005 -- would yield any more elements. There is no clear choice whether we
3006 -- should do this before or after that. Doing it after may make the state
3007 -- machine a bit simpler though.
3008 step gst (ICALSecondInject s1 s2) = do
3009 r <- step2 (adaptState gst) s2
3013 i `seq` return (Skip (ICALFirstInject s1 s i))
3014 Skip s -> return $ Skip (ICALSecondInject s1 s)
3015 Stop -> return $ Skip (ICALFirstOnly s1)
3017 step gst (ICALFirstInject s1 s2 i2) = do
3018 r <- step1 (adaptState gst) s1
3022 i `seq` return (Skip (ICALSecondInner s s2 i i2))
3023 -- i `seq` return (Skip (ICALFirstBuf s s2 i i2))
3024 Skip s -> return $ Skip (ICALFirstInject s s2 i2)
3028 step _ (ICALFirstBuf s1 s2 i1 i2) = do
3031 Yield x i' -> Skip (ICALSecondInner s1 s2 i' i2 x)
3032 Skip i' -> Skip (ICALFirstBuf s1 s2 i' i2)
3035 step _ (ICALSecondInner s1 s2 i1 i2 v) = do
3038 Yield x i' -> Yield x (ICALSecondInner s1 s2 i1 i' v)
3039 Skip i' -> Skip (ICALSecondInner s1 s2 i1 i' v)
3040 Stop -> Skip (ICALFirstResume s1 s2 i1 i2 v)
3043 step _ (ICALSecondInner s1 s2 i1 i2) = do
3046 Yield x i' -> Yield x (ICALSecondInner s1 s2 i1 i')
3047 Skip i' -> Skip (ICALSecondInner s1 s2 i1 i')
3048 Stop -> Skip (ICALFirstInner s1 s2 i1)
3049 -- Stop -> Skip (ICALFirstResume s1 s2 i1 i2)
3052 step _ (ICALFirstResume s1 s2 i1 i2 x) = do
3053 return $ Yield x (ICALFirstInner s1 s2 i1 i2)
3056data InterposeState s1 i1 a =
3058 -- | InterposeFirstYield s1 i1
3059 | InterposeFirstInner s1 i1
3060 | InterposeFirstInject s1
3061 -- | InterposeFirstBuf s1 i1
3062 | InterposeSecondYield s1 i1
3063 -- -- | InterposeSecondYield s1 i1 a
3064 -- -- | InterposeFirstResume s1 i1 a
3066-- Note that this only interposes the pure values, we may run many effects to
3067-- generate those values as some effects may not generate anything (Skip).
3068{-# INLINE_NORMAL interpose #-}
3069interpose :: Monad m => m c -> Unfold m b c -> Stream m b -> Stream m c
3072 (Unfold istep1 inject1) (Stream step1 state1) =
3073 Stream step (InterposeFirst state1)
3077 {-# INLINE_LATE step #-}
3078 step gst (InterposeFirst s1) = do
3079 r <- step1 (adaptState gst) s1
3083 i `seq` return (Skip (InterposeFirstInner s i))
3084 -- i `seq` return (Skip (InterposeFirstYield s i))
3085 Skip s -> return $ Skip (InterposeFirst s)
3089 step _ (InterposeFirstYield s1 i1) = do
3092 Yield x i' -> Yield x (InterposeFirstInner s1 i')
3093 Skip i' -> Skip (InterposeFirstYield s1 i')
3094 Stop -> Skip (InterposeFirst s1)
3097 step _ (InterposeFirstInner s1 i1) = do
3100 Yield x i' -> Yield x (InterposeFirstInner s1 i')
3101 Skip i' -> Skip (InterposeFirstInner s1 i')
3102 Stop -> Skip (InterposeFirstInject s1)
3104 step gst (InterposeFirstInject s1) = do
3105 r <- step1 (adaptState gst) s1
3109 -- i `seq` return (Skip (InterposeFirstBuf s i))
3110 i `seq` return (Skip (InterposeSecondYield s i))
3111 Skip s -> return $ Skip (InterposeFirstInject s)
3115 step _ (InterposeFirstBuf s1 i1) = do
3118 Yield x i' -> Skip (InterposeSecondYield s1 i' x)
3119 Skip i' -> Skip (InterposeFirstBuf s1 i')
3124 step _ (InterposeSecondYield s1 i1 v) = do
3126 return $ Yield r (InterposeFirstResume s1 i1 v)
3128 step _ (InterposeSecondYield s1 i1) = do
3130 return $ Yield r (InterposeFirstInner s1 i1)
3133 step _ (InterposeFirstResume s1 i1 v) = do
3134 return $ Yield v (InterposeFirstInner s1 i1)
3137------------------------------------------------------------------------------
3139------------------------------------------------------------------------------
3141data GbracketState s1 s2 v
3143 | GBracketNormal s1 v
3144 | GBracketException s2
3146-- | The most general bracketing and exception combinator. All other
3147-- combinators can be expressed in terms of this combinator. This can also be
3148-- used for cases which are not covered by the standard combinators.
3152{-# INLINE_NORMAL gbracket_ #-}
3156 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
3157 -> (c -> m d) -- ^ after, on normal stop
3158 -> (c -> e -> Stream m b -> Stream m b) -- ^ on exception
3159 -> (c -> Stream m b) -- ^ stream generator
3161gbracket_ bef exc aft fexc fnormal =
3162 Stream step GBracketInit
3166 {-# INLINE_LATE step #-}
3167 step _ GBracketInit = do
3169 return $ Skip $ GBracketNormal (fnormal r) r
3171 step gst (GBracketNormal (UnStream step1 st) v) = do
3172 res <- exc $ step1 gst st
3174 Right r -> case r of
3176 return $ Yield x (GBracketNormal (Stream step1 s) v)
3177 Skip s -> return $ Skip (GBracketNormal (Stream step1 s) v)
3178 Stop -> aft v >> return Stop
3180 return $ Skip (GBracketException (fexc v e (UnStream step1 st)))
3181 step gst (GBracketException (UnStream step1 st)) = do
3184 Yield x s -> return $ Yield x (GBracketException (Stream step1 s))
3185 Skip s -> return $ Skip (GBracketException (Stream step1 s))
3188------------------------------------------------------------------------------
3190------------------------------------------------------------------------------
3192-- | Make a finalizer from a monadic action @m a@ that can run in IO monad.
3193mkIOFinalizer :: MonadBaseControl IO m => m b -> m (IO ())
3195 mrun <- captureMonadState
3201-- | Run an IO action stored in a finalized IORef.
3202runIORefFinalizerGC :: IORef (Maybe (IO ())) -> IO ()
3203runIORefFinalizerGC ref = do
3204 res <- readIORef ref
3206 Nothing -> return ()
3209-- | Create an IORef holding a finalizer that is called automatically when
3210-- the IORef is garbage collected. The IORef can be written to with a 'Nothing'
3211-- value to deactivate the finalizer.
3213-- Note: The finalizer is always run with the state of the monad captured at
3214-- the time of calling newFinalizedIORef. To run it on garbage collection we
3215-- have no option but to take a snapshot of the monadic state at some point of
3216-- time. For normal case we could run it with the current state of the monad
3217-- but we want to keep both the cases consistent.
3219newFinalizedIORef :: (MonadIO m, MonadBaseControl IO m)
3220 => m a -> m (IORef (Maybe (IO ())))
3221newFinalizedIORef finalizer = do
3222 f <- mkIOFinalizer finalizer
3223 ref <- liftIO $ newIORef $ Just f
3224 _ <- liftIO $ mkWeakIORef ref (runIORefFinalizerGC ref)
3227-- | Run the finalizer stored in an IORef and deactivate it so that it is run
3228-- only once. Note, the action runs with async exceptions masked.
3230runIORefFinalizer :: MonadIO m => IORef (Maybe (IO ())) -> m ()
3231runIORefFinalizer ref = liftIO $ do
3232 res <- readIORef ref
3234 Nothing -> return ()
3236 -- if an async exception comes after writing 'Nothing' then the
3237 -- finalizing action will never be run. We need to do this
3238 -- atomically wrt async exceptions.
3240 writeIORef ref Nothing
3243-- | Run an action clearing the finalizer IORef atomically wrt async
3244-- exceptions. The action is run with async exceptions masked.
3245withIORefFinalizer :: MonadBaseControl IO m
3246 => IORef (Maybe (IO ())) -> m a -> m a
3247withIORefFinalizer ref action = do
3248 control $ \runinio ->
3250 writeIORef ref Nothing
3253------------------------------------------------------------------------------
3255data GbracketIOState s1 s2 v wref
3257 | GBracketIONormal s1 v wref
3258 | GBracketIOException s2
3260-- | The most general bracketing and exception combinator. All other
3261-- combinators can be expressed in terms of this combinator. This can also be
3262-- used for cases which are not covered by the standard combinators.
3264-- | It uses a finalizer to make sure we run the finalizing action
3265-- when the stream is garbage collected.
3269{-# INLINE_NORMAL gbracket #-}
3271 :: (MonadIO m, MonadBaseControl IO m)
3273 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
3274 -> (c -> m d) -- ^ after, on normal stop or GC
3275 -> (c -> e -> Stream m b -> m (Stream m b)) -- ^ on exception
3276 -> (c -> Stream m b) -- ^ stream generator
3278gbracket bef exc aft fexc fnormal =
3279 Stream step GBracketIOInit
3283 -- If the stream is never evaluated the "aft" action will never be
3284 -- called. For that to occur we will need the user of this API to pass a
3285 -- weak pointer to us.
3286 {-# INLINE_LATE step #-}
3287 step _ GBracketIOInit = do
3288 -- We mask asynchronous exceptions to make the execution
3289 -- of 'bef' and the registration of 'aft' atomic.
3290 -- A similar thing is done in the resourcet package: https://git.io/JvKV3
3291 -- Tutorial: https://markkarpov.com/tutorial/exceptions.html
3292 (r, ref) <- liftBaseOp_ mask_ $ do
3294 ref <- newFinalizedIORef (aft r)
3296 return $ Skip $ GBracketIONormal (fnormal r) r ref
3298 step gst (GBracketIONormal (UnStream step1 st) v ref) = do
3299 res <- exc $ step1 gst st
3301 Right r -> case r of
3303 return $ Yield x (GBracketIONormal (Stream step1 s) v ref)
3305 return $ Skip (GBracketIONormal (Stream step1 s) v ref)
3307 runIORefFinalizer ref
3310 -- Clearing of finalizer and running of exception handler must
3311 -- be atomic wrt async exceptions. Otherwise if we have cleared
3312 -- the finalizer and have not run the exception handler then we
3313 -- may leak the resource.
3314 stream <- withIORefFinalizer ref (fexc v e (UnStream step1 st))
3315 return $ Skip (GBracketIOException stream)
3316 step gst (GBracketIOException (UnStream step1 st)) = do
3320 return $ Yield x (GBracketIOException (Stream step1 s))
3321 Skip s -> return $ Skip (GBracketIOException (Stream step1 s))
3324-- | Run a side effect before the stream yields its first element.
3325{-# INLINE_NORMAL before #-}
3326before :: Monad m => m b -> Stream m a -> Stream m a
3327before action (Stream step state) = Stream step' Nothing
3331 {-# INLINE_LATE step' #-}
3332 step' _ Nothing = action >> return (Skip (Just state))
3334 step' gst (Just st) = do
3337 Yield x s -> return $ Yield x (Just s)
3338 Skip s -> return $ Skip (Just s)
3341-- | Run a side effect whenever the stream stops normally.
3342{-# INLINE_NORMAL after_ #-}
3343after_ :: Monad m => m b -> Stream m a -> Stream m a
3344after_ action (Stream step state) = Stream step' state
3348 {-# INLINE_LATE step' #-}
3352 Yield x s -> return $ Yield x s
3353 Skip s -> return $ Skip s
3354 Stop -> action >> return Stop
3356{-# INLINE_NORMAL after #-}
3357after :: (MonadIO m, MonadBaseControl IO m)
3358 => m b -> Stream m a -> Stream m a
3359after action (Stream step state) = Stream step' Nothing
3363 {-# INLINE_LATE step' #-}
3364 step' _ Nothing = do
3365 ref <- newFinalizedIORef action
3366 return $ Skip $ Just (state, ref)
3367 step' gst (Just (st, ref)) = do
3370 Yield x s -> return $ Yield x (Just (s, ref))
3371 Skip s -> return $ Skip (Just (s, ref))
3373 runIORefFinalizer ref
3376-- XXX These combinators are expensive due to the call to
3377-- onException/handle/try on each step. Therefore, when possible, they should
3378-- be called in an outer loop where we perform less iterations. For example, we
3379-- cannot call them on each iteration in a char stream, instead we can call
3380-- them when doing an IO on an array.
3382-- XXX For high performance error checks in busy streams we may need another
3383-- Error constructor in step.
3385-- | Run a side effect whenever the stream aborts due to an exception. The
3386-- exception is not caught, simply rethrown.
3387{-# INLINE_NORMAL onException #-}
3388onException :: MonadCatch m => m b -> Stream m a -> Stream m a
3389onException action str =
3390 gbracket_ (return ()) MC.try return
3391 (\_ (e :: MC.SomeException) _ -> nilM (action >> MC.throwM e))
3394{-# INLINE_NORMAL _onException #-}
3395_onException :: MonadCatch m => m b -> Stream m a -> Stream m a
3396_onException action (Stream step state) = Stream step' state
3400 {-# INLINE_LATE step' #-}
3402 res <- step gst st `MC.onException` action
3404 Yield x s -> return $ Yield x s
3405 Skip s -> return $ Skip s
3408-- XXX bracket is like concatMap, it generates a stream and then flattens it.
3409-- Like concatMap it has 10x worse performance compared to linear fused
3412-- | Run the first action before the stream starts and remember its output,
3413-- generate a stream using the output, run the second action providing the
3414-- remembered value as an argument whenever the stream ends normally or due to
3416{-# INLINE_NORMAL bracket_ #-}
3417bracket_ :: MonadCatch m => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
3418bracket_ bef aft bet =
3419 gbracket_ bef MC.try aft
3420 (\a (e :: SomeException) _ -> nilM (aft a >> MC.throwM e)) bet
3422-- | Run the first action before the stream starts and remember its output,
3423-- generate a stream using the output, run the second action providing the
3424-- remembered value as an argument whenever the stream ends normally or due to
3425-- an exception or gets garbage collected.
3427{-# INLINE_NORMAL bracket #-}
3428bracket :: (MonadAsync m, MonadCatch m)
3429 => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
3430bracket bef aft bet =
3431 gbracket bef MC.try aft
3432 (\a (e :: SomeException) _ -> aft a >> return (nilM (MC.throwM e))) bet
3434data BracketState s v = BracketInit | BracketRun s v
3436{-# INLINE_NORMAL _bracket #-}
3437_bracket :: MonadCatch m => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
3438_bracket bef aft bet = Stream step' BracketInit
3442 {-# INLINE_LATE step' #-}
3443 step' _ BracketInit = bef >>= \x -> return (Skip (BracketRun (bet x) x))
3445 -- NOTE: It is important to use UnStream instead of the Stream pattern
3446 -- here, otherwise we get huge perf degradation, see note in concatMap.
3447 step' gst (BracketRun (UnStream step state) v) = do
3448 -- res <- step gst state `MC.onException` aft v
3449 res <- MC.try $ step gst state
3451 Left (e :: SomeException) -> aft v >> MC.throwM e >> return Stop
3452 Right r -> case r of
3453 Yield x s -> return $ Yield x (BracketRun (Stream step s) v)
3454 Skip s -> return $ Skip (BracketRun (Stream step s) v)
3455 Stop -> aft v >> return Stop
3457-- | Run a side effect whenever the stream stops normally or aborts due to an
3459{-# INLINE finally_ #-}
3460finally_ :: MonadCatch m => m b -> Stream m a -> Stream m a
3461-- finally action xs = after action $ onException action xs
3462finally_ action xs = bracket_ (return ()) (\_ -> action) (const xs)
3464-- | Run a side effect whenever the stream stops normally or aborts due to an
3465-- exception or gets garbage collected.
3467{-# INLINE finally #-}
3468finally :: (MonadAsync m, MonadCatch m) => m b -> Stream m a -> Stream m a
3469finally action xs = bracket (return ()) (\_ -> action) (const xs)
3471-- | When evaluating a stream if an exception occurs, stream
3472-- evaluation aborts and the specified exception handler is run with
3473-- the exception and the Stream which threw the exception as argument.
3475{-# INLINE_NORMAL ghandle #-}
3476ghandle :: (MonadCatch m, Exception e)
3477 => (e -> Stream m a -> Stream m a) -> Stream m a -> Stream m a
3479 gbracket_ (return ()) MC.try return (\_ -> f) (\_ -> str)
3481-- | When evaluating a stream if an exception occurs, stream evaluation aborts
3482-- and the specified exception handler is run with the exception as argument.
3483{-# INLINE_NORMAL handle #-}
3484handle :: (MonadCatch m, Exception e)
3485 => (e -> Stream m a) -> Stream m a -> Stream m a
3487 gbracket_ (return ()) MC.try return (\_ e _ -> f e) (\_ -> str)
3489{-# INLINE_NORMAL _handle #-}
3490_handle :: (MonadCatch m, Exception e)
3491 => (e -> Stream m a) -> Stream m a -> Stream m a
3492_handle f (Stream step state) = Stream step' (Left state)
3496 {-# INLINE_LATE step' #-}
3497 step' gst (Left st) = do
3498 res <- MC.try $ step gst st
3500 Left e -> return $ Skip $ Right (f e)
3501 Right r -> case r of
3502 Yield x s -> return $ Yield x (Left s)
3503 Skip s -> return $ Skip (Left s)
3506 step' gst (Right (UnStream step1 st)) = do
3509 Yield x s -> return $ Yield x (Right (Stream step1 s))
3510 Skip s -> return $ Skip (Right (Stream step1 s))
3513-------------------------------------------------------------------------------
3514-- General transformation
3515-------------------------------------------------------------------------------
3517{-# INLINE_NORMAL transform #-}
3518transform :: Monad m => Pipe m a b -> Stream m a -> Stream m b
3519transform (Pipe pstep1 pstep2 pstate) (Stream step state) =
3520 Stream step' (Consume pstate, state)
3524 {-# INLINE_LATE step' #-}
3526 step' gst (Consume pst, st) = pst `seq` do
3527 r <- step (adaptState gst) st
3532 Pipe.Yield b pst' -> return $ Yield b (pst', s)
3533 Pipe.Continue pst' -> return $ Skip (pst', s)
3534 Skip s -> return $ Skip (Consume pst, s)
3537 step' _ (Produce pst, st) = pst `seq` do
3540 Pipe.Yield b pst' -> return $ Yield b (pst', st)
3541 Pipe.Continue pst' -> return $ Skip (pst', st)
3543------------------------------------------------------------------------------
3544-- Transformation by Folding (Scans)
3545------------------------------------------------------------------------------
3547------------------------------------------------------------------------------
3549------------------------------------------------------------------------------
3551-- XXX Is a prescan useful, discarding the last step does not sound useful? I
3552-- am not sure about the utility of this function, so this is implemented but
3553-- not exposed. We can expose it if someone provides good reasons why this is
3556-- XXX We have to execute the stream one step ahead to know that we are at the
3557-- last step. The vector implementation of prescan executes the last fold step
3558-- but does not yield the result. This means we have executed the effect but
3559-- discarded value. This does not sound right. In this implementation we are
3560-- not executing the last fold step.
3561{-# INLINE_NORMAL prescanlM' #-}
3562prescanlM' :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
3563prescanlM' f mz (Stream step state) = Stream step' (state, mz)
3565 {-# INLINE_LATE step' #-}
3566 step' gst (st, prev) = do
3567 r <- step (adaptState gst) st
3571 return $ Yield acc (s, f acc x)
3572 Skip s -> return $ Skip (s, prev)
3575{-# INLINE prescanl' #-}
3576prescanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
3577prescanl' f z = prescanlM' (\a b -> return (f a b)) (return z)
3579------------------------------------------------------------------------------
3580-- Monolithic postscans (postscan followed by a map)
3581------------------------------------------------------------------------------
3583-- The performance of a modular postscan followed by a map seems to be
3584-- equivalent to this monolithic scan followed by map therefore we may not need
3585-- this implementation. We just have it for performance comparison and in case
3586-- modular version does not perform well in some situation.
3588{-# INLINE_NORMAL postscanlMx' #-}
3589postscanlMx' :: Monad m
3590 => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> Stream m b
3591postscanlMx' fstep begin done (Stream step state) = do
3592 Stream step' (state, begin)
3594 {-# INLINE_LATE step' #-}
3595 step' gst (st, acc) = do
3596 r <- step (adaptState gst) st
3602 v `seq` y `seq` return (Yield v (s, return y))
3603 Skip s -> return $ Skip (s, acc)
3606{-# INLINE_NORMAL postscanlx' #-}
3607postscanlx' :: Monad m
3608 => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> Stream m b
3609postscanlx' fstep begin done s =
3610 postscanlMx' (\b a -> return (fstep b a)) (return begin) (return . done) s
3612-- XXX do we need consM strict to evaluate the begin value?
3613{-# INLINE scanlMx' #-}
3615 => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> Stream m b
3616scanlMx' fstep begin done s =
3617 (begin >>= \x -> x `seq` done x) `consM` postscanlMx' fstep begin done s
3619{-# INLINE scanlx' #-}
3621 => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> Stream m b
3622scanlx' fstep begin done s =
3623 scanlMx' (\b a -> return (fstep b a)) (return begin) (return . done) s
3625------------------------------------------------------------------------------
3627------------------------------------------------------------------------------
3629{-# INLINE_NORMAL postscanlM' #-}
3630postscanlM' :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
3631postscanlM' fstep begin (Stream step state) =
3632 Stream step' Nothing
3634 {-# INLINE_LATE step' #-}
3635 step' _ Nothing = do
3637 return $ Skip (Just (state, x))
3639 step' gst (Just (st, acc)) = do
3640 r <- step (adaptState gst) st
3644 return $ Yield y (Just (s, y))
3645 Skip s -> return $ Skip (Just (s, acc))
3648{-# INLINE_NORMAL postscanl' #-}
3649postscanl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
3650postscanl' f seed = postscanlM' (\a b -> return (f a b)) (return seed)
3652-- We can possibly have the "done" function as a Maybe to provide an option to
3653-- emit or not emit the accumulator when the stream stops.
3655-- TBD: use a single Yield point
3657{-# INLINE_NORMAL postscanlMAfter' #-}
3658postscanlMAfter' :: Monad m
3659 => (b -> a -> m b) -> m b -> (b -> m b) -> Stream m a -> Stream m b
3660postscanlMAfter' fstep initial done (Stream step1 state1) = do
3661 Stream step (Just (state1, initial))
3665 {-# INLINE_LATE step #-}
3666 step gst (Just (st, acc)) = do
3667 r <- step1 (adaptState gst) st
3672 y `seq` return (Yield y (Just (s, return y)))
3673 Skip s -> return $ Skip $ Just (s, acc)
3674 Stop -> acc >>= done >>= \res -> return (Yield res Nothing)
3675 step _ Nothing = return Stop
3677{-# INLINE_NORMAL postscanlM #-}
3678postscanlM :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
3679postscanlM fstep begin (Stream step state) = Stream step' Nothing
3681 {-# INLINE_LATE step' #-}
3682 step' _ Nothing = do
3684 return $ Skip (Just (state, r))
3686 step' gst (Just (st, acc)) = do
3687 r <- step (adaptState gst) st
3691 return (Yield y (Just (s, y)))
3692 Skip s -> return $ Skip (Just (s, acc))
3695{-# INLINE_NORMAL postscanl #-}
3696postscanl :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
3697postscanl f seed = postscanlM (\a b -> return (f a b)) (return seed)
3699{-# INLINE_NORMAL scanlM' #-}
3700scanlM' :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
3701scanlM' fstep begin (Stream step state) = Stream step' Nothing
3703 {-# INLINE_LATE step' #-}
3704 step' _ Nothing = do
3706 return $ Yield x (Just (state, x))
3707 step' gst (Just (st, acc)) = do
3708 r <- step (adaptState gst) st
3712 return $ Yield y (Just (s, y))
3713 Skip s -> return $ Skip (Just (s, acc))
3716{-# INLINE scanlMAfter' #-}
3717scanlMAfter' :: Monad m
3718 => (b -> a -> m b) -> m b -> (b -> m b) -> Stream m a -> Stream m b
3719scanlMAfter' fstep initial done s =
3720 (initial >>= \x -> x `seq` return x) `consM`
3721 postscanlMAfter' fstep initial done s
3723{-# INLINE scanl' #-}
3724scanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
3725scanl' f seed = scanlM' (\a b -> return (f a b)) (return seed)
3727{-# INLINE_NORMAL scanlM #-}
3728scanlM :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
3729scanlM fstep begin (Stream step state) = Stream step' Nothing
3731 {-# INLINE_LATE step' #-}
3732 step' _ Nothing = do
3734 return $ Yield x (Just (state, x))
3735 step' gst (Just (st, acc)) = do
3736 r <- step (adaptState gst) st
3740 return $ Yield y (Just (s, y))
3741 Skip s -> return $ Skip (Just (s, acc))
3742 Stop -> return $ Stop
3745scanl :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
3746scanl f seed = scanlM (\a b -> return (f a b)) (return seed)
3748{-# INLINE_NORMAL scanl1M #-}
3749scanl1M :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
3750scanl1M fstep (Stream step state) = Stream step' (state, Nothing)
3752 {-# INLINE_LATE step' #-}
3753 step' gst (st, Nothing) = do
3756 Yield x s -> return $ Yield x (s, Just x)
3757 Skip s -> return $ Skip (s, Nothing)
3760 step' gst (st, Just acc) = do
3765 return $ Yield z (s, Just z)
3766 Skip s -> return $ Skip (s, Just acc)
3769{-# INLINE scanl1 #-}
3770scanl1 :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
3771scanl1 f = scanl1M (\x y -> return (f x y))
3773{-# INLINE_NORMAL scanl1M' #-}
3774scanl1M' :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
3775scanl1M' fstep (Stream step state) = Stream step' (state, Nothing)
3777 {-# INLINE_LATE step' #-}
3778 step' gst (st, Nothing) = do
3781 Yield x s -> x `seq` return $ Yield x (s, Just x)
3782 Skip s -> return $ Skip (s, Nothing)
3785 step' gst (st, Just acc) = acc `seq` do
3790 z `seq` return $ Yield z (s, Just z)
3791 Skip s -> return $ Skip (s, Just acc)
3794{-# INLINE scanl1' #-}
3795scanl1' :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
3796scanl1' f = scanl1M' (\x y -> return (f x y))
3798------------------------------------------------------------------------------
3800------------------------------------------------------------------------------
3802data RollingMapState s a = RollingMapInit s | RollingMapGo s a
3804{-# INLINE rollingMapM #-}
3805rollingMapM :: Monad m => (a -> a -> m b) -> Stream m a -> Stream m b
3806rollingMapM f (Stream step1 state1) = Stream step (RollingMapInit state1)
3808 step gst (RollingMapInit st) = do
3809 r <- step1 (adaptState gst) st
3811 Yield x s -> Skip $ RollingMapGo s x
3812 Skip s -> Skip $ RollingMapInit s
3815 step gst (RollingMapGo s1 x1) = do
3816 r <- step1 (adaptState gst) s1
3820 return $ Yield res $ RollingMapGo s x
3821 Skip s -> return $ Skip $ RollingMapGo s x1
3822 Stop -> return $ Stop
3824{-# INLINE rollingMap #-}
3825rollingMap :: Monad m => (a -> a -> b) -> Stream m a -> Stream m b
3826rollingMap f = rollingMapM (\x y -> return $ f x y)
3828------------------------------------------------------------------------------
3829-- Tapping/Distributing
3830------------------------------------------------------------------------------
3833tap :: Monad m => Fold m a b -> Stream m a -> Stream m a
3834tap (Fold fstep initial extract) (Stream step state) = Stream step' Nothing
3838 step' _ Nothing = do
3840 return $ Skip (Just (r, state))
3842 step' gst (Just (acc, st)) = acc `seq` do
3847 return $ Yield x (Just (acc', s))
3848 Skip s -> return $ Skip (Just (acc, s))
3853{-# INLINE_NORMAL tapOffsetEvery #-}
3854tapOffsetEvery :: Monad m
3855 => Int -> Int -> Fold m a b -> Stream m a -> Stream m a
3856tapOffsetEvery offset n (Fold fstep initial extract) (Stream step state) =
3857 Stream step' Nothing
3861 {-# INLINE_LATE step' #-}
3862 step' _ Nothing = do
3864 return $ Skip (Just (r, state, offset `mod` n))
3866 step' gst (Just (acc, st, count)) | count <= 0 = do
3870 !acc' <- fstep acc x
3871 return $ Yield x (Just (acc', s, n - 1))
3872 Skip s -> return $ Skip (Just (acc, s, count))
3877 step' gst (Just (acc, st, count)) = do
3880 Yield x s -> return $ Yield x (Just (acc, s, count - 1))
3881 Skip s -> return $ Skip (Just (acc, s, count))
3886{-# INLINE_NORMAL pollCounts #-}
3890 -> (Stream m Int -> Stream m Int)
3894pollCounts predicate transf fld (Stream step state) = Stream step' Nothing
3897 {-# INLINE_LATE step' #-}
3898 step' _ Nothing = do
3899 -- As long as we are using an "Int" for counts lockfree reads from
3900 -- Var should work correctly on both 32-bit and 64-bit machines.
3901 -- However, an Int on a 32-bit machine may overflow quickly.
3902 countVar <- liftIO $ Prim.newIORef (0 :: Int)
3904 $ void $ runFold fld
3905 $ transf $ fromPrimIORef countVar
3906 return $ Skip (Just (countVar, tid, state))
3908 step' gst (Just (countVar, tid, st)) = do
3912 when (predicate x) $ liftIO $ Prim.modifyIORef' countVar (+ 1)
3913 return $ Yield x (Just (countVar, tid, s))
3914 Skip s -> return $ Skip (Just (countVar, tid, s))
3916 liftIO $ killThread tid
3919{-# INLINE_NORMAL tapRate #-}
3921 (MonadAsync m, MonadCatch m)
3926tapRate samplingRate action (Stream step state) = Stream step' Nothing
3928 {-# NOINLINE loop #-}
3929 loop countVar prev = do
3932 (do liftIO $ threadDelay (round $ samplingRate * 1000000)
3933 i <- liftIO $ Prim.readIORef countVar
3934 let !diff = i - prev
3937 (\(e :: AsyncException) -> do
3938 i <- liftIO $ Prim.readIORef countVar
3939 let !diff = i - prev
3941 throwM (MC.toException e))
3944 {-# INLINE_LATE step' #-}
3945 step' _ Nothing = do
3946 countVar <- liftIO $ Prim.newIORef 0
3947 tid <- fork $ loop countVar 0
3948 ref <- liftIO $ newIORef ()
3949 _ <- liftIO $ mkWeakIORef ref (killThread tid)
3950 return $ Skip (Just (countVar, tid, state, ref))
3952 step' gst (Just (countVar, tid, st, ref)) = do
3956 liftIO $ Prim.modifyIORef' countVar (+ 1)
3957 return $ Yield x (Just (countVar, tid, s, ref))
3958 Skip s -> return $ Skip (Just (countVar, tid, s, ref))
3960 liftIO $ killThread tid
3964-------------------------------------------------------------------------------
3966-------------------------------------------------------------------------------
3968{-# INLINE_NORMAL takeWhileM #-}
3969takeWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
3970takeWhileM f (Stream step state) = Stream step' state
3972 {-# INLINE_LATE step' #-}
3978 return $ if b then Yield x s else Stop
3979 Skip s -> return $ Skip s
3982{-# INLINE takeWhile #-}
3983takeWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
3984takeWhile f = takeWhileM (return . f)
3986{-# INLINE_NORMAL drop #-}
3987drop :: Monad m => Int -> Stream m a -> Stream m a
3988drop n (Stream step state) = Stream step' (state, Just n)
3990 {-# INLINE_LATE step' #-}
3991 step' gst (st, Just i)
3996 Yield _ s -> Skip (s, Just (i - 1))
3997 Skip s -> Skip (s, Just i)
3999 | otherwise = return $ Skip (st, Nothing)
4001 step' gst (st, Nothing) = do
4005 Yield x s -> Yield x (s, Nothing)
4006 Skip s -> Skip (s, Nothing)
4009data DropWhileState s a
4011 | DropWhileYield a s
4014{-# INLINE_NORMAL dropWhileM #-}
4015dropWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
4016dropWhileM f (Stream step state) = Stream step' (DropWhileDrop state)
4018 {-# INLINE_LATE step' #-}
4019 step' gst (DropWhileDrop st) = do
4025 then return $ Skip (DropWhileDrop s)
4026 else return $ Skip (DropWhileYield x s)
4027 Skip s -> return $ Skip (DropWhileDrop s)
4030 step' gst (DropWhileNext st) = do
4033 Yield x s -> return $ Skip (DropWhileYield x s)
4034 Skip s -> return $ Skip (DropWhileNext s)
4037 step' _ (DropWhileYield x st) = return $ Yield x (DropWhileNext st)
4039{-# INLINE dropWhile #-}
4040dropWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
4041dropWhile f = dropWhileM (return . f)
4043{-# INLINE_NORMAL filterM #-}
4044filterM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
4045filterM f (Stream step state) = Stream step' state
4047 {-# INLINE_LATE step' #-}
4056 Skip s -> return $ Skip s
4059{-# INLINE filter #-}
4060filter :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
4061filter f = filterM (return . f)
4063{-# INLINE_NORMAL uniq #-}
4064uniq :: (Eq a, Monad m) => Stream m a -> Stream m a
4065uniq (Stream step state) = Stream step' (Nothing, state)
4067 {-# INLINE_LATE step' #-}
4068 step' gst (Nothing, st) = do
4071 Yield x s -> return $ Yield x (Just x, s)
4072 Skip s -> return $ Skip (Nothing, s)
4074 step' gst (Just x, st) = do
4077 Yield y s | x == y -> return $ Skip (Just x, s)
4078 | otherwise -> return $ Yield y (Just y, s)
4079 Skip s -> return $ Skip (Just x, s)
4082------------------------------------------------------------------------------
4083-- Transformation by Mapping
4084------------------------------------------------------------------------------
4086{-# INLINE_NORMAL sequence #-}
4087sequence :: Monad m => Stream m (m a) -> Stream m a
4088sequence (Stream step state) = Stream step' state
4090 {-# INLINE_LATE step' #-}
4092 r <- step (adaptState gst) st
4094 Yield x s -> x >>= \a -> return (Yield a s)
4095 Skip s -> return $ Skip s
4098------------------------------------------------------------------------------
4100------------------------------------------------------------------------------
4102data LoopState x s = FirstYield s
4103 | InterspersingYield s
4106{-# INLINE_NORMAL intersperseM #-}
4107intersperseM :: Monad m => m a -> Stream m a -> Stream m a
4108intersperseM m (Stream step state) = Stream step' (FirstYield state)
4110 {-# INLINE_LATE step' #-}
4111 step' gst (FirstYield st) = do
4115 Yield x s -> Skip (YieldAndCarry x s)
4116 Skip s -> Skip (FirstYield s)
4119 step' gst (InterspersingYield st) = do
4124 return $ Yield a (YieldAndCarry x s)
4125 Skip s -> return $ Skip $ InterspersingYield s
4128 step' _ (YieldAndCarry x st) = return $ Yield x (InterspersingYield st)
4130{-# INLINE_NORMAL intersperseM_ #-}
4131intersperseM_ :: Monad m => m b -> Stream m a -> Stream m a
4132intersperseM_ m (Stream step1 state1) = Stream step (Left (pure (), state1))
4134 {-# INLINE_LATE step #-}
4135 step gst (Left (eff, st)) = do
4138 Yield x s -> eff >> return (Yield x (Right s))
4139 Skip s -> return $ Skip (Left (eff, s))
4142 step _ (Right st) = return $ Skip $ Left (void m, st)
4147 | SuffixYield a (SuffixState s a)
4149{-# INLINE_NORMAL intersperseSuffix #-}
4150intersperseSuffix :: forall m a. Monad m => m a -> Stream m a -> Stream m a
4151intersperseSuffix action (Stream step state) = Stream step' (SuffixElem state)
4153 {-# INLINE_LATE step' #-}
4154 step' gst (SuffixElem st) = do
4157 Yield x s -> Skip (SuffixYield x (SuffixSuffix s))
4158 Skip s -> Skip (SuffixElem s)
4161 step' _ (SuffixSuffix st) = do
4162 action >>= \r -> return $ Skip (SuffixYield r (SuffixElem st))
4164 step' _ (SuffixYield x next) = return $ Yield x next
4166{-# INLINE_NORMAL intersperseSuffix_ #-}
4167intersperseSuffix_ :: Monad m => m b -> Stream m a -> Stream m a
4168intersperseSuffix_ m (Stream step1 state1) = Stream step (Left state1)
4170 {-# INLINE_LATE step #-}
4171 step gst (Left st) = do
4174 Yield x s -> return $ Yield x (Right s)
4175 Skip s -> return $ Skip $ Left s
4178 step _ (Right st) = m >> return (Skip (Left st))
4180data SuffixSpanState s a
4181 = SuffixSpanElem s Int
4182 | SuffixSpanSuffix s
4183 | SuffixSpanYield a (SuffixSpanState s a)
4187-- | intersperse after every n items
4188{-# INLINE_NORMAL intersperseSuffixBySpan #-}
4189intersperseSuffixBySpan :: forall m a. Monad m
4190 => Int -> m a -> Stream m a -> Stream m a
4191intersperseSuffixBySpan n action (Stream step state) =
4192 Stream step' (SuffixSpanElem state n)
4194 {-# INLINE_LATE step' #-}
4195 step' gst (SuffixSpanElem st i) | i > 0 = do
4198 Yield x s -> Skip (SuffixSpanYield x (SuffixSpanElem s (i - 1)))
4199 Skip s -> Skip (SuffixSpanElem s i)
4200 Stop -> if i == n then Stop else Skip SuffixSpanLast
4201 step' _ (SuffixSpanElem st _) = return $ Skip (SuffixSpanSuffix st)
4203 step' _ (SuffixSpanSuffix st) = do
4204 action >>= \r -> return $ Skip (SuffixSpanYield r (SuffixSpanElem st n))
4206 step' _ (SuffixSpanLast) = do
4207 action >>= \r -> return $ Skip (SuffixSpanYield r SuffixSpanStop)
4209 step' _ (SuffixSpanYield x next) = return $ Yield x next
4211 step' _ (SuffixSpanStop) = return Stop
4213{-# INLINE intersperse #-}
4214intersperse :: Monad m => a -> Stream m a -> Stream m a
4215intersperse a = intersperseM (return a)
4217{-# INLINE_NORMAL insertBy #-}
4218insertBy :: Monad m => (a -> a -> Ordering) -> a -> Stream m a -> Stream m a
4219insertBy cmp a (Stream step state) = Stream step' (state, False, Nothing)
4221 {-# INLINE_LATE step' #-}
4222 step' gst (st, False, _) = do
4225 Yield x s -> case cmp a x of
4226 GT -> return $ Yield x (s, False, Nothing)
4227 _ -> return $ Yield a (s, True, Just x)
4228 Skip s -> return $ Skip (s, False, Nothing)
4229 Stop -> return $ Yield a (st, True, Nothing)
4231 step' _ (_, True, Nothing) = return Stop
4233 step' gst (st, True, Just prev) = do
4236 Yield x s -> return $ Yield prev (s, True, Just x)
4237 Skip s -> return $ Skip (s, True, Just prev)
4238 Stop -> return $ Yield prev (st, True, Nothing)
4240------------------------------------------------------------------------------
4242------------------------------------------------------------------------------
4244{-# INLINE_NORMAL deleteBy #-}
4245deleteBy :: Monad m => (a -> a -> Bool) -> a -> Stream m a -> Stream m a
4246deleteBy eq x (Stream step state) = Stream step' (state, False)
4248 {-# INLINE_LATE step' #-}
4249 step' gst (st, False) = do
4252 Yield y s -> return $
4253 if eq x y then Skip (s, True) else Yield y (s, False)
4254 Skip s -> return $ Skip (s, False)
4257 step' gst (st, True) = do
4260 Yield y s -> return $ Yield y (s, True)
4261 Skip s -> return $ Skip (s, True)
4264------------------------------------------------------------------------------
4265-- Transformation by Map and Filter
4266------------------------------------------------------------------------------
4268-- XXX Will this always fuse properly?
4269{-# INLINE_NORMAL mapMaybe #-}
4270mapMaybe :: Monad m => (a -> Maybe b) -> Stream m a -> Stream m b
4271mapMaybe f = fmap fromJust . filter isJust . map f
4273{-# INLINE_NORMAL mapMaybeM #-}
4274mapMaybeM :: Monad m => (a -> m (Maybe b)) -> Stream m a -> Stream m b
4275mapMaybeM f = fmap fromJust . filter isJust . mapM f
4277------------------------------------------------------------------------------
4279------------------------------------------------------------------------------
4281{-# INLINE_NORMAL indexed #-}
4282indexed :: Monad m => Stream m a -> Stream m (Int, a)
4283indexed (Stream step state) = Stream step' (state, 0)
4285 {-# INLINE_LATE step' #-}
4286 step' gst (st, i) = i `seq` do
4287 r <- step (adaptState gst) st
4289 Yield x s -> return $ Yield (i, x) (s, i+1)
4290 Skip s -> return $ Skip (s, i)
4293{-# INLINE_NORMAL indexedR #-}
4294indexedR :: Monad m => Int -> Stream m a -> Stream m (Int, a)
4295indexedR m (Stream step state) = Stream step' (state, m)
4297 {-# INLINE_LATE step' #-}
4298 step' gst (st, i) = i `seq` do
4299 r <- step (adaptState gst) st
4301 Yield x s -> let i' = i - 1
4302 in return $ Yield (i, x) (s, i')
4303 Skip s -> return $ Skip (s, i)
4306{-# INLINE_NORMAL zipWithM #-}
4308 => (a -> b -> m c) -> Stream m a -> Stream m b -> Stream m c
4309zipWithM f (Stream stepa ta) (Stream stepb tb) = Stream step (ta, tb, Nothing)
4311 {-# INLINE_LATE step #-}
4312 step gst (sa, sb, Nothing) = do
4313 r <- stepa (adaptState gst) sa
4316 Yield x sa' -> Skip (sa', sb, Just x)
4317 Skip sa' -> Skip (sa', sb, Nothing)
4320 step gst (sa, sb, Just x) = do
4321 r <- stepb (adaptState gst) sb
4325 return $ Yield z (sa, sb', Nothing)
4326 Skip sb' -> return $ Skip (sa, sb', Just x)
4329#if __GLASGOW_HASKELL__ >= 801
4330{-# RULES "zipWithM xs xs"
4331 forall f xs. zipWithM @Identity f xs xs = mapM (\x -> f x x) xs #-}
4334{-# INLINE zipWith #-}
4335zipWith :: Monad m => (a -> b -> c) -> Stream m a -> Stream m b -> Stream m c
4336zipWith f = zipWithM (\a b -> return (f a b))
4338------------------------------------------------------------------------------
4340------------------------------------------------------------------------------
4342{-# INLINE_NORMAL mergeByM #-}
4345 => (a -> a -> m Ordering) -> Stream m a -> Stream m a -> Stream m a
4346mergeByM cmp (Stream stepa ta) (Stream stepb tb) =
4347 Stream step (Just ta, Just tb, Nothing, Nothing)
4349 {-# INLINE_LATE step #-}
4351 -- one of the values is missing, and the corresponding stream is running
4352 step gst (Just sa, sb, Nothing, b) = do
4355 Yield a sa' -> Skip (Just sa', sb, Just a, b)
4356 Skip sa' -> Skip (Just sa', sb, Nothing, b)
4357 Stop -> Skip (Nothing, sb, Nothing, b)
4359 step gst (sa, Just sb, a, Nothing) = do
4362 Yield b sb' -> Skip (sa, Just sb', a, Just b)
4363 Skip sb' -> Skip (sa, Just sb', a, Nothing)
4364 Stop -> Skip (sa, Nothing, a, Nothing)
4366 -- both the values are available
4367 step _ (sa, sb, Just a, Just b) = do
4369 return $ case res of
4370 GT -> Yield b (sa, sb, Just a, Nothing)
4371 _ -> Yield a (sa, sb, Nothing, Just b)
4373 -- one of the values is missing, corresponding stream is done
4374 step _ (Nothing, sb, Nothing, Just b) =
4375 return $ Yield b (Nothing, sb, Nothing, Nothing)
4377 step _ (sa, Nothing, Just a, Nothing) =
4378 return $ Yield a (sa, Nothing, Nothing, Nothing)
4380 step _ (Nothing, Nothing, Nothing, Nothing) = return Stop
4382{-# INLINE mergeBy #-}
4385 => (a -> a -> Ordering) -> Stream m a -> Stream m a -> Stream m a
4386mergeBy cmp = mergeByM (\a b -> return $ cmp a b)
4388------------------------------------------------------------------------------
4389-- Transformation comprehensions
4390------------------------------------------------------------------------------
4392{-# INLINE_NORMAL the #-}
4393the :: (Eq a, Monad m) => Stream m a -> m (Maybe a)
4394the (Stream step state) = go state
4397 r <- step defState st
4399 Yield x s -> go' x s
4401 Stop -> return Nothing
4403 r <- step defState st
4405 Yield x s | x == n -> go' n s
4406 | otherwise -> return Nothing
4408 Stop -> return (Just n)
4410{-# INLINE runFold #-}
4411runFold :: (Monad m) => Fold m a b -> Stream m a -> m b
4412runFold (Fold step begin done) = foldlMx' step begin done
4414-------------------------------------------------------------------------------
4415-- Concurrent application and fold
4416-------------------------------------------------------------------------------
4418-- XXX These functions should be moved to Stream/Parallel.hs
4420-- Using StreamD the worker stream producing code can fuse with the code to
4421-- queue output to the SVar giving some perf boost.
4423-- Note that StreamD can only be used in limited situations, specifically, we
4424-- cannot implement joinStreamVarPar using this.
4426-- XXX make sure that the SVar passed is a Parallel style SVar.
4428-- | Fold the supplied stream to the SVar asynchronously using Parallel
4429-- concurrency style.
4430-- {-# INLINE_NORMAL toSVarParallel #-}
4431{-# INLINE toSVarParallel #-}
4432toSVarParallel :: MonadAsync m
4433 => State t m a -> SVar t m a -> Stream m a -> m ()
4434toSVarParallel st sv xs =
4435 if svarInspectMode sv
4439 case getYieldLimit st of
4440 Nothing -> doFork (work Nothing)
4442 (handleChildException sv)
4443 Just _ -> doFork (workLim Nothing)
4445 (handleChildException sv)
4450 {-# NOINLINE work #-}
4451 work info = (runFold (FL.toParallelSVar sv info) xs)
4453 {-# NOINLINE workLim #-}
4454 workLim info = runFold (FL.toParallelSVarLimited sv info) xs
4456 {-# NOINLINE forkWithDiag #-}
4458 -- We do not use workerCount in case of ParallelVar but still there is
4459 -- no harm in maintaining it correctly.
4460 liftIO $ atomicModifyIORefCAS_ (workerCount sv) $ \n -> n + 1
4462 -- This allocation matters when significant number of workers are being
4463 -- sent. We allocate it only when needed. The overhead increases by 4x.
4465 case yieldRateInfo sv of
4466 Nothing -> return Nothing
4467 Just _ -> liftIO $ do
4468 cntRef <- newIORef 0
4469 t <- getTime Monotonic
4470 lat <- newIORef (0, t)
4471 return $ Just WorkerInfo
4472 { workerYieldMax = 0
4473 , workerYieldCount = cntRef
4474 , workerLatencyStart = lat
4477 case getYieldLimit st of
4478 Nothing -> doFork (work winfo)
4480 (handleChildException sv)
4481 Just _ -> doFork (workLim winfo)
4483 (handleChildException sv)
4486{-# INLINE_NORMAL mkParallelD #-}
4487mkParallelD :: MonadAsync m => Stream m a -> Stream m a
4488mkParallelD m = Stream step Nothing
4491 step gst Nothing = do
4492 sv <- newParallelVar StopNone gst
4493 toSVarParallel gst sv m
4494 -- XXX use unfold instead?
4495 return $ Skip $ Just $ fromSVar sv
4497 step gst (Just (UnStream step1 st)) = do
4500 Yield a s -> Yield a (Just $ Stream step1 s)
4501 Skip s -> Skip (Just $ Stream step1 s)
4504-- Compare with mkAsync. mkAsync uses an Async style SVar whereas this uses a
4505-- parallel style SVar for evaluation. Currently, parallel style cannot use
4506-- rate control whereas Async style can use rate control. In async style SVar
4507-- the worker thread terminates when the buffer is full whereas in Parallel
4510-- | Make the stream producer and consumer run concurrently by introducing a
4511-- buffer between them. The producer thread evaluates the input stream until
4512-- the buffer fills, it blocks if the buffer is full until there is space in
4513-- the buffer. The consumer consumes the stream lazily from the buffer.
4517{-# INLINE_NORMAL mkParallel #-}
4518mkParallel :: (K.IsStream t, MonadAsync m) => t m a -> t m a
4519mkParallel = fromStreamD . mkParallelD . toStreamD
4521-- Note: we can use another API with two callbacks stop and yield if we want
4522-- the callback to be able to indicate end of stream.
4524-- | Generates a callback and a stream pair. The callback returned is used to
4525-- queue values to the stream. The stream is infinite, there is no way for the
4526-- callback to indicate that it is done now.
4530{-# INLINE_NORMAL newCallbackStream #-}
4531newCallbackStream :: (K.IsStream t, MonadAsync m) => m ((a -> m ()), t m a)
4532newCallbackStream = do
4533 sv <- newParallelVar StopNone defState
4535 -- XXX Add our own thread-id to the SVar as we can not know the callback's
4536 -- thread-id and the callback is not run in a managed worker. We need to
4537 -- handle this better.
4538 liftIO myThreadId >>= modifyThread sv
4540 let callback a = liftIO $ void $ send sv (ChildYield a)
4541 -- XXX we can return an SVar and then the consumer can unfold from the
4543 return (callback, fromStreamD (fromSVar sv))
4545-------------------------------------------------------------------------------
4547-------------------------------------------------------------------------------
4549-- | Create an SVar with a fold consumer that will fold any elements sent to it
4550-- using the supplied fold function.
4551{-# INLINE newFoldSVar #-}
4552newFoldSVar :: MonadAsync m => State t m a -> Fold m a b -> m (SVar t m a)
4553newFoldSVar stt f = do
4554 -- Buffer size for the SVar is derived from the current state
4555 sv <- newParallelVar StopAny (adaptState stt)
4556 -- Add the producer thread-id to the SVar.
4557 liftIO myThreadId >>= modifyThread sv
4558 void $ doFork (work sv) (svarMrun sv) (handleFoldException sv)
4563 {-# NOINLINE work #-}
4564 work sv = void $ runFold f $ fromProducer sv
4566data TapState sv st = TapInit | Tapping sv st | TapDone st
4568{-# INLINE_NORMAL tapAsync #-}
4569tapAsync :: MonadAsync m => Fold m a b -> Stream m a -> Stream m a
4570tapAsync f (Stream step1 state1) = Stream step TapInit
4574 -- In general, a Stop event would come equipped with the result
4575 -- of the fold. It is not used here but it would be useful in
4576 -- applicative and distribute.
4577 done <- fromConsumer svr
4578 when (not done) $ do
4579 liftIO $ withDiagMVar svr "teeToSVar: waiting to drain"
4580 $ takeMVar (outputDoorBellFromConsumer svr)
4584 liftIO $ sendStop svr Nothing
4585 -- drain/wait until a stop event arrives from the fold.
4588 {-# INLINE_LATE step #-}
4589 step gst TapInit = do
4590 sv <- newFoldSVar gst f
4591 return $ Skip (Tapping sv state1)
4593 step gst (Tapping sv st) = do
4597 done <- pushToFold sv a
4600 -- XXX we do not need to wait synchronously here
4602 return $ Yield a (TapDone s)
4603 else return $ Yield a (Tapping sv s)
4604 Skip s -> return $ Skip (Tapping sv s)
4609 step gst (TapDone st) = do
4612 Yield a s -> Yield a (TapDone s)
4613 Skip s -> Skip (TapDone s)
4616-- XXX Exported from Array again as this fold is specific to Array
4617-- | Take last 'n' elements from the stream and discard the rest.
4619lastN :: (Storable a, MonadIO m) => Int -> Fold m a (Array a)
4621 | n <= 0 = fmap (const mempty) FL.drain
4622 | otherwise = A.unsafeFreeze <$> Fold step initial done
4624 step (Tuple3' rb rh i) a = do
4625 rh1 <- liftIO $ RB.unsafeInsert rb rh a
4626 return $ Tuple3' rb rh1 (i + 1)
4627 initial = fmap (\(a, b) -> Tuple3' a b (0 :: Int)) $ liftIO $ RB.new n
4628 done (Tuple3' rb rh i) = do
4629 arr <- liftIO $ MA.newArray n
4630 foldFunc i rh snoc' arr rb
4631 snoc' b a = liftIO $ MA.unsafeSnoc b a
4633 | i < n = RB.unsafeFoldRingM
4634 | otherwise = RB.unsafeFoldRingFullM
4636------------------------------------------------------------------------------
4638------------------------------------------------------------------------------
4640-- XXX using getTime in the loop can be pretty expensive especially for
4641-- computations where iterations are lightweight. We have the following
4644-- 1) Run a timeout thread updating a flag asynchronously and check that
4645-- flag here, that way we can have a cheap termination check.
4647-- 2) Use COARSE clock to get time with lower resolution but more efficiently.
4649-- 3) Use rdtscp/rdtsc to get time directly from the processor, compute the
4650-- termination value of rdtsc in the beginning and then in each iteration just
4651-- get rdtsc and check if we should terminate.
4655 | TakeByTimeCheck st s
4656 | TakeByTimeYield st s
4658{-# INLINE_NORMAL takeByTime #-}
4659takeByTime :: (MonadIO m, TimeUnit64 t) => t -> Stream m a -> Stream m a
4660takeByTime duration (Stream step1 state1) = Stream step (TakeByTimeInit state1)
4663 lim = toRelTime64 duration
4665 {-# INLINE_LATE step #-}
4666 step _ (TakeByTimeInit _) | lim == 0 = return Stop
4667 step _ (TakeByTimeInit st) = do
4668 t0 <- liftIO $ getTime Monotonic
4669 return $ Skip (TakeByTimeYield st t0)
4670 step _ (TakeByTimeCheck st t0) = do
4671 t <- liftIO $ getTime Monotonic
4673 if diffAbsTime64 t t0 > lim
4675 else Skip (TakeByTimeYield st t0)
4676 step gst (TakeByTimeYield st t0) = do
4679 Yield x s -> Yield x (TakeByTimeCheck s t0)
4680 Skip s -> Skip (TakeByTimeCheck s t0)
4683data DropByTime st s x
4685 | DropByTimeGen st s
4686 | DropByTimeCheck st s x
4687 | DropByTimeYield st
4689{-# INLINE_NORMAL dropByTime #-}
4690dropByTime :: (MonadIO m, TimeUnit64 t) => t -> Stream m a -> Stream m a
4691dropByTime duration (Stream step1 state1) = Stream step (DropByTimeInit state1)
4694 lim = toRelTime64 duration
4696 {-# INLINE_LATE step #-}
4697 step _ (DropByTimeInit st) = do
4698 t0 <- liftIO $ getTime Monotonic
4699 return $ Skip (DropByTimeGen st t0)
4700 step gst (DropByTimeGen st t0) = do
4703 Yield x s -> Skip (DropByTimeCheck s t0 x)
4704 Skip s -> Skip (DropByTimeGen s t0)
4706 step _ (DropByTimeCheck st t0 x) = do
4707 t <- liftIO $ getTime Monotonic
4708 if diffAbsTime64 t t0 <= lim
4709 then return $ Skip $ DropByTimeGen st t0
4710 else return $ Yield x $ DropByTimeYield st
4711 step gst (DropByTimeYield st) = do
4714 Yield x s -> Yield x (DropByTimeYield s)
4715 Skip s -> Skip (DropByTimeYield s)
4718{-# INLINE updateTimeVar #-}
4719updateTimeVar :: Prim.IORef Int64 -> IO ()
4720updateTimeVar timeVar = do
4721 MicroSecond64 t <- fromAbsTime <$> getTime Monotonic
4722 Prim.modifyIORef' timeVar (const t)
4724{-# INLINE updateWithDelay #-}
4725updateWithDelay :: RealFrac a => a -> Prim.IORef Int64 -> IO ()
4726updateWithDelay precision timeVar = do
4727 threadDelay (delayTime precision)
4728 updateTimeVar timeVar
4732 -- Keep the minimum at least a millisecond to avoid high CPU usage
4733 {-# INLINE delayTime #-}
4735 | g' >= fromIntegral (maxBound :: Int) = maxBound
4737 | otherwise = round g'
4741 g' = g * 10 ^ (6 :: Int)
4743-- XXX we should move this to stream generation section of this file. Also, the
4744-- take/drop combinators above should be moved to filtering section.
4746{-# INLINE_NORMAL times #-}
4747times :: MonadAsync m => Double -> Stream m (AbsTime, RelTime64)
4748times g = Stream step Nothing
4752 {-# INLINE_LATE step #-}
4754 -- XXX note that this is safe only on a 64-bit machine. On a 32-bit
4755 -- machine a 64-bit 'Var' cannot be read consistently without a lock
4756 -- while another thread is writing to it.
4757 timeVar <- liftIO $ Prim.newIORef (0 :: Int64)
4758 liftIO $ updateTimeVar timeVar
4759 tid <- forkManaged $ liftIO $ forever (updateWithDelay g timeVar)
4760 a <- liftIO $ Prim.readIORef timeVar
4761 return $ Skip $ Just (timeVar, tid, a)
4763 step _ s@(Just (timeVar, _, t0)) = do
4764 a <- liftIO $ Prim.readIORef timeVar
4765 -- XXX we can perhaps use an AbsTime64 using a 64 bit Int for
4766 -- efficiency. or maybe we can use a representation using Double for
4767 -- floating precision time
4768 return $ Yield (toAbsTime (MicroSecond64 t0),
4769 (toRelTime64 (MicroSecond64 (a - t0)))) s
keysesc to close
- j k
- move down, up
- enter
- open
- esc
- back, or up a level
- 1 2 3 4
- readme, files, commits, refs
- /
- filter
- y
- copy the clone command
- c
- compare the chosen commit with another (or shift-click both)
- ?
- show or hide this list