4-- Module : Streamly.Internal.Data.Stream.StreamD
5-- Copyright : (c) 2018 Composewell Technologies
6-- (c) Roman Leshchinskiy 2008-2010
7-- (c) The University of Glasgow, 2009
10-- Maintainer : streamly@composewell.com
11-- Stability : experimental
14-- Direct style re-implementation of CPS style stream in StreamK module. The
15-- symbol or suffix 'D' in this module denotes the "Direct" style. GHC is able
16-- to INLINE and fuse direct style better, providing better performance than
20-- import qualified Streamly.Internal.Data.Stream.StreamD as D
23-- Some of the functions in this file have been adapted from the vector
24-- library, https://hackage.haskell.org/package/vector.
26module Streamly.Internal.Data.Stream.StreamD
30 , Stream (Stream, UnStream)
46 -- ** Specialized Generation
47 -- | Generate a monadic stream from a seed.
60 , enumerateFromStepIntegral
61 , enumerateFromIntegral
62 , enumerateFromThenIntegral
63 , enumerateFromToIntegral
64 , enumerateFromThenToIntegral
66 , enumerateFromStepNum
69 , enumerateFromToFractional
70 , enumerateFromThenToFractional
76 -- | Transform an input structure into a stream.
77 -- | Direct style stream does not support @fromFoldable@.
111 -- ** Specialized Folds
138 -- ** Flattening nested streams
141 , ConcatMapUState (..)
143 , ConcatUnfoldInterleaveState (..)
144 , concatUnfoldInterleave
145 , concatUnfoldRoundrobin
148 , InterleaveState(..)
153 , roundRobin -- interleaveFair?/ParallelFair
186 -- | Transform a stream into another type.
203 -- ** By folding (scans)
254 , intersperseSuffixBySpan
278 -- * Transformation comprehensions
298 -- * Concurrent Application
307import Control.Concurrent (killThread, myThreadId, takeMVar, threadDelay)
308import Control.Exception
309 (assert, Exception, SomeException, AsyncException, fromException, mask_)
310import Control.Monad (void, when, forever)
311import Control.Monad.Catch (MonadCatch, MonadThrow, throwM)
312import Control.Monad.IO.Class (MonadIO(..))
313import Control.Monad.Trans.Reader (ReaderT)
314import Control.Monad.Trans.State.Strict (StateT)
315import Control.Monad.Trans.Class (MonadTrans(lift))
316import Control.Monad.Trans.Control (MonadBaseControl, liftBaseOp_, control)
317import Data.Bits (shiftR, shiftL, (.|.), (.&.))
318import Data.Functor.Identity (Identity(..))
319import Data.Int (Int64)
320import Data.IORef (newIORef, readIORef, mkWeakIORef, writeIORef, IORef)
321import Data.Maybe (fromJust, isJust, isNothing)
322import Data.Word (Word32)
323import Foreign.Ptr (Ptr)
324import Foreign.Storable (Storable(..))
325import GHC.Exts (SpecConstrAnnotation(..))
326import GHC.Types (SPEC(..))
327import System.Mem (performMajorGC)
328import Fusion.Plugin.Types (Fuse(..))
329import Streamly.Internal.Data.IORef.Prim (Prim)
330import Streamly.Internal.Data.Time.Units
331 (TimeUnit64, toRelTime64, diffAbsTime64, RelTime64)
332import Streamly.Internal.Data.Atomics (atomicModifyIORefCAS_)
333import Streamly.Internal.Data.Array.Storable.Foreign.Types (Array(..))
334import Streamly.Internal.Data.Fold.Types (Fold(..))
335import Streamly.Internal.Data.Parser (ParseError(..))
336import Streamly.Internal.Data.Pipe.Types (Pipe(..), PipeState(..))
337import Streamly.Internal.Data.Time.Clock (Clock(Monotonic), getTime)
338import Streamly.Internal.Data.Time.Units
339 (MicroSecond64(..), fromAbsTime, toAbsTime, AbsTime)
340import Streamly.Internal.Data.Unfold.Types (Unfold(..))
341import Streamly.Internal.Data.Tuple.Strict (Tuple3'(..))
342import Streamly.Internal.Data.Stream.SVar (fromConsumer, pushToFold)
344import qualified Streamly.Internal.Data.IORef.Prim as Prim
345import qualified Streamly.Internal.Data.Pipe.Types as Pipe
346import qualified Streamly.Internal.Data.Array.Storable.Foreign.Types as A
347import qualified Streamly.Internal.Data.Array.Storable.Foreign.Mut.Types as MA
348import qualified Streamly.Internal.Data.Fold as FL
349import qualified Streamly.Memory.Ring as RB
350import qualified Streamly.Internal.Data.Stream.StreamK as K
351import qualified Streamly.Internal.Data.Parser as PR
352import qualified Streamly.Internal.Data.Parser.ParserD as PRD
353import qualified Control.Monad.Catch as MC
354import qualified Control.Monad.Trans.Reader as Reader
355import qualified Control.Monad.Trans.State.Strict as State
356import qualified Prelude
359 ( map, mapM, mapM_, repeat, foldr, last, take, filter
360 , takeWhile, drop, dropWhile, all, any, maximum, minimum, elem
361 , notElem, null, head, tail, zipWith, lookup, foldr1, sequence
362 , (!!), scanl, scanl1, concatMap, replicate, enumFromTo, concat
363 , reverse, iterate, splitAt)
364import Streamly.Internal.Data.Stream.StreamD.Type
365import Streamly.Internal.Data.SVar
367------------------------------------------------------------------------------
369------------------------------------------------------------------------------
371-- | An empty 'Stream'.
372{-# INLINE_NORMAL nil #-}
373nil :: Monad m => Stream m a
374nil = Stream (\_ _ -> return Stop) ()
376-- | An empty 'Stream' with a side effect.
377{-# INLINE_NORMAL nilM #-}
378nilM :: Monad m => m b -> Stream m a
379nilM m = Stream (\_ _ -> m >> return Stop) ()
381{-# INLINE_NORMAL consM #-}
382consM :: Monad m => m a -> Stream m a -> Stream m a
383consM m (Stream step state) = Stream step1 Nothing
385 {-# INLINE_LATE step1 #-}
386 step1 _ Nothing = m >>= \x -> return $ Yield x (Just state)
387 step1 gst (Just st) = do
391 Yield a s -> Yield a (Just s)
392 Skip s -> Skip (Just s)
395-- XXX implement in terms of consM?
396-- cons x = consM (return x)
398-- | Can fuse but has O(n^2) complexity.
399{-# INLINE_NORMAL cons #-}
400cons :: Monad m => a -> Stream m a -> Stream m a
401cons x (Stream step state) = Stream step1 Nothing
403 {-# INLINE_LATE step1 #-}
404 step1 _ Nothing = return $ Yield x (Just state)
405 step1 gst (Just st) = do
409 Yield a s -> Yield a (Just s)
410 Skip s -> Skip (Just s)
413-------------------------------------------------------------------------------
415-------------------------------------------------------------------------------
417-- Does not fuse, has the same performance as the StreamK version.
418{-# INLINE_NORMAL uncons #-}
419uncons :: Monad m => Stream m a -> m (Maybe (a, Stream m a))
420uncons (UnStream step state) = go state
423 r <- step defState st
425 Yield x s -> return $ Just (x, Stream step s)
427 Stop -> return Nothing
429------------------------------------------------------------------------------
430-- Generation by unfold
431------------------------------------------------------------------------------
433{-# INLINE_NORMAL unfoldrM #-}
434unfoldrM :: Monad m => (s -> m (Maybe (a, s))) -> s -> Stream m a
435unfoldrM next state = Stream step state
437 {-# INLINE_LATE step #-}
441 Just (x, s) -> Yield x s
444{-# INLINE_LATE unfoldr #-}
445unfoldr :: Monad m => (s -> Maybe (a, s)) -> s -> Stream m a
446unfoldr f = unfoldrM (return . f)
448-- | Convert an 'Unfold' into a 'Stream' by supplying it a seed.
450{-# INLINE_NORMAL unfold #-}
451unfold :: Monad m => Unfold m a b -> a -> Stream m b
452unfold (Unfold ustep inject) seed = Stream step Nothing
454 {-# INLINE_LATE step #-}
455 step _ Nothing = inject seed >>= return . Skip . Just
456 step _ (Just st) = do
459 Yield x s -> Yield x (Just s)
460 Skip s -> Skip (Just s)
463------------------------------------------------------------------------------
464-- Specialized Generation
465------------------------------------------------------------------------------
467{-# INLINE_NORMAL repeatM #-}
468repeatM :: Monad m => m a -> Stream m a
469repeatM x = Stream (\_ _ -> x >>= \r -> return $ Yield r ()) ()
471{-# INLINE_NORMAL repeat #-}
472repeat :: Monad m => a -> Stream m a
473repeat x = Stream (\_ _ -> return $ Yield x ()) ()
475{-# INLINE_NORMAL iterateM #-}
476iterateM :: Monad m => (a -> m a) -> m a -> Stream m a
477iterateM step = Stream (\_ st -> st >>= \x -> return $ Yield x (step x))
479{-# INLINE_NORMAL iterate #-}
480iterate :: Monad m => (a -> a) -> a -> Stream m a
481iterate step st = iterateM (return . step) (return st)
483{-# INLINE_NORMAL replicateM #-}
484replicateM :: forall m a. Monad m => Int -> m a -> Stream m a
485replicateM n p = Stream step n
487 {-# INLINE_LATE step #-}
489 | i <= 0 = return Stop
492 return $ Yield x (i - 1)
494{-# INLINE_NORMAL replicate #-}
495replicate :: Monad m => Int -> a -> Stream m a
496replicate n x = replicateM n (return x)
498-- This would not work properly for floats, therefore we put an Integral
500-- | Can be used to enumerate unbounded integrals. This does not check for
501-- overflow or underflow for bounded integrals.
502{-# INLINE_NORMAL enumerateFromStepIntegral #-}
503enumerateFromStepIntegral :: (Integral a, Monad m) => a -> a -> Stream m a
504enumerateFromStepIntegral from stride =
505 from `seq` stride `seq` Stream step from
507 {-# INLINE_LATE step #-}
508 step _ !x = return $ Yield x $! (x + stride)
510-- We are assuming that "to" is constrained by the type to be within
512{-# INLINE enumerateFromToIntegral #-}
513enumerateFromToIntegral :: (Monad m, Integral a) => a -> a -> Stream m a
514enumerateFromToIntegral from to =
515 takeWhile (<= to) $ enumerateFromStepIntegral from 1
517{-# INLINE enumerateFromIntegral #-}
518enumerateFromIntegral :: (Monad m, Integral a, Bounded a) => a -> Stream m a
519enumerateFromIntegral from = enumerateFromToIntegral from maxBound
521data EnumState a = EnumInit | EnumYield a a a | EnumStop
523{-# INLINE_NORMAL enumerateFromThenToIntegralUp #-}
524enumerateFromThenToIntegralUp
525 :: (Monad m, Integral a)
526 => a -> a -> a -> Stream m a
527enumerateFromThenToIntegralUp from next to = Stream step EnumInit
529 {-# INLINE_LATE step #-}
535 else Yield from EnumStop
536 else -- from <= next <= to
537 let stride = next - from
538 in Skip $ EnumYield from stride (to - stride)
540 step _ (EnumYield x stride toMinus) =
543 then Yield x EnumStop
544 else Yield x $ EnumYield (x + stride) stride toMinus
546 step _ EnumStop = return Stop
548{-# INLINE_NORMAL enumerateFromThenToIntegralDn #-}
549enumerateFromThenToIntegralDn
550 :: (Monad m, Integral a)
551 => a -> a -> a -> Stream m a
552enumerateFromThenToIntegralDn from next to = Stream step EnumInit
554 {-# INLINE_LATE step #-}
556 return $ if to > next
559 else Yield from EnumStop
560 else -- from >= next >= to
561 let stride = next - from
562 in Skip $ EnumYield from stride (to - stride)
564 step _ (EnumYield x stride toMinus) =
567 then Yield x EnumStop
568 else Yield x $ EnumYield (x + stride) stride toMinus
570 step _ EnumStop = return Stop
572{-# INLINE_NORMAL enumerateFromThenToIntegral #-}
573enumerateFromThenToIntegral
574 :: (Monad m, Integral a)
575 => a -> a -> a -> Stream m a
576enumerateFromThenToIntegral from next to
577 | next >= from = enumerateFromThenToIntegralUp from next to
578 | otherwise = enumerateFromThenToIntegralDn from next to
580{-# INLINE_NORMAL enumerateFromThenIntegral #-}
581enumerateFromThenIntegral
582 :: (Monad m, Integral a, Bounded a)
583 => a -> a -> Stream m a
584enumerateFromThenIntegral from next =
586 then enumerateFromThenToIntegralUp from next maxBound
587 else enumerateFromThenToIntegralDn from next minBound
589-- For floating point numbers if the increment is less than the precision then
590-- it just gets lost. Therefore we cannot always increment it correctly by just
592-- 9007199254740992 + 1 + 1 :: Double => 9.007199254740992e15
593-- 9007199254740992 + 2 :: Double => 9.007199254740994e15
595-- Instead we accumulate the increment counter and compute the increment
596-- every time before adding it to the starting number.
598-- This works for Integrals as well as floating point numbers, but
599-- enumerateFromStepIntegral is faster for integrals.
600{-# INLINE_NORMAL enumerateFromStepNum #-}
601enumerateFromStepNum :: (Monad m, Num a) => a -> a -> Stream m a
602enumerateFromStepNum from stride = Stream step 0
604 {-# INLINE_LATE step #-}
605 step _ !i = return $ (Yield $! (from + i * stride)) $! (i + 1)
607{-# INLINE_NORMAL numFrom #-}
608numFrom :: (Monad m, Num a) => a -> Stream m a
609numFrom from = enumerateFromStepNum from 1
611{-# INLINE_NORMAL numFromThen #-}
612numFromThen :: (Monad m, Num a) => a -> a -> Stream m a
613numFromThen from next = enumerateFromStepNum from (next - from)
615-- We cannot write a general function for Num. The only way to write code
616-- portable between the two is to use a 'Real' constraint and convert between
617-- Fractional and Integral using fromRational which is horribly slow.
618{-# INLINE_NORMAL enumerateFromToFractional #-}
619enumerateFromToFractional
620 :: (Monad m, Fractional a, Ord a)
621 => a -> a -> Stream m a
622enumerateFromToFractional from to =
623 takeWhile (<= to + 1 / 2) $ enumerateFromStepNum from 1
625{-# INLINE_NORMAL enumerateFromThenToFractional #-}
626enumerateFromThenToFractional
627 :: (Monad m, Fractional a, Ord a)
628 => a -> a -> a -> Stream m a
629enumerateFromThenToFractional from next to =
630 takeWhile predicate $ numFromThen from next
632 mid = (next - from) / 2
633 predicate | next >= from = (<= to + mid)
634 | otherwise = (>= to + mid)
636-------------------------------------------------------------------------------
637-- Generation by Conversion
638-------------------------------------------------------------------------------
640{-# INLINE_NORMAL fromIndicesM #-}
641fromIndicesM :: Monad m => (Int -> m a) -> Stream m a
642fromIndicesM gen = Stream step 0
644 {-# INLINE_LATE step #-}
647 return $ Yield x (i + 1)
649{-# INLINE fromIndices #-}
650fromIndices :: Monad m => (Int -> a) -> Stream m a
651fromIndices gen = fromIndicesM (return . gen)
653{-# INLINE_NORMAL generateM #-}
654generateM :: Monad m => Int -> (Int -> m a) -> Stream m a
655generateM n gen = n `seq` Stream step 0
657 {-# INLINE_LATE step #-}
658 step _ i | i < n = do
660 return $ Yield x (i + 1)
661 | otherwise = return Stop
663{-# INLINE generate #-}
664generate :: Monad m => Int -> (Int -> a) -> Stream m a
665generate n gen = generateM n (return . gen)
667-- XXX we need the MonadAsync constraint because of a rewrite rule.
668-- | Convert a list of monadic actions to a 'Stream'
669{-# INLINE_LATE fromListM #-}
670fromListM :: MonadAsync m => [m a] -> Stream m a
671fromListM = Stream step
673 {-# INLINE_LATE step #-}
674 step _ (m:ms) = m >>= \x -> return $ Yield x ms
675 step _ [] = return Stop
677{-# INLINE toStreamD #-}
678toStreamD :: (K.IsStream t, Monad m) => t m a -> Stream m a
679toStreamD = fromStreamK . K.toStream
681{-# INLINE_NORMAL fromPrimIORef #-}
682fromPrimIORef :: (MonadIO m, Prim a) => Prim.IORef a -> Stream m a
683fromPrimIORef var = Stream step ()
685 {-# INLINE_LATE step #-}
686 step _ () = liftIO (Prim.readIORef var) >>= \x -> return $ Yield x ()
688-------------------------------------------------------------------------------
689-- Generation from SVar
690-------------------------------------------------------------------------------
692data FromSVarState t m a =
694 | FromSVarRead (SVar t m a)
695 | FromSVarLoop (SVar t m a) [ChildEvent a]
696 | FromSVarDone (SVar t m a)
698{-# INLINE_NORMAL fromSVar #-}
699fromSVar :: (MonadAsync m) => SVar t m a -> Stream m a
700fromSVar svar = Stream step FromSVarInit
703 {-# INLINE_LATE step #-}
704 step _ FromSVarInit = do
705 ref <- liftIO $ newIORef ()
706 _ <- liftIO $ mkWeakIORef ref hook
707 -- when this copy of svar gets garbage collected "ref" will get
708 -- garbage collected and our GC hook will be called.
709 let sv = svar{svarRef = Just ref}
710 return $ Skip (FromSVarRead sv)
714 {-# NOINLINE hook #-}
716 when (svarInspectMode svar) $ do
717 r <- liftIO $ readIORef (svarStopTime (svarStats svar))
719 printSVar svar "SVar Garbage Collected"
721 -- If there are any SVars referenced by this SVar a GC will prompt
722 -- them to be cleaned up quickly.
723 when (svarInspectMode svar) performMajorGC
725 step _ (FromSVarRead sv) = do
726 list <- readOutputQ sv
727 -- Reversing the output is important to guarantee that we process the
728 -- outputs in the same order as they were generated by the constituent
730 return $ Skip $ FromSVarLoop sv (Prelude.reverse list)
732 step _ (FromSVarLoop sv []) = do
733 done <- postProcess sv
734 return $ Skip $ if done
735 then (FromSVarDone sv)
736 else (FromSVarRead sv)
738 step _ (FromSVarLoop sv (ev : es)) = do
740 ChildYield a -> return $ Yield a (FromSVarLoop sv es)
741 ChildStop tid e -> do
745 stop <- shouldStop tid
748 liftIO (cleanupSVar sv)
749 return $ Skip (FromSVarDone sv)
750 else return $ Skip (FromSVarLoop sv es)
752 case fromException ex of
754 return $ Skip (FromSVarLoop sv es)
755 Nothing -> liftIO (cleanupSVar sv) >> throwM ex
759 case svarStopStyle sv of
760 StopNone -> return False
761 StopAny -> return True
763 sid <- liftIO $ readIORef (svarStopBy sv)
764 return $ if tid == sid then True else False
766 step _ (FromSVarDone sv) = do
767 when (svarInspectMode sv) $ do
768 t <- liftIO $ getTime Monotonic
769 liftIO $ writeIORef (svarStopTime (svarStats sv)) (Just t)
770 liftIO $ printSVar sv "SVar Done"
773-------------------------------------------------------------------------------
774-- Process events received by a fold consumer from a stream producer
775-------------------------------------------------------------------------------
777{-# INLINE_NORMAL fromProducer #-}
778fromProducer :: (MonadAsync m) => SVar t m a -> Stream m a
779fromProducer svar = Stream step (FromSVarRead svar)
782 {-# INLINE_LATE step #-}
783 step _ (FromSVarRead sv) = do
784 list <- readOutputQ sv
785 -- Reversing the output is important to guarantee that we process the
786 -- outputs in the same order as they were generated by the constituent
788 return $ Skip $ FromSVarLoop sv (Prelude.reverse list)
790 step _ (FromSVarLoop sv []) = return $ Skip $ FromSVarRead sv
791 step _ (FromSVarLoop sv (ev : es)) = do
793 ChildYield a -> return $ Yield a (FromSVarLoop sv es)
794 ChildStop tid e -> do
798 sendStopToProducer sv
799 return $ Skip (FromSVarDone sv)
800 Just _ -> error "Bug: fromProducer: received exception"
802 step _ (FromSVarDone sv) = do
803 when (svarInspectMode sv) $ do
804 t <- liftIO $ getTime Monotonic
805 liftIO $ writeIORef (svarStopTime (svarStats sv)) (Just t)
806 liftIO $ printSVar sv "SVar Done"
809 step _ FromSVarInit = undefined
811-------------------------------------------------------------------------------
812-- Hoisting the inner monad
813-------------------------------------------------------------------------------
815{-# INLINE_NORMAL hoist #-}
816hoist :: Monad n => (forall x. m x -> n x) -> Stream m a -> Stream n a
817hoist f (Stream step state) = (Stream step' state)
819 {-# INLINE_LATE step' #-}
821 r <- f $ step (adaptState gst) st
823 Yield x s -> Yield x s
827{-# INLINE generally #-}
828generally :: Monad m => Stream Identity a -> Stream m a
829generally = hoist (return . runIdentity)
831{-# INLINE_NORMAL liftInner #-}
832liftInner :: (Monad m, MonadTrans t, Monad (t m))
833 => Stream m a -> Stream (t m) a
834liftInner (Stream step state) = Stream step' state
836 {-# INLINE_LATE step' #-}
838 r <- lift $ step (adaptState gst) st
840 Yield x s -> Yield x s
844{-# INLINE_NORMAL runReaderT #-}
845runReaderT :: Monad m => m s -> Stream (ReaderT s m) a -> Stream m a
846runReaderT env (Stream step state) = Stream step' (state, env)
848 {-# INLINE_LATE step' #-}
849 step' gst (st, action) = do
851 r <- Reader.runReaderT (step (adaptState gst) st) sv
853 Yield x s -> Yield x (s, return sv)
854 Skip s -> Skip (s, return sv)
857{-# INLINE_NORMAL evalStateT #-}
858evalStateT :: Monad m => m s -> Stream (StateT s m) a -> Stream m a
859evalStateT initial (Stream step state) = Stream step' (state, initial)
861 {-# INLINE_LATE step' #-}
862 step' gst (st, action) = do
864 (r, sv') <- State.runStateT (step (adaptState gst) st) sv
866 Yield x s -> Yield x (s, return sv')
867 Skip s -> Skip (s, return sv')
870{-# INLINE_NORMAL runStateT #-}
871runStateT :: Monad m => m s -> Stream (StateT s m) a -> Stream m (s, a)
872runStateT initial (Stream step state) = Stream step' (state, initial)
874 {-# INLINE_LATE step' #-}
875 step' gst (st, action) = do
877 (r, sv') <- State.runStateT (step (adaptState gst) st) sv
879 Yield x s -> Yield (sv', x) (s, return sv')
880 Skip s -> Skip (s, return sv')
883------------------------------------------------------------------------------
884-- Elimination by Folds
885------------------------------------------------------------------------------
887------------------------------------------------------------------------------
889------------------------------------------------------------------------------
891{-# INLINE_NORMAL foldr1 #-}
892foldr1 :: Monad m => (a -> a -> a) -> Stream m a -> m (Maybe a)
896 Nothing -> return Nothing
897 Just (h, t) -> fmap Just (foldr f h t)
899------------------------------------------------------------------------------
901------------------------------------------------------------------------------
903{-# INLINE_NORMAL foldlT #-}
904foldlT :: (Monad m, Monad (s m), MonadTrans s)
905 => (s m b -> a -> s m b) -> s m b -> Stream m a -> s m b
906foldlT fstep begin (Stream step state) = go SPEC begin state
909 r <- lift $ step defState st
911 Yield x s -> go SPEC (fstep acc x) s
912 Skip s -> go SPEC acc s
915-- Note, this is going to have horrible performance, because of the nature of
916-- the stream type (i.e. direct stream vs CPS). Its only for reference, it is
917-- likely be practically unusable.
918{-# INLINE_NORMAL foldlS #-}
920 => (Stream m b -> a -> Stream m b) -> Stream m b -> Stream m a -> Stream m b
921foldlS fstep begin (Stream step state) = Stream step' (Left (state, begin))
923 step' gst (Left (st, acc)) = do
924 r <- step (adaptState gst) st
926 Yield x s -> Skip (Left (s, fstep acc x))
927 Skip s -> Skip (Left (s, acc))
928 Stop -> Skip (Right acc)
930 step' gst (Right (Stream stp stt)) = do
931 r <- stp (adaptState gst) stt
933 Yield x s -> Yield x (Right (Stream stp s))
934 Skip s -> Skip (Right (Stream stp s))
937------------------------------------------------------------------------------
939------------------------------------------------------------------------------
941-- Inlined definition. Without the inline "serially/parser/take" benchmark
942-- degrades and parseMany does not fuse. Even using "inline" at the callsite
944{-# INLINE splitAt #-}
945splitAt :: Int -> [a] -> ([a],[a])
948 | otherwise = splitAt' n ls
950 splitAt' :: Int -> [a] -> ([a], [a])
951 splitAt' _ [] = ([], [])
952 splitAt' 1 (x:xs) = ([x], xs)
953 splitAt' m (x:xs) = (x:xs', xs'')
955 (xs', xs'') = splitAt' (m - 1) xs
957-- GHC parser does not accept {-# ANN type [] NoSpecConstr #-}, so we need
959{-# ANN type List NoSpecConstr #-}
960newtype List a = List {getList :: [a]}
962-- | Run a 'Parse' over a stream.
963{-# INLINE_NORMAL parselMx' #-}
966 => (s -> a -> m (PR.Step s b))
971parselMx' pstep initial extract (Stream step state) = do
972 initial >>= go SPEC state (List [])
976 -- XXX currently we are using a dumb list based approach for backtracking
977 -- buffer. This can be replaced by a sliding/ring buffer using Data.Array.
978 -- That will allow us more efficient random back and forth movement.
980 go !_ st buf !pst = do
981 r <- step defState st
986 PR.Partial 0 pst1 -> go SPEC s (List []) pst1
987 PR.Partial n pst1 -> do
988 assert (n <= length (x:getList buf)) (return ())
989 let src0 = Prelude.take n (x:getList buf)
990 src = Prelude.reverse src0
991 gobuf SPEC s (List []) (List src) pst1
992 PR.Continue 0 pst1 -> go SPEC s (List (x:getList buf)) pst1
993 PR.Continue n pst1 -> do
994 assert (n <= length (x:getList buf)) (return ())
995 let (src0, buf1) = splitAt n (x:getList buf)
996 src = Prelude.reverse src0
997 gobuf SPEC s (List buf1) (List src) pst1
998 PR.Done _ b -> return b
999 PR.Error err -> throwM $ ParseError err
1000 Skip s -> go SPEC s buf pst
1003 gobuf !_ s buf (List []) !pst = go SPEC s buf pst
1004 gobuf !_ s buf (List (x:xs)) !pst = do
1007 PR.Partial 0 pst1 ->
1008 gobuf SPEC s (List []) (List xs) pst1
1009 PR.Partial n pst1 -> do
1010 assert (n <= length (x:getList buf)) (return ())
1011 let src0 = Prelude.take n (x:getList buf)
1012 src = Prelude.reverse src0 ++ xs
1013 gobuf SPEC s (List []) (List src) pst1
1014 PR.Continue 0 pst1 ->
1015 gobuf SPEC s (List (x:getList buf)) (List xs) pst1
1016 PR.Continue n pst1 -> do
1017 assert (n <= length (x:getList buf)) (return ())
1018 let (src0, buf1) = splitAt n (x:getList buf)
1019 src = Prelude.reverse src0 ++ xs
1020 gobuf SPEC s (List buf1) (List src) pst1
1021 PR.Done _ b -> return b
1022 PR.Error err -> throwM $ ParseError err
1024------------------------------------------------------------------------------
1026------------------------------------------------------------------------------
1028{-# ANN type ParseChunksState Fuse #-}
1029data ParseChunksState x inpBuf st pst =
1030 ParseChunksInit inpBuf st
1031 | ParseChunksInitLeftOver inpBuf
1032 | ParseChunksStream st inpBuf !pst
1033 | ParseChunksBuf inpBuf st inpBuf !pst
1034 | ParseChunksYield x (ParseChunksState x inpBuf st pst)
1036{-# INLINE_NORMAL parseMany #-}
1042parseMany (PRD.Parser pstep initial extract) (Stream step state) =
1043 Stream stepOuter (ParseChunksInit [] state)
1047 {-# INLINE_LATE stepOuter #-}
1048 -- Buffer is empty, get the first element from the stream, initialize the
1049 -- fold and then go to stream processing loop.
1050 stepOuter gst (ParseChunksInit [] st) = do
1051 r <- step (adaptState gst) st
1053 Yield x s -> initial >>= return . Skip . ParseChunksBuf [x] s []
1054 Skip s -> return $ Skip $ ParseChunksInit [] s
1057 -- Buffer is not empty, go to buffered processing loop
1058 stepOuter _ (ParseChunksInit src st) = do
1059 initial >>= return . Skip . ParseChunksBuf src st []
1061 -- XXX we just discard any leftover input at the end
1062 stepOuter _ (ParseChunksInitLeftOver _) = return Stop
1064 -- Buffer is empty, process elements from the stream
1065 stepOuter gst (ParseChunksStream st buf pst) = do
1066 r <- step (adaptState gst) st
1071 PR.Partial 0 pst1 ->
1072 return $ Skip $ ParseChunksStream s [] pst1
1073 PR.Partial n pst1 -> do
1074 assert (n <= length (x:buf)) (return ())
1075 let src0 = Prelude.take n (x:buf)
1076 src = Prelude.reverse src0
1077 return $ Skip $ ParseChunksBuf src s [] pst1
1078 PR.Continue 0 pst1 ->
1079 return $ Skip $ ParseChunksStream s (x:buf) pst1
1080 PR.Continue n pst1 -> do
1081 assert (n <= length (x:buf)) (return ())
1082 let (src0, buf1) = splitAt n (x:buf)
1083 src = Prelude.reverse src0
1084 return $ Skip $ ParseChunksBuf src s buf1 pst1
1087 ParseChunksYield b (ParseChunksInit [] s)
1089 assert (n <= length (x:buf)) (return ())
1090 let src = Prelude.reverse (Prelude.take n (x:buf))
1092 ParseChunksYield b (ParseChunksInit src s)
1093 PR.Error err -> throwM $ ParseError err
1094 Skip s -> return $ Skip $ ParseChunksStream s buf pst
1097 let src = Prelude.reverse buf
1099 ParseChunksYield b (ParseChunksInitLeftOver src)
1101 -- go back to stream processing mode
1102 stepOuter _ (ParseChunksBuf [] s buf pst) =
1103 return $ Skip $ ParseChunksStream s buf pst
1105 -- buffered processing loop
1106 stepOuter _ (ParseChunksBuf (x:xs) s buf pst) = do
1109 PR.Partial 0 pst1 ->
1110 return $ Skip $ ParseChunksBuf xs s [] pst1
1111 PR.Partial n pst1 -> do
1112 assert (n <= length (x:buf)) (return ())
1113 let src0 = Prelude.take n (x:buf)
1114 src = Prelude.reverse src0 ++ xs
1115 return $ Skip $ ParseChunksBuf src s [] pst1
1116 PR.Continue 0 pst1 ->
1117 return $ Skip $ ParseChunksBuf xs s (x:buf) pst1
1118 PR.Continue n pst1 -> do
1119 assert (n <= length (x:buf)) (return ())
1120 let (src0, buf1) = splitAt n (x:buf)
1121 src = Prelude.reverse src0 ++ xs
1122 return $ Skip $ ParseChunksBuf src s buf1 pst1
1124 return $ Skip $ ParseChunksYield b (ParseChunksInit xs s)
1126 assert (n <= length (x:buf)) (return ())
1127 let src = Prelude.reverse (Prelude.take n (x:buf)) ++ xs
1128 return $ Skip $ ParseChunksYield b (ParseChunksInit src s)
1129 PR.Error err -> throwM $ ParseError err
1131 stepOuter _ (ParseChunksYield a next) = return $ Yield a next
1133{-# ANN type ConcatParseState Fuse #-}
1134data ConcatParseState x inpBuf st p =
1135 ConcatParseInit inpBuf st p
1136 | ConcatParseInitLeftOver inpBuf
1137 | ConcatParseStream st inpBuf p
1138 | ConcatParseBuf inpBuf st inpBuf p
1139 | ConcatParseYield x (ConcatParseState x inpBuf st p)
1141{-# INLINE_NORMAL parseIterate #-}
1144 => (b -> PRD.Parser m a b)
1148parseIterate func seed (Stream step state) =
1149 Stream stepOuter (ConcatParseInit [] state (func seed))
1153 {-# INLINE_LATE stepOuter #-}
1154 -- Buffer is empty, go to stream processing loop
1155 stepOuter _ (ConcatParseInit [] st (PRD.Parser pstep initial extract)) = do
1156 initial >>= \r -> return $ Skip $ ConcatParseStream st []
1157 (PRD.Parser pstep (return r) extract)
1159 -- Buffer is not empty, go to buffered processing loop
1160 stepOuter _ (ConcatParseInit src st
1161 (PRD.Parser pstep initial extract)) = do
1162 initial >>= \r -> return $ Skip $ ConcatParseBuf src st []
1163 (PRD.Parser pstep (return r) extract)
1165 -- XXX we just discard any leftover input at the end
1166 stepOuter _ (ConcatParseInitLeftOver _) = return Stop
1168 -- Buffer is empty process elements from the stream
1169 stepOuter gst (ConcatParseStream st buf
1170 p@(PRD.Parser pstep initial extract)) = do
1171 r <- step (adaptState gst) st
1177 PR.Partial 0 pst1 ->
1178 return $ Skip $ ConcatParseStream s []
1179 (PRD.Parser pstep (return pst1) extract)
1180 PR.Partial n pst1 -> do
1181 assert (n <= length (x:buf)) (return ())
1182 let src0 = Prelude.take n (x:buf)
1183 src = Prelude.reverse src0
1184 return $ Skip $ ConcatParseBuf src s []
1185 (PRD.Parser pstep (return pst1) extract)
1186 -- PR.Continue 0 pst1 ->
1187 -- return $ Skip $ ConcatParseStream s (x:buf) pst1
1188 PR.Continue n pst1 -> do
1189 assert (n <= length (x:buf)) (return ())
1190 let (src0, buf1) = splitAt n (x:buf)
1191 src = Prelude.reverse src0
1192 return $ Skip $ ConcatParseBuf src s buf1
1193 (PRD.Parser pstep (return pst1) extract)
1194 -- XXX Specialize for Stop 0 common case?
1196 assert (n <= length (x:buf)) (return ())
1197 let src = Prelude.reverse (Prelude.take n (x:buf))
1199 ConcatParseYield b (ConcatParseInit src s (func b))
1200 PR.Error err -> throwM $ ParseError err
1201 Skip s -> return $ Skip $ ConcatParseStream s buf p
1205 let src = Prelude.reverse buf
1207 ConcatParseYield b (ConcatParseInitLeftOver src)
1209 -- go back to stream processing mode
1210 stepOuter _ (ConcatParseBuf [] s buf p) =
1211 return $ Skip $ ConcatParseStream s buf p
1213 -- buffered processing loop
1214 stepOuter _ (ConcatParseBuf (x:xs) s buf
1215 (PRD.Parser pstep initial extract)) = do
1219 PR.Partial 0 pst1 ->
1220 return $ Skip $ ConcatParseBuf xs s []
1221 (PRD.Parser pstep (return pst1) extract)
1222 PR.Partial n pst1 -> do
1223 assert (n <= length (x:buf)) (return ())
1224 let src0 = Prelude.take n (x:buf)
1225 src = Prelude.reverse src0 ++ xs
1226 return $ Skip $ ConcatParseBuf src s []
1227 (PRD.Parser pstep (return pst1) extract)
1228 -- PR.Continue 0 pst1 -> return $ Skip $ ConcatParseBuf xs s (x:buf) pst1
1229 PR.Continue n pst1 -> do
1230 assert (n <= length (x:buf)) (return ())
1231 let (src0, buf1) = splitAt n (x:buf)
1232 src = Prelude.reverse src0 ++ xs
1233 return $ Skip $ ConcatParseBuf src s buf1
1234 (PRD.Parser pstep (return pst1) extract)
1235 -- XXX Specialize for Stop 0 common case?
1237 assert (n <= length (x:buf)) (return ())
1238 let src = Prelude.reverse (Prelude.take n (x:buf)) ++ xs
1239 return $ Skip $ ConcatParseYield b
1240 (ConcatParseInit src s (func b))
1241 PR.Error err -> throwM $ ParseError err
1243 stepOuter _ (ConcatParseYield a next) = return $ Yield a next
1245------------------------------------------------------------------------------
1247------------------------------------------------------------------------------
1249-- | Run a streaming composition, discard the results.
1250{-# INLINE_LATE drain #-}
1251drain :: Monad m => Stream m a -> m ()
1252-- drain = foldrM (\_ xs -> xs) (return ())
1253drain (Stream step state) = go SPEC state
1256 r <- step defState st
1258 Yield _ s -> go SPEC s
1262{-# INLINE_NORMAL null #-}
1263null :: Monad m => Stream m a -> m Bool
1264null m = foldrM (\_ _ -> return False) (return True) m
1266{-# INLINE_NORMAL head #-}
1267head :: Monad m => Stream m a -> m (Maybe a)
1268head m = foldrM (\x _ -> return (Just x)) (return Nothing) m
1270{-# INLINE_NORMAL headElse #-}
1271headElse :: Monad m => a -> Stream m a -> m a
1272headElse a m = foldrM (\x _ -> return x) (return a) m
1274-- Does not fuse, has the same performance as the StreamK version.
1275{-# INLINE_NORMAL tail #-}
1276tail :: Monad m => Stream m a -> m (Maybe (Stream m a))
1277tail (UnStream step state) = go state
1280 r <- step defState st
1282 Yield _ s -> return (Just $ Stream step s)
1284 Stop -> return Nothing
1286-- XXX will it fuse? need custom impl?
1287{-# INLINE_NORMAL last #-}
1288last :: Monad m => Stream m a -> m (Maybe a)
1289last = foldl' (\_ y -> Just y) Nothing
1291{-# INLINE_NORMAL elem #-}
1292elem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
1293-- elem e m = foldrM (\x xs -> if x == e then return True else xs) (return False) m
1294elem e (Stream step state) = go state
1297 r <- step defState st
1300 | x == e -> return True
1303 Stop -> return False
1305{-# INLINE_NORMAL notElem #-}
1306notElem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
1307notElem e s = fmap not (elem e s)
1309{-# INLINE_NORMAL all #-}
1310all :: Monad m => (a -> Bool) -> Stream m a -> m Bool
1311-- all p m = foldrM (\x xs -> if p x then xs else return False) (return True) m
1312all p (Stream step state) = go state
1315 r <- step defState st
1319 | otherwise -> return False
1323{-# INLINE_NORMAL any #-}
1324any :: Monad m => (a -> Bool) -> Stream m a -> m Bool
1325-- any p m = foldrM (\x xs -> if p x then return True else xs) (return False) m
1326any p (Stream step state) = go state
1329 r <- step defState st
1332 | p x -> return True
1335 Stop -> return False
1337{-# INLINE_NORMAL maximum #-}
1338maximum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
1339maximum (Stream step state) = go Nothing state
1342 r <- step defState st
1344 Yield x s -> go (Just x) s
1345 Skip s -> go Nothing s
1346 Stop -> return Nothing
1347 go (Just acc) st = do
1348 r <- step defState st
1351 | acc <= x -> go (Just x) s
1352 | otherwise -> go (Just acc) s
1353 Skip s -> go (Just acc) s
1354 Stop -> return (Just acc)
1356{-# INLINE_NORMAL maximumBy #-}
1357maximumBy :: Monad m => (a -> a -> Ordering) -> Stream m a -> m (Maybe a)
1358maximumBy cmp (Stream step state) = go Nothing state
1361 r <- step defState st
1363 Yield x s -> go (Just x) s
1364 Skip s -> go Nothing s
1365 Stop -> return Nothing
1366 go (Just acc) st = do
1367 r <- step defState st
1369 Yield x s -> case cmp acc x of
1370 GT -> go (Just acc) s
1372 Skip s -> go (Just acc) s
1373 Stop -> return (Just acc)
1375{-# INLINE_NORMAL minimum #-}
1376minimum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
1377minimum (Stream step state) = go Nothing state
1380 r <- step defState st
1382 Yield x s -> go (Just x) s
1383 Skip s -> go Nothing s
1384 Stop -> return Nothing
1385 go (Just acc) st = do
1386 r <- step defState st
1389 | acc <= x -> go (Just acc) s
1390 | otherwise -> go (Just x) s
1391 Skip s -> go (Just acc) s
1392 Stop -> return (Just acc)
1394{-# INLINE_NORMAL minimumBy #-}
1395minimumBy :: Monad m => (a -> a -> Ordering) -> Stream m a -> m (Maybe a)
1396minimumBy cmp (Stream step state) = go Nothing state
1399 r <- step defState st
1401 Yield x s -> go (Just x) s
1402 Skip s -> go Nothing s
1403 Stop -> return Nothing
1404 go (Just acc) st = do
1405 r <- step defState st
1407 Yield x s -> case cmp acc x of
1409 _ -> go (Just acc) s
1410 Skip s -> go (Just acc) s
1411 Stop -> return (Just acc)
1413{-# INLINE_NORMAL (!!) #-}
1414(!!) :: (Monad m) => Stream m a -> Int -> m (Maybe a)
1415(Stream step state) !! i = go i state
1418 r <- step defState st
1420 Yield x s | n < 0 -> return Nothing
1421 | n == 0 -> return $ Just x
1422 | otherwise -> go (n - 1) s
1424 Stop -> return Nothing
1426{-# INLINE_NORMAL lookup #-}
1427lookup :: (Monad m, Eq a) => a -> Stream m (a, b) -> m (Maybe b)
1428lookup e m = foldrM (\(a, b) xs -> if e == a then return (Just b) else xs)
1431{-# INLINE_NORMAL findM #-}
1432findM :: Monad m => (a -> m Bool) -> Stream m a -> m (Maybe a)
1433findM p m = foldrM (\x xs -> p x >>= \r -> if r then return (Just x) else xs)
1437find :: Monad m => (a -> Bool) -> Stream m a -> m (Maybe a)
1438find p = findM (return . p)
1440{-# INLINE_NORMAL findIndices #-}
1441findIndices :: Monad m => (a -> Bool) -> Stream m a -> Stream m Int
1442findIndices p (Stream step state) = Stream step' (state, 0)
1444 {-# INLINE_LATE step' #-}
1445 step' gst (st, i) = i `seq` do
1446 r <- step (adaptState gst) st
1448 Yield x s -> if p x then Yield i (s, i+1) else Skip (s, i+1)
1449 Skip s -> Skip (s, i)
1452{-# INLINE toListRev #-}
1453toListRev :: Monad m => Stream m a -> m [a]
1454toListRev = foldl' (flip (:)) []
1456-- We can implement reverse as:
1458-- > reverse = foldlS (flip cons) nil
1460-- However, this implementation is unusable because of the horrible performance
1461-- of cons. So we just convert it to a list first and then stream from the
1464-- XXX Maybe we can use an Array instead of a list here?
1465{-# INLINE_NORMAL reverse #-}
1466reverse :: Monad m => Stream m a -> Stream m a
1467reverse m = Stream step Nothing
1469 {-# INLINE_LATE step #-}
1472 return $ Skip (Just xs)
1473 step _ (Just (x:xs)) = return $ Yield x (Just xs)
1474 step _ (Just []) = return Stop
1476-- Much faster reverse for Storables
1477{-# INLINE_NORMAL reverse' #-}
1478reverse' :: forall m a. (MonadIO m, Storable a) => Stream m a -> Stream m a
1480-- This commented implementation copies the whole stream into one single array
1481-- and then streams from that array, this is 3-4x faster than the chunked code
1482-- that follows. Though this could be problematic due to unbounded large
1483-- allocations. We need to figure out why the chunked code is slower and if we
1484-- can optimize the chunked code to work as fast as this one. It may be a
1486import Foreign.ForeignPtr (touchForeignPtr)
1487import Foreign.ForeignPtr.Unsafe (unsafeForeignPtrToPtr)
1488import Foreign.Ptr (Ptr, plusPtr)
1489reverse' m = Stream step Nothing
1491 {-# INLINE_LATE step #-}
1493 arr <- A.fromStreamD m
1494 let p = aEnd arr `plusPtr` negate (sizeOf (undefined :: a))
1495 return $ Skip $ Just (aStart arr, p)
1497 step _ (Just (start, p)) | p < unsafeForeignPtrToPtr start = return Stop
1499 step _ (Just (start, p)) = do
1500 let !x = A.unsafeInlineIO $ do
1502 touchForeignPtr start
1504 next = p `plusPtr` negate (sizeOf (undefined :: a))
1505 return $ Yield x (Just (start, next))
1512 $ A.fromStreamDArraysOf A.defaultChunkSize m
1515------------------------------------------------------------------------------
1516-- Grouping/Splitting
1517------------------------------------------------------------------------------
1519{-# INLINE_NORMAL splitSuffixBy' #-}
1520splitSuffixBy' :: Monad m
1521 => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1522splitSuffixBy' predicate f (Stream step state) =
1523 Stream (stepOuter f) (Just state)
1527 {-# INLINE_LATE stepOuter #-}
1528 stepOuter (Fold fstep initial done) gst (Just st) = do
1529 res <- step (adaptState gst) st
1535 then done acc' >>= \val -> return $ Yield val (Just s)
1538 Skip s -> return $ Skip $ Just s
1544 res <- step (adaptState gst) stt
1549 then done acc' >>= \val -> return $ Yield val (Just s)
1551 Skip s -> go SPEC s acc
1552 Stop -> done acc >>= \val -> return $ Yield val Nothing
1554 stepOuter _ _ Nothing = return Stop
1556{-# INLINE_NORMAL groupsBy #-}
1562groupsBy cmp f (Stream step state) = Stream (stepOuter f) (Just state, Nothing)
1566 {-# INLINE_LATE stepOuter #-}
1567 stepOuter (Fold fstep initial done) gst (Just st, Nothing) = do
1568 res <- step (adaptState gst) st
1575 Skip s -> return $ Skip $ (Just s, Nothing)
1580 go !_ prev stt !acc = do
1581 res <- step (adaptState gst) stt
1588 else done acc >>= \r -> return $ Yield r (Just s, Just x)
1589 Skip s -> go SPEC prev s acc
1590 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1592 stepOuter (Fold fstep initial done) gst (Just st, Just prev) = do
1594 acc' <- fstep acc prev
1599 -- XXX code duplicated from the previous equation
1601 res <- step (adaptState gst) stt
1608 else done acc >>= \r -> return $ Yield r (Just s, Just x)
1609 Skip s -> go SPEC s acc
1610 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1612 stepOuter _ _ (Nothing,_) = return Stop
1614{-# INLINE_NORMAL groupsRollingBy #-}
1615groupsRollingBy :: Monad m
1620groupsRollingBy cmp f (Stream step state) =
1621 Stream (stepOuter f) (Just state, Nothing)
1624 {-# INLINE_LATE stepOuter #-}
1625 stepOuter (Fold fstep initial done) gst (Just st, Nothing) = do
1626 res <- step (adaptState gst) st
1633 Skip s -> return $ Skip $ (Just s, Nothing)
1637 go !_ prev stt !acc = do
1638 res <- step (adaptState gst) stt
1646 done acc >>= \r -> return $ Yield r (Just s, Just x)
1647 Skip s -> go SPEC prev s acc
1648 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1650 stepOuter (Fold fstep initial done) gst (Just st, Just prev') = do
1652 acc' <- fstep acc prev'
1653 go SPEC prev' st acc'
1656 go !_ prevv stt !acc = do
1657 res <- step (adaptState gst) stt
1664 else done acc >>= \r -> return $ Yield r (Just s, Just x)
1665 Skip s -> go SPEC prevv s acc
1666 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1668 stepOuter _ _ (Nothing, _) = return Stop
1670{-# INLINE_NORMAL splitBy #-}
1671splitBy :: Monad m => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1672splitBy predicate f (Stream step state) = Stream (step' f) (Just state)
1676 {-# INLINE_LATE step' #-}
1677 step' (Fold fstep initial done) gst (Just st) = initial >>= go SPEC st
1682 res <- step (adaptState gst) stt
1686 then done acc >>= \r -> return $ Yield r (Just s)
1690 Skip s -> go SPEC s acc
1691 Stop -> done acc >>= \r -> return $ Yield r Nothing
1693 step' _ _ Nothing = return Stop
1695-- XXX requires -funfolding-use-threshold=150 in lines-unlines benchmark
1696{-# INLINE_NORMAL splitSuffixBy #-}
1697splitSuffixBy :: Monad m
1698 => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1699splitSuffixBy predicate f (Stream step state) = Stream (step' f) (Just state)
1703 {-# INLINE_LATE step' #-}
1704 step' (Fold fstep initial done) gst (Just st) = do
1705 res <- step (adaptState gst) st
1710 then done acc >>= \val -> return $ Yield val (Just s)
1715 Skip s -> return $ Skip $ Just s
1721 res <- step (adaptState gst) stt
1725 then done acc >>= \r -> return $ Yield r (Just s)
1729 Skip s -> go SPEC s acc
1730 Stop -> done acc >>= \r -> return $ Yield r Nothing
1732 step' _ _ Nothing = return Stop
1734{-# INLINE_NORMAL wordsBy #-}
1735wordsBy :: Monad m => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1736wordsBy predicate f (Stream step state) = Stream (stepOuter f) (Just state)
1740 {-# INLINE_LATE stepOuter #-}
1741 stepOuter (Fold fstep initial done) gst (Just st) = do
1742 res <- step (adaptState gst) st
1746 then return $ Skip (Just s)
1752 Skip s -> return $ Skip $ Just s
1758 res <- step (adaptState gst) stt
1762 then done acc >>= \r -> return $ Yield r (Just s)
1766 Skip s -> go SPEC s acc
1767 Stop -> done acc >>= \r -> return $ Yield r Nothing
1769 stepOuter _ _ Nothing = return Stop
1771-- String search algorithms:
1772-- http://www-igm.univ-mlv.fr/~lecroq/string/index.html
1775-- TODO can we unify the splitting operations using a splitting configuration
1776-- like in the split package.
1778data SplitStyle = Infix | Suffix | Prefix deriving (Eq, Show)
1780data SplitOptions = SplitOptions
1781 { style :: SplitStyle
1782 , withSep :: Bool -- ^ keep the separators in output
1783 -- , compact :: Bool -- ^ treat multiple consecutive separators as one
1784 -- , trimHead :: Bool -- ^ drop blank at head
1785 -- , trimTail :: Bool -- ^ drop blank at tail
1789{-# ANN type SplitOnSeqState Fuse #-}
1790data SplitOnSeqState rb rh ck w fs s b x =
1792 | SplitOnSeqYield b (SplitOnSeqState rb rh ck w fs s b x)
1796 | SplitOnSeqSingle !fs s x
1798 | SplitOnSeqWordInit s
1799 | SplitOnSeqWordLoop !w s !fs
1800 | SplitOnSeqWordDone Int !fs !w
1802 | SplitOnSeqKRInit Int s rb !rh
1803 | SplitOnSeqKRLoop fs s rb !rh !ck
1804 | SplitOnSeqKRCheck fs s rb !rh
1805 | SplitOnSeqKRDone Int !fs rb !rh
1807{-# INLINE_NORMAL splitOnSeq #-}
1809 :: forall m a b. (MonadIO m, Storable a, Enum a, Eq a)
1814splitOnSeq patArr (Fold fstep initial done) (Stream step state) =
1815 Stream stepOuter SplitOnSeqInit
1819 patLen = A.length patArr
1820 maxIndex = patLen - 1
1821 elemBits = sizeOf (undefined :: a) * 8
1823 -- For word pattern case
1825 wordMask = (1 `shiftL` (elemBits * patLen)) - 1
1828 elemMask = (1 `shiftL` elemBits) - 1
1831 wordPat = wordMask .&. A.foldl' addToWord 0 patArr
1833 addToWord wd a = (wd `shiftL` elemBits) .|. fromIntegral (fromEnum a)
1835 -- For Rabin-Karp search
1836 k = 2891336453 :: Word32
1839 addCksum cksum a = cksum * k + fromIntegral (fromEnum a)
1841 deltaCksum cksum old new =
1842 addCksum cksum new - coeff * fromIntegral (fromEnum old)
1844 -- XXX shall we use a random starting hash or 1 instead of 0?
1845 patHash = A.foldl' addCksum 0 patArr
1847 skip = return . Skip
1849 {-# INLINE_LATE stepOuter #-}
1850 stepOuter _ SplitOnSeqInit =
1852 then return $ Skip $ SplitOnSeqEmpty state
1856 let !pat = A.unsafeIndex patArr 0
1857 return $ Skip $ SplitOnSeqSingle acc state pat
1858 else if sizeOf (undefined :: a) * patLen
1859 <= sizeOf (undefined :: Word)
1860 then return $ Skip $ SplitOnSeqWordInit state
1862 (rb, rhead) <- liftIO $ RB.new patLen
1863 skip $ SplitOnSeqKRInit 0 state rb rhead
1865 stepOuter _ (SplitOnSeqYield x next) = return $ Yield x next
1867 ---------------------------
1869 ---------------------------
1871 stepOuter gst (SplitOnSeqEmpty st) = do
1872 res <- step (adaptState gst) st
1878 skip $ SplitOnSeqYield r (SplitOnSeqEmpty s)
1879 Skip s -> return $ Skip (SplitOnSeqEmpty s)
1886 stepOuter _ SplitOnSeqDone = return Stop
1892 stepOuter gst (SplitOnSeqSingle fs st pat) = do
1893 res <- step (adaptState gst) st
1900 return $ Skip $ SplitOnSeqYield r (SplitOnSeqSingle fs1 s pat)
1903 return $ Skip $ (SplitOnSeqSingle fs1 s pat)
1904 Skip s -> return $ Skip $ SplitOnSeqSingle fs s pat
1907 return $ Skip $ SplitOnSeqYield r SplitOnSeqDone
1909 ---------------------------
1910 -- Short Pattern - Shift Or
1911 ---------------------------
1913 stepOuter _ (SplitOnSeqWordDone 0 fs _) = do
1915 skip $ SplitOnSeqYield r SplitOnSeqDone
1916 stepOuter _ (SplitOnSeqWordDone n fs wrd) = do
1917 let old = elemMask .&. (wrd `shiftR` (elemBits * (n - 1)))
1918 fs1 <- fstep fs (toEnum $ fromIntegral old)
1919 skip $ SplitOnSeqWordDone (n - 1) fs1 wrd
1921 stepOuter gst (SplitOnSeqWordInit st0) =
1927 go !_ !idx !wrd !st = do
1928 res <- step (adaptState gst) st
1931 let wrd1 = addToWord wrd x
1935 if wrd1 .&. wordMask == wordPat
1938 let next = SplitOnSeqWordInit s
1939 skip $ SplitOnSeqYield r next
1940 else skip $ SplitOnSeqWordLoop wrd1 s fs
1941 else go SPEC (idx + 1) wrd1 s
1942 Skip s -> go SPEC idx wrd s
1946 then skip $ SplitOnSeqWordDone idx fs wrd
1949 skip $ SplitOnSeqYield r SplitOnSeqDone
1951 stepOuter gst (SplitOnSeqWordLoop wrd0 st0 fs0) =
1952 go SPEC wrd0 st0 fs0
1957 go !_ !wrd !st !fs = do
1958 res <- step (adaptState gst) st
1961 let wrd1 = addToWord wrd x
1962 old = (wordMask .&. wrd)
1963 `shiftR` (elemBits * (patLen - 1))
1964 fs1 <- fstep fs (toEnum $ fromIntegral old)
1965 if wrd1 .&. wordMask == wordPat
1968 let next = SplitOnSeqWordInit s
1969 skip $ SplitOnSeqYield r next
1970 else go SPEC wrd1 s fs1
1971 Skip s -> go SPEC wrd s fs
1972 Stop -> skip $ SplitOnSeqWordDone patLen fs wrd
1974 -------------------------------
1975 -- General Pattern - Karp Rabin
1976 -------------------------------
1978 stepOuter gst (SplitOnSeqKRInit idx st rb rh) = do
1979 res <- step (adaptState gst) st
1982 rh1 <- liftIO $ RB.unsafeInsert rb rh x
1985 let fold = RB.unsafeFoldRing (RB.ringBound rb)
1986 let !ringHash = fold addCksum 0 rb
1988 if ringHash == patHash
1989 then skip $ SplitOnSeqKRCheck fs s rb rh1
1990 else skip $ SplitOnSeqKRLoop fs s rb rh1 ringHash
1991 else skip $ SplitOnSeqKRInit (idx + 1) s rb rh1
1992 Skip s -> skip $ SplitOnSeqKRInit idx s rb rh
1995 skip $ SplitOnSeqKRDone idx fs rb (RB.startOf rb)
1997 -- XXX The recursive "go" is more efficient than the state based recursion
1998 -- code commented out below. Perhaps its more efficient because of
1999 -- factoring out "rb" outside the loop.
2001 stepOuter gst (SplitOnSeqKRLoop fs0 st0 rb rh0 cksum0) =
2002 go SPEC fs0 st0 rh0 cksum0
2006 go !_ !fs !st !rh !cksum = do
2007 res <- step (adaptState gst) st
2010 old <- liftIO $ peek rh
2011 let cksum1 = deltaCksum cksum old x
2013 rh1 <- liftIO (RB.unsafeInsert rb rh x)
2014 if (cksum1 == patHash)
2015 then skip $ SplitOnSeqKRCheck fs1 s rb rh1
2016 else go SPEC fs1 s rh1 cksum1
2017 Skip s -> go SPEC fs s rh cksum
2018 Stop -> skip $ SplitOnSeqKRDone patLen fs rb rh
2020 -- XXX The following code is 5 times slower compared to the recursive loop
2021 -- based code above. Need to investigate why. One possibility is that the
2022 -- go loop above does not thread around the ring buffer (rb). This code may
2023 -- be causing the state to bloat and getting allocated on each iteration.
2024 -- We can check the cmm/asm code to confirm. If so a good GHC solution to
2025 -- such problem is needed. One way to avoid this could be to use unboxed
2028 stepOuter gst (SplitOnSeqKRLoop fs st rb rh cksum) = do
2029 res <- step (adaptState gst) st
2032 old <- liftIO $ peek rh
2033 let cksum1 = deltaCksum cksum old x
2035 if (cksum1 == patHash)
2038 skip $ SplitOnSeqYield r $ SplitOnSeqKRInit 0 s rb rh
2040 rh1 <- liftIO (RB.unsafeInsert rb rh x)
2041 skip $ SplitOnSeqKRLoop fs1 s rb rh1 cksum1
2042 Skip s -> skip $ SplitOnSeqKRLoop fs s rb rh cksum
2043 Stop -> skip $ SplitOnSeqKRDone patLen fs rb rh
2046 stepOuter _ (SplitOnSeqKRCheck fs st rb rh) = do
2047 if RB.unsafeEqArray rb rh patArr
2050 let next = SplitOnSeqKRInit 0 st rb (RB.startOf rb)
2051 skip $ SplitOnSeqYield r next
2052 else skip $ SplitOnSeqKRLoop fs st rb rh patHash
2053 stepOuter _ (SplitOnSeqKRDone 0 fs _ _) = do
2055 skip $ SplitOnSeqYield r SplitOnSeqDone
2056 stepOuter _ (SplitOnSeqKRDone n fs rb rh) = do
2057 old <- liftIO $ peek rh
2058 let rh1 = RB.advance rb rh
2060 skip $ SplitOnSeqKRDone (n - 1) fs1 rb rh1
2062data SplitOnState1 s a =
2067 | GO_KARP_RABIN s !(RB.Ring a) !(Ptr a)
2070{-# INLINE_NORMAL splitSuffixOn #-}
2072 :: forall m a b. (MonadIO m, Storable a, Enum a, Eq a)
2078splitSuffixOn withSep patArr (Fold fstep initial done)
2079 (Stream step state) =
2080 Stream stepOuter GO_START1
2084 patLen = A.length patArr
2085 maxIndex = patLen - 1
2086 elemBits = sizeOf (undefined :: a) * 8
2088 {-# INLINE_LATE stepOuter #-}
2089 stepOuter _ GO_START1 =
2091 then return $ Skip $ GO_EMPTY_PAT1 state
2094 r <- liftIO $ (A.unsafeIndexIO patArr 0)
2095 return $ Skip $ GO_SINGLE_PAT state r
2096 else if sizeOf (undefined :: a) * patLen
2097 <= sizeOf (undefined :: Word)
2098 then return $ Skip $ GO_SHORT_PAT state
2100 (rb, rhead) <- liftIO $ RB.new patLen
2101 return $ Skip $ GO_KARP_RABIN state rb rhead
2103 stepOuter gst (GO_SINGLE_PAT stt pat) = do
2104 -- This first part is the only difference between splitOn and
2106 -- If the last element is a separator do not issue a blank segment.
2107 res <- step (adaptState gst) stt
2113 acc' <- if withSep then fstep acc x else return acc
2114 done acc' >>= \r -> return $ Yield r (GO_SINGLE_PAT s pat)
2115 else fstep acc x >>= go SPEC s
2116 Skip s -> return $ Skip $ (GO_SINGLE_PAT s pat)
2121 -- This is identical for splitOn and splitSuffixOn
2123 res <- step (adaptState gst) st
2128 acc' <- if withSep then fstep acc x else return acc
2130 return $ Yield r (GO_SINGLE_PAT s pat)
2131 else fstep acc x >>= go SPEC s
2132 Skip s -> go SPEC s acc
2133 Stop -> done acc >>= \r -> return $ Yield r GO_DONE1
2135 stepOuter gst (GO_SHORT_PAT stt) = do
2137 -- Call "initial" only if the stream yields an element, otherwise we
2138 -- may call "initial" but never yield anything. initial may produce a
2139 -- side effect, therefore we will end up doing and discard a side
2144 res <- step (adaptState gst) stt
2148 let wrd' = addToWord wrd x
2149 acc' <- if withSep then fstep acc x else return acc
2152 if wrd' .&. mask == patWord
2153 then done acc' >>= \r -> return $ Yield r (GO_SHORT_PAT s)
2154 else go0 SPEC (idx + 1) wrd' s acc'
2155 else go0 SPEC (idx + 1) wrd' s acc'
2156 Skip s -> return $ Skip (GO_SHORT_PAT s)
2162 mask = (1 `shiftL` (elemBits * patLen)) - 1
2164 addToWord wrd a = (wrd `shiftL` elemBits) .|. fromIntegral (fromEnum a)
2167 patWord = mask .&. A.foldl' addToWord 0 patArr
2169 go0 !_ !idx wrd st !acc = do
2170 res <- step (adaptState gst) st
2173 let wrd' = addToWord wrd x
2174 acc' <- if withSep then fstep acc x else return acc
2177 if wrd' .&. mask == patWord
2180 return $ Yield r (GO_SHORT_PAT s)
2181 else go1 SPEC wrd' s acc'
2182 else go0 SPEC (idx + 1) wrd' s acc'
2183 Skip s -> go0 SPEC idx wrd s acc
2185 if (idx == maxIndex) && (wrd .&. mask == patWord)
2188 acc' <- if idx /= 0 && not withSep
2189 then go2 wrd idx acc
2191 done acc' >>= \r -> return $ Yield r GO_DONE1
2194 go1 !_ wrd st !acc = do
2195 res <- step (adaptState gst) st
2198 let wrd' = addToWord wrd x
2199 old = (mask .&. wrd) `shiftR` (elemBits * (patLen - 1))
2202 else fstep acc (toEnum $ fromIntegral old)
2203 if wrd' .&. mask == patWord
2204 then done acc' >>= \r -> return $ Yield r (GO_SHORT_PAT s)
2205 else go1 SPEC wrd' s acc'
2206 Skip s -> go1 SPEC wrd s acc
2208 -- If the last sequence is a separator do not issue a blank
2210 if wrd .&. mask == patWord
2215 else go2 wrd patLen acc
2216 done acc' >>= \r -> return $ Yield r GO_DONE1
2218 go2 !wrd !n !acc | n > 0 = do
2219 let old = (mask .&. wrd) `shiftR` (elemBits * (n - 1))
2220 fstep acc (toEnum $ fromIntegral old) >>= go2 wrd (n - 1)
2221 go2 _ _ acc = return acc
2223 stepOuter gst (GO_KARP_RABIN stt rb rhead) = do
2225 res <- step (adaptState gst) stt
2229 acc' <- if withSep then fstep acc x else return acc
2230 rh' <- liftIO (RB.unsafeInsert rb rhead x)
2233 let fold = RB.unsafeFoldRing (RB.ringBound rb)
2234 let !ringHash = fold addCksum 0 rb
2235 if ringHash == patHash
2236 then go2 SPEC ringHash rh' s acc'
2237 else go0 SPEC (idx + 1) rh' s acc'
2238 else go0 SPEC (idx + 1) rh' s acc'
2239 Skip s -> return $ Skip (GO_KARP_RABIN s rb rhead)
2244 k = 2891336453 :: Word32
2246 addCksum cksum a = cksum * k + fromIntegral (fromEnum a)
2247 deltaCksum cksum old new =
2248 addCksum cksum new - coeff * fromIntegral (fromEnum old)
2250 -- XXX shall we use a random starting hash or 1 instead of 0?
2251 patHash = A.foldl' addCksum 0 patArr
2254 go0 !_ !idx !rh st !acc = do
2255 res <- step (adaptState gst) st
2258 acc' <- if withSep then fstep acc x else return acc
2259 rh' <- liftIO (RB.unsafeInsert rb rh x)
2262 let fold = RB.unsafeFoldRing (RB.ringBound rb)
2263 let !ringHash = fold addCksum 0 rb
2264 if ringHash == patHash
2265 then go2 SPEC ringHash rh' s acc'
2266 else go1 SPEC ringHash rh' s acc'
2267 else go0 SPEC (idx + 1) rh' s acc'
2268 Skip s -> go0 SPEC idx rh s acc
2270 -- do not issue a blank segment when we end at pattern
2271 if (idx == maxIndex) && RB.unsafeEqArray rb rh patArr
2274 !acc' <- if idx /= 0 && not withSep
2275 then RB.unsafeFoldRingM rh fstep acc rb
2277 done acc' >>= \r -> return $ Yield r GO_DONE1
2279 -- XXX Theoretically this code can do 4 times faster if GHC generates
2280 -- optimal code. If we use just "(cksum' == patHash)" condition it goes
2281 -- 4x faster, as soon as we add the "RB.unsafeEqArray rb v" condition
2282 -- the generated code changes drastically and becomes 4x slower. Need
2283 -- to investigate what is going on with GHC.
2285 go1 !_ !cksum !rh st !acc = do
2286 res <- step (adaptState gst) st
2289 old <- liftIO $ peek rh
2290 let cksum' = deltaCksum cksum old x
2295 if (cksum' == patHash)
2297 rh' <- liftIO (RB.unsafeInsert rb rh x)
2298 go2 SPEC cksum' rh' s acc'
2300 rh' <- liftIO (RB.unsafeInsert rb rh x)
2301 go1 SPEC cksum' rh' s acc'
2302 Skip s -> go1 SPEC cksum rh s acc
2304 if RB.unsafeEqArray rb rh patArr
2309 else RB.unsafeFoldRingFullM rh fstep acc rb
2310 done acc' >>= \r -> return $ Yield r GO_DONE1
2312 go2 !_ !cksum' !rh' s !acc' = do
2313 if RB.unsafeEqArray rb rh' patArr
2316 return $ Yield r (GO_KARP_RABIN s rb rhead)
2317 else go1 SPEC cksum' rh' s acc'
2319 stepOuter gst (GO_EMPTY_PAT1 st) = do
2320 res <- step (adaptState gst) st
2325 done acc' >>= \r -> return $ Yield r (GO_EMPTY_PAT1 s)
2326 Skip s -> return $ Skip (GO_EMPTY_PAT1 s)
2329 stepOuter _ GO_DONE1 = return Stop
2331data SplitState s arr
2333 | SplitBuffering s arr
2334 | SplitSplitting s arr
2335 | SplitYielding arr (SplitState s arr)
2338-- XXX An alternative approach would be to use a partial fold (Fold m a b) to
2339-- split using a splitBy like combinator. The Fold would consume upto the
2340-- separator and return any leftover which can then be fed to the next fold.
2342-- We can revisit this once we have partial folds/parsers.
2344-- | Performs infix separator style splitting.
2345{-# INLINE_NORMAL splitInnerBy #-}
2348 => (f a -> m (f a, Maybe (f a))) -- splitter
2349 -> (f a -> f a -> m (f a)) -- joiner
2352splitInnerBy splitter joiner (Stream step1 state1) =
2353 (Stream step (SplitInitial state1))
2357 {-# INLINE_LATE step #-}
2358 step gst (SplitInitial st) = do
2362 (x1, mx2) <- splitter x
2363 return $ case mx2 of
2364 Nothing -> Skip (SplitBuffering s x1)
2365 Just x2 -> Skip (SplitYielding x1 (SplitSplitting s x2))
2366 Skip s -> return $ Skip (SplitInitial s)
2367 Stop -> return $ Stop
2369 step gst (SplitBuffering st buf) = do
2373 (x1, mx2) <- splitter x
2374 buf' <- joiner buf x1
2375 return $ case mx2 of
2376 Nothing -> Skip (SplitBuffering s buf')
2377 Just x2 -> Skip (SplitYielding buf' (SplitSplitting s x2))
2378 Skip s -> return $ Skip (SplitBuffering s buf)
2379 Stop -> return $ Skip (SplitYielding buf SplitFinishing)
2381 step _ (SplitSplitting st buf) = do
2382 (x1, mx2) <- splitter buf
2383 return $ case mx2 of
2384 Nothing -> Skip $ SplitBuffering st x1
2385 Just x2 -> Skip $ SplitYielding x1 (SplitSplitting st x2)
2387 step _ (SplitYielding x next) = return $ Yield x next
2388 step _ SplitFinishing = return $ Stop
2390-- | Performs infix separator style splitting.
2391{-# INLINE_NORMAL splitInnerBySuffix #-}
2393 :: (Monad m, Eq (f a), Monoid (f a))
2394 => (f a -> m (f a, Maybe (f a))) -- splitter
2395 -> (f a -> f a -> m (f a)) -- joiner
2398splitInnerBySuffix splitter joiner (Stream step1 state1) =
2399 (Stream step (SplitInitial state1))
2403 {-# INLINE_LATE step #-}
2404 step gst (SplitInitial st) = do
2408 (x1, mx2) <- splitter x
2409 return $ case mx2 of
2410 Nothing -> Skip (SplitBuffering s x1)
2411 Just x2 -> Skip (SplitYielding x1 (SplitSplitting s x2))
2412 Skip s -> return $ Skip (SplitInitial s)
2413 Stop -> return $ Stop
2415 step gst (SplitBuffering st buf) = do
2419 (x1, mx2) <- splitter x
2420 buf' <- joiner buf x1
2421 return $ case mx2 of
2422 Nothing -> Skip (SplitBuffering s buf')
2423 Just x2 -> Skip (SplitYielding buf' (SplitSplitting s x2))
2424 Skip s -> return $ Skip (SplitBuffering s buf)
2428 else Skip (SplitYielding buf SplitFinishing)
2430 step _ (SplitSplitting st buf) = do
2431 (x1, mx2) <- splitter buf
2432 return $ case mx2 of
2433 Nothing -> Skip $ SplitBuffering st x1
2434 Just x2 -> Skip $ SplitYielding x1 (SplitSplitting st x2)
2436 step _ (SplitYielding x next) = return $ Yield x next
2437 step _ SplitFinishing = return $ Stop
2439------------------------------------------------------------------------------
2441------------------------------------------------------------------------------
2443{-# INLINE_NORMAL isPrefixOf #-}
2444isPrefixOf :: (Eq a, Monad m) => Stream m a -> Stream m a -> m Bool
2445isPrefixOf (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
2447 go (sa, sb, Nothing) = do
2448 r <- stepa defState sa
2450 Yield x sa' -> go (sa', sb, Just x)
2451 Skip sa' -> go (sa', sb, Nothing)
2454 go (sa, sb, Just x) = do
2455 r <- stepb defState sb
2459 then go (sa, sb', Nothing)
2461 Skip sb' -> go (sa, sb', Just x)
2462 Stop -> return False
2464{-# INLINE_NORMAL isSubsequenceOf #-}
2465isSubsequenceOf :: (Eq a, Monad m) => Stream m a -> Stream m a -> m Bool
2466isSubsequenceOf (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
2468 go (sa, sb, Nothing) = do
2469 r <- stepa defState sa
2471 Yield x sa' -> go (sa', sb, Just x)
2472 Skip sa' -> go (sa', sb, Nothing)
2475 go (sa, sb, Just x) = do
2476 r <- stepb defState sb
2480 then go (sa, sb', Nothing)
2481 else go (sa, sb', Just x)
2482 Skip sb' -> go (sa, sb', Just x)
2483 Stop -> return False
2485{-# INLINE_NORMAL stripPrefix #-}
2488 => Stream m a -> Stream m a -> m (Maybe (Stream m a))
2489stripPrefix (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
2491 go (sa, sb, Nothing) = do
2492 r <- stepa defState sa
2494 Yield x sa' -> go (sa', sb, Just x)
2495 Skip sa' -> go (sa', sb, Nothing)
2496 Stop -> return $ Just (Stream stepb sb)
2498 go (sa, sb, Just x) = do
2499 r <- stepb defState sb
2503 then go (sa, sb', Nothing)
2505 Skip sb' -> go (sa, sb', Just x)
2506 Stop -> return Nothing
2508------------------------------------------------------------------------------
2510------------------------------------------------------------------------------
2512-- | Execute a monadic action for each element of the 'Stream'
2513{-# INLINE_NORMAL mapM_ #-}
2514mapM_ :: Monad m => (a -> m b) -> Stream m a -> m ()
2515mapM_ m = drain . mapM m
2517-------------------------------------------------------------------------------
2518-- Stream transformations using Unfolds
2519-------------------------------------------------------------------------------
2521-- Define a unique structure to use in inspection testing
2522data ConcatMapUState o i =
2524 | ConcatMapUInner o i
2526-- | @concatMapU unfold stream@ uses @unfold@ to map the input stream elements
2527-- to streams and then flattens the generated streams into a single output
2530-- This is like 'concatMap' but uses an unfold with an explicit state to
2531-- generate the stream instead of a 'Stream' type generator. This allows better
2532-- optimization via fusion. This can be many times more efficient than
2535{-# INLINE_NORMAL concatMapU #-}
2536concatMapU :: Monad m => Unfold m a b -> Stream m a -> Stream m b
2537concatMapU (Unfold istep inject) (Stream ostep ost) =
2538 Stream step (ConcatMapUOuter ost)
2540 {-# INLINE_LATE step #-}
2541 step gst (ConcatMapUOuter o) = do
2542 r <- ostep (adaptState gst) o
2546 i `seq` return (Skip (ConcatMapUInner o' i))
2547 Skip o' -> return $ Skip (ConcatMapUOuter o')
2548 Stop -> return $ Stop
2550 step _ (ConcatMapUInner o i) = do
2553 Yield x i' -> Yield x (ConcatMapUInner o i')
2554 Skip i' -> Skip (ConcatMapUInner o i')
2555 Stop -> Skip (ConcatMapUOuter o)
2557data ConcatUnfoldInterleaveState o i =
2558 ConcatUnfoldInterleaveOuter o [i]
2559 | ConcatUnfoldInterleaveInner o [i]
2560 | ConcatUnfoldInterleaveInnerL [i] [i]
2561 | ConcatUnfoldInterleaveInnerR [i] [i]
2563-- XXX use arrays to store state instead of lists.
2564-- XXX In general we can use different scheduling strategies e.g. how to
2565-- schedule the outer vs inner loop or assigning weights to different streams
2566-- or outer and inner loops.
2568-- After a yield, switch to the next stream. Do not switch streams on Skip.
2569-- Yield from outer stream switches to the inner stream.
2571-- There are two choices here, (1) exhaust the outer stream first and then
2572-- start yielding from the inner streams, this is much simpler to implement,
2573-- (2) yield at least one element from an inner stream before going back to
2574-- outer stream and opening the next stream from it.
2576-- Ideally, we need some scheduling bias to inner streams vs outer stream.
2577-- Maybe we can configure the behavior.
2579{-# INLINE_NORMAL concatUnfoldInterleave #-}
2580concatUnfoldInterleave :: Monad m => Unfold m a b -> Stream m a -> Stream m b
2581concatUnfoldInterleave (Unfold istep inject) (Stream ostep ost) =
2582 Stream step (ConcatUnfoldInterleaveOuter ost [])
2584 {-# INLINE_LATE step #-}
2585 step gst (ConcatUnfoldInterleaveOuter o ls) = do
2586 r <- ostep (adaptState gst) o
2590 i `seq` return (Skip (ConcatUnfoldInterleaveInner o' (i : ls)))
2591 Skip o' -> return $ Skip (ConcatUnfoldInterleaveOuter o' ls)
2592 Stop -> return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2594 step _ (ConcatUnfoldInterleaveInner _ []) = undefined
2595 step _ (ConcatUnfoldInterleaveInner o (st:ls)) = do
2598 Yield x s -> Yield x (ConcatUnfoldInterleaveOuter o (s:ls))
2599 Skip s -> Skip (ConcatUnfoldInterleaveInner o (s:ls))
2600 Stop -> Skip (ConcatUnfoldInterleaveOuter o ls)
2602 step _ (ConcatUnfoldInterleaveInnerL [] []) = return Stop
2603 step _ (ConcatUnfoldInterleaveInnerL [] rs) =
2604 return $ Skip (ConcatUnfoldInterleaveInnerR [] rs)
2606 step _ (ConcatUnfoldInterleaveInnerL (st:ls) rs) = do
2609 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerL ls (s:rs))
2610 Skip s -> Skip (ConcatUnfoldInterleaveInnerL (s:ls) rs)
2611 Stop -> Skip (ConcatUnfoldInterleaveInnerL ls rs)
2613 step _ (ConcatUnfoldInterleaveInnerR [] []) = return Stop
2614 step _ (ConcatUnfoldInterleaveInnerR ls []) =
2615 return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2617 step _ (ConcatUnfoldInterleaveInnerR ls (st:rs)) = do
2620 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerR (s:ls) rs)
2621 Skip s -> Skip (ConcatUnfoldInterleaveInnerR ls (s:rs))
2622 Stop -> Skip (ConcatUnfoldInterleaveInnerR ls rs)
2624-- XXX In general we can use different scheduling strategies e.g. how to
2625-- schedule the outer vs inner loop or assigning weights to different streams
2626-- or outer and inner loops.
2628-- This could be inefficient if the tasks are too small.
2630-- Compared to concatUnfoldInterleave this one switches streams on Skips.
2632{-# INLINE_NORMAL concatUnfoldRoundrobin #-}
2633concatUnfoldRoundrobin :: Monad m => Unfold m a b -> Stream m a -> Stream m b
2634concatUnfoldRoundrobin (Unfold istep inject) (Stream ostep ost) =
2635 Stream step (ConcatUnfoldInterleaveOuter ost [])
2637 {-# INLINE_LATE step #-}
2638 step gst (ConcatUnfoldInterleaveOuter o ls) = do
2639 r <- ostep (adaptState gst) o
2643 i `seq` return (Skip (ConcatUnfoldInterleaveInner o' (i : ls)))
2644 Skip o' -> return $ Skip (ConcatUnfoldInterleaveInner o' ls)
2645 Stop -> return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2647 step _ (ConcatUnfoldInterleaveInner o []) =
2648 return $ Skip (ConcatUnfoldInterleaveOuter o [])
2650 step _ (ConcatUnfoldInterleaveInner o (st:ls)) = do
2653 Yield x s -> Yield x (ConcatUnfoldInterleaveOuter o (s:ls))
2654 Skip s -> Skip (ConcatUnfoldInterleaveOuter o (s:ls))
2655 Stop -> Skip (ConcatUnfoldInterleaveOuter o ls)
2657 step _ (ConcatUnfoldInterleaveInnerL [] []) = return Stop
2658 step _ (ConcatUnfoldInterleaveInnerL [] rs) =
2659 return $ Skip (ConcatUnfoldInterleaveInnerR [] rs)
2661 step _ (ConcatUnfoldInterleaveInnerL (st:ls) rs) = do
2664 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerL ls (s:rs))
2665 Skip s -> Skip (ConcatUnfoldInterleaveInnerL ls (s:rs))
2666 Stop -> Skip (ConcatUnfoldInterleaveInnerL ls rs)
2668 step _ (ConcatUnfoldInterleaveInnerR [] []) = return Stop
2669 step _ (ConcatUnfoldInterleaveInnerR ls []) =
2670 return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2672 step _ (ConcatUnfoldInterleaveInnerR ls (st:rs)) = do
2675 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerR (s:ls) rs)
2676 Skip s -> Skip (ConcatUnfoldInterleaveInnerR (s:ls) rs)
2677 Stop -> Skip (ConcatUnfoldInterleaveInnerR ls rs)
2679data AppendState s1 s2 = AppendFirst s1 | AppendSecond s2
2681-- Note that this could be much faster compared to the CPS stream. However, as
2682-- the number of streams being composed increases this may become expensive.
2683-- Need to see where the breaking point is between the two.
2685{-# INLINE_NORMAL append #-}
2686append :: Monad m => Stream m a -> Stream m a -> Stream m a
2687append (Stream step1 state1) (Stream step2 state2) =
2688 Stream step (AppendFirst state1)
2692 {-# INLINE_LATE step #-}
2693 step gst (AppendFirst st) = do
2696 Yield a s -> Yield a (AppendFirst s)
2697 Skip s -> Skip (AppendFirst s)
2698 Stop -> Skip (AppendSecond state2)
2700 step gst (AppendSecond st) = do
2703 Yield a s -> Yield a (AppendSecond s)
2704 Skip s -> Skip (AppendSecond s)
2707data InterleaveState s1 s2 = InterleaveFirst s1 s2 | InterleaveSecond s1 s2
2708 | InterleaveSecondOnly s2 | InterleaveFirstOnly s1
2710{-# INLINE_NORMAL interleave #-}
2711interleave :: Monad m => Stream m a -> Stream m a -> Stream m a
2712interleave (Stream step1 state1) (Stream step2 state2) =
2713 Stream step (InterleaveFirst state1 state2)
2717 {-# INLINE_LATE step #-}
2718 step gst (InterleaveFirst st1 st2) = do
2721 Yield a s -> Yield a (InterleaveSecond s st2)
2722 Skip s -> Skip (InterleaveFirst s st2)
2723 Stop -> Skip (InterleaveSecondOnly st2)
2725 step gst (InterleaveSecond st1 st2) = do
2728 Yield a s -> Yield a (InterleaveFirst st1 s)
2729 Skip s -> Skip (InterleaveSecond st1 s)
2730 Stop -> Skip (InterleaveFirstOnly st1)
2732 step gst (InterleaveFirstOnly st1) = do
2735 Yield a s -> Yield a (InterleaveFirstOnly s)
2736 Skip s -> Skip (InterleaveFirstOnly s)
2739 step gst (InterleaveSecondOnly st2) = do
2742 Yield a s -> Yield a (InterleaveSecondOnly s)
2743 Skip s -> Skip (InterleaveSecondOnly s)
2746{-# INLINE_NORMAL interleaveMin #-}
2747interleaveMin :: Monad m => Stream m a -> Stream m a -> Stream m a
2748interleaveMin (Stream step1 state1) (Stream step2 state2) =
2749 Stream step (InterleaveFirst state1 state2)
2753 {-# INLINE_LATE step #-}
2754 step gst (InterleaveFirst st1 st2) = do
2757 Yield a s -> Yield a (InterleaveSecond s st2)
2758 Skip s -> Skip (InterleaveFirst s st2)
2761 step gst (InterleaveSecond st1 st2) = do
2764 Yield a s -> Yield a (InterleaveFirst st1 s)
2765 Skip s -> Skip (InterleaveSecond st1 s)
2768 step _ (InterleaveFirstOnly _) = undefined
2769 step _ (InterleaveSecondOnly _) = undefined
2771{-# INLINE_NORMAL interleaveSuffix #-}
2772interleaveSuffix :: Monad m => Stream m a -> Stream m a -> Stream m a
2773interleaveSuffix (Stream step1 state1) (Stream step2 state2) =
2774 Stream step (InterleaveFirst state1 state2)
2778 {-# INLINE_LATE step #-}
2779 step gst (InterleaveFirst st1 st2) = do
2782 Yield a s -> Yield a (InterleaveSecond s st2)
2783 Skip s -> Skip (InterleaveFirst s st2)
2786 step gst (InterleaveSecond st1 st2) = do
2789 Yield a s -> Yield a (InterleaveFirst st1 s)
2790 Skip s -> Skip (InterleaveSecond st1 s)
2791 Stop -> Skip (InterleaveFirstOnly st1)
2793 step gst (InterleaveFirstOnly st1) = do
2796 Yield a s -> Yield a (InterleaveFirstOnly s)
2797 Skip s -> Skip (InterleaveFirstOnly s)
2800 step _ (InterleaveSecondOnly _) = undefined
2802data InterleaveInfixState s1 s2 a
2803 = InterleaveInfixFirst s1 s2
2804 | InterleaveInfixSecondBuf s1 s2
2805 | InterleaveInfixSecondYield s1 s2 a
2806 | InterleaveInfixFirstYield s1 s2 a
2807 | InterleaveInfixFirstOnly s1
2809{-# INLINE_NORMAL interleaveInfix #-}
2810interleaveInfix :: Monad m => Stream m a -> Stream m a -> Stream m a
2811interleaveInfix (Stream step1 state1) (Stream step2 state2) =
2812 Stream step (InterleaveInfixFirst state1 state2)
2816 {-# INLINE_LATE step #-}
2817 step gst (InterleaveInfixFirst st1 st2) = do
2820 Yield a s -> Yield a (InterleaveInfixSecondBuf s st2)
2821 Skip s -> Skip (InterleaveInfixFirst s st2)
2824 step gst (InterleaveInfixSecondBuf st1 st2) = do
2827 Yield a s -> Skip (InterleaveInfixSecondYield st1 s a)
2828 Skip s -> Skip (InterleaveInfixSecondBuf st1 s)
2829 Stop -> Skip (InterleaveInfixFirstOnly st1)
2831 step gst (InterleaveInfixSecondYield st1 st2 x) = do
2834 Yield a s -> Yield x (InterleaveInfixFirstYield s st2 a)
2835 Skip s -> Skip (InterleaveInfixSecondYield s st2 x)
2838 step _ (InterleaveInfixFirstYield st1 st2 x) = do
2839 return $ Yield x (InterleaveInfixSecondBuf st1 st2)
2841 step gst (InterleaveInfixFirstOnly st1) = do
2844 Yield a s -> Yield a (InterleaveInfixFirstOnly s)
2845 Skip s -> Skip (InterleaveInfixFirstOnly s)
2848{-# INLINE_NORMAL roundRobin #-}
2849roundRobin :: Monad m => Stream m a -> Stream m a -> Stream m a
2850roundRobin (Stream step1 state1) (Stream step2 state2) =
2851 Stream step (InterleaveFirst state1 state2)
2855 {-# INLINE_LATE step #-}
2856 step gst (InterleaveFirst st1 st2) = do
2859 Yield a s -> Yield a (InterleaveSecond s st2)
2860 Skip s -> Skip (InterleaveSecond s st2)
2861 Stop -> Skip (InterleaveSecondOnly st2)
2863 step gst (InterleaveSecond st1 st2) = do
2866 Yield a s -> Yield a (InterleaveFirst st1 s)
2867 Skip s -> Skip (InterleaveFirst st1 s)
2868 Stop -> Skip (InterleaveFirstOnly st1)
2870 step gst (InterleaveSecondOnly st2) = do
2873 Yield a s -> Yield a (InterleaveSecondOnly s)
2874 Skip s -> Skip (InterleaveSecondOnly s)
2877 step gst (InterleaveFirstOnly st1) = do
2880 Yield a s -> Yield a (InterleaveFirstOnly s)
2881 Skip s -> Skip (InterleaveFirstOnly s)
2884data ICUState s1 s2 i1 i2 =
2889 | ICUFirstInner s1 s2 i1
2890 | ICUSecondInner s1 s2 i2
2891 | ICUFirstOnlyInner s1 i1
2892 | ICUSecondOnlyInner s2 i2
2894-- | Interleave streams (full streams, not the elements) unfolded from two
2895-- input streams and concat. Stop when the first stream stops. If the second
2896-- stream ends before the first one then first stream still keeps running alone
2897-- without any interleaving with the second stream.
2899-- [a1, a2, ... an] [b1, b2 ...]
2900-- => [streamA1, streamA2, ... streamAn] [streamB1, streamB2, ...]
2901-- => [streamA1, streamB1, streamA2...StreamAn, streamBn]
2902-- => [a11, a12, ...a1j, b11, b12, ...b1k, a21, a22, ...]
2904{-# INLINE_NORMAL gintercalateSuffix #-}
2907 => Unfold m a c -> Stream m a -> Unfold m b c -> Stream m b -> Stream m c
2909 (Unfold istep1 inject1) (Stream step1 state1)
2910 (Unfold istep2 inject2) (Stream step2 state2) =
2911 Stream step (ICUFirst state1 state2)
2915 {-# INLINE_LATE step #-}
2916 step gst (ICUFirst s1 s2) = do
2917 r <- step1 (adaptState gst) s1
2921 i `seq` return (Skip (ICUFirstInner s s2 i))
2922 Skip s -> return $ Skip (ICUFirst s s2)
2925 step gst (ICUFirstOnly s1) = do
2926 r <- step1 (adaptState gst) s1
2930 i `seq` return (Skip (ICUFirstOnlyInner s i))
2931 Skip s -> return $ Skip (ICUFirstOnly s)
2934 step _ (ICUFirstInner s1 s2 i1) = do
2937 Yield x i' -> Yield x (ICUFirstInner s1 s2 i')
2938 Skip i' -> Skip (ICUFirstInner s1 s2 i')
2939 Stop -> Skip (ICUSecond s1 s2)
2941 step _ (ICUFirstOnlyInner s1 i1) = do
2944 Yield x i' -> Yield x (ICUFirstOnlyInner s1 i')
2945 Skip i' -> Skip (ICUFirstOnlyInner s1 i')
2946 Stop -> Skip (ICUFirstOnly s1)
2948 step gst (ICUSecond s1 s2) = do
2949 r <- step2 (adaptState gst) s2
2953 i `seq` return (Skip (ICUSecondInner s1 s i))
2954 Skip s -> return $ Skip (ICUSecond s1 s)
2955 Stop -> return $ Skip (ICUFirstOnly s1)
2957 step _ (ICUSecondInner s1 s2 i2) = do
2960 Yield x i' -> Yield x (ICUSecondInner s1 s2 i')
2961 Skip i' -> Skip (ICUSecondInner s1 s2 i')
2962 Stop -> Skip (ICUFirst s1 s2)
2964 step _ (ICUSecondOnly _s2) = undefined
2965 step _ (ICUSecondOnlyInner _s2 _i2) = undefined
2967data InterposeSuffixState s1 i1 =
2968 InterposeSuffixFirst s1
2969 -- | InterposeSuffixFirstYield s1 i1
2970 | InterposeSuffixFirstInner s1 i1
2971 | InterposeSuffixSecond s1
2973-- Note that if an unfolded layer turns out to be nil we still emit the
2974-- separator effect. An alternate behavior could be to emit the separator
2975-- effect only if at least one element has been yielded by the unfolding.
2976-- However, that becomes a bit complicated, so we have chosen the former
2978{-# INLINE_NORMAL interposeSuffix #-}
2981 => m c -> Unfold m b c -> Stream m b -> Stream m c
2984 (Unfold istep1 inject1) (Stream step1 state1) =
2985 Stream step (InterposeSuffixFirst state1)
2989 {-# INLINE_LATE step #-}
2990 step gst (InterposeSuffixFirst s1) = do
2991 r <- step1 (adaptState gst) s1
2995 i `seq` return (Skip (InterposeSuffixFirstInner s i))
2996 -- i `seq` return (Skip (InterposeSuffixFirstYield s i))
2997 Skip s -> return $ Skip (InterposeSuffixFirst s)
3001 step _ (InterposeSuffixFirstYield s1 i1) = do
3004 Yield x i' -> Yield x (InterposeSuffixFirstInner s1 i')
3005 Skip i' -> Skip (InterposeSuffixFirstYield s1 i')
3006 Stop -> Skip (InterposeSuffixFirst s1)
3009 step _ (InterposeSuffixFirstInner s1 i1) = do
3012 Yield x i' -> Yield x (InterposeSuffixFirstInner s1 i')
3013 Skip i' -> Skip (InterposeSuffixFirstInner s1 i')
3014 Stop -> Skip (InterposeSuffixSecond s1)
3016 step _ (InterposeSuffixSecond s1) = do
3018 return $ Yield r (InterposeSuffixFirst s1)
3020data ICALState s1 s2 i1 i2 a =
3022 -- | ICALFirstYield s1 s2 i1
3023 | ICALFirstInner s1 s2 i1
3025 | ICALFirstOnlyInner s1 i1
3026 | ICALSecondInject s1 s2
3027 | ICALFirstInject s1 s2 i2
3028 -- | ICALFirstBuf s1 s2 i1 i2
3029 | ICALSecondInner s1 s2 i1 i2
3030 -- -- | ICALSecondInner s1 s2 i1 i2 a
3031 -- -- | ICALFirstResume s1 s2 i1 i2 a
3033-- | Interleave streams (full streams, not the elements) unfolded from two
3034-- input streams and concat. Stop when the first stream stops. If the second
3035-- stream ends before the first one then first stream still keeps running alone
3036-- without any interleaving with the second stream.
3038-- [a1, a2, ... an] [b1, b2 ...]
3039-- => [streamA1, streamA2, ... streamAn] [streamB1, streamB2, ...]
3040-- => [streamA1, streamB1, streamA2...StreamAn, streamBn]
3041-- => [a11, a12, ...a1j, b11, b12, ...b1k, a21, a22, ...]
3043{-# INLINE_NORMAL gintercalate #-}
3046 => Unfold m a c -> Stream m a -> Unfold m b c -> Stream m b -> Stream m c
3048 (Unfold istep1 inject1) (Stream step1 state1)
3049 (Unfold istep2 inject2) (Stream step2 state2) =
3050 Stream step (ICALFirst state1 state2)
3054 {-# INLINE_LATE step #-}
3055 step gst (ICALFirst s1 s2) = do
3056 r <- step1 (adaptState gst) s1
3060 i `seq` return (Skip (ICALFirstInner s s2 i))
3061 -- i `seq` return (Skip (ICALFirstYield s s2 i))
3062 Skip s -> return $ Skip (ICALFirst s s2)
3066 step _ (ICALFirstYield s1 s2 i1) = do
3069 Yield x i' -> Yield x (ICALFirstInner s1 s2 i')
3070 Skip i' -> Skip (ICALFirstYield s1 s2 i')
3071 Stop -> Skip (ICALFirst s1 s2)
3074 step _ (ICALFirstInner s1 s2 i1) = do
3077 Yield x i' -> Yield x (ICALFirstInner s1 s2 i')
3078 Skip i' -> Skip (ICALFirstInner s1 s2 i')
3079 Stop -> Skip (ICALSecondInject s1 s2)
3081 step gst (ICALFirstOnly s1) = do
3082 r <- step1 (adaptState gst) s1
3086 i `seq` return (Skip (ICALFirstOnlyInner s i))
3087 Skip s -> return $ Skip (ICALFirstOnly s)
3090 step _ (ICALFirstOnlyInner s1 i1) = do
3093 Yield x i' -> Yield x (ICALFirstOnlyInner s1 i')
3094 Skip i' -> Skip (ICALFirstOnlyInner s1 i')
3095 Stop -> Skip (ICALFirstOnly s1)
3097 -- We inject the second stream even before checking if the first stream
3098 -- would yield any more elements. There is no clear choice whether we
3099 -- should do this before or after that. Doing it after may make the state
3100 -- machine a bit simpler though.
3101 step gst (ICALSecondInject s1 s2) = do
3102 r <- step2 (adaptState gst) s2
3106 i `seq` return (Skip (ICALFirstInject s1 s i))
3107 Skip s -> return $ Skip (ICALSecondInject s1 s)
3108 Stop -> return $ Skip (ICALFirstOnly s1)
3110 step gst (ICALFirstInject s1 s2 i2) = do
3111 r <- step1 (adaptState gst) s1
3115 i `seq` return (Skip (ICALSecondInner s s2 i i2))
3116 -- i `seq` return (Skip (ICALFirstBuf s s2 i i2))
3117 Skip s -> return $ Skip (ICALFirstInject s s2 i2)
3121 step _ (ICALFirstBuf s1 s2 i1 i2) = do
3124 Yield x i' -> Skip (ICALSecondInner s1 s2 i' i2 x)
3125 Skip i' -> Skip (ICALFirstBuf s1 s2 i' i2)
3128 step _ (ICALSecondInner s1 s2 i1 i2 v) = do
3131 Yield x i' -> Yield x (ICALSecondInner s1 s2 i1 i' v)
3132 Skip i' -> Skip (ICALSecondInner s1 s2 i1 i' v)
3133 Stop -> Skip (ICALFirstResume s1 s2 i1 i2 v)
3136 step _ (ICALSecondInner s1 s2 i1 i2) = do
3139 Yield x i' -> Yield x (ICALSecondInner s1 s2 i1 i')
3140 Skip i' -> Skip (ICALSecondInner s1 s2 i1 i')
3141 Stop -> Skip (ICALFirstInner s1 s2 i1)
3142 -- Stop -> Skip (ICALFirstResume s1 s2 i1 i2)
3145 step _ (ICALFirstResume s1 s2 i1 i2 x) = do
3146 return $ Yield x (ICALFirstInner s1 s2 i1 i2)
3149data InterposeState s1 i1 a =
3151 -- | InterposeFirstYield s1 i1
3152 | InterposeFirstInner s1 i1
3153 | InterposeFirstInject s1
3154 -- | InterposeFirstBuf s1 i1
3155 | InterposeSecondYield s1 i1
3156 -- -- | InterposeSecondYield s1 i1 a
3157 -- -- | InterposeFirstResume s1 i1 a
3159-- Note that this only interposes the pure values, we may run many effects to
3160-- generate those values as some effects may not generate anything (Skip).
3161{-# INLINE_NORMAL interpose #-}
3162interpose :: Monad m => m c -> Unfold m b c -> Stream m b -> Stream m c
3165 (Unfold istep1 inject1) (Stream step1 state1) =
3166 Stream step (InterposeFirst state1)
3170 {-# INLINE_LATE step #-}
3171 step gst (InterposeFirst s1) = do
3172 r <- step1 (adaptState gst) s1
3176 i `seq` return (Skip (InterposeFirstInner s i))
3177 -- i `seq` return (Skip (InterposeFirstYield s i))
3178 Skip s -> return $ Skip (InterposeFirst s)
3182 step _ (InterposeFirstYield s1 i1) = do
3185 Yield x i' -> Yield x (InterposeFirstInner s1 i')
3186 Skip i' -> Skip (InterposeFirstYield s1 i')
3187 Stop -> Skip (InterposeFirst s1)
3190 step _ (InterposeFirstInner s1 i1) = do
3193 Yield x i' -> Yield x (InterposeFirstInner s1 i')
3194 Skip i' -> Skip (InterposeFirstInner s1 i')
3195 Stop -> Skip (InterposeFirstInject s1)
3197 step gst (InterposeFirstInject s1) = do
3198 r <- step1 (adaptState gst) s1
3202 -- i `seq` return (Skip (InterposeFirstBuf s i))
3203 i `seq` return (Skip (InterposeSecondYield s i))
3204 Skip s -> return $ Skip (InterposeFirstInject s)
3208 step _ (InterposeFirstBuf s1 i1) = do
3211 Yield x i' -> Skip (InterposeSecondYield s1 i' x)
3212 Skip i' -> Skip (InterposeFirstBuf s1 i')
3217 step _ (InterposeSecondYield s1 i1 v) = do
3219 return $ Yield r (InterposeFirstResume s1 i1 v)
3221 step _ (InterposeSecondYield s1 i1) = do
3223 return $ Yield r (InterposeFirstInner s1 i1)
3226 step _ (InterposeFirstResume s1 i1 v) = do
3227 return $ Yield v (InterposeFirstInner s1 i1)
3230------------------------------------------------------------------------------
3232------------------------------------------------------------------------------
3234data GbracketState s1 s2 v
3236 | GBracketNormal s1 v
3237 | GBracketException s2
3239-- | The most general bracketing and exception combinator. All other
3240-- combinators can be expressed in terms of this combinator. This can also be
3241-- used for cases which are not covered by the standard combinators.
3245{-# INLINE_NORMAL gbracket_ #-}
3249 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
3250 -> (c -> m d) -- ^ after, on normal stop
3251 -> (c -> e -> Stream m b -> Stream m b) -- ^ on exception
3252 -> (c -> Stream m b) -- ^ stream generator
3254gbracket_ bef exc aft fexc fnormal =
3255 Stream step GBracketInit
3259 {-# INLINE_LATE step #-}
3260 step _ GBracketInit = do
3262 return $ Skip $ GBracketNormal (fnormal r) r
3264 step gst (GBracketNormal (UnStream step1 st) v) = do
3265 res <- exc $ step1 gst st
3267 Right r -> case r of
3269 return $ Yield x (GBracketNormal (Stream step1 s) v)
3270 Skip s -> return $ Skip (GBracketNormal (Stream step1 s) v)
3271 Stop -> aft v >> return Stop
3273 return $ Skip (GBracketException (fexc v e (UnStream step1 st)))
3274 step gst (GBracketException (UnStream step1 st)) = do
3277 Yield x s -> return $ Yield x (GBracketException (Stream step1 s))
3278 Skip s -> return $ Skip (GBracketException (Stream step1 s))
3281------------------------------------------------------------------------------
3283------------------------------------------------------------------------------
3285-- | Make a finalizer from a monadic action @m a@ that can run in IO monad.
3286mkIOFinalizer :: MonadBaseControl IO m => m b -> m (IO ())
3288 mrun <- captureMonadState
3294-- | Run an IO action stored in a finalized IORef.
3295runIORefFinalizerGC :: IORef (Maybe (IO ())) -> IO ()
3296runIORefFinalizerGC ref = do
3297 res <- readIORef ref
3299 Nothing -> return ()
3302-- | Create an IORef holding a finalizer that is called automatically when
3303-- the IORef is garbage collected. The IORef can be written to with a 'Nothing'
3304-- value to deactivate the finalizer.
3306-- Note: The finalizer is always run with the state of the monad captured at
3307-- the time of calling newFinalizedIORef. To run it on garbage collection we
3308-- have no option but to take a snapshot of the monadic state at some point of
3309-- time. For normal case we could run it with the current state of the monad
3310-- but we want to keep both the cases consistent.
3312newFinalizedIORef :: (MonadIO m, MonadBaseControl IO m)
3313 => m a -> m (IORef (Maybe (IO ())))
3314newFinalizedIORef finalizer = do
3315 f <- mkIOFinalizer finalizer
3316 ref <- liftIO $ newIORef $ Just f
3317 _ <- liftIO $ mkWeakIORef ref (runIORefFinalizerGC ref)
3320-- | Run the finalizer stored in an IORef and deactivate it so that it is run
3321-- only once. Note, the action runs with async exceptions masked.
3323runIORefFinalizer :: MonadIO m => IORef (Maybe (IO ())) -> m ()
3324runIORefFinalizer ref = liftIO $ do
3325 res <- readIORef ref
3327 Nothing -> return ()
3329 -- if an async exception comes after writing 'Nothing' then the
3330 -- finalizing action will never be run. We need to do this
3331 -- atomically wrt async exceptions.
3333 writeIORef ref Nothing
3336-- | Run an action clearing the finalizer IORef atomically wrt async
3337-- exceptions. The action is run with async exceptions masked.
3338withIORefFinalizer :: MonadBaseControl IO m
3339 => IORef (Maybe (IO ())) -> m a -> m a
3340withIORefFinalizer ref action = do
3341 control $ \runinio ->
3343 writeIORef ref Nothing
3346------------------------------------------------------------------------------
3348data GbracketIOState s1 s2 v wref
3350 | GBracketIONormal s1 v wref
3351 | GBracketIOException s2
3353-- | The most general bracketing and exception combinator. All other
3354-- combinators can be expressed in terms of this combinator. This can also be
3355-- used for cases which are not covered by the standard combinators.
3357-- | It uses a finalizer to make sure we run the finalizing action
3358-- when the stream is garbage collected.
3362{-# INLINE_NORMAL gbracket #-}
3364 :: (MonadIO m, MonadBaseControl IO m)
3366 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
3367 -> (c -> m d) -- ^ after, on normal stop or GC
3368 -> (c -> e -> Stream m b -> m (Stream m b)) -- ^ on exception
3369 -> (c -> Stream m b) -- ^ stream generator
3371gbracket bef exc aft fexc fnormal =
3372 Stream step GBracketIOInit
3376 -- If the stream is never evaluated the "aft" action will never be
3377 -- called. For that to occur we will need the user of this API to pass a
3378 -- weak pointer to us.
3379 {-# INLINE_LATE step #-}
3380 step _ GBracketIOInit = do
3381 -- We mask asynchronous exceptions to make the execution
3382 -- of 'bef' and the registration of 'aft' atomic.
3383 -- A similar thing is done in the resourcet package: https://git.io/JvKV3
3384 -- Tutorial: https://markkarpov.com/tutorial/exceptions.html
3385 (r, ref) <- liftBaseOp_ mask_ $ do
3387 ref <- newFinalizedIORef (aft r)
3389 return $ Skip $ GBracketIONormal (fnormal r) r ref
3391 step gst (GBracketIONormal (UnStream step1 st) v ref) = do
3392 res <- exc $ step1 gst st
3394 Right r -> case r of
3396 return $ Yield x (GBracketIONormal (Stream step1 s) v ref)
3398 return $ Skip (GBracketIONormal (Stream step1 s) v ref)
3400 runIORefFinalizer ref
3403 -- Clearing of finalizer and running of exception handler must
3404 -- be atomic wrt async exceptions. Otherwise if we have cleared
3405 -- the finalizer and have not run the exception handler then we
3406 -- may leak the resource.
3407 stream <- withIORefFinalizer ref (fexc v e (UnStream step1 st))
3408 return $ Skip (GBracketIOException stream)
3409 step gst (GBracketIOException (UnStream step1 st)) = do
3413 return $ Yield x (GBracketIOException (Stream step1 s))
3414 Skip s -> return $ Skip (GBracketIOException (Stream step1 s))
3417-- | Run a side effect before the stream yields its first element.
3418{-# INLINE_NORMAL before #-}
3419before :: Monad m => m b -> Stream m a -> Stream m a
3420before action (Stream step state) = Stream step' Nothing
3424 {-# INLINE_LATE step' #-}
3425 step' _ Nothing = action >> return (Skip (Just state))
3427 step' gst (Just st) = do
3430 Yield x s -> return $ Yield x (Just s)
3431 Skip s -> return $ Skip (Just s)
3434-- | Run a side effect whenever the stream stops normally.
3435{-# INLINE_NORMAL after_ #-}
3436after_ :: Monad m => m b -> Stream m a -> Stream m a
3437after_ action (Stream step state) = Stream step' state
3441 {-# INLINE_LATE step' #-}
3445 Yield x s -> return $ Yield x s
3446 Skip s -> return $ Skip s
3447 Stop -> action >> return Stop
3449{-# INLINE_NORMAL after #-}
3450after :: (MonadIO m, MonadBaseControl IO m)
3451 => m b -> Stream m a -> Stream m a
3452after action (Stream step state) = Stream step' Nothing
3456 {-# INLINE_LATE step' #-}
3457 step' _ Nothing = do
3458 ref <- newFinalizedIORef action
3459 return $ Skip $ Just (state, ref)
3460 step' gst (Just (st, ref)) = do
3463 Yield x s -> return $ Yield x (Just (s, ref))
3464 Skip s -> return $ Skip (Just (s, ref))
3466 runIORefFinalizer ref
3469-- XXX These combinators are expensive due to the call to
3470-- onException/handle/try on each step. Therefore, when possible, they should
3471-- be called in an outer loop where we perform less iterations. For example, we
3472-- cannot call them on each iteration in a char stream, instead we can call
3473-- them when doing an IO on an array.
3475-- XXX For high performance error checks in busy streams we may need another
3476-- Error constructor in step.
3478-- | Run a side effect whenever the stream aborts due to an exception. The
3479-- exception is not caught, simply rethrown.
3480{-# INLINE_NORMAL onException #-}
3481onException :: MonadCatch m => m b -> Stream m a -> Stream m a
3482onException action str =
3483 gbracket_ (return ()) MC.try return
3484 (\_ (e :: MC.SomeException) _ -> nilM (action >> MC.throwM e))
3487{-# INLINE_NORMAL _onException #-}
3488_onException :: MonadCatch m => m b -> Stream m a -> Stream m a
3489_onException action (Stream step state) = Stream step' state
3493 {-# INLINE_LATE step' #-}
3495 res <- step gst st `MC.onException` action
3497 Yield x s -> return $ Yield x s
3498 Skip s -> return $ Skip s
3501-- XXX bracket is like concatMap, it generates a stream and then flattens it.
3502-- Like concatMap it has 10x worse performance compared to linear fused
3505-- | Run the first action before the stream starts and remember its output,
3506-- generate a stream using the output, run the second action providing the
3507-- remembered value as an argument whenever the stream ends normally or due to
3509{-# INLINE_NORMAL bracket_ #-}
3510bracket_ :: MonadCatch m => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
3511bracket_ bef aft bet =
3512 gbracket_ bef MC.try aft
3513 (\a (e :: SomeException) _ -> nilM (aft a >> MC.throwM e)) bet
3515-- | Run the first action before the stream starts and remember its output,
3516-- generate a stream using the output, run the second action providing the
3517-- remembered value as an argument whenever the stream ends normally or due to
3518-- an exception or gets garbage collected.
3520{-# INLINE_NORMAL bracket #-}
3521bracket :: (MonadAsync m, MonadCatch m)
3522 => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
3523bracket bef aft bet =
3524 gbracket bef MC.try aft
3525 (\a (e :: SomeException) _ -> aft a >> return (nilM (MC.throwM e))) bet
3527data BracketState s v = BracketInit | BracketRun s v
3529{-# INLINE_NORMAL _bracket #-}
3530_bracket :: MonadCatch m => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
3531_bracket bef aft bet = Stream step' BracketInit
3535 {-# INLINE_LATE step' #-}
3536 step' _ BracketInit = bef >>= \x -> return (Skip (BracketRun (bet x) x))
3538 -- NOTE: It is important to use UnStream instead of the Stream pattern
3539 -- here, otherwise we get huge perf degradation, see note in concatMap.
3540 step' gst (BracketRun (UnStream step state) v) = do
3541 -- res <- step gst state `MC.onException` aft v
3542 res <- MC.try $ step gst state
3544 Left (e :: SomeException) -> aft v >> MC.throwM e >> return Stop
3545 Right r -> case r of
3546 Yield x s -> return $ Yield x (BracketRun (Stream step s) v)
3547 Skip s -> return $ Skip (BracketRun (Stream step s) v)
3548 Stop -> aft v >> return Stop
3550-- | Run a side effect whenever the stream stops normally or aborts due to an
3552{-# INLINE finally_ #-}
3553finally_ :: MonadCatch m => m b -> Stream m a -> Stream m a
3554-- finally action xs = after action $ onException action xs
3555finally_ action xs = bracket_ (return ()) (\_ -> action) (const xs)
3557-- | Run a side effect whenever the stream stops normally or aborts due to an
3558-- exception or gets garbage collected.
3560{-# INLINE finally #-}
3561finally :: (MonadAsync m, MonadCatch m) => m b -> Stream m a -> Stream m a
3562finally action xs = bracket (return ()) (\_ -> action) (const xs)
3564-- | When evaluating a stream if an exception occurs, stream
3565-- evaluation aborts and the specified exception handler is run with
3566-- the exception and the Stream which threw the exception as argument.
3568{-# INLINE_NORMAL ghandle #-}
3569ghandle :: (MonadCatch m, Exception e)
3570 => (e -> Stream m a -> Stream m a) -> Stream m a -> Stream m a
3572 gbracket_ (return ()) MC.try return (\_ -> f) (\_ -> str)
3574-- | When evaluating a stream if an exception occurs, stream evaluation aborts
3575-- and the specified exception handler is run with the exception as argument.
3576{-# INLINE_NORMAL handle #-}
3577handle :: (MonadCatch m, Exception e)
3578 => (e -> Stream m a) -> Stream m a -> Stream m a
3580 gbracket_ (return ()) MC.try return (\_ e _ -> f e) (\_ -> str)
3582{-# INLINE_NORMAL _handle #-}
3583_handle :: (MonadCatch m, Exception e)
3584 => (e -> Stream m a) -> Stream m a -> Stream m a
3585_handle f (Stream step state) = Stream step' (Left state)
3589 {-# INLINE_LATE step' #-}
3590 step' gst (Left st) = do
3591 res <- MC.try $ step gst st
3593 Left e -> return $ Skip $ Right (f e)
3594 Right r -> case r of
3595 Yield x s -> return $ Yield x (Left s)
3596 Skip s -> return $ Skip (Left s)
3599 step' gst (Right (UnStream step1 st)) = do
3602 Yield x s -> return $ Yield x (Right (Stream step1 s))
3603 Skip s -> return $ Skip (Right (Stream step1 s))
3606-------------------------------------------------------------------------------
3607-- General transformation
3608-------------------------------------------------------------------------------
3610{-# INLINE_NORMAL transform #-}
3611transform :: Monad m => Pipe m a b -> Stream m a -> Stream m b
3612transform (Pipe pstep1 pstep2 pstate) (Stream step state) =
3613 Stream step' (Consume pstate, state)
3617 {-# INLINE_LATE step' #-}
3619 step' gst (Consume pst, st) = pst `seq` do
3620 r <- step (adaptState gst) st
3625 Pipe.Yield b pst' -> return $ Yield b (pst', s)
3626 Pipe.Continue pst' -> return $ Skip (pst', s)
3627 Skip s -> return $ Skip (Consume pst, s)
3630 step' _ (Produce pst, st) = pst `seq` do
3633 Pipe.Yield b pst' -> return $ Yield b (pst', st)
3634 Pipe.Continue pst' -> return $ Skip (pst', st)
3636------------------------------------------------------------------------------
3637-- Transformation by Folding (Scans)
3638------------------------------------------------------------------------------
3640------------------------------------------------------------------------------
3642------------------------------------------------------------------------------
3644-- XXX Is a prescan useful, discarding the last step does not sound useful? I
3645-- am not sure about the utility of this function, so this is implemented but
3646-- not exposed. We can expose it if someone provides good reasons why this is
3649-- XXX We have to execute the stream one step ahead to know that we are at the
3650-- last step. The vector implementation of prescan executes the last fold step
3651-- but does not yield the result. This means we have executed the effect but
3652-- discarded value. This does not sound right. In this implementation we are
3653-- not executing the last fold step.
3654{-# INLINE_NORMAL prescanlM' #-}
3655prescanlM' :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
3656prescanlM' f mz (Stream step state) = Stream step' (state, mz)
3658 {-# INLINE_LATE step' #-}
3659 step' gst (st, prev) = do
3660 r <- step (adaptState gst) st
3664 return $ Yield acc (s, f acc x)
3665 Skip s -> return $ Skip (s, prev)
3668{-# INLINE prescanl' #-}
3669prescanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
3670prescanl' f z = prescanlM' (\a b -> return (f a b)) (return z)
3672------------------------------------------------------------------------------
3673-- Monolithic postscans (postscan followed by a map)
3674------------------------------------------------------------------------------
3676-- The performance of a modular postscan followed by a map seems to be
3677-- equivalent to this monolithic scan followed by map therefore we may not need
3678-- this implementation. We just have it for performance comparison and in case
3679-- modular version does not perform well in some situation.
3681{-# INLINE_NORMAL postscanlMx' #-}
3682postscanlMx' :: Monad m
3683 => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> Stream m b
3684postscanlMx' fstep begin done (Stream step state) = do
3685 Stream step' (state, begin)
3687 {-# INLINE_LATE step' #-}
3688 step' gst (st, acc) = do
3689 r <- step (adaptState gst) st
3695 v `seq` y `seq` return (Yield v (s, return y))
3696 Skip s -> return $ Skip (s, acc)
3699{-# INLINE_NORMAL postscanlx' #-}
3700postscanlx' :: Monad m
3701 => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> Stream m b
3702postscanlx' fstep begin done s =
3703 postscanlMx' (\b a -> return (fstep b a)) (return begin) (return . done) s
3705-- XXX do we need consM strict to evaluate the begin value?
3706{-# INLINE scanlMx' #-}
3708 => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> Stream m b
3709scanlMx' fstep begin done s =
3710 (begin >>= \x -> x `seq` done x) `consM` postscanlMx' fstep begin done s
3712{-# INLINE scanlx' #-}
3714 => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> Stream m b
3715scanlx' fstep begin done s =
3716 scanlMx' (\b a -> return (fstep b a)) (return begin) (return . done) s
3718------------------------------------------------------------------------------
3720------------------------------------------------------------------------------
3722{-# INLINE_NORMAL postscanlM' #-}
3723postscanlM' :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
3724postscanlM' fstep begin (Stream step state) =
3725 Stream step' Nothing
3727 {-# INLINE_LATE step' #-}
3728 step' _ Nothing = do
3730 return $ Skip (Just (state, x))
3732 step' gst (Just (st, acc)) = do
3733 r <- step (adaptState gst) st
3737 return $ Yield y (Just (s, y))
3738 Skip s -> return $ Skip (Just (s, acc))
3741{-# INLINE_NORMAL postscanl' #-}
3742postscanl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
3743postscanl' f seed = postscanlM' (\a b -> return (f a b)) (return seed)
3745-- We can possibly have the "done" function as a Maybe to provide an option to
3746-- emit or not emit the accumulator when the stream stops.
3748-- TBD: use a single Yield point
3750{-# INLINE_NORMAL postscanlMAfter' #-}
3751postscanlMAfter' :: Monad m
3752 => (b -> a -> m b) -> m b -> (b -> m b) -> Stream m a -> Stream m b
3753postscanlMAfter' fstep initial done (Stream step1 state1) = do
3754 Stream step (Just (state1, initial))
3758 {-# INLINE_LATE step #-}
3759 step gst (Just (st, acc)) = do
3760 r <- step1 (adaptState gst) st
3765 y `seq` return (Yield y (Just (s, return y)))
3766 Skip s -> return $ Skip $ Just (s, acc)
3767 Stop -> acc >>= done >>= \res -> return (Yield res Nothing)
3768 step _ Nothing = return Stop
3770{-# INLINE_NORMAL postscanlM #-}
3771postscanlM :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
3772postscanlM fstep begin (Stream step state) = Stream step' Nothing
3774 {-# INLINE_LATE step' #-}
3775 step' _ Nothing = do
3777 return $ Skip (Just (state, r))
3779 step' gst (Just (st, acc)) = do
3780 r <- step (adaptState gst) st
3784 return (Yield y (Just (s, y)))
3785 Skip s -> return $ Skip (Just (s, acc))
3788{-# INLINE_NORMAL postscanl #-}
3789postscanl :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
3790postscanl f seed = postscanlM (\a b -> return (f a b)) (return seed)
3792{-# INLINE_NORMAL scanlM' #-}
3793scanlM' :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
3794scanlM' fstep begin (Stream step state) = Stream step' Nothing
3796 {-# INLINE_LATE step' #-}
3797 step' _ Nothing = do
3799 return $ Yield x (Just (state, x))
3800 step' gst (Just (st, acc)) = do
3801 r <- step (adaptState gst) st
3805 return $ Yield y (Just (s, y))
3806 Skip s -> return $ Skip (Just (s, acc))
3809{-# INLINE scanlMAfter' #-}
3810scanlMAfter' :: Monad m
3811 => (b -> a -> m b) -> m b -> (b -> m b) -> Stream m a -> Stream m b
3812scanlMAfter' fstep initial done s =
3813 (initial >>= \x -> x `seq` return x) `consM`
3814 postscanlMAfter' fstep initial done s
3816{-# INLINE scanl' #-}
3817scanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
3818scanl' f seed = scanlM' (\a b -> return (f a b)) (return seed)
3820{-# INLINE_NORMAL scanlM #-}
3821scanlM :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
3822scanlM fstep begin (Stream step state) = Stream step' Nothing
3824 {-# INLINE_LATE step' #-}
3825 step' _ Nothing = do
3827 return $ Yield x (Just (state, x))
3828 step' gst (Just (st, acc)) = do
3829 r <- step (adaptState gst) st
3833 return $ Yield y (Just (s, y))
3834 Skip s -> return $ Skip (Just (s, acc))
3835 Stop -> return $ Stop
3838scanl :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
3839scanl f seed = scanlM (\a b -> return (f a b)) (return seed)
3841{-# INLINE_NORMAL scanl1M #-}
3842scanl1M :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
3843scanl1M fstep (Stream step state) = Stream step' (state, Nothing)
3845 {-# INLINE_LATE step' #-}
3846 step' gst (st, Nothing) = do
3849 Yield x s -> return $ Yield x (s, Just x)
3850 Skip s -> return $ Skip (s, Nothing)
3853 step' gst (st, Just acc) = do
3858 return $ Yield z (s, Just z)
3859 Skip s -> return $ Skip (s, Just acc)
3862{-# INLINE scanl1 #-}
3863scanl1 :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
3864scanl1 f = scanl1M (\x y -> return (f x y))
3866{-# INLINE_NORMAL scanl1M' #-}
3867scanl1M' :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
3868scanl1M' fstep (Stream step state) = Stream step' (state, Nothing)
3870 {-# INLINE_LATE step' #-}
3871 step' gst (st, Nothing) = do
3874 Yield x s -> x `seq` return $ Yield x (s, Just x)
3875 Skip s -> return $ Skip (s, Nothing)
3878 step' gst (st, Just acc) = acc `seq` do
3883 z `seq` return $ Yield z (s, Just z)
3884 Skip s -> return $ Skip (s, Just acc)
3887{-# INLINE scanl1' #-}
3888scanl1' :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
3889scanl1' f = scanl1M' (\x y -> return (f x y))
3891------------------------------------------------------------------------------
3893------------------------------------------------------------------------------
3895data RollingMapState s a = RollingMapInit s | RollingMapGo s a
3897{-# INLINE rollingMapM #-}
3898rollingMapM :: Monad m => (a -> a -> m b) -> Stream m a -> Stream m b
3899rollingMapM f (Stream step1 state1) = Stream step (RollingMapInit state1)
3901 step gst (RollingMapInit st) = do
3902 r <- step1 (adaptState gst) st
3904 Yield x s -> Skip $ RollingMapGo s x
3905 Skip s -> Skip $ RollingMapInit s
3908 step gst (RollingMapGo s1 x1) = do
3909 r <- step1 (adaptState gst) s1
3913 return $ Yield res $ RollingMapGo s x
3914 Skip s -> return $ Skip $ RollingMapGo s x1
3915 Stop -> return $ Stop
3917{-# INLINE rollingMap #-}
3918rollingMap :: Monad m => (a -> a -> b) -> Stream m a -> Stream m b
3919rollingMap f = rollingMapM (\x y -> return $ f x y)
3921------------------------------------------------------------------------------
3922-- Tapping/Distributing
3923------------------------------------------------------------------------------
3926tap :: Monad m => Fold m a b -> Stream m a -> Stream m a
3927tap (Fold fstep initial extract) (Stream step state) = Stream step' Nothing
3931 step' _ Nothing = do
3933 return $ Skip (Just (r, state))
3935 step' gst (Just (acc, st)) = acc `seq` do
3940 return $ Yield x (Just (acc', s))
3941 Skip s -> return $ Skip (Just (acc, s))
3946{-# INLINE_NORMAL tapOffsetEvery #-}
3947tapOffsetEvery :: Monad m
3948 => Int -> Int -> Fold m a b -> Stream m a -> Stream m a
3949tapOffsetEvery offset n (Fold fstep initial extract) (Stream step state) =
3950 Stream step' Nothing
3954 {-# INLINE_LATE step' #-}
3955 step' _ Nothing = do
3957 return $ Skip (Just (r, state, offset `mod` n))
3959 step' gst (Just (acc, st, count)) | count <= 0 = do
3963 !acc' <- fstep acc x
3964 return $ Yield x (Just (acc', s, n - 1))
3965 Skip s -> return $ Skip (Just (acc, s, count))
3970 step' gst (Just (acc, st, count)) = do
3973 Yield x s -> return $ Yield x (Just (acc, s, count - 1))
3974 Skip s -> return $ Skip (Just (acc, s, count))
3979{-# INLINE_NORMAL pollCounts #-}
3983 -> (Stream m Int -> Stream m Int)
3987pollCounts predicate transf fld (Stream step state) = Stream step' Nothing
3990 {-# INLINE_LATE step' #-}
3991 step' _ Nothing = do
3992 -- As long as we are using an "Int" for counts lockfree reads from
3993 -- Var should work correctly on both 32-bit and 64-bit machines.
3994 -- However, an Int on a 32-bit machine may overflow quickly.
3995 countVar <- liftIO $ Prim.newIORef (0 :: Int)
3997 $ void $ runFold fld
3998 $ transf $ fromPrimIORef countVar
3999 return $ Skip (Just (countVar, tid, state))
4001 step' gst (Just (countVar, tid, st)) = do
4005 when (predicate x) $ liftIO $ Prim.modifyIORef' countVar (+ 1)
4006 return $ Yield x (Just (countVar, tid, s))
4007 Skip s -> return $ Skip (Just (countVar, tid, s))
4009 liftIO $ killThread tid
4012{-# INLINE_NORMAL tapRate #-}
4014 (MonadAsync m, MonadCatch m)
4019tapRate samplingRate action (Stream step state) = Stream step' Nothing
4021 {-# NOINLINE loop #-}
4022 loop countVar prev = do
4025 (do liftIO $ threadDelay (round $ samplingRate * 1000000)
4026 i <- liftIO $ Prim.readIORef countVar
4027 let !diff = i - prev
4030 (\(e :: AsyncException) -> do
4031 i <- liftIO $ Prim.readIORef countVar
4032 let !diff = i - prev
4034 throwM (MC.toException e))
4037 {-# INLINE_LATE step' #-}
4038 step' _ Nothing = do
4039 countVar <- liftIO $ Prim.newIORef 0
4040 tid <- fork $ loop countVar 0
4041 ref <- liftIO $ newIORef ()
4042 _ <- liftIO $ mkWeakIORef ref (killThread tid)
4043 return $ Skip (Just (countVar, tid, state, ref))
4045 step' gst (Just (countVar, tid, st, ref)) = do
4049 liftIO $ Prim.modifyIORef' countVar (+ 1)
4050 return $ Yield x (Just (countVar, tid, s, ref))
4051 Skip s -> return $ Skip (Just (countVar, tid, s, ref))
4053 liftIO $ killThread tid
4057-------------------------------------------------------------------------------
4059-------------------------------------------------------------------------------
4061{-# INLINE_NORMAL takeWhileM #-}
4062takeWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
4063takeWhileM f (Stream step state) = Stream step' state
4065 {-# INLINE_LATE step' #-}
4071 return $ if b then Yield x s else Stop
4072 Skip s -> return $ Skip s
4075{-# INLINE takeWhile #-}
4076takeWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
4077takeWhile f = takeWhileM (return . f)
4079{-# INLINE_NORMAL drop #-}
4080drop :: Monad m => Int -> Stream m a -> Stream m a
4081drop n (Stream step state) = Stream step' (state, Just n)
4083 {-# INLINE_LATE step' #-}
4084 step' gst (st, Just i)
4089 Yield _ s -> Skip (s, Just (i - 1))
4090 Skip s -> Skip (s, Just i)
4092 | otherwise = return $ Skip (st, Nothing)
4094 step' gst (st, Nothing) = do
4098 Yield x s -> Yield x (s, Nothing)
4099 Skip s -> Skip (s, Nothing)
4102data DropWhileState s a
4104 | DropWhileYield a s
4107{-# INLINE_NORMAL dropWhileM #-}
4108dropWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
4109dropWhileM f (Stream step state) = Stream step' (DropWhileDrop state)
4111 {-# INLINE_LATE step' #-}
4112 step' gst (DropWhileDrop st) = do
4118 then return $ Skip (DropWhileDrop s)
4119 else return $ Skip (DropWhileYield x s)
4120 Skip s -> return $ Skip (DropWhileDrop s)
4123 step' gst (DropWhileNext st) = do
4126 Yield x s -> return $ Skip (DropWhileYield x s)
4127 Skip s -> return $ Skip (DropWhileNext s)
4130 step' _ (DropWhileYield x st) = return $ Yield x (DropWhileNext st)
4132{-# INLINE dropWhile #-}
4133dropWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
4134dropWhile f = dropWhileM (return . f)
4136{-# INLINE_NORMAL filterM #-}
4137filterM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
4138filterM f (Stream step state) = Stream step' state
4140 {-# INLINE_LATE step' #-}
4149 Skip s -> return $ Skip s
4152{-# INLINE filter #-}
4153filter :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
4154filter f = filterM (return . f)
4156{-# INLINE_NORMAL uniq #-}
4157uniq :: (Eq a, Monad m) => Stream m a -> Stream m a
4158uniq (Stream step state) = Stream step' (Nothing, state)
4160 {-# INLINE_LATE step' #-}
4161 step' gst (Nothing, st) = do
4164 Yield x s -> return $ Yield x (Just x, s)
4165 Skip s -> return $ Skip (Nothing, s)
4167 step' gst (Just x, st) = do
4170 Yield y s | x == y -> return $ Skip (Just x, s)
4171 | otherwise -> return $ Yield y (Just y, s)
4172 Skip s -> return $ Skip (Just x, s)
4175------------------------------------------------------------------------------
4176-- Transformation by Mapping
4177------------------------------------------------------------------------------
4179{-# INLINE_NORMAL sequence #-}
4180sequence :: Monad m => Stream m (m a) -> Stream m a
4181sequence (Stream step state) = Stream step' state
4183 {-# INLINE_LATE step' #-}
4185 r <- step (adaptState gst) st
4187 Yield x s -> x >>= \a -> return (Yield a s)
4188 Skip s -> return $ Skip s
4191------------------------------------------------------------------------------
4193------------------------------------------------------------------------------
4195data LoopState x s = FirstYield s
4196 | InterspersingYield s
4199{-# INLINE_NORMAL intersperseM #-}
4200intersperseM :: Monad m => m a -> Stream m a -> Stream m a
4201intersperseM m (Stream step state) = Stream step' (FirstYield state)
4203 {-# INLINE_LATE step' #-}
4204 step' gst (FirstYield st) = do
4208 Yield x s -> Skip (YieldAndCarry x s)
4209 Skip s -> Skip (FirstYield s)
4212 step' gst (InterspersingYield st) = do
4217 return $ Yield a (YieldAndCarry x s)
4218 Skip s -> return $ Skip $ InterspersingYield s
4221 step' _ (YieldAndCarry x st) = return $ Yield x (InterspersingYield st)
4223{-# INLINE_NORMAL intersperseM_ #-}
4224intersperseM_ :: Monad m => m b -> Stream m a -> Stream m a
4225intersperseM_ m (Stream step1 state1) = Stream step (Left (pure (), state1))
4227 {-# INLINE_LATE step #-}
4228 step gst (Left (eff, st)) = do
4231 Yield x s -> eff >> return (Yield x (Right s))
4232 Skip s -> return $ Skip (Left (eff, s))
4235 step _ (Right st) = return $ Skip $ Left (void m, st)
4240 | SuffixYield a (SuffixState s a)
4242{-# INLINE_NORMAL intersperseSuffix #-}
4243intersperseSuffix :: forall m a. Monad m => m a -> Stream m a -> Stream m a
4244intersperseSuffix action (Stream step state) = Stream step' (SuffixElem state)
4246 {-# INLINE_LATE step' #-}
4247 step' gst (SuffixElem st) = do
4250 Yield x s -> Skip (SuffixYield x (SuffixSuffix s))
4251 Skip s -> Skip (SuffixElem s)
4254 step' _ (SuffixSuffix st) = do
4255 action >>= \r -> return $ Skip (SuffixYield r (SuffixElem st))
4257 step' _ (SuffixYield x next) = return $ Yield x next
4259{-# INLINE_NORMAL intersperseSuffix_ #-}
4260intersperseSuffix_ :: Monad m => m b -> Stream m a -> Stream m a
4261intersperseSuffix_ m (Stream step1 state1) = Stream step (Left state1)
4263 {-# INLINE_LATE step #-}
4264 step gst (Left st) = do
4267 Yield x s -> return $ Yield x (Right s)
4268 Skip s -> return $ Skip $ Left s
4271 step _ (Right st) = m >> return (Skip (Left st))
4273data SuffixSpanState s a
4274 = SuffixSpanElem s Int
4275 | SuffixSpanSuffix s
4276 | SuffixSpanYield a (SuffixSpanState s a)
4280-- | intersperse after every n items
4281{-# INLINE_NORMAL intersperseSuffixBySpan #-}
4282intersperseSuffixBySpan :: forall m a. Monad m
4283 => Int -> m a -> Stream m a -> Stream m a
4284intersperseSuffixBySpan n action (Stream step state) =
4285 Stream step' (SuffixSpanElem state n)
4287 {-# INLINE_LATE step' #-}
4288 step' gst (SuffixSpanElem st i) | i > 0 = do
4291 Yield x s -> Skip (SuffixSpanYield x (SuffixSpanElem s (i - 1)))
4292 Skip s -> Skip (SuffixSpanElem s i)
4293 Stop -> if i == n then Stop else Skip SuffixSpanLast
4294 step' _ (SuffixSpanElem st _) = return $ Skip (SuffixSpanSuffix st)
4296 step' _ (SuffixSpanSuffix st) = do
4297 action >>= \r -> return $ Skip (SuffixSpanYield r (SuffixSpanElem st n))
4299 step' _ (SuffixSpanLast) = do
4300 action >>= \r -> return $ Skip (SuffixSpanYield r SuffixSpanStop)
4302 step' _ (SuffixSpanYield x next) = return $ Yield x next
4304 step' _ (SuffixSpanStop) = return Stop
4306{-# INLINE intersperse #-}
4307intersperse :: Monad m => a -> Stream m a -> Stream m a
4308intersperse a = intersperseM (return a)
4310{-# INLINE_NORMAL insertBy #-}
4311insertBy :: Monad m => (a -> a -> Ordering) -> a -> Stream m a -> Stream m a
4312insertBy cmp a (Stream step state) = Stream step' (state, False, Nothing)
4314 {-# INLINE_LATE step' #-}
4315 step' gst (st, False, _) = do
4318 Yield x s -> case cmp a x of
4319 GT -> return $ Yield x (s, False, Nothing)
4320 _ -> return $ Yield a (s, True, Just x)
4321 Skip s -> return $ Skip (s, False, Nothing)
4322 Stop -> return $ Yield a (st, True, Nothing)
4324 step' _ (_, True, Nothing) = return Stop
4326 step' gst (st, True, Just prev) = do
4329 Yield x s -> return $ Yield prev (s, True, Just x)
4330 Skip s -> return $ Skip (s, True, Just prev)
4331 Stop -> return $ Yield prev (st, True, Nothing)
4333------------------------------------------------------------------------------
4335------------------------------------------------------------------------------
4337{-# INLINE_NORMAL deleteBy #-}
4338deleteBy :: Monad m => (a -> a -> Bool) -> a -> Stream m a -> Stream m a
4339deleteBy eq x (Stream step state) = Stream step' (state, False)
4341 {-# INLINE_LATE step' #-}
4342 step' gst (st, False) = do
4345 Yield y s -> return $
4346 if eq x y then Skip (s, True) else Yield y (s, False)
4347 Skip s -> return $ Skip (s, False)
4350 step' gst (st, True) = do
4353 Yield y s -> return $ Yield y (s, True)
4354 Skip s -> return $ Skip (s, True)
4357------------------------------------------------------------------------------
4358-- Transformation by Map and Filter
4359------------------------------------------------------------------------------
4361-- XXX Will this always fuse properly?
4362{-# INLINE_NORMAL mapMaybe #-}
4363mapMaybe :: Monad m => (a -> Maybe b) -> Stream m a -> Stream m b
4364mapMaybe f = fmap fromJust . filter isJust . map f
4366{-# INLINE_NORMAL mapMaybeM #-}
4367mapMaybeM :: Monad m => (a -> m (Maybe b)) -> Stream m a -> Stream m b
4368mapMaybeM f = fmap fromJust . filter isJust . mapM f
4370------------------------------------------------------------------------------
4372------------------------------------------------------------------------------
4374{-# INLINE_NORMAL indexed #-}
4375indexed :: Monad m => Stream m a -> Stream m (Int, a)
4376indexed (Stream step state) = Stream step' (state, 0)
4378 {-# INLINE_LATE step' #-}
4379 step' gst (st, i) = i `seq` do
4380 r <- step (adaptState gst) st
4382 Yield x s -> return $ Yield (i, x) (s, i+1)
4383 Skip s -> return $ Skip (s, i)
4386{-# INLINE_NORMAL indexedR #-}
4387indexedR :: Monad m => Int -> Stream m a -> Stream m (Int, a)
4388indexedR m (Stream step state) = Stream step' (state, m)
4390 {-# INLINE_LATE step' #-}
4391 step' gst (st, i) = i `seq` do
4392 r <- step (adaptState gst) st
4394 Yield x s -> let i' = i - 1
4395 in return $ Yield (i, x) (s, i')
4396 Skip s -> return $ Skip (s, i)
4399{-# INLINE_NORMAL zipWithM #-}
4401 => (a -> b -> m c) -> Stream m a -> Stream m b -> Stream m c
4402zipWithM f (Stream stepa ta) (Stream stepb tb) = Stream step (ta, tb, Nothing)
4404 {-# INLINE_LATE step #-}
4405 step gst (sa, sb, Nothing) = do
4406 r <- stepa (adaptState gst) sa
4409 Yield x sa' -> Skip (sa', sb, Just x)
4410 Skip sa' -> Skip (sa', sb, Nothing)
4413 step gst (sa, sb, Just x) = do
4414 r <- stepb (adaptState gst) sb
4418 return $ Yield z (sa, sb', Nothing)
4419 Skip sb' -> return $ Skip (sa, sb', Just x)
4422#if __GLASGOW_HASKELL__ >= 801
4423{-# RULES "zipWithM xs xs"
4424 forall f xs. zipWithM @Identity f xs xs = mapM (\x -> f x x) xs #-}
4427{-# INLINE zipWith #-}
4428zipWith :: Monad m => (a -> b -> c) -> Stream m a -> Stream m b -> Stream m c
4429zipWith f = zipWithM (\a b -> return (f a b))
4431------------------------------------------------------------------------------
4433------------------------------------------------------------------------------
4435{-# INLINE_NORMAL mergeByM #-}
4438 => (a -> a -> m Ordering) -> Stream m a -> Stream m a -> Stream m a
4439mergeByM cmp (Stream stepa ta) (Stream stepb tb) =
4440 Stream step (Just ta, Just tb, Nothing, Nothing)
4442 {-# INLINE_LATE step #-}
4444 -- one of the values is missing, and the corresponding stream is running
4445 step gst (Just sa, sb, Nothing, b) = do
4448 Yield a sa' -> Skip (Just sa', sb, Just a, b)
4449 Skip sa' -> Skip (Just sa', sb, Nothing, b)
4450 Stop -> Skip (Nothing, sb, Nothing, b)
4452 step gst (sa, Just sb, a, Nothing) = do
4455 Yield b sb' -> Skip (sa, Just sb', a, Just b)
4456 Skip sb' -> Skip (sa, Just sb', a, Nothing)
4457 Stop -> Skip (sa, Nothing, a, Nothing)
4459 -- both the values are available
4460 step _ (sa, sb, Just a, Just b) = do
4462 return $ case res of
4463 GT -> Yield b (sa, sb, Just a, Nothing)
4464 _ -> Yield a (sa, sb, Nothing, Just b)
4466 -- one of the values is missing, corresponding stream is done
4467 step _ (Nothing, sb, Nothing, Just b) =
4468 return $ Yield b (Nothing, sb, Nothing, Nothing)
4470 step _ (sa, Nothing, Just a, Nothing) =
4471 return $ Yield a (sa, Nothing, Nothing, Nothing)
4473 step _ (Nothing, Nothing, Nothing, Nothing) = return Stop
4475{-# INLINE mergeBy #-}
4478 => (a -> a -> Ordering) -> Stream m a -> Stream m a -> Stream m a
4479mergeBy cmp = mergeByM (\a b -> return $ cmp a b)
4481------------------------------------------------------------------------------
4482-- Transformation comprehensions
4483------------------------------------------------------------------------------
4485{-# INLINE_NORMAL the #-}
4486the :: (Eq a, Monad m) => Stream m a -> m (Maybe a)
4487the (Stream step state) = go state
4490 r <- step defState st
4492 Yield x s -> go' x s
4494 Stop -> return Nothing
4496 r <- step defState st
4498 Yield x s | x == n -> go' n s
4499 | otherwise -> return Nothing
4501 Stop -> return (Just n)
4503{-# INLINE runFold #-}
4504runFold :: (Monad m) => Fold m a b -> Stream m a -> m b
4505runFold (Fold step begin done) = foldlMx' step begin done
4507-------------------------------------------------------------------------------
4508-- Concurrent application and fold
4509-------------------------------------------------------------------------------
4511-- XXX These functions should be moved to Stream/Parallel.hs
4513-- Using StreamD the worker stream producing code can fuse with the code to
4514-- queue output to the SVar giving some perf boost.
4516-- Note that StreamD can only be used in limited situations, specifically, we
4517-- cannot implement joinStreamVarPar using this.
4519-- XXX make sure that the SVar passed is a Parallel style SVar.
4521-- | Fold the supplied stream to the SVar asynchronously using Parallel
4522-- concurrency style.
4523-- {-# INLINE_NORMAL toSVarParallel #-}
4524{-# INLINE toSVarParallel #-}
4525toSVarParallel :: MonadAsync m
4526 => State t m a -> SVar t m a -> Stream m a -> m ()
4527toSVarParallel st sv xs =
4528 if svarInspectMode sv
4532 case getYieldLimit st of
4533 Nothing -> doFork (work Nothing)
4535 (handleChildException sv)
4536 Just _ -> doFork (workLim Nothing)
4538 (handleChildException sv)
4543 {-# NOINLINE work #-}
4544 work info = (runFold (FL.toParallelSVar sv info) xs)
4546 {-# NOINLINE workLim #-}
4547 workLim info = runFold (FL.toParallelSVarLimited sv info) xs
4549 {-# NOINLINE forkWithDiag #-}
4551 -- We do not use workerCount in case of ParallelVar but still there is
4552 -- no harm in maintaining it correctly.
4553 liftIO $ atomicModifyIORefCAS_ (workerCount sv) $ \n -> n + 1
4555 -- This allocation matters when significant number of workers are being
4556 -- sent. We allocate it only when needed. The overhead increases by 4x.
4558 case yieldRateInfo sv of
4559 Nothing -> return Nothing
4560 Just _ -> liftIO $ do
4561 cntRef <- newIORef 0
4562 t <- getTime Monotonic
4563 lat <- newIORef (0, t)
4564 return $ Just WorkerInfo
4565 { workerYieldMax = 0
4566 , workerYieldCount = cntRef
4567 , workerLatencyStart = lat
4570 case getYieldLimit st of
4571 Nothing -> doFork (work winfo)
4573 (handleChildException sv)
4574 Just _ -> doFork (workLim winfo)
4576 (handleChildException sv)
4579{-# INLINE_NORMAL mkParallelD #-}
4580mkParallelD :: MonadAsync m => Stream m a -> Stream m a
4581mkParallelD m = Stream step Nothing
4584 step gst Nothing = do
4585 sv <- newParallelVar StopNone gst
4586 toSVarParallel gst sv m
4587 -- XXX use unfold instead?
4588 return $ Skip $ Just $ fromSVar sv
4590 step gst (Just (UnStream step1 st)) = do
4593 Yield a s -> Yield a (Just $ Stream step1 s)
4594 Skip s -> Skip (Just $ Stream step1 s)
4597-- Compare with mkAsync. mkAsync uses an Async style SVar whereas this uses a
4598-- parallel style SVar for evaluation. Currently, parallel style cannot use
4599-- rate control whereas Async style can use rate control. In async style SVar
4600-- the worker thread terminates when the buffer is full whereas in Parallel
4603-- | Make the stream producer and consumer run concurrently by introducing a
4604-- buffer between them. The producer thread evaluates the input stream until
4605-- the buffer fills, it blocks if the buffer is full until there is space in
4606-- the buffer. The consumer consumes the stream lazily from the buffer.
4610{-# INLINE_NORMAL mkParallel #-}
4611mkParallel :: (K.IsStream t, MonadAsync m) => t m a -> t m a
4612mkParallel = fromStreamD . mkParallelD . toStreamD
4614-- Note: we can use another API with two callbacks stop and yield if we want
4615-- the callback to be able to indicate end of stream.
4617-- | Generates a callback and a stream pair. The callback returned is used to
4618-- queue values to the stream. The stream is infinite, there is no way for the
4619-- callback to indicate that it is done now.
4623{-# INLINE_NORMAL newCallbackStream #-}
4624newCallbackStream :: (K.IsStream t, MonadAsync m) => m ((a -> m ()), t m a)
4625newCallbackStream = do
4626 sv <- newParallelVar StopNone defState
4628 -- XXX Add our own thread-id to the SVar as we can not know the callback's
4629 -- thread-id and the callback is not run in a managed worker. We need to
4630 -- handle this better.
4631 liftIO myThreadId >>= modifyThread sv
4633 let callback a = liftIO $ void $ send sv (ChildYield a)
4634 -- XXX we can return an SVar and then the consumer can unfold from the
4636 return (callback, fromStreamD (fromSVar sv))
4638-------------------------------------------------------------------------------
4640-------------------------------------------------------------------------------
4642-- | Create an SVar with a fold consumer that will fold any elements sent to it
4643-- using the supplied fold function.
4644{-# INLINE newFoldSVar #-}
4645newFoldSVar :: MonadAsync m => State t m a -> Fold m a b -> m (SVar t m a)
4646newFoldSVar stt f = do
4647 -- Buffer size for the SVar is derived from the current state
4648 sv <- newParallelVar StopAny (adaptState stt)
4649 -- Add the producer thread-id to the SVar.
4650 liftIO myThreadId >>= modifyThread sv
4651 void $ doFork (work sv) (svarMrun sv) (handleFoldException sv)
4656 {-# NOINLINE work #-}
4657 work sv = void $ runFold f $ fromProducer sv
4659data TapState sv st = TapInit | Tapping sv st | TapDone st
4661{-# INLINE_NORMAL tapAsync #-}
4662tapAsync :: MonadAsync m => Fold m a b -> Stream m a -> Stream m a
4663tapAsync f (Stream step1 state1) = Stream step TapInit
4667 -- In general, a Stop event would come equipped with the result
4668 -- of the fold. It is not used here but it would be useful in
4669 -- applicative and distribute.
4670 done <- fromConsumer svr
4671 when (not done) $ do
4672 liftIO $ withDiagMVar svr "teeToSVar: waiting to drain"
4673 $ takeMVar (outputDoorBellFromConsumer svr)
4677 liftIO $ sendStop svr Nothing
4678 -- drain/wait until a stop event arrives from the fold.
4681 {-# INLINE_LATE step #-}
4682 step gst TapInit = do
4683 sv <- newFoldSVar gst f
4684 return $ Skip (Tapping sv state1)
4686 step gst (Tapping sv st) = do
4690 done <- pushToFold sv a
4693 -- XXX we do not need to wait synchronously here
4695 return $ Yield a (TapDone s)
4696 else return $ Yield a (Tapping sv s)
4697 Skip s -> return $ Skip (Tapping sv s)
4702 step gst (TapDone st) = do
4705 Yield a s -> Yield a (TapDone s)
4706 Skip s -> Skip (TapDone s)
4709-- XXX Exported from Array again as this fold is specific to Array
4710-- | Take last 'n' elements from the stream and discard the rest.
4712lastN :: (Storable a, MonadIO m) => Int -> Fold m a (Array a)
4714 | n <= 0 = fmap (const mempty) FL.drain
4715 | otherwise = A.unsafeFreeze <$> Fold step initial done
4717 step (Tuple3' rb rh i) a = do
4718 rh1 <- liftIO $ RB.unsafeInsert rb rh a
4719 return $ Tuple3' rb rh1 (i + 1)
4720 initial = fmap (\(a, b) -> Tuple3' a b (0 :: Int)) $ liftIO $ RB.new n
4721 done (Tuple3' rb rh i) = do
4722 arr <- liftIO $ MA.newArray n
4723 foldFunc i rh snoc' arr rb
4724 snoc' b a = liftIO $ MA.unsafeSnoc b a
4726 | i < n = RB.unsafeFoldRingM
4727 | otherwise = RB.unsafeFoldRingFullM
4729------------------------------------------------------------------------------
4731------------------------------------------------------------------------------
4733-- XXX using getTime in the loop can be pretty expensive especially for
4734-- computations where iterations are lightweight. We have the following
4737-- 1) Run a timeout thread updating a flag asynchronously and check that
4738-- flag here, that way we can have a cheap termination check.
4740-- 2) Use COARSE clock to get time with lower resolution but more efficiently.
4742-- 3) Use rdtscp/rdtsc to get time directly from the processor, compute the
4743-- termination value of rdtsc in the beginning and then in each iteration just
4744-- get rdtsc and check if we should terminate.
4748 | TakeByTimeCheck st s
4749 | TakeByTimeYield st s
4751{-# INLINE_NORMAL takeByTime #-}
4752takeByTime :: (MonadIO m, TimeUnit64 t) => t -> Stream m a -> Stream m a
4753takeByTime duration (Stream step1 state1) = Stream step (TakeByTimeInit state1)
4756 lim = toRelTime64 duration
4758 {-# INLINE_LATE step #-}
4759 step _ (TakeByTimeInit _) | lim == 0 = return Stop
4760 step _ (TakeByTimeInit st) = do
4761 t0 <- liftIO $ getTime Monotonic
4762 return $ Skip (TakeByTimeYield st t0)
4763 step _ (TakeByTimeCheck st t0) = do
4764 t <- liftIO $ getTime Monotonic
4766 if diffAbsTime64 t t0 > lim
4768 else Skip (TakeByTimeYield st t0)
4769 step gst (TakeByTimeYield st t0) = do
4772 Yield x s -> Yield x (TakeByTimeCheck s t0)
4773 Skip s -> Skip (TakeByTimeCheck s t0)
4776data DropByTime st s x
4778 | DropByTimeGen st s
4779 | DropByTimeCheck st s x
4780 | DropByTimeYield st
4782{-# INLINE_NORMAL dropByTime #-}
4783dropByTime :: (MonadIO m, TimeUnit64 t) => t -> Stream m a -> Stream m a
4784dropByTime duration (Stream step1 state1) = Stream step (DropByTimeInit state1)
4787 lim = toRelTime64 duration
4789 {-# INLINE_LATE step #-}
4790 step _ (DropByTimeInit st) = do
4791 t0 <- liftIO $ getTime Monotonic
4792 return $ Skip (DropByTimeGen st t0)
4793 step gst (DropByTimeGen st t0) = do
4796 Yield x s -> Skip (DropByTimeCheck s t0 x)
4797 Skip s -> Skip (DropByTimeGen s t0)
4799 step _ (DropByTimeCheck st t0 x) = do
4800 t <- liftIO $ getTime Monotonic
4801 if diffAbsTime64 t t0 <= lim
4802 then return $ Skip $ DropByTimeGen st t0
4803 else return $ Yield x $ DropByTimeYield st
4804 step gst (DropByTimeYield st) = do
4807 Yield x s -> Yield x (DropByTimeYield s)
4808 Skip s -> Skip (DropByTimeYield s)
4811{-# INLINE updateTimeVar #-}
4812updateTimeVar :: Prim.IORef Int64 -> IO ()
4813updateTimeVar timeVar = do
4814 MicroSecond64 t <- fromAbsTime <$> getTime Monotonic
4815 Prim.modifyIORef' timeVar (const t)
4817{-# INLINE updateWithDelay #-}
4818updateWithDelay :: RealFrac a => a -> Prim.IORef Int64 -> IO ()
4819updateWithDelay precision timeVar = do
4820 threadDelay (delayTime precision)
4821 updateTimeVar timeVar
4825 -- Keep the minimum at least a millisecond to avoid high CPU usage
4826 {-# INLINE delayTime #-}
4828 | g' >= fromIntegral (maxBound :: Int) = maxBound
4830 | otherwise = round g'
4834 g' = g * 10 ^ (6 :: Int)
4836-- XXX we should move this to stream generation section of this file. Also, the
4837-- take/drop combinators above should be moved to filtering section.
4839{-# INLINE_NORMAL times #-}
4840times :: MonadAsync m => Double -> Stream m (AbsTime, RelTime64)
4841times g = Stream step Nothing
4845 {-# INLINE_LATE step #-}
4847 -- XXX note that this is safe only on a 64-bit machine. On a 32-bit
4848 -- machine a 64-bit 'Var' cannot be read consistently without a lock
4849 -- while another thread is writing to it.
4850 timeVar <- liftIO $ Prim.newIORef (0 :: Int64)
4851 liftIO $ updateTimeVar timeVar
4852 tid <- forkManaged $ liftIO $ forever (updateWithDelay g timeVar)
4853 a <- liftIO $ Prim.readIORef timeVar
4854 return $ Skip $ Just (timeVar, tid, a)
4856 step _ s@(Just (timeVar, _, t0)) = do
4857 a <- liftIO $ Prim.readIORef timeVar
4858 -- XXX we can perhaps use an AbsTime64 using a 64 bit Int for
4859 -- efficiency. or maybe we can use a representation using Double for
4860 -- floating precision time
4861 return $ Yield (toAbsTime (MicroSecond64 t0),
4862 (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