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
1085 -- XXX Specialize for Stop 0 common case?
1087 assert (n <= length (x:buf)) (return ())
1088 let src = Prelude.reverse (Prelude.take n (x:buf))
1090 ParseChunksYield b (ParseChunksInit src s)
1091 PR.Error err -> throwM $ ParseError err
1092 Skip s -> return $ Skip $ ParseChunksStream s buf pst
1095 let src = Prelude.reverse buf
1097 ParseChunksYield b (ParseChunksInitLeftOver src)
1099 -- go back to stream processing mode
1100 stepOuter _ (ParseChunksBuf [] s buf pst) =
1101 return $ Skip $ ParseChunksStream s buf pst
1103 -- buffered processing loop
1104 stepOuter _ (ParseChunksBuf (x:xs) s buf pst) = do
1107 PR.Partial 0 pst1 ->
1108 return $ Skip $ ParseChunksBuf xs s [] pst1
1109 PR.Partial n pst1 -> do
1110 assert (n <= length (x:buf)) (return ())
1111 let src0 = Prelude.take n (x:buf)
1112 src = Prelude.reverse src0 ++ xs
1113 return $ Skip $ ParseChunksBuf src s [] pst1
1114 -- PR.Continue 0 pst1 -> return $ Skip $ ParseChunksBuf xs s (x:buf) pst1
1115 PR.Continue n pst1 -> do
1116 assert (n <= length (x:buf)) (return ())
1117 let (src0, buf1) = splitAt n (x:buf)
1118 src = Prelude.reverse src0 ++ xs
1119 return $ Skip $ ParseChunksBuf src s buf1 pst1
1120 -- XXX Specialize for Stop 0 common case?
1122 assert (n <= length (x:buf)) (return ())
1123 let src = Prelude.reverse (Prelude.take n (x:buf)) ++ xs
1124 return $ Skip $ ParseChunksYield b (ParseChunksInit src s)
1125 PR.Error err -> throwM $ ParseError err
1127 stepOuter _ (ParseChunksYield a next) = return $ Yield a next
1129{-# ANN type ConcatParseState Fuse #-}
1130data ConcatParseState x inpBuf st p =
1131 ConcatParseInit inpBuf st p
1132 | ConcatParseInitLeftOver inpBuf
1133 | ConcatParseStream st inpBuf p
1134 | ConcatParseBuf inpBuf st inpBuf p
1135 | ConcatParseYield x (ConcatParseState x inpBuf st p)
1137{-# INLINE_NORMAL parseIterate #-}
1140 => (b -> PRD.Parser m a b)
1144parseIterate func seed (Stream step state) =
1145 Stream stepOuter (ConcatParseInit [] state (func seed))
1149 {-# INLINE_LATE stepOuter #-}
1150 -- Buffer is empty, go to stream processing loop
1151 stepOuter _ (ConcatParseInit [] st (PRD.Parser pstep initial extract)) = do
1152 initial >>= \r -> return $ Skip $ ConcatParseStream st []
1153 (PRD.Parser pstep (return r) extract)
1155 -- Buffer is not empty, go to buffered processing loop
1156 stepOuter _ (ConcatParseInit src st
1157 (PRD.Parser pstep initial extract)) = do
1158 initial >>= \r -> return $ Skip $ ConcatParseBuf src st []
1159 (PRD.Parser pstep (return r) extract)
1161 -- XXX we just discard any leftover input at the end
1162 stepOuter _ (ConcatParseInitLeftOver _) = return Stop
1164 -- Buffer is empty process elements from the stream
1165 stepOuter gst (ConcatParseStream st buf
1166 p@(PRD.Parser pstep initial extract)) = do
1167 r <- step (adaptState gst) st
1173 PR.Partial 0 pst1 ->
1174 return $ Skip $ ConcatParseStream s []
1175 (PRD.Parser pstep (return pst1) extract)
1176 PR.Partial n pst1 -> do
1177 assert (n <= length (x:buf)) (return ())
1178 let src0 = Prelude.take n (x:buf)
1179 src = Prelude.reverse src0
1180 return $ Skip $ ConcatParseBuf src s []
1181 (PRD.Parser pstep (return pst1) extract)
1182 -- PR.Continue 0 pst1 ->
1183 -- return $ Skip $ ConcatParseStream s (x:buf) pst1
1184 PR.Continue n pst1 -> do
1185 assert (n <= length (x:buf)) (return ())
1186 let (src0, buf1) = splitAt n (x:buf)
1187 src = Prelude.reverse src0
1188 return $ Skip $ ConcatParseBuf src s buf1
1189 (PRD.Parser pstep (return pst1) extract)
1190 -- XXX Specialize for Stop 0 common case?
1192 assert (n <= length (x:buf)) (return ())
1193 let src = Prelude.reverse (Prelude.take n (x:buf))
1195 ConcatParseYield b (ConcatParseInit src s (func b))
1196 PR.Error err -> throwM $ ParseError err
1197 Skip s -> return $ Skip $ ConcatParseStream s buf p
1201 let src = Prelude.reverse buf
1203 ConcatParseYield b (ConcatParseInitLeftOver src)
1205 -- go back to stream processing mode
1206 stepOuter _ (ConcatParseBuf [] s buf p) =
1207 return $ Skip $ ConcatParseStream s buf p
1209 -- buffered processing loop
1210 stepOuter _ (ConcatParseBuf (x:xs) s buf
1211 (PRD.Parser pstep initial extract)) = do
1215 PR.Partial 0 pst1 ->
1216 return $ Skip $ ConcatParseBuf xs s []
1217 (PRD.Parser pstep (return pst1) extract)
1218 PR.Partial n pst1 -> do
1219 assert (n <= length (x:buf)) (return ())
1220 let src0 = Prelude.take n (x:buf)
1221 src = Prelude.reverse src0 ++ xs
1222 return $ Skip $ ConcatParseBuf src s []
1223 (PRD.Parser pstep (return pst1) extract)
1224 -- PR.Continue 0 pst1 -> return $ Skip $ ConcatParseBuf xs s (x:buf) pst1
1225 PR.Continue n pst1 -> do
1226 assert (n <= length (x:buf)) (return ())
1227 let (src0, buf1) = splitAt n (x:buf)
1228 src = Prelude.reverse src0 ++ xs
1229 return $ Skip $ ConcatParseBuf src s buf1
1230 (PRD.Parser pstep (return pst1) extract)
1231 -- XXX Specialize for Stop 0 common case?
1233 assert (n <= length (x:buf)) (return ())
1234 let src = Prelude.reverse (Prelude.take n (x:buf)) ++ xs
1235 return $ Skip $ ConcatParseYield b
1236 (ConcatParseInit src s (func b))
1237 PR.Error err -> throwM $ ParseError err
1239 stepOuter _ (ConcatParseYield a next) = return $ Yield a next
1241------------------------------------------------------------------------------
1243------------------------------------------------------------------------------
1245-- | Run a streaming composition, discard the results.
1246{-# INLINE_LATE drain #-}
1247drain :: Monad m => Stream m a -> m ()
1248-- drain = foldrM (\_ xs -> xs) (return ())
1249drain (Stream step state) = go SPEC state
1252 r <- step defState st
1254 Yield _ s -> go SPEC s
1258{-# INLINE_NORMAL null #-}
1259null :: Monad m => Stream m a -> m Bool
1260null m = foldrM (\_ _ -> return False) (return True) m
1262{-# INLINE_NORMAL head #-}
1263head :: Monad m => Stream m a -> m (Maybe a)
1264head m = foldrM (\x _ -> return (Just x)) (return Nothing) m
1266{-# INLINE_NORMAL headElse #-}
1267headElse :: Monad m => a -> Stream m a -> m a
1268headElse a m = foldrM (\x _ -> return x) (return a) m
1270-- Does not fuse, has the same performance as the StreamK version.
1271{-# INLINE_NORMAL tail #-}
1272tail :: Monad m => Stream m a -> m (Maybe (Stream m a))
1273tail (UnStream step state) = go state
1276 r <- step defState st
1278 Yield _ s -> return (Just $ Stream step s)
1280 Stop -> return Nothing
1282-- XXX will it fuse? need custom impl?
1283{-# INLINE_NORMAL last #-}
1284last :: Monad m => Stream m a -> m (Maybe a)
1285last = foldl' (\_ y -> Just y) Nothing
1287{-# INLINE_NORMAL elem #-}
1288elem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
1289-- elem e m = foldrM (\x xs -> if x == e then return True else xs) (return False) m
1290elem e (Stream step state) = go state
1293 r <- step defState st
1296 | x == e -> return True
1299 Stop -> return False
1301{-# INLINE_NORMAL notElem #-}
1302notElem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
1303notElem e s = fmap not (elem e s)
1305{-# INLINE_NORMAL all #-}
1306all :: Monad m => (a -> Bool) -> Stream m a -> m Bool
1307-- all p m = foldrM (\x xs -> if p x then xs else return False) (return True) m
1308all p (Stream step state) = go state
1311 r <- step defState st
1315 | otherwise -> return False
1319{-# INLINE_NORMAL any #-}
1320any :: Monad m => (a -> Bool) -> Stream m a -> m Bool
1321-- any p m = foldrM (\x xs -> if p x then return True else xs) (return False) m
1322any p (Stream step state) = go state
1325 r <- step defState st
1328 | p x -> return True
1331 Stop -> return False
1333{-# INLINE_NORMAL maximum #-}
1334maximum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
1335maximum (Stream step state) = go Nothing state
1338 r <- step defState st
1340 Yield x s -> go (Just x) s
1341 Skip s -> go Nothing s
1342 Stop -> return Nothing
1343 go (Just acc) st = do
1344 r <- step defState st
1347 | acc <= x -> go (Just x) s
1348 | otherwise -> go (Just acc) s
1349 Skip s -> go (Just acc) s
1350 Stop -> return (Just acc)
1352{-# INLINE_NORMAL maximumBy #-}
1353maximumBy :: Monad m => (a -> a -> Ordering) -> Stream m a -> m (Maybe a)
1354maximumBy cmp (Stream step state) = go Nothing state
1357 r <- step defState st
1359 Yield x s -> go (Just x) s
1360 Skip s -> go Nothing s
1361 Stop -> return Nothing
1362 go (Just acc) st = do
1363 r <- step defState st
1365 Yield x s -> case cmp acc x of
1366 GT -> go (Just acc) s
1368 Skip s -> go (Just acc) s
1369 Stop -> return (Just acc)
1371{-# INLINE_NORMAL minimum #-}
1372minimum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
1373minimum (Stream step state) = go Nothing state
1376 r <- step defState st
1378 Yield x s -> go (Just x) s
1379 Skip s -> go Nothing s
1380 Stop -> return Nothing
1381 go (Just acc) st = do
1382 r <- step defState st
1385 | acc <= x -> go (Just acc) s
1386 | otherwise -> go (Just x) s
1387 Skip s -> go (Just acc) s
1388 Stop -> return (Just acc)
1390{-# INLINE_NORMAL minimumBy #-}
1391minimumBy :: Monad m => (a -> a -> Ordering) -> Stream m a -> m (Maybe a)
1392minimumBy cmp (Stream step state) = go Nothing state
1395 r <- step defState st
1397 Yield x s -> go (Just x) s
1398 Skip s -> go Nothing s
1399 Stop -> return Nothing
1400 go (Just acc) st = do
1401 r <- step defState st
1403 Yield x s -> case cmp acc x of
1405 _ -> go (Just acc) s
1406 Skip s -> go (Just acc) s
1407 Stop -> return (Just acc)
1409{-# INLINE_NORMAL (!!) #-}
1410(!!) :: (Monad m) => Stream m a -> Int -> m (Maybe a)
1411(Stream step state) !! i = go i state
1414 r <- step defState st
1416 Yield x s | n < 0 -> return Nothing
1417 | n == 0 -> return $ Just x
1418 | otherwise -> go (n - 1) s
1420 Stop -> return Nothing
1422{-# INLINE_NORMAL lookup #-}
1423lookup :: (Monad m, Eq a) => a -> Stream m (a, b) -> m (Maybe b)
1424lookup e m = foldrM (\(a, b) xs -> if e == a then return (Just b) else xs)
1427{-# INLINE_NORMAL findM #-}
1428findM :: Monad m => (a -> m Bool) -> Stream m a -> m (Maybe a)
1429findM p m = foldrM (\x xs -> p x >>= \r -> if r then return (Just x) else xs)
1433find :: Monad m => (a -> Bool) -> Stream m a -> m (Maybe a)
1434find p = findM (return . p)
1436{-# INLINE_NORMAL findIndices #-}
1437findIndices :: Monad m => (a -> Bool) -> Stream m a -> Stream m Int
1438findIndices p (Stream step state) = Stream step' (state, 0)
1440 {-# INLINE_LATE step' #-}
1441 step' gst (st, i) = i `seq` do
1442 r <- step (adaptState gst) st
1444 Yield x s -> if p x then Yield i (s, i+1) else Skip (s, i+1)
1445 Skip s -> Skip (s, i)
1448{-# INLINE toListRev #-}
1449toListRev :: Monad m => Stream m a -> m [a]
1450toListRev = foldl' (flip (:)) []
1452-- We can implement reverse as:
1454-- > reverse = foldlS (flip cons) nil
1456-- However, this implementation is unusable because of the horrible performance
1457-- of cons. So we just convert it to a list first and then stream from the
1460-- XXX Maybe we can use an Array instead of a list here?
1461{-# INLINE_NORMAL reverse #-}
1462reverse :: Monad m => Stream m a -> Stream m a
1463reverse m = Stream step Nothing
1465 {-# INLINE_LATE step #-}
1468 return $ Skip (Just xs)
1469 step _ (Just (x:xs)) = return $ Yield x (Just xs)
1470 step _ (Just []) = return Stop
1472-- Much faster reverse for Storables
1473{-# INLINE_NORMAL reverse' #-}
1474reverse' :: forall m a. (MonadIO m, Storable a) => Stream m a -> Stream m a
1476-- This commented implementation copies the whole stream into one single array
1477-- and then streams from that array, this is 3-4x faster than the chunked code
1478-- that follows. Though this could be problematic due to unbounded large
1479-- allocations. We need to figure out why the chunked code is slower and if we
1480-- can optimize the chunked code to work as fast as this one. It may be a
1482import Foreign.ForeignPtr (touchForeignPtr)
1483import Foreign.ForeignPtr.Unsafe (unsafeForeignPtrToPtr)
1484import Foreign.Ptr (Ptr, plusPtr)
1485reverse' m = Stream step Nothing
1487 {-# INLINE_LATE step #-}
1489 arr <- A.fromStreamD m
1490 let p = aEnd arr `plusPtr` negate (sizeOf (undefined :: a))
1491 return $ Skip $ Just (aStart arr, p)
1493 step _ (Just (start, p)) | p < unsafeForeignPtrToPtr start = return Stop
1495 step _ (Just (start, p)) = do
1496 let !x = A.unsafeInlineIO $ do
1498 touchForeignPtr start
1500 next = p `plusPtr` negate (sizeOf (undefined :: a))
1501 return $ Yield x (Just (start, next))
1508 $ A.fromStreamDArraysOf A.defaultChunkSize m
1511------------------------------------------------------------------------------
1512-- Grouping/Splitting
1513------------------------------------------------------------------------------
1515{-# INLINE_NORMAL splitSuffixBy' #-}
1516splitSuffixBy' :: Monad m
1517 => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1518splitSuffixBy' predicate f (Stream step state) =
1519 Stream (stepOuter f) (Just state)
1523 {-# INLINE_LATE stepOuter #-}
1524 stepOuter (Fold fstep initial done) gst (Just st) = do
1525 res <- step (adaptState gst) st
1531 then done acc' >>= \val -> return $ Yield val (Just s)
1534 Skip s -> return $ Skip $ Just s
1540 res <- step (adaptState gst) stt
1545 then done acc' >>= \val -> return $ Yield val (Just s)
1547 Skip s -> go SPEC s acc
1548 Stop -> done acc >>= \val -> return $ Yield val Nothing
1550 stepOuter _ _ Nothing = return Stop
1552{-# INLINE_NORMAL groupsBy #-}
1558groupsBy cmp f (Stream step state) = Stream (stepOuter f) (Just state, Nothing)
1562 {-# INLINE_LATE stepOuter #-}
1563 stepOuter (Fold fstep initial done) gst (Just st, Nothing) = do
1564 res <- step (adaptState gst) st
1571 Skip s -> return $ Skip $ (Just s, Nothing)
1576 go !_ prev stt !acc = do
1577 res <- step (adaptState gst) stt
1584 else done acc >>= \r -> return $ Yield r (Just s, Just x)
1585 Skip s -> go SPEC prev s acc
1586 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1588 stepOuter (Fold fstep initial done) gst (Just st, Just prev) = do
1590 acc' <- fstep acc prev
1595 -- XXX code duplicated from the previous equation
1597 res <- step (adaptState gst) stt
1604 else done acc >>= \r -> return $ Yield r (Just s, Just x)
1605 Skip s -> go SPEC s acc
1606 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1608 stepOuter _ _ (Nothing,_) = return Stop
1610{-# INLINE_NORMAL groupsRollingBy #-}
1611groupsRollingBy :: Monad m
1616groupsRollingBy cmp f (Stream step state) =
1617 Stream (stepOuter f) (Just state, Nothing)
1620 {-# INLINE_LATE stepOuter #-}
1621 stepOuter (Fold fstep initial done) gst (Just st, Nothing) = do
1622 res <- step (adaptState gst) st
1629 Skip s -> return $ Skip $ (Just s, Nothing)
1633 go !_ prev stt !acc = do
1634 res <- step (adaptState gst) stt
1642 done acc >>= \r -> return $ Yield r (Just s, Just x)
1643 Skip s -> go SPEC prev s acc
1644 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1646 stepOuter (Fold fstep initial done) gst (Just st, Just prev') = do
1648 acc' <- fstep acc prev'
1649 go SPEC prev' st acc'
1652 go !_ prevv stt !acc = do
1653 res <- step (adaptState gst) stt
1660 else done acc >>= \r -> return $ Yield r (Just s, Just x)
1661 Skip s -> go SPEC prevv s acc
1662 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1664 stepOuter _ _ (Nothing, _) = return Stop
1666{-# INLINE_NORMAL splitBy #-}
1667splitBy :: Monad m => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1668splitBy predicate f (Stream step state) = Stream (step' f) (Just state)
1672 {-# INLINE_LATE step' #-}
1673 step' (Fold fstep initial done) gst (Just st) = initial >>= go SPEC st
1678 res <- step (adaptState gst) stt
1682 then done acc >>= \r -> return $ Yield r (Just s)
1686 Skip s -> go SPEC s acc
1687 Stop -> done acc >>= \r -> return $ Yield r Nothing
1689 step' _ _ Nothing = return Stop
1691-- XXX requires -funfolding-use-threshold=150 in lines-unlines benchmark
1692{-# INLINE_NORMAL splitSuffixBy #-}
1693splitSuffixBy :: Monad m
1694 => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1695splitSuffixBy predicate f (Stream step state) = Stream (step' f) (Just state)
1699 {-# INLINE_LATE step' #-}
1700 step' (Fold fstep initial done) gst (Just st) = do
1701 res <- step (adaptState gst) st
1706 then done acc >>= \val -> return $ Yield val (Just s)
1711 Skip s -> return $ Skip $ Just s
1717 res <- step (adaptState gst) stt
1721 then done acc >>= \r -> return $ Yield r (Just s)
1725 Skip s -> go SPEC s acc
1726 Stop -> done acc >>= \r -> return $ Yield r Nothing
1728 step' _ _ Nothing = return Stop
1730{-# INLINE_NORMAL wordsBy #-}
1731wordsBy :: Monad m => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1732wordsBy predicate f (Stream step state) = Stream (stepOuter f) (Just state)
1736 {-# INLINE_LATE stepOuter #-}
1737 stepOuter (Fold fstep initial done) gst (Just st) = do
1738 res <- step (adaptState gst) st
1742 then return $ Skip (Just s)
1748 Skip s -> return $ Skip $ Just s
1754 res <- step (adaptState gst) stt
1758 then done acc >>= \r -> return $ Yield r (Just s)
1762 Skip s -> go SPEC s acc
1763 Stop -> done acc >>= \r -> return $ Yield r Nothing
1765 stepOuter _ _ Nothing = return Stop
1767-- String search algorithms:
1768-- http://www-igm.univ-mlv.fr/~lecroq/string/index.html
1771-- TODO can we unify the splitting operations using a splitting configuration
1772-- like in the split package.
1774data SplitStyle = Infix | Suffix | Prefix deriving (Eq, Show)
1776data SplitOptions = SplitOptions
1777 { style :: SplitStyle
1778 , withSep :: Bool -- ^ keep the separators in output
1779 -- , compact :: Bool -- ^ treat multiple consecutive separators as one
1780 -- , trimHead :: Bool -- ^ drop blank at head
1781 -- , trimTail :: Bool -- ^ drop blank at tail
1785data SplitOnState s a =
1790 | GO_KARP_RABIN s !(RB.Ring a) !(Ptr a)
1793{-# INLINE_NORMAL splitOn #-}
1795 :: forall m a b. (MonadIO m, Storable a, Enum a, Eq a)
1800splitOn patArr (Fold fstep initial done) (Stream step state) =
1801 Stream stepOuter GO_START
1805 patLen = A.length patArr
1806 maxIndex = patLen - 1
1807 elemBits = sizeOf (undefined :: a) * 8
1809 {-# INLINE_LATE stepOuter #-}
1810 stepOuter _ GO_START =
1812 then return $ Skip $ GO_EMPTY_PAT state
1815 r <- liftIO $ (A.unsafeIndexIO patArr 0)
1816 return $ Skip $ GO_SINGLE_PAT state r
1817 else if sizeOf (undefined :: a) * patLen
1818 <= sizeOf (undefined :: Word)
1819 then return $ Skip $ GO_SHORT_PAT state
1821 (rb, rhead) <- liftIO $ RB.new patLen
1822 return $ Skip $ GO_KARP_RABIN state rb rhead
1824 stepOuter gst (GO_SINGLE_PAT stt pat) = initial >>= go SPEC stt
1829 res <- step (adaptState gst) st
1835 return $ Yield r (GO_SINGLE_PAT s pat)
1836 else fstep acc x >>= go SPEC s
1837 Skip s -> go SPEC s acc
1838 Stop -> done acc >>= \r -> return $ Yield r GO_DONE
1840 stepOuter gst (GO_SHORT_PAT stt) = initial >>= go0 SPEC 0 (0 :: Word) stt
1845 mask = (1 `shiftL` (elemBits * patLen)) - 1
1847 addToWord wrd a = (wrd `shiftL` elemBits) .|. fromIntegral (fromEnum a)
1850 patWord = mask .&. A.foldl' addToWord 0 patArr
1852 go0 !_ !idx wrd st !acc = do
1853 res <- step (adaptState gst) st
1856 let wrd' = addToWord wrd x
1859 if wrd' .&. mask == patWord
1862 return $ Yield r (GO_SHORT_PAT s)
1863 else go1 SPEC wrd' s acc
1864 else go0 SPEC (idx + 1) wrd' s acc
1865 Skip s -> go0 SPEC idx wrd s acc
1868 then go2 wrd idx acc
1870 done acc' >>= \r -> return $ Yield r GO_DONE
1873 go1 !_ wrd st !acc = do
1874 res <- step (adaptState gst) st
1877 let wrd' = addToWord wrd x
1878 old = (mask .&. wrd) `shiftR` (elemBits * (patLen - 1))
1879 acc' <- fstep acc (toEnum $ fromIntegral old)
1880 if wrd' .&. mask == patWord
1881 then done acc' >>= \r -> return $ Yield r (GO_SHORT_PAT s)
1882 else go1 SPEC wrd' s acc'
1883 Skip s -> go1 SPEC wrd s acc
1885 acc' <- go2 wrd patLen acc
1886 done acc' >>= \r -> return $ Yield r GO_DONE
1888 go2 !wrd !n !acc | n > 0 = do
1889 let old = (mask .&. wrd) `shiftR` (elemBits * (n - 1))
1890 fstep acc (toEnum $ fromIntegral old) >>= go2 wrd (n - 1)
1891 go2 _ _ acc = return acc
1893 stepOuter gst (GO_KARP_RABIN stt rb rhead) = do
1894 initial >>= go0 SPEC 0 rhead stt
1898 k = 2891336453 :: Word32
1900 addCksum cksum a = cksum * k + fromIntegral (fromEnum a)
1901 deltaCksum cksum old new =
1902 addCksum cksum new - coeff * fromIntegral (fromEnum old)
1904 -- XXX shall we use a random starting hash or 1 instead of 0?
1905 patHash = A.foldl' addCksum 0 patArr
1908 go0 !_ !idx !rh st !acc = do
1909 res <- step (adaptState gst) st
1912 rh' <- liftIO $ RB.unsafeInsert rb rh x
1915 let fold = RB.unsafeFoldRing (RB.ringBound rb)
1916 let !ringHash = fold addCksum 0 rb
1917 if ringHash == patHash
1918 then go2 SPEC ringHash rh' s acc
1919 else go1 SPEC ringHash rh' s acc
1920 else go0 SPEC (idx + 1) rh' s acc
1921 Skip s -> go0 SPEC idx rh s acc
1923 !acc' <- if idx /= 0
1924 then RB.unsafeFoldRingM rh fstep acc rb
1926 done acc' >>= \r -> return $ Yield r GO_DONE
1928 -- XXX Theoretically this code can do 4 times faster if GHC generates
1929 -- optimal code. If we use just "(cksum' == patHash)" condition it goes
1930 -- 4x faster, as soon as we add the "RB.unsafeEqArray rb v" condition
1931 -- the generated code changes drastically and becomes 4x slower. Need
1932 -- to investigate what is going on with GHC.
1934 go1 !_ !cksum !rh st !acc = do
1935 res <- step (adaptState gst) st
1938 old <- liftIO $ peek rh
1939 let cksum' = deltaCksum cksum old x
1940 acc' <- fstep acc old
1942 if (cksum' == patHash)
1944 rh' <- liftIO (RB.unsafeInsert rb rh x)
1945 go2 SPEC cksum' rh' s acc'
1947 rh' <- liftIO (RB.unsafeInsert rb rh x)
1948 go1 SPEC cksum' rh' s acc'
1949 Skip s -> go1 SPEC cksum rh s acc
1951 acc' <- RB.unsafeFoldRingFullM rh fstep acc rb
1952 done acc' >>= \r -> return $ Yield r GO_DONE
1954 go2 !_ !cksum' !rh' s !acc' = do
1955 if RB.unsafeEqArray rb rh' patArr
1958 return $ Yield r (GO_KARP_RABIN s rb rhead)
1959 else go1 SPEC cksum' rh' s acc'
1961 stepOuter gst (GO_EMPTY_PAT st) = do
1962 res <- step (adaptState gst) st
1967 done acc' >>= \r -> return $ Yield r (GO_EMPTY_PAT s)
1968 Skip s -> return $ Skip (GO_EMPTY_PAT s)
1971 stepOuter _ GO_DONE = return Stop
1973{-# INLINE_NORMAL splitSuffixOn #-}
1975 :: forall m a b. (MonadIO m, Storable a, Enum a, Eq a)
1981splitSuffixOn withSep patArr (Fold fstep initial done)
1982 (Stream step state) =
1983 Stream stepOuter GO_START
1987 patLen = A.length patArr
1988 maxIndex = patLen - 1
1989 elemBits = sizeOf (undefined :: a) * 8
1991 {-# INLINE_LATE stepOuter #-}
1992 stepOuter _ GO_START =
1994 then return $ Skip $ GO_EMPTY_PAT state
1997 r <- liftIO $ (A.unsafeIndexIO patArr 0)
1998 return $ Skip $ GO_SINGLE_PAT state r
1999 else if sizeOf (undefined :: a) * patLen
2000 <= sizeOf (undefined :: Word)
2001 then return $ Skip $ GO_SHORT_PAT state
2003 (rb, rhead) <- liftIO $ RB.new patLen
2004 return $ Skip $ GO_KARP_RABIN state rb rhead
2006 stepOuter gst (GO_SINGLE_PAT stt pat) = do
2007 -- This first part is the only difference between splitOn and
2009 -- If the last element is a separator do not issue a blank segment.
2010 res <- step (adaptState gst) stt
2016 acc' <- if withSep then fstep acc x else return acc
2017 done acc' >>= \r -> return $ Yield r (GO_SINGLE_PAT s pat)
2018 else fstep acc x >>= go SPEC s
2019 Skip s -> return $ Skip $ (GO_SINGLE_PAT s pat)
2024 -- This is identical for splitOn and splitSuffixOn
2026 res <- step (adaptState gst) st
2031 acc' <- if withSep then fstep acc x else return acc
2033 return $ Yield r (GO_SINGLE_PAT s pat)
2034 else fstep acc x >>= go SPEC s
2035 Skip s -> go SPEC s acc
2036 Stop -> done acc >>= \r -> return $ Yield r GO_DONE
2038 stepOuter gst (GO_SHORT_PAT stt) = do
2040 -- Call "initial" only if the stream yields an element, otherwise we
2041 -- may call "initial" but never yield anything. initial may produce a
2042 -- side effect, therefore we will end up doing and discard a side
2047 res <- step (adaptState gst) stt
2051 let wrd' = addToWord wrd x
2052 acc' <- if withSep then fstep acc x else return acc
2055 if wrd' .&. mask == patWord
2056 then done acc' >>= \r -> return $ Yield r (GO_SHORT_PAT s)
2057 else go0 SPEC (idx + 1) wrd' s acc'
2058 else go0 SPEC (idx + 1) wrd' s acc'
2059 Skip s -> return $ Skip (GO_SHORT_PAT s)
2065 mask = (1 `shiftL` (elemBits * patLen)) - 1
2067 addToWord wrd a = (wrd `shiftL` elemBits) .|. fromIntegral (fromEnum a)
2070 patWord = mask .&. A.foldl' addToWord 0 patArr
2072 go0 !_ !idx wrd st !acc = do
2073 res <- step (adaptState gst) st
2076 let wrd' = addToWord wrd x
2077 acc' <- if withSep then fstep acc x else return acc
2080 if wrd' .&. mask == patWord
2083 return $ Yield r (GO_SHORT_PAT s)
2084 else go1 SPEC wrd' s acc'
2085 else go0 SPEC (idx + 1) wrd' s acc'
2086 Skip s -> go0 SPEC idx wrd s acc
2088 if (idx == maxIndex) && (wrd .&. mask == patWord)
2091 acc' <- if idx /= 0 && not withSep
2092 then go2 wrd idx acc
2094 done acc' >>= \r -> return $ Yield r GO_DONE
2097 go1 !_ wrd st !acc = do
2098 res <- step (adaptState gst) st
2101 let wrd' = addToWord wrd x
2102 old = (mask .&. wrd) `shiftR` (elemBits * (patLen - 1))
2105 else fstep acc (toEnum $ fromIntegral old)
2106 if wrd' .&. mask == patWord
2107 then done acc' >>= \r -> return $ Yield r (GO_SHORT_PAT s)
2108 else go1 SPEC wrd' s acc'
2109 Skip s -> go1 SPEC wrd s acc
2111 -- If the last sequence is a separator do not issue a blank
2113 if wrd .&. mask == patWord
2118 else go2 wrd patLen acc
2119 done acc' >>= \r -> return $ Yield r GO_DONE
2121 go2 !wrd !n !acc | n > 0 = do
2122 let old = (mask .&. wrd) `shiftR` (elemBits * (n - 1))
2123 fstep acc (toEnum $ fromIntegral old) >>= go2 wrd (n - 1)
2124 go2 _ _ acc = return acc
2126 stepOuter gst (GO_KARP_RABIN stt rb rhead) = do
2128 res <- step (adaptState gst) stt
2132 acc' <- if withSep then fstep acc x else return acc
2133 rh' <- liftIO (RB.unsafeInsert rb rhead x)
2136 let fold = RB.unsafeFoldRing (RB.ringBound rb)
2137 let !ringHash = fold addCksum 0 rb
2138 if ringHash == patHash
2139 then go2 SPEC ringHash rh' s acc'
2140 else go0 SPEC (idx + 1) rh' s acc'
2141 else go0 SPEC (idx + 1) rh' s acc'
2142 Skip s -> return $ Skip (GO_KARP_RABIN s rb rhead)
2147 k = 2891336453 :: Word32
2149 addCksum cksum a = cksum * k + fromIntegral (fromEnum a)
2150 deltaCksum cksum old new =
2151 addCksum cksum new - coeff * fromIntegral (fromEnum old)
2153 -- XXX shall we use a random starting hash or 1 instead of 0?
2154 patHash = A.foldl' addCksum 0 patArr
2157 go0 !_ !idx !rh st !acc = do
2158 res <- step (adaptState gst) st
2161 acc' <- if withSep then fstep acc x else return acc
2162 rh' <- liftIO (RB.unsafeInsert rb rh x)
2165 let fold = RB.unsafeFoldRing (RB.ringBound rb)
2166 let !ringHash = fold addCksum 0 rb
2167 if ringHash == patHash
2168 then go2 SPEC ringHash rh' s acc'
2169 else go1 SPEC ringHash rh' s acc'
2170 else go0 SPEC (idx + 1) rh' s acc'
2171 Skip s -> go0 SPEC idx rh s acc
2173 -- do not issue a blank segment when we end at pattern
2174 if (idx == maxIndex) && RB.unsafeEqArray rb rh patArr
2177 !acc' <- if idx /= 0 && not withSep
2178 then RB.unsafeFoldRingM rh fstep acc rb
2180 done acc' >>= \r -> return $ Yield r GO_DONE
2182 -- XXX Theoretically this code can do 4 times faster if GHC generates
2183 -- optimal code. If we use just "(cksum' == patHash)" condition it goes
2184 -- 4x faster, as soon as we add the "RB.unsafeEqArray rb v" condition
2185 -- the generated code changes drastically and becomes 4x slower. Need
2186 -- to investigate what is going on with GHC.
2188 go1 !_ !cksum !rh st !acc = do
2189 res <- step (adaptState gst) st
2192 old <- liftIO $ peek rh
2193 let cksum' = deltaCksum cksum old x
2198 if (cksum' == patHash)
2200 rh' <- liftIO (RB.unsafeInsert rb rh x)
2201 go2 SPEC cksum' rh' s acc'
2203 rh' <- liftIO (RB.unsafeInsert rb rh x)
2204 go1 SPEC cksum' rh' s acc'
2205 Skip s -> go1 SPEC cksum rh s acc
2207 if RB.unsafeEqArray rb rh patArr
2212 else RB.unsafeFoldRingFullM rh fstep acc rb
2213 done acc' >>= \r -> return $ Yield r GO_DONE
2215 go2 !_ !cksum' !rh' s !acc' = do
2216 if RB.unsafeEqArray rb rh' patArr
2219 return $ Yield r (GO_KARP_RABIN s rb rhead)
2220 else go1 SPEC cksum' rh' s acc'
2222 stepOuter gst (GO_EMPTY_PAT st) = do
2223 res <- step (adaptState gst) st
2228 done acc' >>= \r -> return $ Yield r (GO_EMPTY_PAT s)
2229 Skip s -> return $ Skip (GO_EMPTY_PAT s)
2232 stepOuter _ GO_DONE = return Stop
2234data SplitState s arr
2236 | SplitBuffering s arr
2237 | SplitSplitting s arr
2238 | SplitYielding arr (SplitState s arr)
2241-- XXX An alternative approach would be to use a partial fold (Fold m a b) to
2242-- split using a splitBy like combinator. The Fold would consume upto the
2243-- separator and return any leftover which can then be fed to the next fold.
2245-- We can revisit this once we have partial folds/parsers.
2247-- | Performs infix separator style splitting.
2248{-# INLINE_NORMAL splitInnerBy #-}
2251 => (f a -> m (f a, Maybe (f a))) -- splitter
2252 -> (f a -> f a -> m (f a)) -- joiner
2255splitInnerBy splitter joiner (Stream step1 state1) =
2256 (Stream step (SplitInitial state1))
2260 {-# INLINE_LATE step #-}
2261 step gst (SplitInitial st) = do
2265 (x1, mx2) <- splitter x
2266 return $ case mx2 of
2267 Nothing -> Skip (SplitBuffering s x1)
2268 Just x2 -> Skip (SplitYielding x1 (SplitSplitting s x2))
2269 Skip s -> return $ Skip (SplitInitial s)
2270 Stop -> return $ Stop
2272 step gst (SplitBuffering st buf) = do
2276 (x1, mx2) <- splitter x
2277 buf' <- joiner buf x1
2278 return $ case mx2 of
2279 Nothing -> Skip (SplitBuffering s buf')
2280 Just x2 -> Skip (SplitYielding buf' (SplitSplitting s x2))
2281 Skip s -> return $ Skip (SplitBuffering s buf)
2282 Stop -> return $ Skip (SplitYielding buf SplitFinishing)
2284 step _ (SplitSplitting st buf) = do
2285 (x1, mx2) <- splitter buf
2286 return $ case mx2 of
2287 Nothing -> Skip $ SplitBuffering st x1
2288 Just x2 -> Skip $ SplitYielding x1 (SplitSplitting st x2)
2290 step _ (SplitYielding x next) = return $ Yield x next
2291 step _ SplitFinishing = return $ Stop
2293-- | Performs infix separator style splitting.
2294{-# INLINE_NORMAL splitInnerBySuffix #-}
2296 :: (Monad m, Eq (f a), Monoid (f a))
2297 => (f a -> m (f a, Maybe (f a))) -- splitter
2298 -> (f a -> f a -> m (f a)) -- joiner
2301splitInnerBySuffix splitter joiner (Stream step1 state1) =
2302 (Stream step (SplitInitial state1))
2306 {-# INLINE_LATE step #-}
2307 step gst (SplitInitial st) = do
2311 (x1, mx2) <- splitter x
2312 return $ case mx2 of
2313 Nothing -> Skip (SplitBuffering s x1)
2314 Just x2 -> Skip (SplitYielding x1 (SplitSplitting s x2))
2315 Skip s -> return $ Skip (SplitInitial s)
2316 Stop -> return $ Stop
2318 step gst (SplitBuffering st buf) = do
2322 (x1, mx2) <- splitter x
2323 buf' <- joiner buf x1
2324 return $ case mx2 of
2325 Nothing -> Skip (SplitBuffering s buf')
2326 Just x2 -> Skip (SplitYielding buf' (SplitSplitting s x2))
2327 Skip s -> return $ Skip (SplitBuffering s buf)
2331 else Skip (SplitYielding buf SplitFinishing)
2333 step _ (SplitSplitting st buf) = do
2334 (x1, mx2) <- splitter buf
2335 return $ case mx2 of
2336 Nothing -> Skip $ SplitBuffering st x1
2337 Just x2 -> Skip $ SplitYielding x1 (SplitSplitting st x2)
2339 step _ (SplitYielding x next) = return $ Yield x next
2340 step _ SplitFinishing = return $ Stop
2342------------------------------------------------------------------------------
2344------------------------------------------------------------------------------
2346{-# INLINE_NORMAL isPrefixOf #-}
2347isPrefixOf :: (Eq a, Monad m) => Stream m a -> Stream m a -> m Bool
2348isPrefixOf (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
2350 go (sa, sb, Nothing) = do
2351 r <- stepa defState sa
2353 Yield x sa' -> go (sa', sb, Just x)
2354 Skip sa' -> go (sa', sb, Nothing)
2357 go (sa, sb, Just x) = do
2358 r <- stepb defState sb
2362 then go (sa, sb', Nothing)
2364 Skip sb' -> go (sa, sb', Just x)
2365 Stop -> return False
2367{-# INLINE_NORMAL isSubsequenceOf #-}
2368isSubsequenceOf :: (Eq a, Monad m) => Stream m a -> Stream m a -> m Bool
2369isSubsequenceOf (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
2371 go (sa, sb, Nothing) = do
2372 r <- stepa defState sa
2374 Yield x sa' -> go (sa', sb, Just x)
2375 Skip sa' -> go (sa', sb, Nothing)
2378 go (sa, sb, Just x) = do
2379 r <- stepb defState sb
2383 then go (sa, sb', Nothing)
2384 else go (sa, sb', Just x)
2385 Skip sb' -> go (sa, sb', Just x)
2386 Stop -> return False
2388{-# INLINE_NORMAL stripPrefix #-}
2391 => Stream m a -> Stream m a -> m (Maybe (Stream m a))
2392stripPrefix (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
2394 go (sa, sb, Nothing) = do
2395 r <- stepa defState sa
2397 Yield x sa' -> go (sa', sb, Just x)
2398 Skip sa' -> go (sa', sb, Nothing)
2399 Stop -> return $ Just (Stream stepb sb)
2401 go (sa, sb, Just x) = do
2402 r <- stepb defState sb
2406 then go (sa, sb', Nothing)
2408 Skip sb' -> go (sa, sb', Just x)
2409 Stop -> return Nothing
2411------------------------------------------------------------------------------
2413------------------------------------------------------------------------------
2415-- | Execute a monadic action for each element of the 'Stream'
2416{-# INLINE_NORMAL mapM_ #-}
2417mapM_ :: Monad m => (a -> m b) -> Stream m a -> m ()
2418mapM_ m = drain . mapM m
2420-------------------------------------------------------------------------------
2421-- Stream transformations using Unfolds
2422-------------------------------------------------------------------------------
2424-- Define a unique structure to use in inspection testing
2425data ConcatMapUState o i =
2427 | ConcatMapUInner o i
2429-- | @concatMapU unfold stream@ uses @unfold@ to map the input stream elements
2430-- to streams and then flattens the generated streams into a single output
2433-- This is like 'concatMap' but uses an unfold with an explicit state to
2434-- generate the stream instead of a 'Stream' type generator. This allows better
2435-- optimization via fusion. This can be many times more efficient than
2438{-# INLINE_NORMAL concatMapU #-}
2439concatMapU :: Monad m => Unfold m a b -> Stream m a -> Stream m b
2440concatMapU (Unfold istep inject) (Stream ostep ost) =
2441 Stream step (ConcatMapUOuter ost)
2443 {-# INLINE_LATE step #-}
2444 step gst (ConcatMapUOuter o) = do
2445 r <- ostep (adaptState gst) o
2449 i `seq` return (Skip (ConcatMapUInner o' i))
2450 Skip o' -> return $ Skip (ConcatMapUOuter o')
2451 Stop -> return $ Stop
2453 step _ (ConcatMapUInner o i) = do
2456 Yield x i' -> Yield x (ConcatMapUInner o i')
2457 Skip i' -> Skip (ConcatMapUInner o i')
2458 Stop -> Skip (ConcatMapUOuter o)
2460data ConcatUnfoldInterleaveState o i =
2461 ConcatUnfoldInterleaveOuter o [i]
2462 | ConcatUnfoldInterleaveInner o [i]
2463 | ConcatUnfoldInterleaveInnerL [i] [i]
2464 | ConcatUnfoldInterleaveInnerR [i] [i]
2466-- XXX use arrays to store state instead of lists.
2467-- XXX In general we can use different scheduling strategies e.g. how to
2468-- schedule the outer vs inner loop or assigning weights to different streams
2469-- or outer and inner loops.
2471-- After a yield, switch to the next stream. Do not switch streams on Skip.
2472-- Yield from outer stream switches to the inner stream.
2474-- There are two choices here, (1) exhaust the outer stream first and then
2475-- start yielding from the inner streams, this is much simpler to implement,
2476-- (2) yield at least one element from an inner stream before going back to
2477-- outer stream and opening the next stream from it.
2479-- Ideally, we need some scheduling bias to inner streams vs outer stream.
2480-- Maybe we can configure the behavior.
2482{-# INLINE_NORMAL concatUnfoldInterleave #-}
2483concatUnfoldInterleave :: Monad m => Unfold m a b -> Stream m a -> Stream m b
2484concatUnfoldInterleave (Unfold istep inject) (Stream ostep ost) =
2485 Stream step (ConcatUnfoldInterleaveOuter ost [])
2487 {-# INLINE_LATE step #-}
2488 step gst (ConcatUnfoldInterleaveOuter o ls) = do
2489 r <- ostep (adaptState gst) o
2493 i `seq` return (Skip (ConcatUnfoldInterleaveInner o' (i : ls)))
2494 Skip o' -> return $ Skip (ConcatUnfoldInterleaveOuter o' ls)
2495 Stop -> return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2497 step _ (ConcatUnfoldInterleaveInner _ []) = undefined
2498 step _ (ConcatUnfoldInterleaveInner o (st:ls)) = do
2501 Yield x s -> Yield x (ConcatUnfoldInterleaveOuter o (s:ls))
2502 Skip s -> Skip (ConcatUnfoldInterleaveInner o (s:ls))
2503 Stop -> Skip (ConcatUnfoldInterleaveOuter o ls)
2505 step _ (ConcatUnfoldInterleaveInnerL [] []) = return Stop
2506 step _ (ConcatUnfoldInterleaveInnerL [] rs) =
2507 return $ Skip (ConcatUnfoldInterleaveInnerR [] rs)
2509 step _ (ConcatUnfoldInterleaveInnerL (st:ls) rs) = do
2512 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerL ls (s:rs))
2513 Skip s -> Skip (ConcatUnfoldInterleaveInnerL (s:ls) rs)
2514 Stop -> Skip (ConcatUnfoldInterleaveInnerL ls rs)
2516 step _ (ConcatUnfoldInterleaveInnerR [] []) = return Stop
2517 step _ (ConcatUnfoldInterleaveInnerR ls []) =
2518 return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2520 step _ (ConcatUnfoldInterleaveInnerR ls (st:rs)) = do
2523 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerR (s:ls) rs)
2524 Skip s -> Skip (ConcatUnfoldInterleaveInnerR ls (s:rs))
2525 Stop -> Skip (ConcatUnfoldInterleaveInnerR ls rs)
2527-- XXX In general we can use different scheduling strategies e.g. how to
2528-- schedule the outer vs inner loop or assigning weights to different streams
2529-- or outer and inner loops.
2531-- This could be inefficient if the tasks are too small.
2533-- Compared to concatUnfoldInterleave this one switches streams on Skips.
2535{-# INLINE_NORMAL concatUnfoldRoundrobin #-}
2536concatUnfoldRoundrobin :: Monad m => Unfold m a b -> Stream m a -> Stream m b
2537concatUnfoldRoundrobin (Unfold istep inject) (Stream ostep ost) =
2538 Stream step (ConcatUnfoldInterleaveOuter ost [])
2540 {-# INLINE_LATE step #-}
2541 step gst (ConcatUnfoldInterleaveOuter o ls) = do
2542 r <- ostep (adaptState gst) o
2546 i `seq` return (Skip (ConcatUnfoldInterleaveInner o' (i : ls)))
2547 Skip o' -> return $ Skip (ConcatUnfoldInterleaveInner o' ls)
2548 Stop -> return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2550 step _ (ConcatUnfoldInterleaveInner o []) =
2551 return $ Skip (ConcatUnfoldInterleaveOuter o [])
2553 step _ (ConcatUnfoldInterleaveInner o (st:ls)) = do
2556 Yield x s -> Yield x (ConcatUnfoldInterleaveOuter o (s:ls))
2557 Skip s -> Skip (ConcatUnfoldInterleaveOuter o (s:ls))
2558 Stop -> Skip (ConcatUnfoldInterleaveOuter o ls)
2560 step _ (ConcatUnfoldInterleaveInnerL [] []) = return Stop
2561 step _ (ConcatUnfoldInterleaveInnerL [] rs) =
2562 return $ Skip (ConcatUnfoldInterleaveInnerR [] rs)
2564 step _ (ConcatUnfoldInterleaveInnerL (st:ls) rs) = do
2567 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerL ls (s:rs))
2568 Skip s -> Skip (ConcatUnfoldInterleaveInnerL ls (s:rs))
2569 Stop -> Skip (ConcatUnfoldInterleaveInnerL ls rs)
2571 step _ (ConcatUnfoldInterleaveInnerR [] []) = return Stop
2572 step _ (ConcatUnfoldInterleaveInnerR ls []) =
2573 return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2575 step _ (ConcatUnfoldInterleaveInnerR ls (st:rs)) = do
2578 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerR (s:ls) rs)
2579 Skip s -> Skip (ConcatUnfoldInterleaveInnerR (s:ls) rs)
2580 Stop -> Skip (ConcatUnfoldInterleaveInnerR ls rs)
2582data AppendState s1 s2 = AppendFirst s1 | AppendSecond s2
2584-- Note that this could be much faster compared to the CPS stream. However, as
2585-- the number of streams being composed increases this may become expensive.
2586-- Need to see where the breaking point is between the two.
2588{-# INLINE_NORMAL append #-}
2589append :: Monad m => Stream m a -> Stream m a -> Stream m a
2590append (Stream step1 state1) (Stream step2 state2) =
2591 Stream step (AppendFirst state1)
2595 {-# INLINE_LATE step #-}
2596 step gst (AppendFirst st) = do
2599 Yield a s -> Yield a (AppendFirst s)
2600 Skip s -> Skip (AppendFirst s)
2601 Stop -> Skip (AppendSecond state2)
2603 step gst (AppendSecond st) = do
2606 Yield a s -> Yield a (AppendSecond s)
2607 Skip s -> Skip (AppendSecond s)
2610data InterleaveState s1 s2 = InterleaveFirst s1 s2 | InterleaveSecond s1 s2
2611 | InterleaveSecondOnly s2 | InterleaveFirstOnly s1
2613{-# INLINE_NORMAL interleave #-}
2614interleave :: Monad m => Stream m a -> Stream m a -> Stream m a
2615interleave (Stream step1 state1) (Stream step2 state2) =
2616 Stream step (InterleaveFirst state1 state2)
2620 {-# INLINE_LATE step #-}
2621 step gst (InterleaveFirst st1 st2) = do
2624 Yield a s -> Yield a (InterleaveSecond s st2)
2625 Skip s -> Skip (InterleaveFirst s st2)
2626 Stop -> Skip (InterleaveSecondOnly st2)
2628 step gst (InterleaveSecond st1 st2) = do
2631 Yield a s -> Yield a (InterleaveFirst st1 s)
2632 Skip s -> Skip (InterleaveSecond st1 s)
2633 Stop -> Skip (InterleaveFirstOnly st1)
2635 step gst (InterleaveFirstOnly st1) = do
2638 Yield a s -> Yield a (InterleaveFirstOnly s)
2639 Skip s -> Skip (InterleaveFirstOnly s)
2642 step gst (InterleaveSecondOnly st2) = do
2645 Yield a s -> Yield a (InterleaveSecondOnly s)
2646 Skip s -> Skip (InterleaveSecondOnly s)
2649{-# INLINE_NORMAL interleaveMin #-}
2650interleaveMin :: Monad m => Stream m a -> Stream m a -> Stream m a
2651interleaveMin (Stream step1 state1) (Stream step2 state2) =
2652 Stream step (InterleaveFirst state1 state2)
2656 {-# INLINE_LATE step #-}
2657 step gst (InterleaveFirst st1 st2) = do
2660 Yield a s -> Yield a (InterleaveSecond s st2)
2661 Skip s -> Skip (InterleaveFirst s st2)
2664 step gst (InterleaveSecond st1 st2) = do
2667 Yield a s -> Yield a (InterleaveFirst st1 s)
2668 Skip s -> Skip (InterleaveSecond st1 s)
2671 step _ (InterleaveFirstOnly _) = undefined
2672 step _ (InterleaveSecondOnly _) = undefined
2674{-# INLINE_NORMAL interleaveSuffix #-}
2675interleaveSuffix :: Monad m => Stream m a -> Stream m a -> Stream m a
2676interleaveSuffix (Stream step1 state1) (Stream step2 state2) =
2677 Stream step (InterleaveFirst state1 state2)
2681 {-# INLINE_LATE step #-}
2682 step gst (InterleaveFirst st1 st2) = do
2685 Yield a s -> Yield a (InterleaveSecond s st2)
2686 Skip s -> Skip (InterleaveFirst s st2)
2689 step gst (InterleaveSecond st1 st2) = do
2692 Yield a s -> Yield a (InterleaveFirst st1 s)
2693 Skip s -> Skip (InterleaveSecond st1 s)
2694 Stop -> Skip (InterleaveFirstOnly st1)
2696 step gst (InterleaveFirstOnly st1) = do
2699 Yield a s -> Yield a (InterleaveFirstOnly s)
2700 Skip s -> Skip (InterleaveFirstOnly s)
2703 step _ (InterleaveSecondOnly _) = undefined
2705data InterleaveInfixState s1 s2 a
2706 = InterleaveInfixFirst s1 s2
2707 | InterleaveInfixSecondBuf s1 s2
2708 | InterleaveInfixSecondYield s1 s2 a
2709 | InterleaveInfixFirstYield s1 s2 a
2710 | InterleaveInfixFirstOnly s1
2712{-# INLINE_NORMAL interleaveInfix #-}
2713interleaveInfix :: Monad m => Stream m a -> Stream m a -> Stream m a
2714interleaveInfix (Stream step1 state1) (Stream step2 state2) =
2715 Stream step (InterleaveInfixFirst state1 state2)
2719 {-# INLINE_LATE step #-}
2720 step gst (InterleaveInfixFirst st1 st2) = do
2723 Yield a s -> Yield a (InterleaveInfixSecondBuf s st2)
2724 Skip s -> Skip (InterleaveInfixFirst s st2)
2727 step gst (InterleaveInfixSecondBuf st1 st2) = do
2730 Yield a s -> Skip (InterleaveInfixSecondYield st1 s a)
2731 Skip s -> Skip (InterleaveInfixSecondBuf st1 s)
2732 Stop -> Skip (InterleaveInfixFirstOnly st1)
2734 step gst (InterleaveInfixSecondYield st1 st2 x) = do
2737 Yield a s -> Yield x (InterleaveInfixFirstYield s st2 a)
2738 Skip s -> Skip (InterleaveInfixSecondYield s st2 x)
2741 step _ (InterleaveInfixFirstYield st1 st2 x) = do
2742 return $ Yield x (InterleaveInfixSecondBuf st1 st2)
2744 step gst (InterleaveInfixFirstOnly st1) = do
2747 Yield a s -> Yield a (InterleaveInfixFirstOnly s)
2748 Skip s -> Skip (InterleaveInfixFirstOnly s)
2751{-# INLINE_NORMAL roundRobin #-}
2752roundRobin :: Monad m => Stream m a -> Stream m a -> Stream m a
2753roundRobin (Stream step1 state1) (Stream step2 state2) =
2754 Stream step (InterleaveFirst state1 state2)
2758 {-# INLINE_LATE step #-}
2759 step gst (InterleaveFirst st1 st2) = do
2762 Yield a s -> Yield a (InterleaveSecond s st2)
2763 Skip s -> Skip (InterleaveSecond s st2)
2764 Stop -> Skip (InterleaveSecondOnly st2)
2766 step gst (InterleaveSecond st1 st2) = do
2769 Yield a s -> Yield a (InterleaveFirst st1 s)
2770 Skip s -> Skip (InterleaveFirst st1 s)
2771 Stop -> Skip (InterleaveFirstOnly st1)
2773 step gst (InterleaveSecondOnly st2) = do
2776 Yield a s -> Yield a (InterleaveSecondOnly s)
2777 Skip s -> Skip (InterleaveSecondOnly s)
2780 step gst (InterleaveFirstOnly st1) = do
2783 Yield a s -> Yield a (InterleaveFirstOnly s)
2784 Skip s -> Skip (InterleaveFirstOnly s)
2787data ICUState s1 s2 i1 i2 =
2792 | ICUFirstInner s1 s2 i1
2793 | ICUSecondInner s1 s2 i2
2794 | ICUFirstOnlyInner s1 i1
2795 | ICUSecondOnlyInner s2 i2
2797-- | Interleave streams (full streams, not the elements) unfolded from two
2798-- input streams and concat. Stop when the first stream stops. If the second
2799-- stream ends before the first one then first stream still keeps running alone
2800-- without any interleaving with the second stream.
2802-- [a1, a2, ... an] [b1, b2 ...]
2803-- => [streamA1, streamA2, ... streamAn] [streamB1, streamB2, ...]
2804-- => [streamA1, streamB1, streamA2...StreamAn, streamBn]
2805-- => [a11, a12, ...a1j, b11, b12, ...b1k, a21, a22, ...]
2807{-# INLINE_NORMAL gintercalateSuffix #-}
2810 => Unfold m a c -> Stream m a -> Unfold m b c -> Stream m b -> Stream m c
2812 (Unfold istep1 inject1) (Stream step1 state1)
2813 (Unfold istep2 inject2) (Stream step2 state2) =
2814 Stream step (ICUFirst state1 state2)
2818 {-# INLINE_LATE step #-}
2819 step gst (ICUFirst s1 s2) = do
2820 r <- step1 (adaptState gst) s1
2824 i `seq` return (Skip (ICUFirstInner s s2 i))
2825 Skip s -> return $ Skip (ICUFirst s s2)
2828 step gst (ICUFirstOnly s1) = do
2829 r <- step1 (adaptState gst) s1
2833 i `seq` return (Skip (ICUFirstOnlyInner s i))
2834 Skip s -> return $ Skip (ICUFirstOnly s)
2837 step _ (ICUFirstInner s1 s2 i1) = do
2840 Yield x i' -> Yield x (ICUFirstInner s1 s2 i')
2841 Skip i' -> Skip (ICUFirstInner s1 s2 i')
2842 Stop -> Skip (ICUSecond s1 s2)
2844 step _ (ICUFirstOnlyInner s1 i1) = do
2847 Yield x i' -> Yield x (ICUFirstOnlyInner s1 i')
2848 Skip i' -> Skip (ICUFirstOnlyInner s1 i')
2849 Stop -> Skip (ICUFirstOnly s1)
2851 step gst (ICUSecond s1 s2) = do
2852 r <- step2 (adaptState gst) s2
2856 i `seq` return (Skip (ICUSecondInner s1 s i))
2857 Skip s -> return $ Skip (ICUSecond s1 s)
2858 Stop -> return $ Skip (ICUFirstOnly s1)
2860 step _ (ICUSecondInner s1 s2 i2) = do
2863 Yield x i' -> Yield x (ICUSecondInner s1 s2 i')
2864 Skip i' -> Skip (ICUSecondInner s1 s2 i')
2865 Stop -> Skip (ICUFirst s1 s2)
2867 step _ (ICUSecondOnly _s2) = undefined
2868 step _ (ICUSecondOnlyInner _s2 _i2) = undefined
2870data InterposeSuffixState s1 i1 =
2871 InterposeSuffixFirst s1
2872 -- | InterposeSuffixFirstYield s1 i1
2873 | InterposeSuffixFirstInner s1 i1
2874 | InterposeSuffixSecond s1
2876-- Note that if an unfolded layer turns out to be nil we still emit the
2877-- separator effect. An alternate behavior could be to emit the separator
2878-- effect only if at least one element has been yielded by the unfolding.
2879-- However, that becomes a bit complicated, so we have chosen the former
2881{-# INLINE_NORMAL interposeSuffix #-}
2884 => m c -> Unfold m b c -> Stream m b -> Stream m c
2887 (Unfold istep1 inject1) (Stream step1 state1) =
2888 Stream step (InterposeSuffixFirst state1)
2892 {-# INLINE_LATE step #-}
2893 step gst (InterposeSuffixFirst s1) = do
2894 r <- step1 (adaptState gst) s1
2898 i `seq` return (Skip (InterposeSuffixFirstInner s i))
2899 -- i `seq` return (Skip (InterposeSuffixFirstYield s i))
2900 Skip s -> return $ Skip (InterposeSuffixFirst s)
2904 step _ (InterposeSuffixFirstYield s1 i1) = do
2907 Yield x i' -> Yield x (InterposeSuffixFirstInner s1 i')
2908 Skip i' -> Skip (InterposeSuffixFirstYield s1 i')
2909 Stop -> Skip (InterposeSuffixFirst s1)
2912 step _ (InterposeSuffixFirstInner s1 i1) = do
2915 Yield x i' -> Yield x (InterposeSuffixFirstInner s1 i')
2916 Skip i' -> Skip (InterposeSuffixFirstInner s1 i')
2917 Stop -> Skip (InterposeSuffixSecond s1)
2919 step _ (InterposeSuffixSecond s1) = do
2921 return $ Yield r (InterposeSuffixFirst s1)
2923data ICALState s1 s2 i1 i2 a =
2925 -- | ICALFirstYield s1 s2 i1
2926 | ICALFirstInner s1 s2 i1
2928 | ICALFirstOnlyInner s1 i1
2929 | ICALSecondInject s1 s2
2930 | ICALFirstInject s1 s2 i2
2931 -- | ICALFirstBuf s1 s2 i1 i2
2932 | ICALSecondInner s1 s2 i1 i2
2933 -- -- | ICALSecondInner s1 s2 i1 i2 a
2934 -- -- | ICALFirstResume s1 s2 i1 i2 a
2936-- | Interleave streams (full streams, not the elements) unfolded from two
2937-- input streams and concat. Stop when the first stream stops. If the second
2938-- stream ends before the first one then first stream still keeps running alone
2939-- without any interleaving with the second stream.
2941-- [a1, a2, ... an] [b1, b2 ...]
2942-- => [streamA1, streamA2, ... streamAn] [streamB1, streamB2, ...]
2943-- => [streamA1, streamB1, streamA2...StreamAn, streamBn]
2944-- => [a11, a12, ...a1j, b11, b12, ...b1k, a21, a22, ...]
2946{-# INLINE_NORMAL gintercalate #-}
2949 => Unfold m a c -> Stream m a -> Unfold m b c -> Stream m b -> Stream m c
2951 (Unfold istep1 inject1) (Stream step1 state1)
2952 (Unfold istep2 inject2) (Stream step2 state2) =
2953 Stream step (ICALFirst state1 state2)
2957 {-# INLINE_LATE step #-}
2958 step gst (ICALFirst s1 s2) = do
2959 r <- step1 (adaptState gst) s1
2963 i `seq` return (Skip (ICALFirstInner s s2 i))
2964 -- i `seq` return (Skip (ICALFirstYield s s2 i))
2965 Skip s -> return $ Skip (ICALFirst s s2)
2969 step _ (ICALFirstYield s1 s2 i1) = do
2972 Yield x i' -> Yield x (ICALFirstInner s1 s2 i')
2973 Skip i' -> Skip (ICALFirstYield s1 s2 i')
2974 Stop -> Skip (ICALFirst s1 s2)
2977 step _ (ICALFirstInner s1 s2 i1) = do
2980 Yield x i' -> Yield x (ICALFirstInner s1 s2 i')
2981 Skip i' -> Skip (ICALFirstInner s1 s2 i')
2982 Stop -> Skip (ICALSecondInject s1 s2)
2984 step gst (ICALFirstOnly s1) = do
2985 r <- step1 (adaptState gst) s1
2989 i `seq` return (Skip (ICALFirstOnlyInner s i))
2990 Skip s -> return $ Skip (ICALFirstOnly s)
2993 step _ (ICALFirstOnlyInner s1 i1) = do
2996 Yield x i' -> Yield x (ICALFirstOnlyInner s1 i')
2997 Skip i' -> Skip (ICALFirstOnlyInner s1 i')
2998 Stop -> Skip (ICALFirstOnly s1)
3000 -- We inject the second stream even before checking if the first stream
3001 -- would yield any more elements. There is no clear choice whether we
3002 -- should do this before or after that. Doing it after may make the state
3003 -- machine a bit simpler though.
3004 step gst (ICALSecondInject s1 s2) = do
3005 r <- step2 (adaptState gst) s2
3009 i `seq` return (Skip (ICALFirstInject s1 s i))
3010 Skip s -> return $ Skip (ICALSecondInject s1 s)
3011 Stop -> return $ Skip (ICALFirstOnly s1)
3013 step gst (ICALFirstInject s1 s2 i2) = do
3014 r <- step1 (adaptState gst) s1
3018 i `seq` return (Skip (ICALSecondInner s s2 i i2))
3019 -- i `seq` return (Skip (ICALFirstBuf s s2 i i2))
3020 Skip s -> return $ Skip (ICALFirstInject s s2 i2)
3024 step _ (ICALFirstBuf s1 s2 i1 i2) = do
3027 Yield x i' -> Skip (ICALSecondInner s1 s2 i' i2 x)
3028 Skip i' -> Skip (ICALFirstBuf s1 s2 i' i2)
3031 step _ (ICALSecondInner s1 s2 i1 i2 v) = do
3034 Yield x i' -> Yield x (ICALSecondInner s1 s2 i1 i' v)
3035 Skip i' -> Skip (ICALSecondInner s1 s2 i1 i' v)
3036 Stop -> Skip (ICALFirstResume s1 s2 i1 i2 v)
3039 step _ (ICALSecondInner s1 s2 i1 i2) = do
3042 Yield x i' -> Yield x (ICALSecondInner s1 s2 i1 i')
3043 Skip i' -> Skip (ICALSecondInner s1 s2 i1 i')
3044 Stop -> Skip (ICALFirstInner s1 s2 i1)
3045 -- Stop -> Skip (ICALFirstResume s1 s2 i1 i2)
3048 step _ (ICALFirstResume s1 s2 i1 i2 x) = do
3049 return $ Yield x (ICALFirstInner s1 s2 i1 i2)
3052data InterposeState s1 i1 a =
3054 -- | InterposeFirstYield s1 i1
3055 | InterposeFirstInner s1 i1
3056 | InterposeFirstInject s1
3057 -- | InterposeFirstBuf s1 i1
3058 | InterposeSecondYield s1 i1
3059 -- -- | InterposeSecondYield s1 i1 a
3060 -- -- | InterposeFirstResume s1 i1 a
3062-- Note that this only interposes the pure values, we may run many effects to
3063-- generate those values as some effects may not generate anything (Skip).
3064{-# INLINE_NORMAL interpose #-}
3065interpose :: Monad m => m c -> Unfold m b c -> Stream m b -> Stream m c
3068 (Unfold istep1 inject1) (Stream step1 state1) =
3069 Stream step (InterposeFirst state1)
3073 {-# INLINE_LATE step #-}
3074 step gst (InterposeFirst s1) = do
3075 r <- step1 (adaptState gst) s1
3079 i `seq` return (Skip (InterposeFirstInner s i))
3080 -- i `seq` return (Skip (InterposeFirstYield s i))
3081 Skip s -> return $ Skip (InterposeFirst s)
3085 step _ (InterposeFirstYield s1 i1) = do
3088 Yield x i' -> Yield x (InterposeFirstInner s1 i')
3089 Skip i' -> Skip (InterposeFirstYield s1 i')
3090 Stop -> Skip (InterposeFirst s1)
3093 step _ (InterposeFirstInner s1 i1) = do
3096 Yield x i' -> Yield x (InterposeFirstInner s1 i')
3097 Skip i' -> Skip (InterposeFirstInner s1 i')
3098 Stop -> Skip (InterposeFirstInject s1)
3100 step gst (InterposeFirstInject s1) = do
3101 r <- step1 (adaptState gst) s1
3105 -- i `seq` return (Skip (InterposeFirstBuf s i))
3106 i `seq` return (Skip (InterposeSecondYield s i))
3107 Skip s -> return $ Skip (InterposeFirstInject s)
3111 step _ (InterposeFirstBuf s1 i1) = do
3114 Yield x i' -> Skip (InterposeSecondYield s1 i' x)
3115 Skip i' -> Skip (InterposeFirstBuf s1 i')
3120 step _ (InterposeSecondYield s1 i1 v) = do
3122 return $ Yield r (InterposeFirstResume s1 i1 v)
3124 step _ (InterposeSecondYield s1 i1) = do
3126 return $ Yield r (InterposeFirstInner s1 i1)
3129 step _ (InterposeFirstResume s1 i1 v) = do
3130 return $ Yield v (InterposeFirstInner s1 i1)
3133------------------------------------------------------------------------------
3135------------------------------------------------------------------------------
3137data GbracketState s1 s2 v
3139 | GBracketNormal s1 v
3140 | GBracketException s2
3142-- | The most general bracketing and exception combinator. All other
3143-- combinators can be expressed in terms of this combinator. This can also be
3144-- used for cases which are not covered by the standard combinators.
3148{-# INLINE_NORMAL gbracket_ #-}
3152 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
3153 -> (c -> m d) -- ^ after, on normal stop
3154 -> (c -> e -> Stream m b -> Stream m b) -- ^ on exception
3155 -> (c -> Stream m b) -- ^ stream generator
3157gbracket_ bef exc aft fexc fnormal =
3158 Stream step GBracketInit
3162 {-# INLINE_LATE step #-}
3163 step _ GBracketInit = do
3165 return $ Skip $ GBracketNormal (fnormal r) r
3167 step gst (GBracketNormal (UnStream step1 st) v) = do
3168 res <- exc $ step1 gst st
3170 Right r -> case r of
3172 return $ Yield x (GBracketNormal (Stream step1 s) v)
3173 Skip s -> return $ Skip (GBracketNormal (Stream step1 s) v)
3174 Stop -> aft v >> return Stop
3176 return $ Skip (GBracketException (fexc v e (UnStream step1 st)))
3177 step gst (GBracketException (UnStream step1 st)) = do
3180 Yield x s -> return $ Yield x (GBracketException (Stream step1 s))
3181 Skip s -> return $ Skip (GBracketException (Stream step1 s))
3184------------------------------------------------------------------------------
3186------------------------------------------------------------------------------
3188-- | Make a finalizer from a monadic action @m a@ that can run in IO monad.
3189mkIOFinalizer :: MonadBaseControl IO m => m b -> m (IO ())
3191 mrun <- captureMonadState
3197-- | Run an IO action stored in a finalized IORef.
3198runIORefFinalizerGC :: IORef (Maybe (IO ())) -> IO ()
3199runIORefFinalizerGC ref = do
3200 res <- readIORef ref
3202 Nothing -> return ()
3205-- | Create an IORef holding a finalizer that is called automatically when
3206-- the IORef is garbage collected. The IORef can be written to with a 'Nothing'
3207-- value to deactivate the finalizer.
3209-- Note: The finalizer is always run with the state of the monad captured at
3210-- the time of calling newFinalizedIORef. To run it on garbage collection we
3211-- have no option but to take a snapshot of the monadic state at some point of
3212-- time. For normal case we could run it with the current state of the monad
3213-- but we want to keep both the cases consistent.
3215newFinalizedIORef :: (MonadIO m, MonadBaseControl IO m)
3216 => m a -> m (IORef (Maybe (IO ())))
3217newFinalizedIORef finalizer = do
3218 f <- mkIOFinalizer finalizer
3219 ref <- liftIO $ newIORef $ Just f
3220 _ <- liftIO $ mkWeakIORef ref (runIORefFinalizerGC ref)
3223-- | Run the finalizer stored in an IORef and deactivate it so that it is run
3224-- only once. Note, the action runs with async exceptions masked.
3226runIORefFinalizer :: MonadIO m => IORef (Maybe (IO ())) -> m ()
3227runIORefFinalizer ref = liftIO $ do
3228 res <- readIORef ref
3230 Nothing -> return ()
3232 -- if an async exception comes after writing 'Nothing' then the
3233 -- finalizing action will never be run. We need to do this
3234 -- atomically wrt async exceptions.
3236 writeIORef ref Nothing
3239-- | Run an action clearing the finalizer IORef atomically wrt async
3240-- exceptions. The action is run with async exceptions masked.
3241withIORefFinalizer :: MonadBaseControl IO m
3242 => IORef (Maybe (IO ())) -> m a -> m a
3243withIORefFinalizer ref action = do
3244 control $ \runinio ->
3246 writeIORef ref Nothing
3249------------------------------------------------------------------------------
3251data GbracketIOState s1 s2 v wref
3253 | GBracketIONormal s1 v wref
3254 | GBracketIOException s2
3256-- | The most general bracketing and exception combinator. All other
3257-- combinators can be expressed in terms of this combinator. This can also be
3258-- used for cases which are not covered by the standard combinators.
3260-- | It uses a finalizer to make sure we run the finalizing action
3261-- when the stream is garbage collected.
3265{-# INLINE_NORMAL gbracket #-}
3267 :: (MonadIO m, MonadBaseControl IO m)
3269 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
3270 -> (c -> m d) -- ^ after, on normal stop or GC
3271 -> (c -> e -> Stream m b -> m (Stream m b)) -- ^ on exception
3272 -> (c -> Stream m b) -- ^ stream generator
3274gbracket bef exc aft fexc fnormal =
3275 Stream step GBracketIOInit
3279 -- If the stream is never evaluated the "aft" action will never be
3280 -- called. For that to occur we will need the user of this API to pass a
3281 -- weak pointer to us.
3282 {-# INLINE_LATE step #-}
3283 step _ GBracketIOInit = do
3284 -- We mask asynchronous exceptions to make the execution
3285 -- of 'bef' and the registration of 'aft' atomic.
3286 -- A similar thing is done in the resourcet package: https://git.io/JvKV3
3287 -- Tutorial: https://markkarpov.com/tutorial/exceptions.html
3288 (r, ref) <- liftBaseOp_ mask_ $ do
3290 ref <- newFinalizedIORef (aft r)
3292 return $ Skip $ GBracketIONormal (fnormal r) r ref
3294 step gst (GBracketIONormal (UnStream step1 st) v ref) = do
3295 res <- exc $ step1 gst st
3297 Right r -> case r of
3299 return $ Yield x (GBracketIONormal (Stream step1 s) v ref)
3301 return $ Skip (GBracketIONormal (Stream step1 s) v ref)
3303 runIORefFinalizer ref
3306 -- Clearing of finalizer and running of exception handler must
3307 -- be atomic wrt async exceptions. Otherwise if we have cleared
3308 -- the finalizer and have not run the exception handler then we
3309 -- may leak the resource.
3310 stream <- withIORefFinalizer ref (fexc v e (UnStream step1 st))
3311 return $ Skip (GBracketIOException stream)
3312 step gst (GBracketIOException (UnStream step1 st)) = do
3316 return $ Yield x (GBracketIOException (Stream step1 s))
3317 Skip s -> return $ Skip (GBracketIOException (Stream step1 s))
3320-- | Run a side effect before the stream yields its first element.
3321{-# INLINE_NORMAL before #-}
3322before :: Monad m => m b -> Stream m a -> Stream m a
3323before action (Stream step state) = Stream step' Nothing
3327 {-# INLINE_LATE step' #-}
3328 step' _ Nothing = action >> return (Skip (Just state))
3330 step' gst (Just st) = do
3333 Yield x s -> return $ Yield x (Just s)
3334 Skip s -> return $ Skip (Just s)
3337-- | Run a side effect whenever the stream stops normally.
3338{-# INLINE_NORMAL after_ #-}
3339after_ :: Monad m => m b -> Stream m a -> Stream m a
3340after_ action (Stream step state) = Stream step' state
3344 {-# INLINE_LATE step' #-}
3348 Yield x s -> return $ Yield x s
3349 Skip s -> return $ Skip s
3350 Stop -> action >> return Stop
3352{-# INLINE_NORMAL after #-}
3353after :: (MonadIO m, MonadBaseControl IO m)
3354 => m b -> Stream m a -> Stream m a
3355after action (Stream step state) = Stream step' Nothing
3359 {-# INLINE_LATE step' #-}
3360 step' _ Nothing = do
3361 ref <- newFinalizedIORef action
3362 return $ Skip $ Just (state, ref)
3363 step' gst (Just (st, ref)) = do
3366 Yield x s -> return $ Yield x (Just (s, ref))
3367 Skip s -> return $ Skip (Just (s, ref))
3369 runIORefFinalizer ref
3372-- XXX These combinators are expensive due to the call to
3373-- onException/handle/try on each step. Therefore, when possible, they should
3374-- be called in an outer loop where we perform less iterations. For example, we
3375-- cannot call them on each iteration in a char stream, instead we can call
3376-- them when doing an IO on an array.
3378-- XXX For high performance error checks in busy streams we may need another
3379-- Error constructor in step.
3381-- | Run a side effect whenever the stream aborts due to an exception. The
3382-- exception is not caught, simply rethrown.
3383{-# INLINE_NORMAL onException #-}
3384onException :: MonadCatch m => m b -> Stream m a -> Stream m a
3385onException action str =
3386 gbracket_ (return ()) MC.try return
3387 (\_ (e :: MC.SomeException) _ -> nilM (action >> MC.throwM e))
3390{-# INLINE_NORMAL _onException #-}
3391_onException :: MonadCatch m => m b -> Stream m a -> Stream m a
3392_onException action (Stream step state) = Stream step' state
3396 {-# INLINE_LATE step' #-}
3398 res <- step gst st `MC.onException` action
3400 Yield x s -> return $ Yield x s
3401 Skip s -> return $ Skip s
3404-- XXX bracket is like concatMap, it generates a stream and then flattens it.
3405-- Like concatMap it has 10x worse performance compared to linear fused
3408-- | Run the first action before the stream starts and remember its output,
3409-- generate a stream using the output, run the second action providing the
3410-- remembered value as an argument whenever the stream ends normally or due to
3412{-# INLINE_NORMAL bracket_ #-}
3413bracket_ :: MonadCatch m => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
3414bracket_ bef aft bet =
3415 gbracket_ bef MC.try aft
3416 (\a (e :: SomeException) _ -> nilM (aft a >> MC.throwM e)) bet
3418-- | Run the first action before the stream starts and remember its output,
3419-- generate a stream using the output, run the second action providing the
3420-- remembered value as an argument whenever the stream ends normally or due to
3421-- an exception or gets garbage collected.
3423{-# INLINE_NORMAL bracket #-}
3424bracket :: (MonadAsync m, MonadCatch m)
3425 => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
3426bracket bef aft bet =
3427 gbracket bef MC.try aft
3428 (\a (e :: SomeException) _ -> aft a >> return (nilM (MC.throwM e))) bet
3430data BracketState s v = BracketInit | BracketRun s v
3432{-# INLINE_NORMAL _bracket #-}
3433_bracket :: MonadCatch m => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
3434_bracket bef aft bet = Stream step' BracketInit
3438 {-# INLINE_LATE step' #-}
3439 step' _ BracketInit = bef >>= \x -> return (Skip (BracketRun (bet x) x))
3441 -- NOTE: It is important to use UnStream instead of the Stream pattern
3442 -- here, otherwise we get huge perf degradation, see note in concatMap.
3443 step' gst (BracketRun (UnStream step state) v) = do
3444 -- res <- step gst state `MC.onException` aft v
3445 res <- MC.try $ step gst state
3447 Left (e :: SomeException) -> aft v >> MC.throwM e >> return Stop
3448 Right r -> case r of
3449 Yield x s -> return $ Yield x (BracketRun (Stream step s) v)
3450 Skip s -> return $ Skip (BracketRun (Stream step s) v)
3451 Stop -> aft v >> return Stop
3453-- | Run a side effect whenever the stream stops normally or aborts due to an
3455{-# INLINE finally_ #-}
3456finally_ :: MonadCatch m => m b -> Stream m a -> Stream m a
3457-- finally action xs = after action $ onException action xs
3458finally_ action xs = bracket_ (return ()) (\_ -> action) (const xs)
3460-- | Run a side effect whenever the stream stops normally or aborts due to an
3461-- exception or gets garbage collected.
3463{-# INLINE finally #-}
3464finally :: (MonadAsync m, MonadCatch m) => m b -> Stream m a -> Stream m a
3465finally action xs = bracket (return ()) (\_ -> action) (const xs)
3467-- | When evaluating a stream if an exception occurs, stream
3468-- evaluation aborts and the specified exception handler is run with
3469-- the exception and the Stream which threw the exception as argument.
3471{-# INLINE_NORMAL ghandle #-}
3472ghandle :: (MonadCatch m, Exception e)
3473 => (e -> Stream m a -> Stream m a) -> Stream m a -> Stream m a
3475 gbracket_ (return ()) MC.try return (\_ -> f) (\_ -> str)
3477-- | When evaluating a stream if an exception occurs, stream evaluation aborts
3478-- and the specified exception handler is run with the exception as argument.
3479{-# INLINE_NORMAL handle #-}
3480handle :: (MonadCatch m, Exception e)
3481 => (e -> Stream m a) -> Stream m a -> Stream m a
3483 gbracket_ (return ()) MC.try return (\_ e _ -> f e) (\_ -> str)
3485{-# INLINE_NORMAL _handle #-}
3486_handle :: (MonadCatch m, Exception e)
3487 => (e -> Stream m a) -> Stream m a -> Stream m a
3488_handle f (Stream step state) = Stream step' (Left state)
3492 {-# INLINE_LATE step' #-}
3493 step' gst (Left st) = do
3494 res <- MC.try $ step gst st
3496 Left e -> return $ Skip $ Right (f e)
3497 Right r -> case r of
3498 Yield x s -> return $ Yield x (Left s)
3499 Skip s -> return $ Skip (Left s)
3502 step' gst (Right (UnStream step1 st)) = do
3505 Yield x s -> return $ Yield x (Right (Stream step1 s))
3506 Skip s -> return $ Skip (Right (Stream step1 s))
3509-------------------------------------------------------------------------------
3510-- General transformation
3511-------------------------------------------------------------------------------
3513{-# INLINE_NORMAL transform #-}
3514transform :: Monad m => Pipe m a b -> Stream m a -> Stream m b
3515transform (Pipe pstep1 pstep2 pstate) (Stream step state) =
3516 Stream step' (Consume pstate, state)
3520 {-# INLINE_LATE step' #-}
3522 step' gst (Consume pst, st) = pst `seq` do
3523 r <- step (adaptState gst) st
3528 Pipe.Yield b pst' -> return $ Yield b (pst', s)
3529 Pipe.Continue pst' -> return $ Skip (pst', s)
3530 Skip s -> return $ Skip (Consume pst, s)
3533 step' _ (Produce pst, st) = pst `seq` do
3536 Pipe.Yield b pst' -> return $ Yield b (pst', st)
3537 Pipe.Continue pst' -> return $ Skip (pst', st)
3539------------------------------------------------------------------------------
3540-- Transformation by Folding (Scans)
3541------------------------------------------------------------------------------
3543------------------------------------------------------------------------------
3545------------------------------------------------------------------------------
3547-- XXX Is a prescan useful, discarding the last step does not sound useful? I
3548-- am not sure about the utility of this function, so this is implemented but
3549-- not exposed. We can expose it if someone provides good reasons why this is
3552-- XXX We have to execute the stream one step ahead to know that we are at the
3553-- last step. The vector implementation of prescan executes the last fold step
3554-- but does not yield the result. This means we have executed the effect but
3555-- discarded value. This does not sound right. In this implementation we are
3556-- not executing the last fold step.
3557{-# INLINE_NORMAL prescanlM' #-}
3558prescanlM' :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
3559prescanlM' f mz (Stream step state) = Stream step' (state, mz)
3561 {-# INLINE_LATE step' #-}
3562 step' gst (st, prev) = do
3563 r <- step (adaptState gst) st
3567 return $ Yield acc (s, f acc x)
3568 Skip s -> return $ Skip (s, prev)
3571{-# INLINE prescanl' #-}
3572prescanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
3573prescanl' f z = prescanlM' (\a b -> return (f a b)) (return z)
3575------------------------------------------------------------------------------
3576-- Monolithic postscans (postscan followed by a map)
3577------------------------------------------------------------------------------
3579-- The performance of a modular postscan followed by a map seems to be
3580-- equivalent to this monolithic scan followed by map therefore we may not need
3581-- this implementation. We just have it for performance comparison and in case
3582-- modular version does not perform well in some situation.
3584{-# INLINE_NORMAL postscanlMx' #-}
3585postscanlMx' :: Monad m
3586 => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> Stream m b
3587postscanlMx' fstep begin done (Stream step state) = do
3588 Stream step' (state, begin)
3590 {-# INLINE_LATE step' #-}
3591 step' gst (st, acc) = do
3592 r <- step (adaptState gst) st
3598 v `seq` y `seq` return (Yield v (s, return y))
3599 Skip s -> return $ Skip (s, acc)
3602{-# INLINE_NORMAL postscanlx' #-}
3603postscanlx' :: Monad m
3604 => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> Stream m b
3605postscanlx' fstep begin done s =
3606 postscanlMx' (\b a -> return (fstep b a)) (return begin) (return . done) s
3608-- XXX do we need consM strict to evaluate the begin value?
3609{-# INLINE scanlMx' #-}
3611 => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> Stream m b
3612scanlMx' fstep begin done s =
3613 (begin >>= \x -> x `seq` done x) `consM` postscanlMx' fstep begin done s
3615{-# INLINE scanlx' #-}
3617 => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> Stream m b
3618scanlx' fstep begin done s =
3619 scanlMx' (\b a -> return (fstep b a)) (return begin) (return . done) s
3621------------------------------------------------------------------------------
3623------------------------------------------------------------------------------
3625{-# INLINE_NORMAL postscanlM' #-}
3626postscanlM' :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
3627postscanlM' fstep begin (Stream step state) =
3628 Stream step' Nothing
3630 {-# INLINE_LATE step' #-}
3631 step' _ Nothing = do
3633 return $ Skip (Just (state, x))
3635 step' gst (Just (st, acc)) = do
3636 r <- step (adaptState gst) st
3640 return $ Yield y (Just (s, y))
3641 Skip s -> return $ Skip (Just (s, acc))
3644{-# INLINE_NORMAL postscanl' #-}
3645postscanl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
3646postscanl' f seed = postscanlM' (\a b -> return (f a b)) (return seed)
3648-- We can possibly have the "done" function as a Maybe to provide an option to
3649-- emit or not emit the accumulator when the stream stops.
3651-- TBD: use a single Yield point
3653{-# INLINE_NORMAL postscanlMAfter' #-}
3654postscanlMAfter' :: Monad m
3655 => (b -> a -> m b) -> m b -> (b -> m b) -> Stream m a -> Stream m b
3656postscanlMAfter' fstep initial done (Stream step1 state1) = do
3657 Stream step (Just (state1, initial))
3661 {-# INLINE_LATE step #-}
3662 step gst (Just (st, acc)) = do
3663 r <- step1 (adaptState gst) st
3668 y `seq` return (Yield y (Just (s, return y)))
3669 Skip s -> return $ Skip $ Just (s, acc)
3670 Stop -> acc >>= done >>= \res -> return (Yield res Nothing)
3671 step _ Nothing = return Stop
3673{-# INLINE_NORMAL postscanlM #-}
3674postscanlM :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
3675postscanlM fstep begin (Stream step state) = Stream step' Nothing
3677 {-# INLINE_LATE step' #-}
3678 step' _ Nothing = do
3680 return $ Skip (Just (state, r))
3682 step' gst (Just (st, acc)) = do
3683 r <- step (adaptState gst) st
3687 return (Yield y (Just (s, y)))
3688 Skip s -> return $ Skip (Just (s, acc))
3691{-# INLINE_NORMAL postscanl #-}
3692postscanl :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
3693postscanl f seed = postscanlM (\a b -> return (f a b)) (return seed)
3695{-# INLINE_NORMAL scanlM' #-}
3696scanlM' :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
3697scanlM' fstep begin (Stream step state) = Stream step' Nothing
3699 {-# INLINE_LATE step' #-}
3700 step' _ Nothing = do
3702 return $ Yield x (Just (state, x))
3703 step' gst (Just (st, acc)) = do
3704 r <- step (adaptState gst) st
3708 return $ Yield y (Just (s, y))
3709 Skip s -> return $ Skip (Just (s, acc))
3712{-# INLINE scanlMAfter' #-}
3713scanlMAfter' :: Monad m
3714 => (b -> a -> m b) -> m b -> (b -> m b) -> Stream m a -> Stream m b
3715scanlMAfter' fstep initial done s =
3716 (initial >>= \x -> x `seq` return x) `consM`
3717 postscanlMAfter' fstep initial done s
3719{-# INLINE scanl' #-}
3720scanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
3721scanl' f seed = scanlM' (\a b -> return (f a b)) (return seed)
3723{-# INLINE_NORMAL scanlM #-}
3724scanlM :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
3725scanlM fstep begin (Stream step state) = Stream step' Nothing
3727 {-# INLINE_LATE step' #-}
3728 step' _ Nothing = do
3730 return $ Yield x (Just (state, x))
3731 step' gst (Just (st, acc)) = do
3732 r <- step (adaptState gst) st
3736 return $ Yield y (Just (s, y))
3737 Skip s -> return $ Skip (Just (s, acc))
3738 Stop -> return $ Stop
3741scanl :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
3742scanl f seed = scanlM (\a b -> return (f a b)) (return seed)
3744{-# INLINE_NORMAL scanl1M #-}
3745scanl1M :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
3746scanl1M fstep (Stream step state) = Stream step' (state, Nothing)
3748 {-# INLINE_LATE step' #-}
3749 step' gst (st, Nothing) = do
3752 Yield x s -> return $ Yield x (s, Just x)
3753 Skip s -> return $ Skip (s, Nothing)
3756 step' gst (st, Just acc) = do
3761 return $ Yield z (s, Just z)
3762 Skip s -> return $ Skip (s, Just acc)
3765{-# INLINE scanl1 #-}
3766scanl1 :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
3767scanl1 f = scanl1M (\x y -> return (f x y))
3769{-# INLINE_NORMAL scanl1M' #-}
3770scanl1M' :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
3771scanl1M' fstep (Stream step state) = Stream step' (state, Nothing)
3773 {-# INLINE_LATE step' #-}
3774 step' gst (st, Nothing) = do
3777 Yield x s -> x `seq` return $ Yield x (s, Just x)
3778 Skip s -> return $ Skip (s, Nothing)
3781 step' gst (st, Just acc) = acc `seq` do
3786 z `seq` return $ Yield z (s, Just z)
3787 Skip s -> return $ Skip (s, Just acc)
3790{-# INLINE scanl1' #-}
3791scanl1' :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
3792scanl1' f = scanl1M' (\x y -> return (f x y))
3794------------------------------------------------------------------------------
3796------------------------------------------------------------------------------
3798data RollingMapState s a = RollingMapInit s | RollingMapGo s a
3800{-# INLINE rollingMapM #-}
3801rollingMapM :: Monad m => (a -> a -> m b) -> Stream m a -> Stream m b
3802rollingMapM f (Stream step1 state1) = Stream step (RollingMapInit state1)
3804 step gst (RollingMapInit st) = do
3805 r <- step1 (adaptState gst) st
3807 Yield x s -> Skip $ RollingMapGo s x
3808 Skip s -> Skip $ RollingMapInit s
3811 step gst (RollingMapGo s1 x1) = do
3812 r <- step1 (adaptState gst) s1
3816 return $ Yield res $ RollingMapGo s x
3817 Skip s -> return $ Skip $ RollingMapGo s x1
3818 Stop -> return $ Stop
3820{-# INLINE rollingMap #-}
3821rollingMap :: Monad m => (a -> a -> b) -> Stream m a -> Stream m b
3822rollingMap f = rollingMapM (\x y -> return $ f x y)
3824------------------------------------------------------------------------------
3825-- Tapping/Distributing
3826------------------------------------------------------------------------------
3829tap :: Monad m => Fold m a b -> Stream m a -> Stream m a
3830tap (Fold fstep initial extract) (Stream step state) = Stream step' Nothing
3834 step' _ Nothing = do
3836 return $ Skip (Just (r, state))
3838 step' gst (Just (acc, st)) = acc `seq` do
3843 return $ Yield x (Just (acc', s))
3844 Skip s -> return $ Skip (Just (acc, s))
3849{-# INLINE_NORMAL tapOffsetEvery #-}
3850tapOffsetEvery :: Monad m
3851 => Int -> Int -> Fold m a b -> Stream m a -> Stream m a
3852tapOffsetEvery offset n (Fold fstep initial extract) (Stream step state) =
3853 Stream step' Nothing
3857 {-# INLINE_LATE step' #-}
3858 step' _ Nothing = do
3860 return $ Skip (Just (r, state, offset `mod` n))
3862 step' gst (Just (acc, st, count)) | count <= 0 = do
3866 !acc' <- fstep acc x
3867 return $ Yield x (Just (acc', s, n - 1))
3868 Skip s -> return $ Skip (Just (acc, s, count))
3873 step' gst (Just (acc, st, count)) = do
3876 Yield x s -> return $ Yield x (Just (acc, s, count - 1))
3877 Skip s -> return $ Skip (Just (acc, s, count))
3882{-# INLINE_NORMAL pollCounts #-}
3886 -> (Stream m Int -> Stream m Int)
3890pollCounts predicate transf fld (Stream step state) = Stream step' Nothing
3893 {-# INLINE_LATE step' #-}
3894 step' _ Nothing = do
3895 -- As long as we are using an "Int" for counts lockfree reads from
3896 -- Var should work correctly on both 32-bit and 64-bit machines.
3897 -- However, an Int on a 32-bit machine may overflow quickly.
3898 countVar <- liftIO $ Prim.newIORef (0 :: Int)
3900 $ void $ runFold fld
3901 $ transf $ fromPrimIORef countVar
3902 return $ Skip (Just (countVar, tid, state))
3904 step' gst (Just (countVar, tid, st)) = do
3908 when (predicate x) $ liftIO $ Prim.modifyIORef' countVar (+ 1)
3909 return $ Yield x (Just (countVar, tid, s))
3910 Skip s -> return $ Skip (Just (countVar, tid, s))
3912 liftIO $ killThread tid
3915{-# INLINE_NORMAL tapRate #-}
3917 (MonadAsync m, MonadCatch m)
3922tapRate samplingRate action (Stream step state) = Stream step' Nothing
3924 {-# NOINLINE loop #-}
3925 loop countVar prev = do
3928 (do liftIO $ threadDelay (round $ samplingRate * 1000000)
3929 i <- liftIO $ Prim.readIORef countVar
3930 let !diff = i - prev
3933 (\(e :: AsyncException) -> do
3934 i <- liftIO $ Prim.readIORef countVar
3935 let !diff = i - prev
3937 throwM (MC.toException e))
3940 {-# INLINE_LATE step' #-}
3941 step' _ Nothing = do
3942 countVar <- liftIO $ Prim.newIORef 0
3943 tid <- fork $ loop countVar 0
3944 ref <- liftIO $ newIORef ()
3945 _ <- liftIO $ mkWeakIORef ref (killThread tid)
3946 return $ Skip (Just (countVar, tid, state, ref))
3948 step' gst (Just (countVar, tid, st, ref)) = do
3952 liftIO $ Prim.modifyIORef' countVar (+ 1)
3953 return $ Yield x (Just (countVar, tid, s, ref))
3954 Skip s -> return $ Skip (Just (countVar, tid, s, ref))
3956 liftIO $ killThread tid
3960-------------------------------------------------------------------------------
3962-------------------------------------------------------------------------------
3964{-# INLINE_NORMAL takeWhileM #-}
3965takeWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
3966takeWhileM f (Stream step state) = Stream step' state
3968 {-# INLINE_LATE step' #-}
3974 return $ if b then Yield x s else Stop
3975 Skip s -> return $ Skip s
3978{-# INLINE takeWhile #-}
3979takeWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
3980takeWhile f = takeWhileM (return . f)
3982{-# INLINE_NORMAL drop #-}
3983drop :: Monad m => Int -> Stream m a -> Stream m a
3984drop n (Stream step state) = Stream step' (state, Just n)
3986 {-# INLINE_LATE step' #-}
3987 step' gst (st, Just i)
3992 Yield _ s -> Skip (s, Just (i - 1))
3993 Skip s -> Skip (s, Just i)
3995 | otherwise = return $ Skip (st, Nothing)
3997 step' gst (st, Nothing) = do
4001 Yield x s -> Yield x (s, Nothing)
4002 Skip s -> Skip (s, Nothing)
4005data DropWhileState s a
4007 | DropWhileYield a s
4010{-# INLINE_NORMAL dropWhileM #-}
4011dropWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
4012dropWhileM f (Stream step state) = Stream step' (DropWhileDrop state)
4014 {-# INLINE_LATE step' #-}
4015 step' gst (DropWhileDrop st) = do
4021 then return $ Skip (DropWhileDrop s)
4022 else return $ Skip (DropWhileYield x s)
4023 Skip s -> return $ Skip (DropWhileDrop s)
4026 step' gst (DropWhileNext st) = do
4029 Yield x s -> return $ Skip (DropWhileYield x s)
4030 Skip s -> return $ Skip (DropWhileNext s)
4033 step' _ (DropWhileYield x st) = return $ Yield x (DropWhileNext st)
4035{-# INLINE dropWhile #-}
4036dropWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
4037dropWhile f = dropWhileM (return . f)
4039{-# INLINE_NORMAL filterM #-}
4040filterM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
4041filterM f (Stream step state) = Stream step' state
4043 {-# INLINE_LATE step' #-}
4052 Skip s -> return $ Skip s
4055{-# INLINE filter #-}
4056filter :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
4057filter f = filterM (return . f)
4059{-# INLINE_NORMAL uniq #-}
4060uniq :: (Eq a, Monad m) => Stream m a -> Stream m a
4061uniq (Stream step state) = Stream step' (Nothing, state)
4063 {-# INLINE_LATE step' #-}
4064 step' gst (Nothing, st) = do
4067 Yield x s -> return $ Yield x (Just x, s)
4068 Skip s -> return $ Skip (Nothing, s)
4070 step' gst (Just x, st) = do
4073 Yield y s | x == y -> return $ Skip (Just x, s)
4074 | otherwise -> return $ Yield y (Just y, s)
4075 Skip s -> return $ Skip (Just x, s)
4078------------------------------------------------------------------------------
4079-- Transformation by Mapping
4080------------------------------------------------------------------------------
4082{-# INLINE_NORMAL sequence #-}
4083sequence :: Monad m => Stream m (m a) -> Stream m a
4084sequence (Stream step state) = Stream step' state
4086 {-# INLINE_LATE step' #-}
4088 r <- step (adaptState gst) st
4090 Yield x s -> x >>= \a -> return (Yield a s)
4091 Skip s -> return $ Skip s
4094------------------------------------------------------------------------------
4096------------------------------------------------------------------------------
4098data LoopState x s = FirstYield s
4099 | InterspersingYield s
4102{-# INLINE_NORMAL intersperseM #-}
4103intersperseM :: Monad m => m a -> Stream m a -> Stream m a
4104intersperseM m (Stream step state) = Stream step' (FirstYield state)
4106 {-# INLINE_LATE step' #-}
4107 step' gst (FirstYield st) = do
4111 Yield x s -> Skip (YieldAndCarry x s)
4112 Skip s -> Skip (FirstYield s)
4115 step' gst (InterspersingYield st) = do
4120 return $ Yield a (YieldAndCarry x s)
4121 Skip s -> return $ Skip $ InterspersingYield s
4124 step' _ (YieldAndCarry x st) = return $ Yield x (InterspersingYield st)
4126{-# INLINE_NORMAL intersperseM_ #-}
4127intersperseM_ :: Monad m => m b -> Stream m a -> Stream m a
4128intersperseM_ m (Stream step1 state1) = Stream step (Left (pure (), state1))
4130 {-# INLINE_LATE step #-}
4131 step gst (Left (eff, st)) = do
4134 Yield x s -> eff >> return (Yield x (Right s))
4135 Skip s -> return $ Skip (Left (eff, s))
4138 step _ (Right st) = return $ Skip $ Left (void m, st)
4143 | SuffixYield a (SuffixState s a)
4145{-# INLINE_NORMAL intersperseSuffix #-}
4146intersperseSuffix :: forall m a. Monad m => m a -> Stream m a -> Stream m a
4147intersperseSuffix action (Stream step state) = Stream step' (SuffixElem state)
4149 {-# INLINE_LATE step' #-}
4150 step' gst (SuffixElem st) = do
4153 Yield x s -> Skip (SuffixYield x (SuffixSuffix s))
4154 Skip s -> Skip (SuffixElem s)
4157 step' _ (SuffixSuffix st) = do
4158 action >>= \r -> return $ Skip (SuffixYield r (SuffixElem st))
4160 step' _ (SuffixYield x next) = return $ Yield x next
4162{-# INLINE_NORMAL intersperseSuffix_ #-}
4163intersperseSuffix_ :: Monad m => m b -> Stream m a -> Stream m a
4164intersperseSuffix_ m (Stream step1 state1) = Stream step (Left state1)
4166 {-# INLINE_LATE step #-}
4167 step gst (Left st) = do
4170 Yield x s -> return $ Yield x (Right s)
4171 Skip s -> return $ Skip $ Left s
4174 step _ (Right st) = m >> return (Skip (Left st))
4176data SuffixSpanState s a
4177 = SuffixSpanElem s Int
4178 | SuffixSpanSuffix s
4179 | SuffixSpanYield a (SuffixSpanState s a)
4183-- | intersperse after every n items
4184{-# INLINE_NORMAL intersperseSuffixBySpan #-}
4185intersperseSuffixBySpan :: forall m a. Monad m
4186 => Int -> m a -> Stream m a -> Stream m a
4187intersperseSuffixBySpan n action (Stream step state) =
4188 Stream step' (SuffixSpanElem state n)
4190 {-# INLINE_LATE step' #-}
4191 step' gst (SuffixSpanElem st i) | i > 0 = do
4194 Yield x s -> Skip (SuffixSpanYield x (SuffixSpanElem s (i - 1)))
4195 Skip s -> Skip (SuffixSpanElem s i)
4196 Stop -> if i == n then Stop else Skip SuffixSpanLast
4197 step' _ (SuffixSpanElem st _) = return $ Skip (SuffixSpanSuffix st)
4199 step' _ (SuffixSpanSuffix st) = do
4200 action >>= \r -> return $ Skip (SuffixSpanYield r (SuffixSpanElem st n))
4202 step' _ (SuffixSpanLast) = do
4203 action >>= \r -> return $ Skip (SuffixSpanYield r SuffixSpanStop)
4205 step' _ (SuffixSpanYield x next) = return $ Yield x next
4207 step' _ (SuffixSpanStop) = return Stop
4209{-# INLINE intersperse #-}
4210intersperse :: Monad m => a -> Stream m a -> Stream m a
4211intersperse a = intersperseM (return a)
4213{-# INLINE_NORMAL insertBy #-}
4214insertBy :: Monad m => (a -> a -> Ordering) -> a -> Stream m a -> Stream m a
4215insertBy cmp a (Stream step state) = Stream step' (state, False, Nothing)
4217 {-# INLINE_LATE step' #-}
4218 step' gst (st, False, _) = do
4221 Yield x s -> case cmp a x of
4222 GT -> return $ Yield x (s, False, Nothing)
4223 _ -> return $ Yield a (s, True, Just x)
4224 Skip s -> return $ Skip (s, False, Nothing)
4225 Stop -> return $ Yield a (st, True, Nothing)
4227 step' _ (_, True, Nothing) = return Stop
4229 step' gst (st, True, Just prev) = do
4232 Yield x s -> return $ Yield prev (s, True, Just x)
4233 Skip s -> return $ Skip (s, True, Just prev)
4234 Stop -> return $ Yield prev (st, True, Nothing)
4236------------------------------------------------------------------------------
4238------------------------------------------------------------------------------
4240{-# INLINE_NORMAL deleteBy #-}
4241deleteBy :: Monad m => (a -> a -> Bool) -> a -> Stream m a -> Stream m a
4242deleteBy eq x (Stream step state) = Stream step' (state, False)
4244 {-# INLINE_LATE step' #-}
4245 step' gst (st, False) = do
4248 Yield y s -> return $
4249 if eq x y then Skip (s, True) else Yield y (s, False)
4250 Skip s -> return $ Skip (s, False)
4253 step' gst (st, True) = do
4256 Yield y s -> return $ Yield y (s, True)
4257 Skip s -> return $ Skip (s, True)
4260------------------------------------------------------------------------------
4261-- Transformation by Map and Filter
4262------------------------------------------------------------------------------
4264-- XXX Will this always fuse properly?
4265{-# INLINE_NORMAL mapMaybe #-}
4266mapMaybe :: Monad m => (a -> Maybe b) -> Stream m a -> Stream m b
4267mapMaybe f = fmap fromJust . filter isJust . map f
4269{-# INLINE_NORMAL mapMaybeM #-}
4270mapMaybeM :: Monad m => (a -> m (Maybe b)) -> Stream m a -> Stream m b
4271mapMaybeM f = fmap fromJust . filter isJust . mapM f
4273------------------------------------------------------------------------------
4275------------------------------------------------------------------------------
4277{-# INLINE_NORMAL indexed #-}
4278indexed :: Monad m => Stream m a -> Stream m (Int, a)
4279indexed (Stream step state) = Stream step' (state, 0)
4281 {-# INLINE_LATE step' #-}
4282 step' gst (st, i) = i `seq` do
4283 r <- step (adaptState gst) st
4285 Yield x s -> return $ Yield (i, x) (s, i+1)
4286 Skip s -> return $ Skip (s, i)
4289{-# INLINE_NORMAL indexedR #-}
4290indexedR :: Monad m => Int -> Stream m a -> Stream m (Int, a)
4291indexedR m (Stream step state) = Stream step' (state, m)
4293 {-# INLINE_LATE step' #-}
4294 step' gst (st, i) = i `seq` do
4295 r <- step (adaptState gst) st
4297 Yield x s -> let i' = i - 1
4298 in return $ Yield (i, x) (s, i')
4299 Skip s -> return $ Skip (s, i)
4302{-# INLINE_NORMAL zipWithM #-}
4304 => (a -> b -> m c) -> Stream m a -> Stream m b -> Stream m c
4305zipWithM f (Stream stepa ta) (Stream stepb tb) = Stream step (ta, tb, Nothing)
4307 {-# INLINE_LATE step #-}
4308 step gst (sa, sb, Nothing) = do
4309 r <- stepa (adaptState gst) sa
4312 Yield x sa' -> Skip (sa', sb, Just x)
4313 Skip sa' -> Skip (sa', sb, Nothing)
4316 step gst (sa, sb, Just x) = do
4317 r <- stepb (adaptState gst) sb
4321 return $ Yield z (sa, sb', Nothing)
4322 Skip sb' -> return $ Skip (sa, sb', Just x)
4325#if __GLASGOW_HASKELL__ >= 801
4326{-# RULES "zipWithM xs xs"
4327 forall f xs. zipWithM @Identity f xs xs = mapM (\x -> f x x) xs #-}
4330{-# INLINE zipWith #-}
4331zipWith :: Monad m => (a -> b -> c) -> Stream m a -> Stream m b -> Stream m c
4332zipWith f = zipWithM (\a b -> return (f a b))
4334------------------------------------------------------------------------------
4336------------------------------------------------------------------------------
4338{-# INLINE_NORMAL mergeByM #-}
4341 => (a -> a -> m Ordering) -> Stream m a -> Stream m a -> Stream m a
4342mergeByM cmp (Stream stepa ta) (Stream stepb tb) =
4343 Stream step (Just ta, Just tb, Nothing, Nothing)
4345 {-# INLINE_LATE step #-}
4347 -- one of the values is missing, and the corresponding stream is running
4348 step gst (Just sa, sb, Nothing, b) = do
4351 Yield a sa' -> Skip (Just sa', sb, Just a, b)
4352 Skip sa' -> Skip (Just sa', sb, Nothing, b)
4353 Stop -> Skip (Nothing, sb, Nothing, b)
4355 step gst (sa, Just sb, a, Nothing) = do
4358 Yield b sb' -> Skip (sa, Just sb', a, Just b)
4359 Skip sb' -> Skip (sa, Just sb', a, Nothing)
4360 Stop -> Skip (sa, Nothing, a, Nothing)
4362 -- both the values are available
4363 step _ (sa, sb, Just a, Just b) = do
4365 return $ case res of
4366 GT -> Yield b (sa, sb, Just a, Nothing)
4367 _ -> Yield a (sa, sb, Nothing, Just b)
4369 -- one of the values is missing, corresponding stream is done
4370 step _ (Nothing, sb, Nothing, Just b) =
4371 return $ Yield b (Nothing, sb, Nothing, Nothing)
4373 step _ (sa, Nothing, Just a, Nothing) =
4374 return $ Yield a (sa, Nothing, Nothing, Nothing)
4376 step _ (Nothing, Nothing, Nothing, Nothing) = return Stop
4378{-# INLINE mergeBy #-}
4381 => (a -> a -> Ordering) -> Stream m a -> Stream m a -> Stream m a
4382mergeBy cmp = mergeByM (\a b -> return $ cmp a b)
4384------------------------------------------------------------------------------
4385-- Transformation comprehensions
4386------------------------------------------------------------------------------
4388{-# INLINE_NORMAL the #-}
4389the :: (Eq a, Monad m) => Stream m a -> m (Maybe a)
4390the (Stream step state) = go state
4393 r <- step defState st
4395 Yield x s -> go' x s
4397 Stop -> return Nothing
4399 r <- step defState st
4401 Yield x s | x == n -> go' n s
4402 | otherwise -> return Nothing
4404 Stop -> return (Just n)
4406{-# INLINE runFold #-}
4407runFold :: (Monad m) => Fold m a b -> Stream m a -> m b
4408runFold (Fold step begin done) = foldlMx' step begin done
4410-------------------------------------------------------------------------------
4411-- Concurrent application and fold
4412-------------------------------------------------------------------------------
4414-- XXX These functions should be moved to Stream/Parallel.hs
4416-- Using StreamD the worker stream producing code can fuse with the code to
4417-- queue output to the SVar giving some perf boost.
4419-- Note that StreamD can only be used in limited situations, specifically, we
4420-- cannot implement joinStreamVarPar using this.
4422-- XXX make sure that the SVar passed is a Parallel style SVar.
4424-- | Fold the supplied stream to the SVar asynchronously using Parallel
4425-- concurrency style.
4426-- {-# INLINE_NORMAL toSVarParallel #-}
4427{-# INLINE toSVarParallel #-}
4428toSVarParallel :: MonadAsync m
4429 => State t m a -> SVar t m a -> Stream m a -> m ()
4430toSVarParallel st sv xs =
4431 if svarInspectMode sv
4435 case getYieldLimit st of
4436 Nothing -> doFork (work Nothing)
4438 (handleChildException sv)
4439 Just _ -> doFork (workLim Nothing)
4441 (handleChildException sv)
4446 {-# NOINLINE work #-}
4447 work info = (runFold (FL.toParallelSVar sv info) xs)
4449 {-# NOINLINE workLim #-}
4450 workLim info = runFold (FL.toParallelSVarLimited sv info) xs
4452 {-# NOINLINE forkWithDiag #-}
4454 -- We do not use workerCount in case of ParallelVar but still there is
4455 -- no harm in maintaining it correctly.
4456 liftIO $ atomicModifyIORefCAS_ (workerCount sv) $ \n -> n + 1
4458 -- This allocation matters when significant number of workers are being
4459 -- sent. We allocate it only when needed. The overhead increases by 4x.
4461 case yieldRateInfo sv of
4462 Nothing -> return Nothing
4463 Just _ -> liftIO $ do
4464 cntRef <- newIORef 0
4465 t <- getTime Monotonic
4466 lat <- newIORef (0, t)
4467 return $ Just WorkerInfo
4468 { workerYieldMax = 0
4469 , workerYieldCount = cntRef
4470 , workerLatencyStart = lat
4473 case getYieldLimit st of
4474 Nothing -> doFork (work winfo)
4476 (handleChildException sv)
4477 Just _ -> doFork (workLim winfo)
4479 (handleChildException sv)
4482{-# INLINE_NORMAL mkParallelD #-}
4483mkParallelD :: MonadAsync m => Stream m a -> Stream m a
4484mkParallelD m = Stream step Nothing
4487 step gst Nothing = do
4488 sv <- newParallelVar StopNone gst
4489 toSVarParallel gst sv m
4490 -- XXX use unfold instead?
4491 return $ Skip $ Just $ fromSVar sv
4493 step gst (Just (UnStream step1 st)) = do
4496 Yield a s -> Yield a (Just $ Stream step1 s)
4497 Skip s -> Skip (Just $ Stream step1 s)
4500-- Compare with mkAsync. mkAsync uses an Async style SVar whereas this uses a
4501-- parallel style SVar for evaluation. Currently, parallel style cannot use
4502-- rate control whereas Async style can use rate control. In async style SVar
4503-- the worker thread terminates when the buffer is full whereas in Parallel
4506-- | Make the stream producer and consumer run concurrently by introducing a
4507-- buffer between them. The producer thread evaluates the input stream until
4508-- the buffer fills, it blocks if the buffer is full until there is space in
4509-- the buffer. The consumer consumes the stream lazily from the buffer.
4513{-# INLINE_NORMAL mkParallel #-}
4514mkParallel :: (K.IsStream t, MonadAsync m) => t m a -> t m a
4515mkParallel = fromStreamD . mkParallelD . toStreamD
4517-- Note: we can use another API with two callbacks stop and yield if we want
4518-- the callback to be able to indicate end of stream.
4520-- | Generates a callback and a stream pair. The callback returned is used to
4521-- queue values to the stream. The stream is infinite, there is no way for the
4522-- callback to indicate that it is done now.
4526{-# INLINE_NORMAL newCallbackStream #-}
4527newCallbackStream :: (K.IsStream t, MonadAsync m) => m ((a -> m ()), t m a)
4528newCallbackStream = do
4529 sv <- newParallelVar StopNone defState
4531 -- XXX Add our own thread-id to the SVar as we can not know the callback's
4532 -- thread-id and the callback is not run in a managed worker. We need to
4533 -- handle this better.
4534 liftIO myThreadId >>= modifyThread sv
4536 let callback a = liftIO $ void $ send sv (ChildYield a)
4537 -- XXX we can return an SVar and then the consumer can unfold from the
4539 return (callback, fromStreamD (fromSVar sv))
4541-------------------------------------------------------------------------------
4543-------------------------------------------------------------------------------
4545-- | Create an SVar with a fold consumer that will fold any elements sent to it
4546-- using the supplied fold function.
4547{-# INLINE newFoldSVar #-}
4548newFoldSVar :: MonadAsync m => State t m a -> Fold m a b -> m (SVar t m a)
4549newFoldSVar stt f = do
4550 -- Buffer size for the SVar is derived from the current state
4551 sv <- newParallelVar StopAny (adaptState stt)
4552 -- Add the producer thread-id to the SVar.
4553 liftIO myThreadId >>= modifyThread sv
4554 void $ doFork (work sv) (svarMrun sv) (handleFoldException sv)
4559 {-# NOINLINE work #-}
4560 work sv = void $ runFold f $ fromProducer sv
4562data TapState sv st = TapInit | Tapping sv st | TapDone st
4564{-# INLINE_NORMAL tapAsync #-}
4565tapAsync :: MonadAsync m => Fold m a b -> Stream m a -> Stream m a
4566tapAsync f (Stream step1 state1) = Stream step TapInit
4570 -- In general, a Stop event would come equipped with the result
4571 -- of the fold. It is not used here but it would be useful in
4572 -- applicative and distribute.
4573 done <- fromConsumer svr
4574 when (not done) $ do
4575 liftIO $ withDiagMVar svr "teeToSVar: waiting to drain"
4576 $ takeMVar (outputDoorBellFromConsumer svr)
4580 liftIO $ sendStop svr Nothing
4581 -- drain/wait until a stop event arrives from the fold.
4584 {-# INLINE_LATE step #-}
4585 step gst TapInit = do
4586 sv <- newFoldSVar gst f
4587 return $ Skip (Tapping sv state1)
4589 step gst (Tapping sv st) = do
4593 done <- pushToFold sv a
4596 -- XXX we do not need to wait synchronously here
4598 return $ Yield a (TapDone s)
4599 else return $ Yield a (Tapping sv s)
4600 Skip s -> return $ Skip (Tapping sv s)
4605 step gst (TapDone st) = do
4608 Yield a s -> Yield a (TapDone s)
4609 Skip s -> Skip (TapDone s)
4612-- XXX Exported from Array again as this fold is specific to Array
4613-- | Take last 'n' elements from the stream and discard the rest.
4615lastN :: (Storable a, MonadIO m) => Int -> Fold m a (Array a)
4617 | n <= 0 = fmap (const mempty) FL.drain
4618 | otherwise = A.unsafeFreeze <$> Fold step initial done
4620 step (Tuple3' rb rh i) a = do
4621 rh1 <- liftIO $ RB.unsafeInsert rb rh a
4622 return $ Tuple3' rb rh1 (i + 1)
4623 initial = fmap (\(a, b) -> Tuple3' a b (0 :: Int)) $ liftIO $ RB.new n
4624 done (Tuple3' rb rh i) = do
4625 arr <- liftIO $ MA.newArray n
4626 foldFunc i rh snoc' arr rb
4627 snoc' b a = liftIO $ MA.unsafeSnoc b a
4629 | i < n = RB.unsafeFoldRingM
4630 | otherwise = RB.unsafeFoldRingFullM
4632------------------------------------------------------------------------------
4634------------------------------------------------------------------------------
4636-- XXX using getTime in the loop can be pretty expensive especially for
4637-- computations where iterations are lightweight. We have the following
4640-- 1) Run a timeout thread updating a flag asynchronously and check that
4641-- flag here, that way we can have a cheap termination check.
4643-- 2) Use COARSE clock to get time with lower resolution but more efficiently.
4645-- 3) Use rdtscp/rdtsc to get time directly from the processor, compute the
4646-- termination value of rdtsc in the beginning and then in each iteration just
4647-- get rdtsc and check if we should terminate.
4651 | TakeByTimeCheck st s
4652 | TakeByTimeYield st s
4654{-# INLINE_NORMAL takeByTime #-}
4655takeByTime :: (MonadIO m, TimeUnit64 t) => t -> Stream m a -> Stream m a
4656takeByTime duration (Stream step1 state1) = Stream step (TakeByTimeInit state1)
4659 lim = toRelTime64 duration
4661 {-# INLINE_LATE step #-}
4662 step _ (TakeByTimeInit _) | lim == 0 = return Stop
4663 step _ (TakeByTimeInit st) = do
4664 t0 <- liftIO $ getTime Monotonic
4665 return $ Skip (TakeByTimeYield st t0)
4666 step _ (TakeByTimeCheck st t0) = do
4667 t <- liftIO $ getTime Monotonic
4669 if diffAbsTime64 t t0 > lim
4671 else Skip (TakeByTimeYield st t0)
4672 step gst (TakeByTimeYield st t0) = do
4675 Yield x s -> Yield x (TakeByTimeCheck s t0)
4676 Skip s -> Skip (TakeByTimeCheck s t0)
4679data DropByTime st s x
4681 | DropByTimeGen st s
4682 | DropByTimeCheck st s x
4683 | DropByTimeYield st
4685{-# INLINE_NORMAL dropByTime #-}
4686dropByTime :: (MonadIO m, TimeUnit64 t) => t -> Stream m a -> Stream m a
4687dropByTime duration (Stream step1 state1) = Stream step (DropByTimeInit state1)
4690 lim = toRelTime64 duration
4692 {-# INLINE_LATE step #-}
4693 step _ (DropByTimeInit st) = do
4694 t0 <- liftIO $ getTime Monotonic
4695 return $ Skip (DropByTimeGen st t0)
4696 step gst (DropByTimeGen st t0) = do
4699 Yield x s -> Skip (DropByTimeCheck s t0 x)
4700 Skip s -> Skip (DropByTimeGen s t0)
4702 step _ (DropByTimeCheck st t0 x) = do
4703 t <- liftIO $ getTime Monotonic
4704 if diffAbsTime64 t t0 <= lim
4705 then return $ Skip $ DropByTimeGen st t0
4706 else return $ Yield x $ DropByTimeYield st
4707 step gst (DropByTimeYield st) = do
4710 Yield x s -> Yield x (DropByTimeYield s)
4711 Skip s -> Skip (DropByTimeYield s)
4714{-# INLINE updateTimeVar #-}
4715updateTimeVar :: Prim.IORef Int64 -> IO ()
4716updateTimeVar timeVar = do
4717 MicroSecond64 t <- fromAbsTime <$> getTime Monotonic
4718 Prim.modifyIORef' timeVar (const t)
4720{-# INLINE updateWithDelay #-}
4721updateWithDelay :: RealFrac a => a -> Prim.IORef Int64 -> IO ()
4722updateWithDelay precision timeVar = do
4723 threadDelay (delayTime precision)
4724 updateTimeVar timeVar
4728 -- Keep the minimum at least a millisecond to avoid high CPU usage
4729 {-# INLINE delayTime #-}
4731 | g' >= fromIntegral (maxBound :: Int) = maxBound
4733 | otherwise = round g'
4737 g' = g * 10 ^ (6 :: Int)
4739-- XXX we should move this to stream generation section of this file. Also, the
4740-- take/drop combinators above should be moved to filtering section.
4742{-# INLINE_NORMAL times #-}
4743times :: MonadAsync m => Double -> Stream m (AbsTime, RelTime64)
4744times g = Stream step Nothing
4748 {-# INLINE_LATE step #-}
4750 -- XXX note that this is safe only on a 64-bit machine. On a 32-bit
4751 -- machine a 64-bit 'Var' cannot be read consistently without a lock
4752 -- while another thread is writing to it.
4753 timeVar <- liftIO $ Prim.newIORef (0 :: Int64)
4754 liftIO $ updateTimeVar timeVar
4755 tid <- forkManaged $ liftIO $ forever (updateWithDelay g timeVar)
4756 a <- liftIO $ Prim.readIORef timeVar
4757 return $ Skip $ Just (timeVar, tid, a)
4759 step _ s@(Just (timeVar, _, t0)) = do
4760 a <- liftIO $ Prim.readIORef timeVar
4761 -- XXX we can perhaps use an AbsTime64 using a 64 bit Int for
4762 -- efficiency. or maybe we can use a representation using Double for
4763 -- floating precision time
4764 return $ Yield (toAbsTime (MicroSecond64 t0),
4765 (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