4-- Module : Streamly.Internal.Data.Stream.StreamD
5-- Copyright : (c) 2018 Harendra Kumar
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.Types (SPEC(..))
326import System.Mem (performMajorGC)
327import Fusion.Plugin.Types (Fuse(..))
328import Streamly.Internal.Data.IORef.Prim (Prim)
329import Streamly.Internal.Data.Time.Units
330 (TimeUnit64, toRelTime64, diffAbsTime64, RelTime64)
331import Streamly.Internal.Data.Atomics (atomicModifyIORefCAS_)
332import Streamly.Internal.Data.Array.Storable.Foreign.Types (Array(..))
333import Streamly.Internal.Data.Fold.Types (Fold(..))
334import Streamly.Internal.Data.Parser (ParseError(..))
335import Streamly.Internal.Data.Pipe.Types (Pipe(..), PipeState(..))
336import Streamly.Internal.Data.Time.Clock (Clock(Monotonic), getTime)
337import Streamly.Internal.Data.Time.Units
338 (MicroSecond64(..), fromAbsTime, toAbsTime, AbsTime)
339import Streamly.Internal.Data.Unfold.Types (Unfold(..))
340import Streamly.Internal.Data.Tuple.Strict (Tuple3'(..))
341import Streamly.Internal.Data.Stream.SVar (fromConsumer, pushToFold)
343import qualified Streamly.Internal.Data.IORef.Prim as Prim
344import qualified Streamly.Internal.Data.Pipe.Types as Pipe
345import qualified Streamly.Internal.Data.Array.Storable.Foreign.Types as A
346import qualified Streamly.Internal.Data.Array.Storable.Foreign.Mut.Types as MA
347import qualified Streamly.Internal.Data.Fold as FL
348import qualified Streamly.Memory.Ring as RB
349import qualified Streamly.Internal.Data.Stream.StreamK as K
350import qualified Streamly.Internal.Data.Parser as PR
351import qualified Streamly.Internal.Data.Parser.ParserD as PRD
352import qualified Control.Monad.Catch as MC
353import qualified Control.Monad.Trans.Reader as Reader
354import qualified Control.Monad.Trans.State.Strict as State
355import qualified Prelude
358 ( map, mapM, mapM_, repeat, foldr, last, take, filter
359 , takeWhile, drop, dropWhile, all, any, maximum, minimum, elem
360 , notElem, null, head, tail, zipWith, lookup, foldr1, sequence
361 , (!!), scanl, scanl1, concatMap, replicate, enumFromTo, concat
362 , reverse, iterate, splitAt)
363import Streamly.Internal.Data.Stream.StreamD.Type
364import Streamly.Internal.Data.SVar
366------------------------------------------------------------------------------
368------------------------------------------------------------------------------
370-- | An empty 'Stream'.
371{-# INLINE_NORMAL nil #-}
372nil :: Monad m => Stream m a
373nil = Stream (\_ _ -> return Stop) ()
375-- | An empty 'Stream' with a side effect.
376{-# INLINE_NORMAL nilM #-}
377nilM :: Monad m => m b -> Stream m a
378nilM m = Stream (\_ _ -> m >> return Stop) ()
380{-# INLINE_NORMAL consM #-}
381consM :: Monad m => m a -> Stream m a -> Stream m a
382consM m (Stream step state) = Stream step1 Nothing
384 {-# INLINE_LATE step1 #-}
385 step1 _ Nothing = m >>= \x -> return $ Yield x (Just state)
386 step1 gst (Just st) = do
390 Yield a s -> Yield a (Just s)
391 Skip s -> Skip (Just s)
394-- XXX implement in terms of consM?
395-- cons x = consM (return x)
397-- | Can fuse but has O(n^2) complexity.
398{-# INLINE_NORMAL cons #-}
399cons :: Monad m => a -> Stream m a -> Stream m a
400cons x (Stream step state) = Stream step1 Nothing
402 {-# INLINE_LATE step1 #-}
403 step1 _ Nothing = return $ Yield x (Just state)
404 step1 gst (Just st) = do
408 Yield a s -> Yield a (Just s)
409 Skip s -> Skip (Just s)
412-------------------------------------------------------------------------------
414-------------------------------------------------------------------------------
416-- Does not fuse, has the same performance as the StreamK version.
417{-# INLINE_NORMAL uncons #-}
418uncons :: Monad m => Stream m a -> m (Maybe (a, Stream m a))
419uncons (UnStream step state) = go state
422 r <- step defState st
424 Yield x s -> return $ Just (x, Stream step s)
426 Stop -> return Nothing
428------------------------------------------------------------------------------
429-- Generation by unfold
430------------------------------------------------------------------------------
432{-# INLINE_NORMAL unfoldrM #-}
433unfoldrM :: Monad m => (s -> m (Maybe (a, s))) -> s -> Stream m a
434unfoldrM next state = Stream step state
436 {-# INLINE_LATE step #-}
440 Just (x, s) -> Yield x s
443{-# INLINE_LATE unfoldr #-}
444unfoldr :: Monad m => (s -> Maybe (a, s)) -> s -> Stream m a
445unfoldr f = unfoldrM (return . f)
447-- | Convert an 'Unfold' into a 'Stream' by supplying it a seed.
449{-# INLINE_NORMAL unfold #-}
450unfold :: Monad m => Unfold m a b -> a -> Stream m b
451unfold (Unfold ustep inject) seed = Stream step Nothing
453 {-# INLINE_LATE step #-}
454 step _ Nothing = inject seed >>= return . Skip . Just
455 step _ (Just st) = do
458 Yield x s -> Yield x (Just s)
459 Skip s -> Skip (Just s)
462------------------------------------------------------------------------------
463-- Specialized Generation
464------------------------------------------------------------------------------
466{-# INLINE_NORMAL repeatM #-}
467repeatM :: Monad m => m a -> Stream m a
468repeatM x = Stream (\_ _ -> x >>= \r -> return $ Yield r ()) ()
470{-# INLINE_NORMAL repeat #-}
471repeat :: Monad m => a -> Stream m a
472repeat x = Stream (\_ _ -> return $ Yield x ()) ()
474{-# INLINE_NORMAL iterateM #-}
475iterateM :: Monad m => (a -> m a) -> m a -> Stream m a
476iterateM step = Stream (\_ st -> st >>= \x -> return $ Yield x (step x))
478{-# INLINE_NORMAL iterate #-}
479iterate :: Monad m => (a -> a) -> a -> Stream m a
480iterate step st = iterateM (return . step) (return st)
482{-# INLINE_NORMAL replicateM #-}
483replicateM :: forall m a. Monad m => Int -> m a -> Stream m a
484replicateM n p = Stream step n
486 {-# INLINE_LATE step #-}
488 | i <= 0 = return Stop
491 return $ Yield x (i - 1)
493{-# INLINE_NORMAL replicate #-}
494replicate :: Monad m => Int -> a -> Stream m a
495replicate n x = replicateM n (return x)
497-- This would not work properly for floats, therefore we put an Integral
499-- | Can be used to enumerate unbounded integrals. This does not check for
500-- overflow or underflow for bounded integrals.
501{-# INLINE_NORMAL enumerateFromStepIntegral #-}
502enumerateFromStepIntegral :: (Integral a, Monad m) => a -> a -> Stream m a
503enumerateFromStepIntegral from stride =
504 from `seq` stride `seq` Stream step from
506 {-# INLINE_LATE step #-}
507 step _ !x = return $ Yield x $! (x + stride)
509-- We are assuming that "to" is constrained by the type to be within
511{-# INLINE enumerateFromToIntegral #-}
512enumerateFromToIntegral :: (Monad m, Integral a) => a -> a -> Stream m a
513enumerateFromToIntegral from to =
514 takeWhile (<= to) $ enumerateFromStepIntegral from 1
516{-# INLINE enumerateFromIntegral #-}
517enumerateFromIntegral :: (Monad m, Integral a, Bounded a) => a -> Stream m a
518enumerateFromIntegral from = enumerateFromToIntegral from maxBound
520data EnumState a = EnumInit | EnumYield a a a | EnumStop
522{-# INLINE_NORMAL enumerateFromThenToIntegralUp #-}
523enumerateFromThenToIntegralUp
524 :: (Monad m, Integral a)
525 => a -> a -> a -> Stream m a
526enumerateFromThenToIntegralUp from next to = Stream step EnumInit
528 {-# INLINE_LATE step #-}
534 else Yield from EnumStop
535 else -- from <= next <= to
536 let stride = next - from
537 in Skip $ EnumYield from stride (to - stride)
539 step _ (EnumYield x stride toMinus) =
542 then Yield x EnumStop
543 else Yield x $ EnumYield (x + stride) stride toMinus
545 step _ EnumStop = return Stop
547{-# INLINE_NORMAL enumerateFromThenToIntegralDn #-}
548enumerateFromThenToIntegralDn
549 :: (Monad m, Integral a)
550 => a -> a -> a -> Stream m a
551enumerateFromThenToIntegralDn from next to = Stream step EnumInit
553 {-# INLINE_LATE step #-}
555 return $ if to > next
558 else Yield from EnumStop
559 else -- from >= next >= to
560 let stride = next - from
561 in Skip $ EnumYield from stride (to - stride)
563 step _ (EnumYield x stride toMinus) =
566 then Yield x EnumStop
567 else Yield x $ EnumYield (x + stride) stride toMinus
569 step _ EnumStop = return Stop
571{-# INLINE_NORMAL enumerateFromThenToIntegral #-}
572enumerateFromThenToIntegral
573 :: (Monad m, Integral a)
574 => a -> a -> a -> Stream m a
575enumerateFromThenToIntegral from next to
576 | next >= from = enumerateFromThenToIntegralUp from next to
577 | otherwise = enumerateFromThenToIntegralDn from next to
579{-# INLINE_NORMAL enumerateFromThenIntegral #-}
580enumerateFromThenIntegral
581 :: (Monad m, Integral a, Bounded a)
582 => a -> a -> Stream m a
583enumerateFromThenIntegral from next =
585 then enumerateFromThenToIntegralUp from next maxBound
586 else enumerateFromThenToIntegralDn from next minBound
588-- For floating point numbers if the increment is less than the precision then
589-- it just gets lost. Therefore we cannot always increment it correctly by just
591-- 9007199254740992 + 1 + 1 :: Double => 9.007199254740992e15
592-- 9007199254740992 + 2 :: Double => 9.007199254740994e15
594-- Instead we accumulate the increment counter and compute the increment
595-- every time before adding it to the starting number.
597-- This works for Integrals as well as floating point numbers, but
598-- enumerateFromStepIntegral is faster for integrals.
599{-# INLINE_NORMAL enumerateFromStepNum #-}
600enumerateFromStepNum :: (Monad m, Num a) => a -> a -> Stream m a
601enumerateFromStepNum from stride = Stream step 0
603 {-# INLINE_LATE step #-}
604 step _ !i = return $ (Yield $! (from + i * stride)) $! (i + 1)
606{-# INLINE_NORMAL numFrom #-}
607numFrom :: (Monad m, Num a) => a -> Stream m a
608numFrom from = enumerateFromStepNum from 1
610{-# INLINE_NORMAL numFromThen #-}
611numFromThen :: (Monad m, Num a) => a -> a -> Stream m a
612numFromThen from next = enumerateFromStepNum from (next - from)
614-- We cannot write a general function for Num. The only way to write code
615-- portable between the two is to use a 'Real' constraint and convert between
616-- Fractional and Integral using fromRational which is horribly slow.
617{-# INLINE_NORMAL enumerateFromToFractional #-}
618enumerateFromToFractional
619 :: (Monad m, Fractional a, Ord a)
620 => a -> a -> Stream m a
621enumerateFromToFractional from to =
622 takeWhile (<= to + 1 / 2) $ enumerateFromStepNum from 1
624{-# INLINE_NORMAL enumerateFromThenToFractional #-}
625enumerateFromThenToFractional
626 :: (Monad m, Fractional a, Ord a)
627 => a -> a -> a -> Stream m a
628enumerateFromThenToFractional from next to =
629 takeWhile predicate $ numFromThen from next
631 mid = (next - from) / 2
632 predicate | next >= from = (<= to + mid)
633 | otherwise = (>= to + mid)
635-------------------------------------------------------------------------------
636-- Generation by Conversion
637-------------------------------------------------------------------------------
639{-# INLINE_NORMAL fromIndicesM #-}
640fromIndicesM :: Monad m => (Int -> m a) -> Stream m a
641fromIndicesM gen = Stream step 0
643 {-# INLINE_LATE step #-}
646 return $ Yield x (i + 1)
648{-# INLINE fromIndices #-}
649fromIndices :: Monad m => (Int -> a) -> Stream m a
650fromIndices gen = fromIndicesM (return . gen)
652{-# INLINE_NORMAL generateM #-}
653generateM :: Monad m => Int -> (Int -> m a) -> Stream m a
654generateM n gen = n `seq` Stream step 0
656 {-# INLINE_LATE step #-}
657 step _ i | i < n = do
659 return $ Yield x (i + 1)
660 | otherwise = return Stop
662{-# INLINE generate #-}
663generate :: Monad m => Int -> (Int -> a) -> Stream m a
664generate n gen = generateM n (return . gen)
666-- XXX we need the MonadAsync constraint because of a rewrite rule.
667-- | Convert a list of monadic actions to a 'Stream'
668{-# INLINE_LATE fromListM #-}
669fromListM :: MonadAsync m => [m a] -> Stream m a
670fromListM = Stream step
672 {-# INLINE_LATE step #-}
673 step _ (m:ms) = m >>= \x -> return $ Yield x ms
674 step _ [] = return Stop
676{-# INLINE toStreamD #-}
677toStreamD :: (K.IsStream t, Monad m) => t m a -> Stream m a
678toStreamD = fromStreamK . K.toStream
680{-# INLINE_NORMAL fromPrimIORef #-}
681fromPrimIORef :: (MonadIO m, Prim a) => Prim.IORef a -> Stream m a
682fromPrimIORef var = Stream step ()
684 {-# INLINE_LATE step #-}
685 step _ () = liftIO (Prim.readIORef var) >>= \x -> return $ Yield x ()
687-------------------------------------------------------------------------------
688-- Generation from SVar
689-------------------------------------------------------------------------------
691data FromSVarState t m a =
693 | FromSVarRead (SVar t m a)
694 | FromSVarLoop (SVar t m a) [ChildEvent a]
695 | FromSVarDone (SVar t m a)
697{-# INLINE_NORMAL fromSVar #-}
698fromSVar :: (MonadAsync m) => SVar t m a -> Stream m a
699fromSVar svar = Stream step FromSVarInit
702 {-# INLINE_LATE step #-}
703 step _ FromSVarInit = do
704 ref <- liftIO $ newIORef ()
705 _ <- liftIO $ mkWeakIORef ref hook
706 -- when this copy of svar gets garbage collected "ref" will get
707 -- garbage collected and our GC hook will be called.
708 let sv = svar{svarRef = Just ref}
709 return $ Skip (FromSVarRead sv)
713 {-# NOINLINE hook #-}
715 when (svarInspectMode svar) $ do
716 r <- liftIO $ readIORef (svarStopTime (svarStats svar))
718 printSVar svar "SVar Garbage Collected"
720 -- If there are any SVars referenced by this SVar a GC will prompt
721 -- them to be cleaned up quickly.
722 when (svarInspectMode svar) performMajorGC
724 step _ (FromSVarRead sv) = do
725 list <- readOutputQ sv
726 -- Reversing the output is important to guarantee that we process the
727 -- outputs in the same order as they were generated by the constituent
729 return $ Skip $ FromSVarLoop sv (Prelude.reverse list)
731 step _ (FromSVarLoop sv []) = do
732 done <- postProcess sv
733 return $ Skip $ if done
734 then (FromSVarDone sv)
735 else (FromSVarRead sv)
737 step _ (FromSVarLoop sv (ev : es)) = do
739 ChildYield a -> return $ Yield a (FromSVarLoop sv es)
740 ChildStop tid e -> do
744 stop <- shouldStop tid
747 liftIO (cleanupSVar sv)
748 return $ Skip (FromSVarDone sv)
749 else return $ Skip (FromSVarLoop sv es)
751 case fromException ex of
753 return $ Skip (FromSVarLoop sv es)
754 Nothing -> liftIO (cleanupSVar sv) >> throwM ex
758 case svarStopStyle sv of
759 StopNone -> return False
760 StopAny -> return True
762 sid <- liftIO $ readIORef (svarStopBy sv)
763 return $ if tid == sid then True else False
765 step _ (FromSVarDone sv) = do
766 when (svarInspectMode sv) $ do
767 t <- liftIO $ getTime Monotonic
768 liftIO $ writeIORef (svarStopTime (svarStats sv)) (Just t)
769 liftIO $ printSVar sv "SVar Done"
772-------------------------------------------------------------------------------
773-- Process events received by a fold consumer from a stream producer
774-------------------------------------------------------------------------------
776{-# INLINE_NORMAL fromProducer #-}
777fromProducer :: (MonadAsync m) => SVar t m a -> Stream m a
778fromProducer svar = Stream step (FromSVarRead svar)
781 {-# INLINE_LATE step #-}
782 step _ (FromSVarRead sv) = do
783 list <- readOutputQ sv
784 -- Reversing the output is important to guarantee that we process the
785 -- outputs in the same order as they were generated by the constituent
787 return $ Skip $ FromSVarLoop sv (Prelude.reverse list)
789 step _ (FromSVarLoop sv []) = return $ Skip $ FromSVarRead sv
790 step _ (FromSVarLoop sv (ev : es)) = do
792 ChildYield a -> return $ Yield a (FromSVarLoop sv es)
793 ChildStop tid e -> do
797 sendStopToProducer sv
798 return $ Skip (FromSVarDone sv)
799 Just _ -> error "Bug: fromProducer: received exception"
801 step _ (FromSVarDone sv) = do
802 when (svarInspectMode sv) $ do
803 t <- liftIO $ getTime Monotonic
804 liftIO $ writeIORef (svarStopTime (svarStats sv)) (Just t)
805 liftIO $ printSVar sv "SVar Done"
808 step _ FromSVarInit = undefined
810-------------------------------------------------------------------------------
811-- Hoisting the inner monad
812-------------------------------------------------------------------------------
814{-# INLINE_NORMAL hoist #-}
815hoist :: Monad n => (forall x. m x -> n x) -> Stream m a -> Stream n a
816hoist f (Stream step state) = (Stream step' state)
818 {-# INLINE_LATE step' #-}
820 r <- f $ step (adaptState gst) st
822 Yield x s -> Yield x s
826{-# INLINE generally #-}
827generally :: Monad m => Stream Identity a -> Stream m a
828generally = hoist (return . runIdentity)
830{-# INLINE_NORMAL liftInner #-}
831liftInner :: (Monad m, MonadTrans t, Monad (t m))
832 => Stream m a -> Stream (t m) a
833liftInner (Stream step state) = Stream step' state
835 {-# INLINE_LATE step' #-}
837 r <- lift $ step (adaptState gst) st
839 Yield x s -> Yield x s
843{-# INLINE_NORMAL runReaderT #-}
844runReaderT :: Monad m => m s -> Stream (ReaderT s m) a -> Stream m a
845runReaderT env (Stream step state) = Stream step' (state, env)
847 {-# INLINE_LATE step' #-}
848 step' gst (st, action) = do
850 r <- Reader.runReaderT (step (adaptState gst) st) sv
852 Yield x s -> Yield x (s, return sv)
853 Skip s -> Skip (s, return sv)
856{-# INLINE_NORMAL evalStateT #-}
857evalStateT :: Monad m => m s -> Stream (StateT s m) a -> Stream m a
858evalStateT initial (Stream step state) = Stream step' (state, initial)
860 {-# INLINE_LATE step' #-}
861 step' gst (st, action) = do
863 (r, sv') <- State.runStateT (step (adaptState gst) st) sv
865 Yield x s -> Yield x (s, return sv')
866 Skip s -> Skip (s, return sv')
869{-# INLINE_NORMAL runStateT #-}
870runStateT :: Monad m => m s -> Stream (StateT s m) a -> Stream m (s, a)
871runStateT initial (Stream step state) = Stream step' (state, initial)
873 {-# INLINE_LATE step' #-}
874 step' gst (st, action) = do
876 (r, sv') <- State.runStateT (step (adaptState gst) st) sv
878 Yield x s -> Yield (sv', x) (s, return sv')
879 Skip s -> Skip (s, return sv')
882------------------------------------------------------------------------------
883-- Elimination by Folds
884------------------------------------------------------------------------------
886------------------------------------------------------------------------------
888------------------------------------------------------------------------------
890{-# INLINE_NORMAL foldr1 #-}
891foldr1 :: Monad m => (a -> a -> a) -> Stream m a -> m (Maybe a)
895 Nothing -> return Nothing
896 Just (h, t) -> fmap Just (foldr f h t)
898------------------------------------------------------------------------------
900------------------------------------------------------------------------------
902{-# INLINE_NORMAL foldlT #-}
903foldlT :: (Monad m, Monad (s m), MonadTrans s)
904 => (s m b -> a -> s m b) -> s m b -> Stream m a -> s m b
905foldlT fstep begin (Stream step state) = go SPEC begin state
908 r <- lift $ step defState st
910 Yield x s -> go SPEC (fstep acc x) s
911 Skip s -> go SPEC acc s
914-- Note, this is going to have horrible performance, because of the nature of
915-- the stream type (i.e. direct stream vs CPS). Its only for reference, it is
916-- likely be practically unusable.
917{-# INLINE_NORMAL foldlS #-}
919 => (Stream m b -> a -> Stream m b) -> Stream m b -> Stream m a -> Stream m b
920foldlS fstep begin (Stream step state) = Stream step' (Left (state, begin))
922 step' gst (Left (st, acc)) = do
923 r <- step (adaptState gst) st
925 Yield x s -> Skip (Left (s, fstep acc x))
926 Skip s -> Skip (Left (s, acc))
927 Stop -> Skip (Right acc)
929 step' gst (Right (Stream stp stt)) = do
930 r <- stp (adaptState gst) stt
932 Yield x s -> Yield x (Right (Stream stp s))
933 Skip s -> Skip (Right (Stream stp s))
936------------------------------------------------------------------------------
938------------------------------------------------------------------------------
940-- Inlined definition. Without the inline "serially/parser/take" benchmark
941-- degrades and parseMany does not fuse. Even using "inline" at the callsite
943{-# INLINE splitAt #-}
944splitAt :: Int -> [a] -> ([a],[a])
947 | otherwise = splitAt' n ls
949 splitAt' :: Int -> [a] -> ([a], [a])
950 splitAt' _ [] = ([], [])
951 splitAt' 1 (x:xs) = ([x], xs)
952 splitAt' m (x:xs) = (x:xs', xs'')
954 (xs', xs'') = splitAt' (m - 1) xs
956-- | Run a 'Parse' over a stream.
957{-# INLINE_NORMAL parselMx' #-}
960 => (s -> a -> m (PR.Step s b))
965parselMx' pstep initial extract (Stream step state) = do
966 initial >>= go SPEC state []
970 -- XXX currently we are using a dumb list based approach for backtracking
971 -- buffer. This can be replaced by a sliding/ring buffer using Data.Array.
972 -- That will allow us more efficient random back and forth movement.
974 go !_ st buf !pst = do
975 r <- step defState st
980 PR.Partial 0 pst1 -> go SPEC s [] pst1
981 PR.Partial n pst1 -> do
982 assert (n <= length (x:buf)) (return ())
983 let src0 = Prelude.take n (x:buf)
984 src = Prelude.reverse src0
985 gobuf SPEC s [] src pst1
986 PR.Continue 0 pst1 -> go SPEC s (x:buf) pst1
987 PR.Continue n pst1 -> do
988 assert (n <= length (x:buf)) (return ())
989 let (src0, buf1) = splitAt n (x:buf)
990 src = Prelude.reverse src0
991 gobuf SPEC s buf1 src pst1
992 PR.Done _ b -> return b
993 PR.Error err -> throwM $ ParseError err
994 Skip s -> go SPEC s buf pst
997 gobuf !_ s buf [] !pst = go SPEC s buf pst
998 gobuf !_ s buf (x:xs) !pst = do
1001 PR.Partial 0 pst1 ->
1002 gobuf SPEC s [] xs pst1
1003 PR.Partial n pst1 -> do
1004 assert (n <= length (x:buf)) (return ())
1005 let src0 = Prelude.take n (x:buf)
1006 src = Prelude.reverse src0 ++ xs
1007 gobuf SPEC s [] src pst1
1008 PR.Continue 0 pst1 -> gobuf SPEC s (x:buf) xs pst1
1009 PR.Continue n pst1 -> do
1010 assert (n <= length (x:buf)) (return ())
1011 let (src0, buf1) = splitAt n (x:buf)
1012 src = Prelude.reverse src0 ++ xs
1013 gobuf SPEC s buf1 src pst1
1014 PR.Done _ b -> return b
1015 PR.Error err -> throwM $ ParseError err
1017------------------------------------------------------------------------------
1019------------------------------------------------------------------------------
1021{-# ANN type ParseChunksState Fuse #-}
1022data ParseChunksState x inpBuf st pst =
1023 ParseChunksInit inpBuf st
1024 | ParseChunksInitLeftOver inpBuf
1025 | ParseChunksStream st inpBuf pst
1026 | ParseChunksBuf inpBuf st inpBuf pst
1027 | ParseChunksYield x (ParseChunksState x inpBuf st pst)
1029{-# INLINE_NORMAL parseMany #-}
1035parseMany (PRD.Parser pstep initial extract) (Stream step state) =
1036 Stream stepOuter (ParseChunksInit [] state)
1040 {-# INLINE_LATE stepOuter #-}
1041 -- Buffer is empty, get the first element from the stream, initialize the
1042 -- fold and then go to stream processing loop.
1043 stepOuter gst (ParseChunksInit [] st) = do
1044 r <- step (adaptState gst) st
1046 Yield x s -> initial >>= return . Skip . ParseChunksBuf [x] s []
1047 Skip s -> return $ Skip $ ParseChunksInit [] s
1050 -- Buffer is not empty, go to buffered processing loop
1051 stepOuter _ (ParseChunksInit src st) = do
1052 initial >>= return . Skip . ParseChunksBuf src st []
1054 -- XXX we just discard any leftover input at the end
1055 stepOuter _ (ParseChunksInitLeftOver _) = return Stop
1057 -- Buffer is empty, process elements from the stream
1058 stepOuter gst (ParseChunksStream st buf pst) = do
1059 r <- step (adaptState gst) st
1064 PR.Partial 0 pst1 ->
1065 return $ Skip $ ParseChunksStream s [] pst1
1066 PR.Partial n pst1 -> do
1067 assert (n <= length (x:buf)) (return ())
1068 let src0 = Prelude.take n (x:buf)
1069 src = Prelude.reverse src0
1070 return $ Skip $ ParseChunksBuf src s [] pst1
1071 -- PR.Continue 0 pst1 ->
1072 -- return $ Skip $ ParseChunksStream s (x:buf) pst1
1073 PR.Continue n pst1 -> do
1074 assert (n <= length (x:buf)) (return ())
1075 let (src0, buf1) = splitAt n (x:buf)
1076 src = Prelude.reverse src0
1077 return $ Skip $ ParseChunksBuf src s buf1 pst1
1078 -- XXX Specialize for Stop 0 common case?
1080 assert (n <= length (x:buf)) (return ())
1081 let src = Prelude.reverse (Prelude.take n (x:buf))
1083 ParseChunksYield b (ParseChunksInit src s)
1084 PR.Error err -> throwM $ ParseError err
1085 Skip s -> return $ Skip $ ParseChunksStream s buf pst
1088 let src = Prelude.reverse buf
1090 ParseChunksYield b (ParseChunksInitLeftOver src)
1092 -- go back to stream processing mode
1093 stepOuter _ (ParseChunksBuf [] s buf pst) =
1094 return $ Skip $ ParseChunksStream s buf pst
1096 -- buffered processing loop
1097 stepOuter _ (ParseChunksBuf (x:xs) s buf pst) = do
1100 PR.Partial 0 pst1 ->
1101 return $ Skip $ ParseChunksBuf xs s [] pst1
1102 PR.Partial n pst1 -> do
1103 assert (n <= length (x:buf)) (return ())
1104 let src0 = Prelude.take n (x:buf)
1105 src = Prelude.reverse src0 ++ xs
1106 return $ Skip $ ParseChunksBuf src s [] pst1
1107 -- PR.Continue 0 pst1 -> return $ Skip $ ParseChunksBuf xs s (x:buf) pst1
1108 PR.Continue n pst1 -> do
1109 assert (n <= length (x:buf)) (return ())
1110 let (src0, buf1) = splitAt n (x:buf)
1111 src = Prelude.reverse src0 ++ xs
1112 return $ Skip $ ParseChunksBuf src s buf1 pst1
1113 -- XXX Specialize for Stop 0 common case?
1115 assert (n <= length (x:buf)) (return ())
1116 let src = Prelude.reverse (Prelude.take n (x:buf)) ++ xs
1117 return $ Skip $ ParseChunksYield b (ParseChunksInit src s)
1118 PR.Error err -> throwM $ ParseError err
1120 stepOuter _ (ParseChunksYield a next) = return $ Yield a next
1122{-# ANN type ConcatParseState Fuse #-}
1123data ConcatParseState x inpBuf st p =
1124 ConcatParseInit inpBuf st p
1125 | ConcatParseInitLeftOver inpBuf
1126 | ConcatParseStream st inpBuf p
1127 | ConcatParseBuf inpBuf st inpBuf p
1128 | ConcatParseYield x (ConcatParseState x inpBuf st p)
1130{-# INLINE_NORMAL parseIterate #-}
1133 => (b -> PRD.Parser m a b)
1137parseIterate func seed (Stream step state) =
1138 Stream stepOuter (ConcatParseInit [] state (func seed))
1142 {-# INLINE_LATE stepOuter #-}
1143 -- Buffer is empty, go to stream processing loop
1144 stepOuter _ (ConcatParseInit [] st (PRD.Parser pstep initial extract)) = do
1145 initial >>= \r -> return $ Skip $ ConcatParseStream st []
1146 (PRD.Parser pstep (return r) extract)
1148 -- Buffer is not empty, go to buffered processing loop
1149 stepOuter _ (ConcatParseInit src st
1150 (PRD.Parser pstep initial extract)) = do
1151 initial >>= \r -> return $ Skip $ ConcatParseBuf src st []
1152 (PRD.Parser pstep (return r) extract)
1154 -- XXX we just discard any leftover input at the end
1155 stepOuter _ (ConcatParseInitLeftOver _) = return Stop
1157 -- Buffer is empty process elements from the stream
1158 stepOuter gst (ConcatParseStream st buf
1159 p@(PRD.Parser pstep initial extract)) = do
1160 r <- step (adaptState gst) st
1166 PR.Partial 0 pst1 ->
1167 return $ Skip $ ConcatParseStream s []
1168 (PRD.Parser pstep (return pst1) extract)
1169 PR.Partial n pst1 -> do
1170 assert (n <= length (x:buf)) (return ())
1171 let src0 = Prelude.take n (x:buf)
1172 src = Prelude.reverse src0
1173 return $ Skip $ ConcatParseBuf src s []
1174 (PRD.Parser pstep (return pst1) extract)
1175 -- PR.Continue 0 pst1 ->
1176 -- return $ Skip $ ConcatParseStream s (x:buf) pst1
1177 PR.Continue n pst1 -> do
1178 assert (n <= length (x:buf)) (return ())
1179 let (src0, buf1) = splitAt n (x:buf)
1180 src = Prelude.reverse src0
1181 return $ Skip $ ConcatParseBuf src s buf1
1182 (PRD.Parser pstep (return pst1) extract)
1183 -- XXX Specialize for Stop 0 common case?
1185 assert (n <= length (x:buf)) (return ())
1186 let src = Prelude.reverse (Prelude.take n (x:buf))
1188 ConcatParseYield b (ConcatParseInit src s (func b))
1189 PR.Error err -> throwM $ ParseError err
1190 Skip s -> return $ Skip $ ConcatParseStream s buf p
1194 let src = Prelude.reverse buf
1196 ConcatParseYield b (ConcatParseInitLeftOver src)
1198 -- go back to stream processing mode
1199 stepOuter _ (ConcatParseBuf [] s buf p) =
1200 return $ Skip $ ConcatParseStream s buf p
1202 -- buffered processing loop
1203 stepOuter _ (ConcatParseBuf (x:xs) s buf
1204 (PRD.Parser pstep initial extract)) = do
1208 PR.Partial 0 pst1 ->
1209 return $ Skip $ ConcatParseBuf xs s []
1210 (PRD.Parser pstep (return pst1) extract)
1211 PR.Partial n pst1 -> do
1212 assert (n <= length (x:buf)) (return ())
1213 let src0 = Prelude.take n (x:buf)
1214 src = Prelude.reverse src0 ++ xs
1215 return $ Skip $ ConcatParseBuf src s []
1216 (PRD.Parser pstep (return pst1) extract)
1217 -- PR.Continue 0 pst1 -> return $ Skip $ ConcatParseBuf xs s (x:buf) pst1
1218 PR.Continue n pst1 -> do
1219 assert (n <= length (x:buf)) (return ())
1220 let (src0, buf1) = splitAt n (x:buf)
1221 src = Prelude.reverse src0 ++ xs
1222 return $ Skip $ ConcatParseBuf src s buf1
1223 (PRD.Parser pstep (return pst1) extract)
1224 -- XXX Specialize for Stop 0 common case?
1226 assert (n <= length (x:buf)) (return ())
1227 let src = Prelude.reverse (Prelude.take n (x:buf)) ++ xs
1228 return $ Skip $ ConcatParseYield b
1229 (ConcatParseInit src s (func b))
1230 PR.Error err -> throwM $ ParseError err
1232 stepOuter _ (ConcatParseYield a next) = return $ Yield a next
1234------------------------------------------------------------------------------
1236------------------------------------------------------------------------------
1238-- | Run a streaming composition, discard the results.
1239{-# INLINE_LATE drain #-}
1240drain :: Monad m => Stream m a -> m ()
1241-- drain = foldrM (\_ xs -> xs) (return ())
1242drain (Stream step state) = go SPEC state
1245 r <- step defState st
1247 Yield _ s -> go SPEC s
1251{-# INLINE_NORMAL null #-}
1252null :: Monad m => Stream m a -> m Bool
1253null m = foldrM (\_ _ -> return False) (return True) m
1255{-# INLINE_NORMAL head #-}
1256head :: Monad m => Stream m a -> m (Maybe a)
1257head m = foldrM (\x _ -> return (Just x)) (return Nothing) m
1259{-# INLINE_NORMAL headElse #-}
1260headElse :: Monad m => a -> Stream m a -> m a
1261headElse a m = foldrM (\x _ -> return x) (return a) m
1263-- Does not fuse, has the same performance as the StreamK version.
1264{-# INLINE_NORMAL tail #-}
1265tail :: Monad m => Stream m a -> m (Maybe (Stream m a))
1266tail (UnStream step state) = go state
1269 r <- step defState st
1271 Yield _ s -> return (Just $ Stream step s)
1273 Stop -> return Nothing
1275-- XXX will it fuse? need custom impl?
1276{-# INLINE_NORMAL last #-}
1277last :: Monad m => Stream m a -> m (Maybe a)
1278last = foldl' (\_ y -> Just y) Nothing
1280{-# INLINE_NORMAL elem #-}
1281elem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
1282-- elem e m = foldrM (\x xs -> if x == e then return True else xs) (return False) m
1283elem e (Stream step state) = go state
1286 r <- step defState st
1289 | x == e -> return True
1292 Stop -> return False
1294{-# INLINE_NORMAL notElem #-}
1295notElem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
1296notElem e s = fmap not (elem e s)
1298{-# INLINE_NORMAL all #-}
1299all :: Monad m => (a -> Bool) -> Stream m a -> m Bool
1300-- all p m = foldrM (\x xs -> if p x then xs else return False) (return True) m
1301all p (Stream step state) = go state
1304 r <- step defState st
1308 | otherwise -> return False
1312{-# INLINE_NORMAL any #-}
1313any :: Monad m => (a -> Bool) -> Stream m a -> m Bool
1314-- any p m = foldrM (\x xs -> if p x then return True else xs) (return False) m
1315any p (Stream step state) = go state
1318 r <- step defState st
1321 | p x -> return True
1324 Stop -> return False
1326{-# INLINE_NORMAL maximum #-}
1327maximum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
1328maximum (Stream step state) = go Nothing state
1331 r <- step defState st
1333 Yield x s -> go (Just x) s
1334 Skip s -> go Nothing s
1335 Stop -> return Nothing
1336 go (Just acc) st = do
1337 r <- step defState st
1340 | acc <= x -> go (Just x) s
1341 | otherwise -> go (Just acc) s
1342 Skip s -> go (Just acc) s
1343 Stop -> return (Just acc)
1345{-# INLINE_NORMAL maximumBy #-}
1346maximumBy :: Monad m => (a -> a -> Ordering) -> Stream m a -> m (Maybe a)
1347maximumBy cmp (Stream step state) = go Nothing state
1350 r <- step defState st
1352 Yield x s -> go (Just x) s
1353 Skip s -> go Nothing s
1354 Stop -> return Nothing
1355 go (Just acc) st = do
1356 r <- step defState st
1358 Yield x s -> case cmp acc x of
1359 GT -> go (Just acc) s
1361 Skip s -> go (Just acc) s
1362 Stop -> return (Just acc)
1364{-# INLINE_NORMAL minimum #-}
1365minimum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
1366minimum (Stream step state) = go Nothing state
1369 r <- step defState st
1371 Yield x s -> go (Just x) s
1372 Skip s -> go Nothing s
1373 Stop -> return Nothing
1374 go (Just acc) st = do
1375 r <- step defState st
1378 | acc <= x -> go (Just acc) s
1379 | otherwise -> go (Just x) s
1380 Skip s -> go (Just acc) s
1381 Stop -> return (Just acc)
1383{-# INLINE_NORMAL minimumBy #-}
1384minimumBy :: Monad m => (a -> a -> Ordering) -> Stream m a -> m (Maybe a)
1385minimumBy cmp (Stream step state) = go Nothing state
1388 r <- step defState st
1390 Yield x s -> go (Just x) s
1391 Skip s -> go Nothing s
1392 Stop -> return Nothing
1393 go (Just acc) st = do
1394 r <- step defState st
1396 Yield x s -> case cmp acc x of
1398 _ -> go (Just acc) s
1399 Skip s -> go (Just acc) s
1400 Stop -> return (Just acc)
1402{-# INLINE_NORMAL (!!) #-}
1403(!!) :: (Monad m) => Stream m a -> Int -> m (Maybe a)
1404(Stream step state) !! i = go i state
1407 r <- step defState st
1409 Yield x s | n < 0 -> return Nothing
1410 | n == 0 -> return $ Just x
1411 | otherwise -> go (n - 1) s
1413 Stop -> return Nothing
1415{-# INLINE_NORMAL lookup #-}
1416lookup :: (Monad m, Eq a) => a -> Stream m (a, b) -> m (Maybe b)
1417lookup e m = foldrM (\(a, b) xs -> if e == a then return (Just b) else xs)
1420{-# INLINE_NORMAL findM #-}
1421findM :: Monad m => (a -> m Bool) -> Stream m a -> m (Maybe a)
1422findM p m = foldrM (\x xs -> p x >>= \r -> if r then return (Just x) else xs)
1426find :: Monad m => (a -> Bool) -> Stream m a -> m (Maybe a)
1427find p = findM (return . p)
1429{-# INLINE_NORMAL findIndices #-}
1430findIndices :: Monad m => (a -> Bool) -> Stream m a -> Stream m Int
1431findIndices p (Stream step state) = Stream step' (state, 0)
1433 {-# INLINE_LATE step' #-}
1434 step' gst (st, i) = i `seq` do
1435 r <- step (adaptState gst) st
1437 Yield x s -> if p x then Yield i (s, i+1) else Skip (s, i+1)
1438 Skip s -> Skip (s, i)
1441{-# INLINE toListRev #-}
1442toListRev :: Monad m => Stream m a -> m [a]
1443toListRev = foldl' (flip (:)) []
1445-- We can implement reverse as:
1447-- > reverse = foldlS (flip cons) nil
1449-- However, this implementation is unusable because of the horrible performance
1450-- of cons. So we just convert it to a list first and then stream from the
1453-- XXX Maybe we can use an Array instead of a list here?
1454{-# INLINE_NORMAL reverse #-}
1455reverse :: Monad m => Stream m a -> Stream m a
1456reverse m = Stream step Nothing
1458 {-# INLINE_LATE step #-}
1461 return $ Skip (Just xs)
1462 step _ (Just (x:xs)) = return $ Yield x (Just xs)
1463 step _ (Just []) = return Stop
1465-- Much faster reverse for Storables
1466{-# INLINE_NORMAL reverse' #-}
1467reverse' :: forall m a. (MonadIO m, Storable a) => Stream m a -> Stream m a
1469-- This commented implementation copies the whole stream into one single array
1470-- and then streams from that array, this is 3-4x faster than the chunked code
1471-- that follows. Though this could be problematic due to unbounded large
1472-- allocations. We need to figure out why the chunked code is slower and if we
1473-- can optimize the chunked code to work as fast as this one. It may be a
1475import Foreign.ForeignPtr (touchForeignPtr)
1476import Foreign.ForeignPtr.Unsafe (unsafeForeignPtrToPtr)
1477import Foreign.Ptr (Ptr, plusPtr)
1478reverse' m = Stream step Nothing
1480 {-# INLINE_LATE step #-}
1482 arr <- A.fromStreamD m
1483 let p = aEnd arr `plusPtr` negate (sizeOf (undefined :: a))
1484 return $ Skip $ Just (aStart arr, p)
1486 step _ (Just (start, p)) | p < unsafeForeignPtrToPtr start = return Stop
1488 step _ (Just (start, p)) = do
1489 let !x = A.unsafeInlineIO $ do
1491 touchForeignPtr start
1493 next = p `plusPtr` negate (sizeOf (undefined :: a))
1494 return $ Yield x (Just (start, next))
1501 $ A.fromStreamDArraysOf A.defaultChunkSize m
1504------------------------------------------------------------------------------
1505-- Grouping/Splitting
1506------------------------------------------------------------------------------
1508{-# INLINE_NORMAL splitSuffixBy' #-}
1509splitSuffixBy' :: Monad m
1510 => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1511splitSuffixBy' predicate f (Stream step state) =
1512 Stream (stepOuter f) (Just state)
1516 {-# INLINE_LATE stepOuter #-}
1517 stepOuter (Fold fstep initial done) gst (Just st) = do
1518 res <- step (adaptState gst) st
1524 then done acc' >>= \val -> return $ Yield val (Just s)
1527 Skip s -> return $ Skip $ Just s
1533 res <- step (adaptState gst) stt
1538 then done acc' >>= \val -> return $ Yield val (Just s)
1540 Skip s -> go SPEC s acc
1541 Stop -> done acc >>= \val -> return $ Yield val Nothing
1543 stepOuter _ _ Nothing = return Stop
1545{-# INLINE_NORMAL groupsBy #-}
1551groupsBy cmp f (Stream step state) = Stream (stepOuter f) (Just state, Nothing)
1555 {-# INLINE_LATE stepOuter #-}
1556 stepOuter (Fold fstep initial done) gst (Just st, Nothing) = do
1557 res <- step (adaptState gst) st
1564 Skip s -> return $ Skip $ (Just s, Nothing)
1569 go !_ prev stt !acc = do
1570 res <- step (adaptState gst) stt
1577 else done acc >>= \r -> return $ Yield r (Just s, Just x)
1578 Skip s -> go SPEC prev s acc
1579 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1581 stepOuter (Fold fstep initial done) gst (Just st, Just prev) = do
1583 acc' <- fstep acc prev
1588 -- XXX code duplicated from the previous equation
1590 res <- step (adaptState gst) stt
1597 else done acc >>= \r -> return $ Yield r (Just s, Just x)
1598 Skip s -> go SPEC s acc
1599 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1601 stepOuter _ _ (Nothing,_) = return Stop
1603{-# INLINE_NORMAL groupsRollingBy #-}
1604groupsRollingBy :: Monad m
1609groupsRollingBy cmp f (Stream step state) =
1610 Stream (stepOuter f) (Just state, Nothing)
1613 {-# INLINE_LATE stepOuter #-}
1614 stepOuter (Fold fstep initial done) gst (Just st, Nothing) = do
1615 res <- step (adaptState gst) st
1622 Skip s -> return $ Skip $ (Just s, Nothing)
1626 go !_ prev stt !acc = do
1627 res <- step (adaptState gst) stt
1635 done acc >>= \r -> return $ Yield r (Just s, Just x)
1636 Skip s -> go SPEC prev s acc
1637 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1639 stepOuter (Fold fstep initial done) gst (Just st, Just prev') = do
1641 acc' <- fstep acc prev'
1642 go SPEC prev' st acc'
1645 go !_ prevv stt !acc = do
1646 res <- step (adaptState gst) stt
1653 else done acc >>= \r -> return $ Yield r (Just s, Just x)
1654 Skip s -> go SPEC prevv s acc
1655 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1657 stepOuter _ _ (Nothing, _) = return Stop
1659{-# INLINE_NORMAL splitBy #-}
1660splitBy :: Monad m => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1661splitBy predicate f (Stream step state) = Stream (step' f) (Just state)
1665 {-# INLINE_LATE step' #-}
1666 step' (Fold fstep initial done) gst (Just st) = initial >>= go SPEC st
1671 res <- step (adaptState gst) stt
1675 then done acc >>= \r -> return $ Yield r (Just s)
1679 Skip s -> go SPEC s acc
1680 Stop -> done acc >>= \r -> return $ Yield r Nothing
1682 step' _ _ Nothing = return Stop
1684-- XXX requires -funfolding-use-threshold=150 in lines-unlines benchmark
1685{-# INLINE_NORMAL splitSuffixBy #-}
1686splitSuffixBy :: Monad m
1687 => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1688splitSuffixBy predicate f (Stream step state) = Stream (step' f) (Just state)
1692 {-# INLINE_LATE step' #-}
1693 step' (Fold fstep initial done) gst (Just st) = do
1694 res <- step (adaptState gst) st
1699 then done acc >>= \val -> return $ Yield val (Just s)
1704 Skip s -> return $ Skip $ Just s
1710 res <- step (adaptState gst) stt
1714 then done acc >>= \r -> return $ Yield r (Just s)
1718 Skip s -> go SPEC s acc
1719 Stop -> done acc >>= \r -> return $ Yield r Nothing
1721 step' _ _ Nothing = return Stop
1723{-# INLINE_NORMAL wordsBy #-}
1724wordsBy :: Monad m => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1725wordsBy predicate f (Stream step state) = Stream (stepOuter f) (Just state)
1729 {-# INLINE_LATE stepOuter #-}
1730 stepOuter (Fold fstep initial done) gst (Just st) = do
1731 res <- step (adaptState gst) st
1735 then return $ Skip (Just s)
1741 Skip s -> return $ Skip $ Just s
1747 res <- step (adaptState gst) stt
1751 then done acc >>= \r -> return $ Yield r (Just s)
1755 Skip s -> go SPEC s acc
1756 Stop -> done acc >>= \r -> return $ Yield r Nothing
1758 stepOuter _ _ Nothing = return Stop
1760-- String search algorithms:
1761-- http://www-igm.univ-mlv.fr/~lecroq/string/index.html
1764-- TODO can we unify the splitting operations using a splitting configuration
1765-- like in the split package.
1767data SplitStyle = Infix | Suffix | Prefix deriving (Eq, Show)
1769data SplitOptions = SplitOptions
1770 { style :: SplitStyle
1771 , withSep :: Bool -- ^ keep the separators in output
1772 -- , compact :: Bool -- ^ treat multiple consecutive separators as one
1773 -- , trimHead :: Bool -- ^ drop blank at head
1774 -- , trimTail :: Bool -- ^ drop blank at tail
1778data SplitOnState s a =
1783 | GO_KARP_RABIN s !(RB.Ring a) !(Ptr a)
1786{-# INLINE_NORMAL splitOn #-}
1788 :: forall m a b. (MonadIO m, Storable a, Enum a, Eq a)
1793splitOn patArr (Fold fstep initial done) (Stream step state) =
1794 Stream stepOuter GO_START
1798 patLen = A.length patArr
1799 maxIndex = patLen - 1
1800 elemBits = sizeOf (undefined :: a) * 8
1802 {-# INLINE_LATE stepOuter #-}
1803 stepOuter _ GO_START =
1805 then return $ Skip $ GO_EMPTY_PAT state
1808 r <- liftIO $ (A.unsafeIndexIO patArr 0)
1809 return $ Skip $ GO_SINGLE_PAT state r
1810 else if sizeOf (undefined :: a) * patLen
1811 <= sizeOf (undefined :: Word)
1812 then return $ Skip $ GO_SHORT_PAT state
1814 (rb, rhead) <- liftIO $ RB.new patLen
1815 return $ Skip $ GO_KARP_RABIN state rb rhead
1817 stepOuter gst (GO_SINGLE_PAT stt pat) = initial >>= go SPEC stt
1822 res <- step (adaptState gst) st
1828 return $ Yield r (GO_SINGLE_PAT s pat)
1829 else fstep acc x >>= go SPEC s
1830 Skip s -> go SPEC s acc
1831 Stop -> done acc >>= \r -> return $ Yield r GO_DONE
1833 stepOuter gst (GO_SHORT_PAT stt) = initial >>= go0 SPEC 0 (0 :: Word) stt
1838 mask = (1 `shiftL` (elemBits * patLen)) - 1
1840 addToWord wrd a = (wrd `shiftL` elemBits) .|. fromIntegral (fromEnum a)
1843 patWord = mask .&. A.foldl' addToWord 0 patArr
1845 go0 !_ !idx wrd st !acc = do
1846 res <- step (adaptState gst) st
1849 let wrd' = addToWord wrd x
1852 if wrd' .&. mask == patWord
1855 return $ Yield r (GO_SHORT_PAT s)
1856 else go1 SPEC wrd' s acc
1857 else go0 SPEC (idx + 1) wrd' s acc
1858 Skip s -> go0 SPEC idx wrd s acc
1861 then go2 wrd idx acc
1863 done acc' >>= \r -> return $ Yield r GO_DONE
1866 go1 !_ wrd st !acc = do
1867 res <- step (adaptState gst) st
1870 let wrd' = addToWord wrd x
1871 old = (mask .&. wrd) `shiftR` (elemBits * (patLen - 1))
1872 acc' <- fstep acc (toEnum $ fromIntegral old)
1873 if wrd' .&. mask == patWord
1874 then done acc' >>= \r -> return $ Yield r (GO_SHORT_PAT s)
1875 else go1 SPEC wrd' s acc'
1876 Skip s -> go1 SPEC wrd s acc
1878 acc' <- go2 wrd patLen acc
1879 done acc' >>= \r -> return $ Yield r GO_DONE
1881 go2 !wrd !n !acc | n > 0 = do
1882 let old = (mask .&. wrd) `shiftR` (elemBits * (n - 1))
1883 fstep acc (toEnum $ fromIntegral old) >>= go2 wrd (n - 1)
1884 go2 _ _ acc = return acc
1886 stepOuter gst (GO_KARP_RABIN stt rb rhead) = do
1887 initial >>= go0 SPEC 0 rhead stt
1891 k = 2891336453 :: Word32
1893 addCksum cksum a = cksum * k + fromIntegral (fromEnum a)
1894 deltaCksum cksum old new =
1895 addCksum cksum new - coeff * fromIntegral (fromEnum old)
1897 -- XXX shall we use a random starting hash or 1 instead of 0?
1898 patHash = A.foldl' addCksum 0 patArr
1901 go0 !_ !idx !rh st !acc = do
1902 res <- step (adaptState gst) st
1905 rh' <- liftIO $ RB.unsafeInsert rb rh x
1908 let fold = RB.unsafeFoldRing (RB.ringBound rb)
1909 let !ringHash = fold addCksum 0 rb
1910 if ringHash == patHash
1911 then go2 SPEC ringHash rh' s acc
1912 else go1 SPEC ringHash rh' s acc
1913 else go0 SPEC (idx + 1) rh' s acc
1914 Skip s -> go0 SPEC idx rh s acc
1916 !acc' <- if idx /= 0
1917 then RB.unsafeFoldRingM rh fstep acc rb
1919 done acc' >>= \r -> return $ Yield r GO_DONE
1921 -- XXX Theoretically this code can do 4 times faster if GHC generates
1922 -- optimal code. If we use just "(cksum' == patHash)" condition it goes
1923 -- 4x faster, as soon as we add the "RB.unsafeEqArray rb v" condition
1924 -- the generated code changes drastically and becomes 4x slower. Need
1925 -- to investigate what is going on with GHC.
1927 go1 !_ !cksum !rh st !acc = do
1928 res <- step (adaptState gst) st
1931 old <- liftIO $ peek rh
1932 let cksum' = deltaCksum cksum old x
1933 acc' <- fstep acc old
1935 if (cksum' == patHash)
1937 rh' <- liftIO (RB.unsafeInsert rb rh x)
1938 go2 SPEC cksum' rh' s acc'
1940 rh' <- liftIO (RB.unsafeInsert rb rh x)
1941 go1 SPEC cksum' rh' s acc'
1942 Skip s -> go1 SPEC cksum rh s acc
1944 acc' <- RB.unsafeFoldRingFullM rh fstep acc rb
1945 done acc' >>= \r -> return $ Yield r GO_DONE
1947 go2 !_ !cksum' !rh' s !acc' = do
1948 if RB.unsafeEqArray rb rh' patArr
1951 return $ Yield r (GO_KARP_RABIN s rb rhead)
1952 else go1 SPEC cksum' rh' s acc'
1954 stepOuter gst (GO_EMPTY_PAT st) = do
1955 res <- step (adaptState gst) st
1960 done acc' >>= \r -> return $ Yield r (GO_EMPTY_PAT s)
1961 Skip s -> return $ Skip (GO_EMPTY_PAT s)
1964 stepOuter _ GO_DONE = return Stop
1966{-# INLINE_NORMAL splitSuffixOn #-}
1968 :: forall m a b. (MonadIO m, Storable a, Enum a, Eq a)
1974splitSuffixOn withSep patArr (Fold fstep initial done)
1975 (Stream step state) =
1976 Stream stepOuter GO_START
1980 patLen = A.length patArr
1981 maxIndex = patLen - 1
1982 elemBits = sizeOf (undefined :: a) * 8
1984 {-# INLINE_LATE stepOuter #-}
1985 stepOuter _ GO_START =
1987 then return $ Skip $ GO_EMPTY_PAT state
1990 r <- liftIO $ (A.unsafeIndexIO patArr 0)
1991 return $ Skip $ GO_SINGLE_PAT state r
1992 else if sizeOf (undefined :: a) * patLen
1993 <= sizeOf (undefined :: Word)
1994 then return $ Skip $ GO_SHORT_PAT state
1996 (rb, rhead) <- liftIO $ RB.new patLen
1997 return $ Skip $ GO_KARP_RABIN state rb rhead
1999 stepOuter gst (GO_SINGLE_PAT stt pat) = do
2000 -- This first part is the only difference between splitOn and
2002 -- If the last element is a separator do not issue a blank segment.
2003 res <- step (adaptState gst) stt
2009 acc' <- if withSep then fstep acc x else return acc
2010 done acc' >>= \r -> return $ Yield r (GO_SINGLE_PAT s pat)
2011 else fstep acc x >>= go SPEC s
2012 Skip s -> return $ Skip $ (GO_SINGLE_PAT s pat)
2017 -- This is identical for splitOn and splitSuffixOn
2019 res <- step (adaptState gst) st
2024 acc' <- if withSep then fstep acc x else return acc
2026 return $ Yield r (GO_SINGLE_PAT s pat)
2027 else fstep acc x >>= go SPEC s
2028 Skip s -> go SPEC s acc
2029 Stop -> done acc >>= \r -> return $ Yield r GO_DONE
2031 stepOuter gst (GO_SHORT_PAT stt) = do
2033 -- Call "initial" only if the stream yields an element, otherwise we
2034 -- may call "initial" but never yield anything. initial may produce a
2035 -- side effect, therefore we will end up doing and discard a side
2040 res <- step (adaptState gst) stt
2044 let wrd' = addToWord wrd x
2045 acc' <- if withSep then fstep acc x else return acc
2048 if wrd' .&. mask == patWord
2049 then done acc' >>= \r -> return $ Yield r (GO_SHORT_PAT s)
2050 else go0 SPEC (idx + 1) wrd' s acc'
2051 else go0 SPEC (idx + 1) wrd' s acc'
2052 Skip s -> return $ Skip (GO_SHORT_PAT s)
2058 mask = (1 `shiftL` (elemBits * patLen)) - 1
2060 addToWord wrd a = (wrd `shiftL` elemBits) .|. fromIntegral (fromEnum a)
2063 patWord = mask .&. A.foldl' addToWord 0 patArr
2065 go0 !_ !idx wrd st !acc = do
2066 res <- step (adaptState gst) st
2069 let wrd' = addToWord wrd x
2070 acc' <- if withSep then fstep acc x else return acc
2073 if wrd' .&. mask == patWord
2076 return $ Yield r (GO_SHORT_PAT s)
2077 else go1 SPEC wrd' s acc'
2078 else go0 SPEC (idx + 1) wrd' s acc'
2079 Skip s -> go0 SPEC idx wrd s acc
2081 if (idx == maxIndex) && (wrd .&. mask == patWord)
2084 acc' <- if idx /= 0 && not withSep
2085 then go2 wrd idx acc
2087 done acc' >>= \r -> return $ Yield r GO_DONE
2090 go1 !_ wrd st !acc = do
2091 res <- step (adaptState gst) st
2094 let wrd' = addToWord wrd x
2095 old = (mask .&. wrd) `shiftR` (elemBits * (patLen - 1))
2098 else fstep acc (toEnum $ fromIntegral old)
2099 if wrd' .&. mask == patWord
2100 then done acc' >>= \r -> return $ Yield r (GO_SHORT_PAT s)
2101 else go1 SPEC wrd' s acc'
2102 Skip s -> go1 SPEC wrd s acc
2104 -- If the last sequence is a separator do not issue a blank
2106 if wrd .&. mask == patWord
2111 else go2 wrd patLen acc
2112 done acc' >>= \r -> return $ Yield r GO_DONE
2114 go2 !wrd !n !acc | n > 0 = do
2115 let old = (mask .&. wrd) `shiftR` (elemBits * (n - 1))
2116 fstep acc (toEnum $ fromIntegral old) >>= go2 wrd (n - 1)
2117 go2 _ _ acc = return acc
2119 stepOuter gst (GO_KARP_RABIN stt rb rhead) = do
2121 res <- step (adaptState gst) stt
2125 acc' <- if withSep then fstep acc x else return acc
2126 rh' <- liftIO (RB.unsafeInsert rb rhead x)
2129 let fold = RB.unsafeFoldRing (RB.ringBound rb)
2130 let !ringHash = fold addCksum 0 rb
2131 if ringHash == patHash
2132 then go2 SPEC ringHash rh' s acc'
2133 else go0 SPEC (idx + 1) rh' s acc'
2134 else go0 SPEC (idx + 1) rh' s acc'
2135 Skip s -> return $ Skip (GO_KARP_RABIN s rb rhead)
2140 k = 2891336453 :: Word32
2142 addCksum cksum a = cksum * k + fromIntegral (fromEnum a)
2143 deltaCksum cksum old new =
2144 addCksum cksum new - coeff * fromIntegral (fromEnum old)
2146 -- XXX shall we use a random starting hash or 1 instead of 0?
2147 patHash = A.foldl' addCksum 0 patArr
2150 go0 !_ !idx !rh st !acc = do
2151 res <- step (adaptState gst) st
2154 acc' <- if withSep then fstep acc x else return acc
2155 rh' <- liftIO (RB.unsafeInsert rb rh x)
2158 let fold = RB.unsafeFoldRing (RB.ringBound rb)
2159 let !ringHash = fold addCksum 0 rb
2160 if ringHash == patHash
2161 then go2 SPEC ringHash rh' s acc'
2162 else go1 SPEC ringHash rh' s acc'
2163 else go0 SPEC (idx + 1) rh' s acc'
2164 Skip s -> go0 SPEC idx rh s acc
2166 -- do not issue a blank segment when we end at pattern
2167 if (idx == maxIndex) && RB.unsafeEqArray rb rh patArr
2170 !acc' <- if idx /= 0 && not withSep
2171 then RB.unsafeFoldRingM rh fstep acc rb
2173 done acc' >>= \r -> return $ Yield r GO_DONE
2175 -- XXX Theoretically this code can do 4 times faster if GHC generates
2176 -- optimal code. If we use just "(cksum' == patHash)" condition it goes
2177 -- 4x faster, as soon as we add the "RB.unsafeEqArray rb v" condition
2178 -- the generated code changes drastically and becomes 4x slower. Need
2179 -- to investigate what is going on with GHC.
2181 go1 !_ !cksum !rh st !acc = do
2182 res <- step (adaptState gst) st
2185 old <- liftIO $ peek rh
2186 let cksum' = deltaCksum cksum old x
2191 if (cksum' == patHash)
2193 rh' <- liftIO (RB.unsafeInsert rb rh x)
2194 go2 SPEC cksum' rh' s acc'
2196 rh' <- liftIO (RB.unsafeInsert rb rh x)
2197 go1 SPEC cksum' rh' s acc'
2198 Skip s -> go1 SPEC cksum rh s acc
2200 if RB.unsafeEqArray rb rh patArr
2205 else RB.unsafeFoldRingFullM rh fstep acc rb
2206 done acc' >>= \r -> return $ Yield r GO_DONE
2208 go2 !_ !cksum' !rh' s !acc' = do
2209 if RB.unsafeEqArray rb rh' patArr
2212 return $ Yield r (GO_KARP_RABIN s rb rhead)
2213 else go1 SPEC cksum' rh' s acc'
2215 stepOuter gst (GO_EMPTY_PAT st) = do
2216 res <- step (adaptState gst) st
2221 done acc' >>= \r -> return $ Yield r (GO_EMPTY_PAT s)
2222 Skip s -> return $ Skip (GO_EMPTY_PAT s)
2225 stepOuter _ GO_DONE = return Stop
2227data SplitState s arr
2229 | SplitBuffering s arr
2230 | SplitSplitting s arr
2231 | SplitYielding arr (SplitState s arr)
2234-- XXX An alternative approach would be to use a partial fold (Fold m a b) to
2235-- split using a splitBy like combinator. The Fold would consume upto the
2236-- separator and return any leftover which can then be fed to the next fold.
2238-- We can revisit this once we have partial folds/parsers.
2240-- | Performs infix separator style splitting.
2241{-# INLINE_NORMAL splitInnerBy #-}
2244 => (f a -> m (f a, Maybe (f a))) -- splitter
2245 -> (f a -> f a -> m (f a)) -- joiner
2248splitInnerBy splitter joiner (Stream step1 state1) =
2249 (Stream step (SplitInitial state1))
2253 {-# INLINE_LATE step #-}
2254 step gst (SplitInitial st) = do
2258 (x1, mx2) <- splitter x
2259 return $ case mx2 of
2260 Nothing -> Skip (SplitBuffering s x1)
2261 Just x2 -> Skip (SplitYielding x1 (SplitSplitting s x2))
2262 Skip s -> return $ Skip (SplitInitial s)
2263 Stop -> return $ Stop
2265 step gst (SplitBuffering st buf) = do
2269 (x1, mx2) <- splitter x
2270 buf' <- joiner buf x1
2271 return $ case mx2 of
2272 Nothing -> Skip (SplitBuffering s buf')
2273 Just x2 -> Skip (SplitYielding buf' (SplitSplitting s x2))
2274 Skip s -> return $ Skip (SplitBuffering s buf)
2275 Stop -> return $ Skip (SplitYielding buf SplitFinishing)
2277 step _ (SplitSplitting st buf) = do
2278 (x1, mx2) <- splitter buf
2279 return $ case mx2 of
2280 Nothing -> Skip $ SplitBuffering st x1
2281 Just x2 -> Skip $ SplitYielding x1 (SplitSplitting st x2)
2283 step _ (SplitYielding x next) = return $ Yield x next
2284 step _ SplitFinishing = return $ Stop
2286-- | Performs infix separator style splitting.
2287{-# INLINE_NORMAL splitInnerBySuffix #-}
2289 :: (Monad m, Eq (f a), Monoid (f a))
2290 => (f a -> m (f a, Maybe (f a))) -- splitter
2291 -> (f a -> f a -> m (f a)) -- joiner
2294splitInnerBySuffix splitter joiner (Stream step1 state1) =
2295 (Stream step (SplitInitial state1))
2299 {-# INLINE_LATE step #-}
2300 step gst (SplitInitial st) = do
2304 (x1, mx2) <- splitter x
2305 return $ case mx2 of
2306 Nothing -> Skip (SplitBuffering s x1)
2307 Just x2 -> Skip (SplitYielding x1 (SplitSplitting s x2))
2308 Skip s -> return $ Skip (SplitInitial s)
2309 Stop -> return $ Stop
2311 step gst (SplitBuffering st buf) = do
2315 (x1, mx2) <- splitter x
2316 buf' <- joiner buf x1
2317 return $ case mx2 of
2318 Nothing -> Skip (SplitBuffering s buf')
2319 Just x2 -> Skip (SplitYielding buf' (SplitSplitting s x2))
2320 Skip s -> return $ Skip (SplitBuffering s buf)
2324 else Skip (SplitYielding buf SplitFinishing)
2326 step _ (SplitSplitting st buf) = do
2327 (x1, mx2) <- splitter buf
2328 return $ case mx2 of
2329 Nothing -> Skip $ SplitBuffering st x1
2330 Just x2 -> Skip $ SplitYielding x1 (SplitSplitting st x2)
2332 step _ (SplitYielding x next) = return $ Yield x next
2333 step _ SplitFinishing = return $ Stop
2335------------------------------------------------------------------------------
2337------------------------------------------------------------------------------
2339{-# INLINE_NORMAL isPrefixOf #-}
2340isPrefixOf :: (Eq a, Monad m) => Stream m a -> Stream m a -> m Bool
2341isPrefixOf (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
2343 go (sa, sb, Nothing) = do
2344 r <- stepa defState sa
2346 Yield x sa' -> go (sa', sb, Just x)
2347 Skip sa' -> go (sa', sb, Nothing)
2350 go (sa, sb, Just x) = do
2351 r <- stepb defState sb
2355 then go (sa, sb', Nothing)
2357 Skip sb' -> go (sa, sb', Just x)
2358 Stop -> return False
2360{-# INLINE_NORMAL isSubsequenceOf #-}
2361isSubsequenceOf :: (Eq a, Monad m) => Stream m a -> Stream m a -> m Bool
2362isSubsequenceOf (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
2364 go (sa, sb, Nothing) = do
2365 r <- stepa defState sa
2367 Yield x sa' -> go (sa', sb, Just x)
2368 Skip sa' -> go (sa', sb, Nothing)
2371 go (sa, sb, Just x) = do
2372 r <- stepb defState sb
2376 then go (sa, sb', Nothing)
2377 else go (sa, sb', Just x)
2378 Skip sb' -> go (sa, sb', Just x)
2379 Stop -> return False
2381{-# INLINE_NORMAL stripPrefix #-}
2384 => Stream m a -> Stream m a -> m (Maybe (Stream m a))
2385stripPrefix (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
2387 go (sa, sb, Nothing) = do
2388 r <- stepa defState sa
2390 Yield x sa' -> go (sa', sb, Just x)
2391 Skip sa' -> go (sa', sb, Nothing)
2392 Stop -> return $ Just (Stream stepb sb)
2394 go (sa, sb, Just x) = do
2395 r <- stepb defState sb
2399 then go (sa, sb', Nothing)
2401 Skip sb' -> go (sa, sb', Just x)
2402 Stop -> return Nothing
2404------------------------------------------------------------------------------
2406------------------------------------------------------------------------------
2408-- | Execute a monadic action for each element of the 'Stream'
2409{-# INLINE_NORMAL mapM_ #-}
2410mapM_ :: Monad m => (a -> m b) -> Stream m a -> m ()
2411mapM_ m = drain . mapM m
2413-------------------------------------------------------------------------------
2414-- Stream transformations using Unfolds
2415-------------------------------------------------------------------------------
2417-- Define a unique structure to use in inspection testing
2418data ConcatMapUState o i =
2420 | ConcatMapUInner o i
2422-- | @concatMapU unfold stream@ uses @unfold@ to map the input stream elements
2423-- to streams and then flattens the generated streams into a single output
2426-- This is like 'concatMap' but uses an unfold with an explicit state to
2427-- generate the stream instead of a 'Stream' type generator. This allows better
2428-- optimization via fusion. This can be many times more efficient than
2431{-# INLINE_NORMAL concatMapU #-}
2432concatMapU :: Monad m => Unfold m a b -> Stream m a -> Stream m b
2433concatMapU (Unfold istep inject) (Stream ostep ost) =
2434 Stream step (ConcatMapUOuter ost)
2436 {-# INLINE_LATE step #-}
2437 step gst (ConcatMapUOuter o) = do
2438 r <- ostep (adaptState gst) o
2442 i `seq` return (Skip (ConcatMapUInner o' i))
2443 Skip o' -> return $ Skip (ConcatMapUOuter o')
2444 Stop -> return $ Stop
2446 step _ (ConcatMapUInner o i) = do
2449 Yield x i' -> Yield x (ConcatMapUInner o i')
2450 Skip i' -> Skip (ConcatMapUInner o i')
2451 Stop -> Skip (ConcatMapUOuter o)
2453data ConcatUnfoldInterleaveState o i =
2454 ConcatUnfoldInterleaveOuter o [i]
2455 | ConcatUnfoldInterleaveInner o [i]
2456 | ConcatUnfoldInterleaveInnerL [i] [i]
2457 | ConcatUnfoldInterleaveInnerR [i] [i]
2459-- XXX use arrays to store state instead of lists.
2460-- XXX In general we can use different scheduling strategies e.g. how to
2461-- schedule the outer vs inner loop or assigning weights to different streams
2462-- or outer and inner loops.
2464-- After a yield, switch to the next stream. Do not switch streams on Skip.
2465-- Yield from outer stream switches to the inner stream.
2467-- There are two choices here, (1) exhaust the outer stream first and then
2468-- start yielding from the inner streams, this is much simpler to implement,
2469-- (2) yield at least one element from an inner stream before going back to
2470-- outer stream and opening the next stream from it.
2472-- Ideally, we need some scheduling bias to inner streams vs outer stream.
2473-- Maybe we can configure the behavior.
2475{-# INLINE_NORMAL concatUnfoldInterleave #-}
2476concatUnfoldInterleave :: Monad m => Unfold m a b -> Stream m a -> Stream m b
2477concatUnfoldInterleave (Unfold istep inject) (Stream ostep ost) =
2478 Stream step (ConcatUnfoldInterleaveOuter ost [])
2480 {-# INLINE_LATE step #-}
2481 step gst (ConcatUnfoldInterleaveOuter o ls) = do
2482 r <- ostep (adaptState gst) o
2486 i `seq` return (Skip (ConcatUnfoldInterleaveInner o' (i : ls)))
2487 Skip o' -> return $ Skip (ConcatUnfoldInterleaveOuter o' ls)
2488 Stop -> return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2490 step _ (ConcatUnfoldInterleaveInner _ []) = undefined
2491 step _ (ConcatUnfoldInterleaveInner o (st:ls)) = do
2494 Yield x s -> Yield x (ConcatUnfoldInterleaveOuter o (s:ls))
2495 Skip s -> Skip (ConcatUnfoldInterleaveInner o (s:ls))
2496 Stop -> Skip (ConcatUnfoldInterleaveOuter o ls)
2498 step _ (ConcatUnfoldInterleaveInnerL [] []) = return Stop
2499 step _ (ConcatUnfoldInterleaveInnerL [] rs) =
2500 return $ Skip (ConcatUnfoldInterleaveInnerR [] rs)
2502 step _ (ConcatUnfoldInterleaveInnerL (st:ls) rs) = do
2505 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerL ls (s:rs))
2506 Skip s -> Skip (ConcatUnfoldInterleaveInnerL (s:ls) rs)
2507 Stop -> Skip (ConcatUnfoldInterleaveInnerL ls rs)
2509 step _ (ConcatUnfoldInterleaveInnerR [] []) = return Stop
2510 step _ (ConcatUnfoldInterleaveInnerR ls []) =
2511 return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2513 step _ (ConcatUnfoldInterleaveInnerR ls (st:rs)) = do
2516 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerR (s:ls) rs)
2517 Skip s -> Skip (ConcatUnfoldInterleaveInnerR ls (s:rs))
2518 Stop -> Skip (ConcatUnfoldInterleaveInnerR ls rs)
2520-- XXX In general we can use different scheduling strategies e.g. how to
2521-- schedule the outer vs inner loop or assigning weights to different streams
2522-- or outer and inner loops.
2524-- This could be inefficient if the tasks are too small.
2526-- Compared to concatUnfoldInterleave this one switches streams on Skips.
2528{-# INLINE_NORMAL concatUnfoldRoundrobin #-}
2529concatUnfoldRoundrobin :: Monad m => Unfold m a b -> Stream m a -> Stream m b
2530concatUnfoldRoundrobin (Unfold istep inject) (Stream ostep ost) =
2531 Stream step (ConcatUnfoldInterleaveOuter ost [])
2533 {-# INLINE_LATE step #-}
2534 step gst (ConcatUnfoldInterleaveOuter o ls) = do
2535 r <- ostep (adaptState gst) o
2539 i `seq` return (Skip (ConcatUnfoldInterleaveInner o' (i : ls)))
2540 Skip o' -> return $ Skip (ConcatUnfoldInterleaveInner o' ls)
2541 Stop -> return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2543 step _ (ConcatUnfoldInterleaveInner o []) =
2544 return $ Skip (ConcatUnfoldInterleaveOuter o [])
2546 step _ (ConcatUnfoldInterleaveInner o (st:ls)) = do
2549 Yield x s -> Yield x (ConcatUnfoldInterleaveOuter o (s:ls))
2550 Skip s -> Skip (ConcatUnfoldInterleaveOuter o (s:ls))
2551 Stop -> Skip (ConcatUnfoldInterleaveOuter o ls)
2553 step _ (ConcatUnfoldInterleaveInnerL [] []) = return Stop
2554 step _ (ConcatUnfoldInterleaveInnerL [] rs) =
2555 return $ Skip (ConcatUnfoldInterleaveInnerR [] rs)
2557 step _ (ConcatUnfoldInterleaveInnerL (st:ls) rs) = do
2560 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerL ls (s:rs))
2561 Skip s -> Skip (ConcatUnfoldInterleaveInnerL ls (s:rs))
2562 Stop -> Skip (ConcatUnfoldInterleaveInnerL ls rs)
2564 step _ (ConcatUnfoldInterleaveInnerR [] []) = return Stop
2565 step _ (ConcatUnfoldInterleaveInnerR ls []) =
2566 return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2568 step _ (ConcatUnfoldInterleaveInnerR ls (st:rs)) = do
2571 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerR (s:ls) rs)
2572 Skip s -> Skip (ConcatUnfoldInterleaveInnerR (s:ls) rs)
2573 Stop -> Skip (ConcatUnfoldInterleaveInnerR ls rs)
2575data AppendState s1 s2 = AppendFirst s1 | AppendSecond s2
2577-- Note that this could be much faster compared to the CPS stream. However, as
2578-- the number of streams being composed increases this may become expensive.
2579-- Need to see where the breaking point is between the two.
2581{-# INLINE_NORMAL append #-}
2582append :: Monad m => Stream m a -> Stream m a -> Stream m a
2583append (Stream step1 state1) (Stream step2 state2) =
2584 Stream step (AppendFirst state1)
2588 {-# INLINE_LATE step #-}
2589 step gst (AppendFirst st) = do
2592 Yield a s -> Yield a (AppendFirst s)
2593 Skip s -> Skip (AppendFirst s)
2594 Stop -> Skip (AppendSecond state2)
2596 step gst (AppendSecond st) = do
2599 Yield a s -> Yield a (AppendSecond s)
2600 Skip s -> Skip (AppendSecond s)
2603data InterleaveState s1 s2 = InterleaveFirst s1 s2 | InterleaveSecond s1 s2
2604 | InterleaveSecondOnly s2 | InterleaveFirstOnly s1
2606{-# INLINE_NORMAL interleave #-}
2607interleave :: Monad m => Stream m a -> Stream m a -> Stream m a
2608interleave (Stream step1 state1) (Stream step2 state2) =
2609 Stream step (InterleaveFirst state1 state2)
2613 {-# INLINE_LATE step #-}
2614 step gst (InterleaveFirst st1 st2) = do
2617 Yield a s -> Yield a (InterleaveSecond s st2)
2618 Skip s -> Skip (InterleaveFirst s st2)
2619 Stop -> Skip (InterleaveSecondOnly st2)
2621 step gst (InterleaveSecond st1 st2) = do
2624 Yield a s -> Yield a (InterleaveFirst st1 s)
2625 Skip s -> Skip (InterleaveSecond st1 s)
2626 Stop -> Skip (InterleaveFirstOnly st1)
2628 step gst (InterleaveFirstOnly st1) = do
2631 Yield a s -> Yield a (InterleaveFirstOnly s)
2632 Skip s -> Skip (InterleaveFirstOnly s)
2635 step gst (InterleaveSecondOnly st2) = do
2638 Yield a s -> Yield a (InterleaveSecondOnly s)
2639 Skip s -> Skip (InterleaveSecondOnly s)
2642{-# INLINE_NORMAL interleaveMin #-}
2643interleaveMin :: Monad m => Stream m a -> Stream m a -> Stream m a
2644interleaveMin (Stream step1 state1) (Stream step2 state2) =
2645 Stream step (InterleaveFirst state1 state2)
2649 {-# INLINE_LATE step #-}
2650 step gst (InterleaveFirst st1 st2) = do
2653 Yield a s -> Yield a (InterleaveSecond s st2)
2654 Skip s -> Skip (InterleaveFirst s st2)
2657 step gst (InterleaveSecond st1 st2) = do
2660 Yield a s -> Yield a (InterleaveFirst st1 s)
2661 Skip s -> Skip (InterleaveSecond st1 s)
2664 step _ (InterleaveFirstOnly _) = undefined
2665 step _ (InterleaveSecondOnly _) = undefined
2667{-# INLINE_NORMAL interleaveSuffix #-}
2668interleaveSuffix :: Monad m => Stream m a -> Stream m a -> Stream m a
2669interleaveSuffix (Stream step1 state1) (Stream step2 state2) =
2670 Stream step (InterleaveFirst state1 state2)
2674 {-# INLINE_LATE step #-}
2675 step gst (InterleaveFirst st1 st2) = do
2678 Yield a s -> Yield a (InterleaveSecond s st2)
2679 Skip s -> Skip (InterleaveFirst s st2)
2682 step gst (InterleaveSecond st1 st2) = do
2685 Yield a s -> Yield a (InterleaveFirst st1 s)
2686 Skip s -> Skip (InterleaveSecond st1 s)
2687 Stop -> Skip (InterleaveFirstOnly st1)
2689 step gst (InterleaveFirstOnly st1) = do
2692 Yield a s -> Yield a (InterleaveFirstOnly s)
2693 Skip s -> Skip (InterleaveFirstOnly s)
2696 step _ (InterleaveSecondOnly _) = undefined
2698data InterleaveInfixState s1 s2 a
2699 = InterleaveInfixFirst s1 s2
2700 | InterleaveInfixSecondBuf s1 s2
2701 | InterleaveInfixSecondYield s1 s2 a
2702 | InterleaveInfixFirstYield s1 s2 a
2703 | InterleaveInfixFirstOnly s1
2705{-# INLINE_NORMAL interleaveInfix #-}
2706interleaveInfix :: Monad m => Stream m a -> Stream m a -> Stream m a
2707interleaveInfix (Stream step1 state1) (Stream step2 state2) =
2708 Stream step (InterleaveInfixFirst state1 state2)
2712 {-# INLINE_LATE step #-}
2713 step gst (InterleaveInfixFirst st1 st2) = do
2716 Yield a s -> Yield a (InterleaveInfixSecondBuf s st2)
2717 Skip s -> Skip (InterleaveInfixFirst s st2)
2720 step gst (InterleaveInfixSecondBuf st1 st2) = do
2723 Yield a s -> Skip (InterleaveInfixSecondYield st1 s a)
2724 Skip s -> Skip (InterleaveInfixSecondBuf st1 s)
2725 Stop -> Skip (InterleaveInfixFirstOnly st1)
2727 step gst (InterleaveInfixSecondYield st1 st2 x) = do
2730 Yield a s -> Yield x (InterleaveInfixFirstYield s st2 a)
2731 Skip s -> Skip (InterleaveInfixSecondYield s st2 x)
2734 step _ (InterleaveInfixFirstYield st1 st2 x) = do
2735 return $ Yield x (InterleaveInfixSecondBuf st1 st2)
2737 step gst (InterleaveInfixFirstOnly st1) = do
2740 Yield a s -> Yield a (InterleaveInfixFirstOnly s)
2741 Skip s -> Skip (InterleaveInfixFirstOnly s)
2744{-# INLINE_NORMAL roundRobin #-}
2745roundRobin :: Monad m => Stream m a -> Stream m a -> Stream m a
2746roundRobin (Stream step1 state1) (Stream step2 state2) =
2747 Stream step (InterleaveFirst state1 state2)
2751 {-# INLINE_LATE step #-}
2752 step gst (InterleaveFirst st1 st2) = do
2755 Yield a s -> Yield a (InterleaveSecond s st2)
2756 Skip s -> Skip (InterleaveSecond s st2)
2757 Stop -> Skip (InterleaveSecondOnly st2)
2759 step gst (InterleaveSecond st1 st2) = do
2762 Yield a s -> Yield a (InterleaveFirst st1 s)
2763 Skip s -> Skip (InterleaveFirst st1 s)
2764 Stop -> Skip (InterleaveFirstOnly st1)
2766 step gst (InterleaveSecondOnly st2) = do
2769 Yield a s -> Yield a (InterleaveSecondOnly s)
2770 Skip s -> Skip (InterleaveSecondOnly s)
2773 step gst (InterleaveFirstOnly st1) = do
2776 Yield a s -> Yield a (InterleaveFirstOnly s)
2777 Skip s -> Skip (InterleaveFirstOnly s)
2780data ICUState s1 s2 i1 i2 =
2785 | ICUFirstInner s1 s2 i1
2786 | ICUSecondInner s1 s2 i2
2787 | ICUFirstOnlyInner s1 i1
2788 | ICUSecondOnlyInner s2 i2
2790-- | Interleave streams (full streams, not the elements) unfolded from two
2791-- input streams and concat. Stop when the first stream stops. If the second
2792-- stream ends before the first one then first stream still keeps running alone
2793-- without any interleaving with the second stream.
2795-- [a1, a2, ... an] [b1, b2 ...]
2796-- => [streamA1, streamA2, ... streamAn] [streamB1, streamB2, ...]
2797-- => [streamA1, streamB1, streamA2...StreamAn, streamBn]
2798-- => [a11, a12, ...a1j, b11, b12, ...b1k, a21, a22, ...]
2800{-# INLINE_NORMAL gintercalateSuffix #-}
2803 => Unfold m a c -> Stream m a -> Unfold m b c -> Stream m b -> Stream m c
2805 (Unfold istep1 inject1) (Stream step1 state1)
2806 (Unfold istep2 inject2) (Stream step2 state2) =
2807 Stream step (ICUFirst state1 state2)
2811 {-# INLINE_LATE step #-}
2812 step gst (ICUFirst s1 s2) = do
2813 r <- step1 (adaptState gst) s1
2817 i `seq` return (Skip (ICUFirstInner s s2 i))
2818 Skip s -> return $ Skip (ICUFirst s s2)
2821 step gst (ICUFirstOnly s1) = do
2822 r <- step1 (adaptState gst) s1
2826 i `seq` return (Skip (ICUFirstOnlyInner s i))
2827 Skip s -> return $ Skip (ICUFirstOnly s)
2830 step _ (ICUFirstInner s1 s2 i1) = do
2833 Yield x i' -> Yield x (ICUFirstInner s1 s2 i')
2834 Skip i' -> Skip (ICUFirstInner s1 s2 i')
2835 Stop -> Skip (ICUSecond s1 s2)
2837 step _ (ICUFirstOnlyInner s1 i1) = do
2840 Yield x i' -> Yield x (ICUFirstOnlyInner s1 i')
2841 Skip i' -> Skip (ICUFirstOnlyInner s1 i')
2842 Stop -> Skip (ICUFirstOnly s1)
2844 step gst (ICUSecond s1 s2) = do
2845 r <- step2 (adaptState gst) s2
2849 i `seq` return (Skip (ICUSecondInner s1 s i))
2850 Skip s -> return $ Skip (ICUSecond s1 s)
2851 Stop -> return $ Skip (ICUFirstOnly s1)
2853 step _ (ICUSecondInner s1 s2 i2) = do
2856 Yield x i' -> Yield x (ICUSecondInner s1 s2 i')
2857 Skip i' -> Skip (ICUSecondInner s1 s2 i')
2858 Stop -> Skip (ICUFirst s1 s2)
2860 step _ (ICUSecondOnly _s2) = undefined
2861 step _ (ICUSecondOnlyInner _s2 _i2) = undefined
2863data InterposeSuffixState s1 i1 =
2864 InterposeSuffixFirst s1
2865 -- | InterposeSuffixFirstYield s1 i1
2866 | InterposeSuffixFirstInner s1 i1
2867 | InterposeSuffixSecond s1
2869-- Note that if an unfolded layer turns out to be nil we still emit the
2870-- separator effect. An alternate behavior could be to emit the separator
2871-- effect only if at least one element has been yielded by the unfolding.
2872-- However, that becomes a bit complicated, so we have chosen the former
2874{-# INLINE_NORMAL interposeSuffix #-}
2877 => m c -> Unfold m b c -> Stream m b -> Stream m c
2880 (Unfold istep1 inject1) (Stream step1 state1) =
2881 Stream step (InterposeSuffixFirst state1)
2885 {-# INLINE_LATE step #-}
2886 step gst (InterposeSuffixFirst s1) = do
2887 r <- step1 (adaptState gst) s1
2891 i `seq` return (Skip (InterposeSuffixFirstInner s i))
2892 -- i `seq` return (Skip (InterposeSuffixFirstYield s i))
2893 Skip s -> return $ Skip (InterposeSuffixFirst s)
2897 step _ (InterposeSuffixFirstYield s1 i1) = do
2900 Yield x i' -> Yield x (InterposeSuffixFirstInner s1 i')
2901 Skip i' -> Skip (InterposeSuffixFirstYield s1 i')
2902 Stop -> Skip (InterposeSuffixFirst s1)
2905 step _ (InterposeSuffixFirstInner s1 i1) = do
2908 Yield x i' -> Yield x (InterposeSuffixFirstInner s1 i')
2909 Skip i' -> Skip (InterposeSuffixFirstInner s1 i')
2910 Stop -> Skip (InterposeSuffixSecond s1)
2912 step _ (InterposeSuffixSecond s1) = do
2914 return $ Yield r (InterposeSuffixFirst s1)
2916data ICALState s1 s2 i1 i2 a =
2918 -- | ICALFirstYield s1 s2 i1
2919 | ICALFirstInner s1 s2 i1
2921 | ICALFirstOnlyInner s1 i1
2922 | ICALSecondInject s1 s2
2923 | ICALFirstInject s1 s2 i2
2924 -- | ICALFirstBuf s1 s2 i1 i2
2925 | ICALSecondInner s1 s2 i1 i2
2926 -- -- | ICALSecondInner s1 s2 i1 i2 a
2927 -- -- | ICALFirstResume s1 s2 i1 i2 a
2929-- | Interleave streams (full streams, not the elements) unfolded from two
2930-- input streams and concat. Stop when the first stream stops. If the second
2931-- stream ends before the first one then first stream still keeps running alone
2932-- without any interleaving with the second stream.
2934-- [a1, a2, ... an] [b1, b2 ...]
2935-- => [streamA1, streamA2, ... streamAn] [streamB1, streamB2, ...]
2936-- => [streamA1, streamB1, streamA2...StreamAn, streamBn]
2937-- => [a11, a12, ...a1j, b11, b12, ...b1k, a21, a22, ...]
2939{-# INLINE_NORMAL gintercalate #-}
2942 => Unfold m a c -> Stream m a -> Unfold m b c -> Stream m b -> Stream m c
2944 (Unfold istep1 inject1) (Stream step1 state1)
2945 (Unfold istep2 inject2) (Stream step2 state2) =
2946 Stream step (ICALFirst state1 state2)
2950 {-# INLINE_LATE step #-}
2951 step gst (ICALFirst s1 s2) = do
2952 r <- step1 (adaptState gst) s1
2956 i `seq` return (Skip (ICALFirstInner s s2 i))
2957 -- i `seq` return (Skip (ICALFirstYield s s2 i))
2958 Skip s -> return $ Skip (ICALFirst s s2)
2962 step _ (ICALFirstYield s1 s2 i1) = do
2965 Yield x i' -> Yield x (ICALFirstInner s1 s2 i')
2966 Skip i' -> Skip (ICALFirstYield s1 s2 i')
2967 Stop -> Skip (ICALFirst s1 s2)
2970 step _ (ICALFirstInner s1 s2 i1) = do
2973 Yield x i' -> Yield x (ICALFirstInner s1 s2 i')
2974 Skip i' -> Skip (ICALFirstInner s1 s2 i')
2975 Stop -> Skip (ICALSecondInject s1 s2)
2977 step gst (ICALFirstOnly s1) = do
2978 r <- step1 (adaptState gst) s1
2982 i `seq` return (Skip (ICALFirstOnlyInner s i))
2983 Skip s -> return $ Skip (ICALFirstOnly s)
2986 step _ (ICALFirstOnlyInner s1 i1) = do
2989 Yield x i' -> Yield x (ICALFirstOnlyInner s1 i')
2990 Skip i' -> Skip (ICALFirstOnlyInner s1 i')
2991 Stop -> Skip (ICALFirstOnly s1)
2993 -- We inject the second stream even before checking if the first stream
2994 -- would yield any more elements. There is no clear choice whether we
2995 -- should do this before or after that. Doing it after may make the state
2996 -- machine a bit simpler though.
2997 step gst (ICALSecondInject s1 s2) = do
2998 r <- step2 (adaptState gst) s2
3002 i `seq` return (Skip (ICALFirstInject s1 s i))
3003 Skip s -> return $ Skip (ICALSecondInject s1 s)
3004 Stop -> return $ Skip (ICALFirstOnly s1)
3006 step gst (ICALFirstInject s1 s2 i2) = do
3007 r <- step1 (adaptState gst) s1
3011 i `seq` return (Skip (ICALSecondInner s s2 i i2))
3012 -- i `seq` return (Skip (ICALFirstBuf s s2 i i2))
3013 Skip s -> return $ Skip (ICALFirstInject s s2 i2)
3017 step _ (ICALFirstBuf s1 s2 i1 i2) = do
3020 Yield x i' -> Skip (ICALSecondInner s1 s2 i' i2 x)
3021 Skip i' -> Skip (ICALFirstBuf s1 s2 i' i2)
3024 step _ (ICALSecondInner s1 s2 i1 i2 v) = do
3027 Yield x i' -> Yield x (ICALSecondInner s1 s2 i1 i' v)
3028 Skip i' -> Skip (ICALSecondInner s1 s2 i1 i' v)
3029 Stop -> Skip (ICALFirstResume s1 s2 i1 i2 v)
3032 step _ (ICALSecondInner s1 s2 i1 i2) = do
3035 Yield x i' -> Yield x (ICALSecondInner s1 s2 i1 i')
3036 Skip i' -> Skip (ICALSecondInner s1 s2 i1 i')
3037 Stop -> Skip (ICALFirstInner s1 s2 i1)
3038 -- Stop -> Skip (ICALFirstResume s1 s2 i1 i2)
3041 step _ (ICALFirstResume s1 s2 i1 i2 x) = do
3042 return $ Yield x (ICALFirstInner s1 s2 i1 i2)
3045data InterposeState s1 i1 a =
3047 -- | InterposeFirstYield s1 i1
3048 | InterposeFirstInner s1 i1
3049 | InterposeFirstInject s1
3050 -- | InterposeFirstBuf s1 i1
3051 | InterposeSecondYield s1 i1
3052 -- -- | InterposeSecondYield s1 i1 a
3053 -- -- | InterposeFirstResume s1 i1 a
3055-- Note that this only interposes the pure values, we may run many effects to
3056-- generate those values as some effects may not generate anything (Skip).
3057{-# INLINE_NORMAL interpose #-}
3058interpose :: Monad m => m c -> Unfold m b c -> Stream m b -> Stream m c
3061 (Unfold istep1 inject1) (Stream step1 state1) =
3062 Stream step (InterposeFirst state1)
3066 {-# INLINE_LATE step #-}
3067 step gst (InterposeFirst s1) = do
3068 r <- step1 (adaptState gst) s1
3072 i `seq` return (Skip (InterposeFirstInner s i))
3073 -- i `seq` return (Skip (InterposeFirstYield s i))
3074 Skip s -> return $ Skip (InterposeFirst s)
3078 step _ (InterposeFirstYield s1 i1) = do
3081 Yield x i' -> Yield x (InterposeFirstInner s1 i')
3082 Skip i' -> Skip (InterposeFirstYield s1 i')
3083 Stop -> Skip (InterposeFirst s1)
3086 step _ (InterposeFirstInner s1 i1) = do
3089 Yield x i' -> Yield x (InterposeFirstInner s1 i')
3090 Skip i' -> Skip (InterposeFirstInner s1 i')
3091 Stop -> Skip (InterposeFirstInject s1)
3093 step gst (InterposeFirstInject s1) = do
3094 r <- step1 (adaptState gst) s1
3098 -- i `seq` return (Skip (InterposeFirstBuf s i))
3099 i `seq` return (Skip (InterposeSecondYield s i))
3100 Skip s -> return $ Skip (InterposeFirstInject s)
3104 step _ (InterposeFirstBuf s1 i1) = do
3107 Yield x i' -> Skip (InterposeSecondYield s1 i' x)
3108 Skip i' -> Skip (InterposeFirstBuf s1 i')
3113 step _ (InterposeSecondYield s1 i1 v) = do
3115 return $ Yield r (InterposeFirstResume s1 i1 v)
3117 step _ (InterposeSecondYield s1 i1) = do
3119 return $ Yield r (InterposeFirstInner s1 i1)
3122 step _ (InterposeFirstResume s1 i1 v) = do
3123 return $ Yield v (InterposeFirstInner s1 i1)
3126------------------------------------------------------------------------------
3128------------------------------------------------------------------------------
3130data GbracketState s1 s2 v
3132 | GBracketNormal s1 v
3133 | GBracketException s2
3135-- | The most general bracketing and exception combinator. All other
3136-- combinators can be expressed in terms of this combinator. This can also be
3137-- used for cases which are not covered by the standard combinators.
3141{-# INLINE_NORMAL gbracket_ #-}
3145 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
3146 -> (c -> m d) -- ^ after, on normal stop
3147 -> (c -> e -> Stream m b -> Stream m b) -- ^ on exception
3148 -> (c -> Stream m b) -- ^ stream generator
3150gbracket_ bef exc aft fexc fnormal =
3151 Stream step GBracketInit
3155 {-# INLINE_LATE step #-}
3156 step _ GBracketInit = do
3158 return $ Skip $ GBracketNormal (fnormal r) r
3160 step gst (GBracketNormal (UnStream step1 st) v) = do
3161 res <- exc $ step1 gst st
3163 Right r -> case r of
3165 return $ Yield x (GBracketNormal (Stream step1 s) v)
3166 Skip s -> return $ Skip (GBracketNormal (Stream step1 s) v)
3167 Stop -> aft v >> return Stop
3169 return $ Skip (GBracketException (fexc v e (UnStream step1 st)))
3170 step gst (GBracketException (UnStream step1 st)) = do
3173 Yield x s -> return $ Yield x (GBracketException (Stream step1 s))
3174 Skip s -> return $ Skip (GBracketException (Stream step1 s))
3177------------------------------------------------------------------------------
3179------------------------------------------------------------------------------
3181-- | Make a finalizer from a monadic action @m a@ that can run in IO monad.
3182mkIOFinalizer :: MonadBaseControl IO m => m b -> m (IO ())
3184 mrun <- captureMonadState
3190-- | Run an IO action stored in a finalized IORef.
3191runIORefFinalizerGC :: IORef (Maybe (IO ())) -> IO ()
3192runIORefFinalizerGC ref = do
3193 res <- readIORef ref
3195 Nothing -> return ()
3198-- | Create an IORef holding a finalizer that is called automatically when
3199-- the IORef is garbage collected. The IORef can be written to with a 'Nothing'
3200-- value to deactivate the finalizer.
3202-- Note: The finalizer is always run with the state of the monad captured at
3203-- the time of calling newFinalizedIORef. To run it on garbage collection we
3204-- have no option but to take a snapshot of the monadic state at some point of
3205-- time. For normal case we could run it with the current state of the monad
3206-- but we want to keep both the cases consistent.
3208newFinalizedIORef :: (MonadIO m, MonadBaseControl IO m)
3209 => m a -> m (IORef (Maybe (IO ())))
3210newFinalizedIORef finalizer = do
3211 f <- mkIOFinalizer finalizer
3212 ref <- liftIO $ newIORef $ Just f
3213 _ <- liftIO $ mkWeakIORef ref (runIORefFinalizerGC ref)
3216-- | Run the finalizer stored in an IORef and deactivate it so that it is run
3217-- only once. Note, the action runs with async exceptions masked.
3219runIORefFinalizer :: MonadIO m => IORef (Maybe (IO ())) -> m ()
3220runIORefFinalizer ref = liftIO $ do
3221 res <- readIORef ref
3223 Nothing -> return ()
3225 -- if an async exception comes after writing 'Nothing' then the
3226 -- finalizing action will never be run. We need to do this
3227 -- atomically wrt async exceptions.
3229 writeIORef ref Nothing
3232-- | Run an action clearing the finalizer IORef atomically wrt async
3233-- exceptions. The action is run with async exceptions masked.
3234withIORefFinalizer :: MonadBaseControl IO m
3235 => IORef (Maybe (IO ())) -> m a -> m a
3236withIORefFinalizer ref action = do
3237 control $ \runinio ->
3239 writeIORef ref Nothing
3242------------------------------------------------------------------------------
3244data GbracketIOState s1 s2 v wref
3246 | GBracketIONormal s1 v wref
3247 | GBracketIOException s2
3249-- | The most general bracketing and exception combinator. All other
3250-- combinators can be expressed in terms of this combinator. This can also be
3251-- used for cases which are not covered by the standard combinators.
3253-- | It uses a finalizer to make sure we run the finalizing action
3254-- when the stream is garbage collected.
3258{-# INLINE_NORMAL gbracket #-}
3260 :: (MonadIO m, MonadBaseControl IO m)
3262 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
3263 -> (c -> m d) -- ^ after, on normal stop or GC
3264 -> (c -> e -> Stream m b -> m (Stream m b)) -- ^ on exception
3265 -> (c -> Stream m b) -- ^ stream generator
3267gbracket bef exc aft fexc fnormal =
3268 Stream step GBracketIOInit
3272 -- If the stream is never evaluated the "aft" action will never be
3273 -- called. For that to occur we will need the user of this API to pass a
3274 -- weak pointer to us.
3275 {-# INLINE_LATE step #-}
3276 step _ GBracketIOInit = do
3277 -- We mask asynchronous exceptions to make the execution
3278 -- of 'bef' and the registration of 'aft' atomic.
3279 -- A similar thing is done in the resourcet package: https://git.io/JvKV3
3280 -- Tutorial: https://markkarpov.com/tutorial/exceptions.html
3281 (r, ref) <- liftBaseOp_ mask_ $ do
3283 ref <- newFinalizedIORef (aft r)
3285 return $ Skip $ GBracketIONormal (fnormal r) r ref
3287 step gst (GBracketIONormal (UnStream step1 st) v ref) = do
3288 res <- exc $ step1 gst st
3290 Right r -> case r of
3292 return $ Yield x (GBracketIONormal (Stream step1 s) v ref)
3294 return $ Skip (GBracketIONormal (Stream step1 s) v ref)
3296 runIORefFinalizer ref
3299 -- Clearing of finalizer and running of exception handler must
3300 -- be atomic wrt async exceptions. Otherwise if we have cleared
3301 -- the finalizer and have not run the exception handler then we
3302 -- may leak the resource.
3303 stream <- withIORefFinalizer ref (fexc v e (UnStream step1 st))
3304 return $ Skip (GBracketIOException stream)
3305 step gst (GBracketIOException (UnStream step1 st)) = do
3309 return $ Yield x (GBracketIOException (Stream step1 s))
3310 Skip s -> return $ Skip (GBracketIOException (Stream step1 s))
3313-- | Run a side effect before the stream yields its first element.
3314{-# INLINE_NORMAL before #-}
3315before :: Monad m => m b -> Stream m a -> Stream m a
3316before action (Stream step state) = Stream step' Nothing
3320 {-# INLINE_LATE step' #-}
3321 step' _ Nothing = action >> return (Skip (Just state))
3323 step' gst (Just st) = do
3326 Yield x s -> return $ Yield x (Just s)
3327 Skip s -> return $ Skip (Just s)
3330-- | Run a side effect whenever the stream stops normally.
3331{-# INLINE_NORMAL after_ #-}
3332after_ :: Monad m => m b -> Stream m a -> Stream m a
3333after_ action (Stream step state) = Stream step' state
3337 {-# INLINE_LATE step' #-}
3341 Yield x s -> return $ Yield x s
3342 Skip s -> return $ Skip s
3343 Stop -> action >> return Stop
3345{-# INLINE_NORMAL after #-}
3346after :: (MonadIO m, MonadBaseControl IO m)
3347 => m b -> Stream m a -> Stream m a
3348after action (Stream step state) = Stream step' Nothing
3352 {-# INLINE_LATE step' #-}
3353 step' _ Nothing = do
3354 ref <- newFinalizedIORef action
3355 return $ Skip $ Just (state, ref)
3356 step' gst (Just (st, ref)) = do
3359 Yield x s -> return $ Yield x (Just (s, ref))
3360 Skip s -> return $ Skip (Just (s, ref))
3362 runIORefFinalizer ref
3365-- XXX These combinators are expensive due to the call to
3366-- onException/handle/try on each step. Therefore, when possible, they should
3367-- be called in an outer loop where we perform less iterations. For example, we
3368-- cannot call them on each iteration in a char stream, instead we can call
3369-- them when doing an IO on an array.
3371-- XXX For high performance error checks in busy streams we may need another
3372-- Error constructor in step.
3374-- | Run a side effect whenever the stream aborts due to an exception. The
3375-- exception is not caught, simply rethrown.
3376{-# INLINE_NORMAL onException #-}
3377onException :: MonadCatch m => m b -> Stream m a -> Stream m a
3378onException action str =
3379 gbracket_ (return ()) MC.try return
3380 (\_ (e :: MC.SomeException) _ -> nilM (action >> MC.throwM e))
3383{-# INLINE_NORMAL _onException #-}
3384_onException :: MonadCatch m => m b -> Stream m a -> Stream m a
3385_onException action (Stream step state) = Stream step' state
3389 {-# INLINE_LATE step' #-}
3391 res <- step gst st `MC.onException` action
3393 Yield x s -> return $ Yield x s
3394 Skip s -> return $ Skip s
3397-- XXX bracket is like concatMap, it generates a stream and then flattens it.
3398-- Like concatMap it has 10x worse performance compared to linear fused
3401-- | Run the first action before the stream starts and remember its output,
3402-- generate a stream using the output, run the second action providing the
3403-- remembered value as an argument whenever the stream ends normally or due to
3405{-# INLINE_NORMAL bracket_ #-}
3406bracket_ :: MonadCatch m => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
3407bracket_ bef aft bet =
3408 gbracket_ bef MC.try aft
3409 (\a (e :: SomeException) _ -> nilM (aft a >> MC.throwM e)) bet
3411-- | Run the first action before the stream starts and remember its output,
3412-- generate a stream using the output, run the second action providing the
3413-- remembered value as an argument whenever the stream ends normally or due to
3414-- an exception or gets garbage collected.
3416{-# INLINE_NORMAL bracket #-}
3417bracket :: (MonadAsync m, MonadCatch m)
3418 => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
3419bracket bef aft bet =
3420 gbracket bef MC.try aft
3421 (\a (e :: SomeException) _ -> aft a >> return (nilM (MC.throwM e))) bet
3423data BracketState s v = BracketInit | BracketRun s v
3425{-# INLINE_NORMAL _bracket #-}
3426_bracket :: MonadCatch m => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
3427_bracket bef aft bet = Stream step' BracketInit
3431 {-# INLINE_LATE step' #-}
3432 step' _ BracketInit = bef >>= \x -> return (Skip (BracketRun (bet x) x))
3434 -- NOTE: It is important to use UnStream instead of the Stream pattern
3435 -- here, otherwise we get huge perf degradation, see note in concatMap.
3436 step' gst (BracketRun (UnStream step state) v) = do
3437 -- res <- step gst state `MC.onException` aft v
3438 res <- MC.try $ step gst state
3440 Left (e :: SomeException) -> aft v >> MC.throwM e >> return Stop
3441 Right r -> case r of
3442 Yield x s -> return $ Yield x (BracketRun (Stream step s) v)
3443 Skip s -> return $ Skip (BracketRun (Stream step s) v)
3444 Stop -> aft v >> return Stop
3446-- | Run a side effect whenever the stream stops normally or aborts due to an
3448{-# INLINE finally_ #-}
3449finally_ :: MonadCatch m => m b -> Stream m a -> Stream m a
3450-- finally action xs = after action $ onException action xs
3451finally_ action xs = bracket_ (return ()) (\_ -> action) (const xs)
3453-- | Run a side effect whenever the stream stops normally or aborts due to an
3454-- exception or gets garbage collected.
3456{-# INLINE finally #-}
3457finally :: (MonadAsync m, MonadCatch m) => m b -> Stream m a -> Stream m a
3458finally action xs = bracket (return ()) (\_ -> action) (const xs)
3460-- | When evaluating a stream if an exception occurs, stream
3461-- evaluation aborts and the specified exception handler is run with
3462-- the exception and the Stream which threw the exception as argument.
3464{-# INLINE_NORMAL ghandle #-}
3465ghandle :: (MonadCatch m, Exception e)
3466 => (e -> Stream m a -> Stream m a) -> Stream m a -> Stream m a
3468 gbracket_ (return ()) MC.try return (\_ -> f) (\_ -> str)
3470-- | When evaluating a stream if an exception occurs, stream evaluation aborts
3471-- and the specified exception handler is run with the exception as argument.
3472{-# INLINE_NORMAL handle #-}
3473handle :: (MonadCatch m, Exception e)
3474 => (e -> Stream m a) -> Stream m a -> Stream m a
3476 gbracket_ (return ()) MC.try return (\_ e _ -> f e) (\_ -> str)
3478{-# INLINE_NORMAL _handle #-}
3479_handle :: (MonadCatch m, Exception e)
3480 => (e -> Stream m a) -> Stream m a -> Stream m a
3481_handle f (Stream step state) = Stream step' (Left state)
3485 {-# INLINE_LATE step' #-}
3486 step' gst (Left st) = do
3487 res <- MC.try $ step gst st
3489 Left e -> return $ Skip $ Right (f e)
3490 Right r -> case r of
3491 Yield x s -> return $ Yield x (Left s)
3492 Skip s -> return $ Skip (Left s)
3495 step' gst (Right (UnStream step1 st)) = do
3498 Yield x s -> return $ Yield x (Right (Stream step1 s))
3499 Skip s -> return $ Skip (Right (Stream step1 s))
3502-------------------------------------------------------------------------------
3503-- General transformation
3504-------------------------------------------------------------------------------
3506{-# INLINE_NORMAL transform #-}
3507transform :: Monad m => Pipe m a b -> Stream m a -> Stream m b
3508transform (Pipe pstep1 pstep2 pstate) (Stream step state) =
3509 Stream step' (Consume pstate, state)
3513 {-# INLINE_LATE step' #-}
3515 step' gst (Consume pst, st) = pst `seq` do
3516 r <- step (adaptState gst) st
3521 Pipe.Yield b pst' -> return $ Yield b (pst', s)
3522 Pipe.Continue pst' -> return $ Skip (pst', s)
3523 Skip s -> return $ Skip (Consume pst, s)
3526 step' _ (Produce pst, st) = pst `seq` do
3529 Pipe.Yield b pst' -> return $ Yield b (pst', st)
3530 Pipe.Continue pst' -> return $ Skip (pst', st)
3532------------------------------------------------------------------------------
3533-- Transformation by Folding (Scans)
3534------------------------------------------------------------------------------
3536------------------------------------------------------------------------------
3538------------------------------------------------------------------------------
3540-- XXX Is a prescan useful, discarding the last step does not sound useful? I
3541-- am not sure about the utility of this function, so this is implemented but
3542-- not exposed. We can expose it if someone provides good reasons why this is
3545-- XXX We have to execute the stream one step ahead to know that we are at the
3546-- last step. The vector implementation of prescan executes the last fold step
3547-- but does not yield the result. This means we have executed the effect but
3548-- discarded value. This does not sound right. In this implementation we are
3549-- not executing the last fold step.
3550{-# INLINE_NORMAL prescanlM' #-}
3551prescanlM' :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
3552prescanlM' f mz (Stream step state) = Stream step' (state, mz)
3554 {-# INLINE_LATE step' #-}
3555 step' gst (st, prev) = do
3556 r <- step (adaptState gst) st
3560 return $ Yield acc (s, f acc x)
3561 Skip s -> return $ Skip (s, prev)
3564{-# INLINE prescanl' #-}
3565prescanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
3566prescanl' f z = prescanlM' (\a b -> return (f a b)) (return z)
3568------------------------------------------------------------------------------
3569-- Monolithic postscans (postscan followed by a map)
3570------------------------------------------------------------------------------
3572-- The performance of a modular postscan followed by a map seems to be
3573-- equivalent to this monolithic scan followed by map therefore we may not need
3574-- this implementation. We just have it for performance comparison and in case
3575-- modular version does not perform well in some situation.
3577{-# INLINE_NORMAL postscanlMx' #-}
3578postscanlMx' :: Monad m
3579 => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> Stream m b
3580postscanlMx' fstep begin done (Stream step state) = do
3581 Stream step' (state, begin)
3583 {-# INLINE_LATE step' #-}
3584 step' gst (st, acc) = do
3585 r <- step (adaptState gst) st
3591 v `seq` y `seq` return (Yield v (s, return y))
3592 Skip s -> return $ Skip (s, acc)
3595{-# INLINE_NORMAL postscanlx' #-}
3596postscanlx' :: Monad m
3597 => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> Stream m b
3598postscanlx' fstep begin done s =
3599 postscanlMx' (\b a -> return (fstep b a)) (return begin) (return . done) s
3601-- XXX do we need consM strict to evaluate the begin value?
3602{-# INLINE scanlMx' #-}
3604 => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> Stream m b
3605scanlMx' fstep begin done s =
3606 (begin >>= \x -> x `seq` done x) `consM` postscanlMx' fstep begin done s
3608{-# INLINE scanlx' #-}
3610 => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> Stream m b
3611scanlx' fstep begin done s =
3612 scanlMx' (\b a -> return (fstep b a)) (return begin) (return . done) s
3614------------------------------------------------------------------------------
3616------------------------------------------------------------------------------
3618{-# INLINE_NORMAL postscanlM' #-}
3619postscanlM' :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
3620postscanlM' fstep begin (Stream step state) =
3621 Stream step' Nothing
3623 {-# INLINE_LATE step' #-}
3624 step' _ Nothing = do
3626 return $ Skip (Just (state, x))
3628 step' gst (Just (st, acc)) = do
3629 r <- step (adaptState gst) st
3633 return $ Yield y (Just (s, y))
3634 Skip s -> return $ Skip (Just (s, acc))
3637{-# INLINE_NORMAL postscanl' #-}
3638postscanl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
3639postscanl' f seed = postscanlM' (\a b -> return (f a b)) (return seed)
3641-- We can possibly have the "done" function as a Maybe to provide an option to
3642-- emit or not emit the accumulator when the stream stops.
3644-- TBD: use a single Yield point
3646{-# INLINE_NORMAL postscanlMAfter' #-}
3647postscanlMAfter' :: Monad m
3648 => (b -> a -> m b) -> m b -> (b -> m b) -> Stream m a -> Stream m b
3649postscanlMAfter' fstep initial done (Stream step1 state1) = do
3650 Stream step (Just (state1, initial))
3654 {-# INLINE_LATE step #-}
3655 step gst (Just (st, acc)) = do
3656 r <- step1 (adaptState gst) st
3661 y `seq` return (Yield y (Just (s, return y)))
3662 Skip s -> return $ Skip $ Just (s, acc)
3663 Stop -> acc >>= done >>= \res -> return (Yield res Nothing)
3664 step _ Nothing = return Stop
3666{-# INLINE_NORMAL postscanlM #-}
3667postscanlM :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
3668postscanlM fstep begin (Stream step state) = Stream step' Nothing
3670 {-# INLINE_LATE step' #-}
3671 step' _ Nothing = do
3673 return $ Skip (Just (state, r))
3675 step' gst (Just (st, acc)) = do
3676 r <- step (adaptState gst) st
3680 return (Yield y (Just (s, y)))
3681 Skip s -> return $ Skip (Just (s, acc))
3684{-# INLINE_NORMAL postscanl #-}
3685postscanl :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
3686postscanl f seed = postscanlM (\a b -> return (f a b)) (return seed)
3688{-# INLINE_NORMAL scanlM' #-}
3689scanlM' :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
3690scanlM' fstep begin (Stream step state) = Stream step' Nothing
3692 {-# INLINE_LATE step' #-}
3693 step' _ Nothing = do
3695 return $ Yield x (Just (state, x))
3696 step' gst (Just (st, acc)) = do
3697 r <- step (adaptState gst) st
3701 return $ Yield y (Just (s, y))
3702 Skip s -> return $ Skip (Just (s, acc))
3705{-# INLINE scanlMAfter' #-}
3706scanlMAfter' :: Monad m
3707 => (b -> a -> m b) -> m b -> (b -> m b) -> Stream m a -> Stream m b
3708scanlMAfter' fstep initial done s =
3709 (initial >>= \x -> x `seq` return x) `consM`
3710 postscanlMAfter' fstep initial done s
3712{-# INLINE scanl' #-}
3713scanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
3714scanl' f seed = scanlM' (\a b -> return (f a b)) (return seed)
3716{-# INLINE_NORMAL scanlM #-}
3717scanlM :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
3718scanlM fstep begin (Stream step state) = Stream step' Nothing
3720 {-# INLINE_LATE step' #-}
3721 step' _ Nothing = do
3723 return $ Yield x (Just (state, x))
3724 step' gst (Just (st, acc)) = do
3725 r <- step (adaptState gst) st
3729 return $ Yield y (Just (s, y))
3730 Skip s -> return $ Skip (Just (s, acc))
3731 Stop -> return $ Stop
3734scanl :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
3735scanl f seed = scanlM (\a b -> return (f a b)) (return seed)
3737{-# INLINE_NORMAL scanl1M #-}
3738scanl1M :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
3739scanl1M fstep (Stream step state) = Stream step' (state, Nothing)
3741 {-# INLINE_LATE step' #-}
3742 step' gst (st, Nothing) = do
3745 Yield x s -> return $ Yield x (s, Just x)
3746 Skip s -> return $ Skip (s, Nothing)
3749 step' gst (st, Just acc) = do
3754 return $ Yield z (s, Just z)
3755 Skip s -> return $ Skip (s, Just acc)
3758{-# INLINE scanl1 #-}
3759scanl1 :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
3760scanl1 f = scanl1M (\x y -> return (f x y))
3762{-# INLINE_NORMAL scanl1M' #-}
3763scanl1M' :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
3764scanl1M' fstep (Stream step state) = Stream step' (state, Nothing)
3766 {-# INLINE_LATE step' #-}
3767 step' gst (st, Nothing) = do
3770 Yield x s -> x `seq` return $ Yield x (s, Just x)
3771 Skip s -> return $ Skip (s, Nothing)
3774 step' gst (st, Just acc) = acc `seq` do
3779 z `seq` return $ Yield z (s, Just z)
3780 Skip s -> return $ Skip (s, Just acc)
3783{-# INLINE scanl1' #-}
3784scanl1' :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
3785scanl1' f = scanl1M' (\x y -> return (f x y))
3787------------------------------------------------------------------------------
3789------------------------------------------------------------------------------
3791data RollingMapState s a = RollingMapInit s | RollingMapGo s a
3793{-# INLINE rollingMapM #-}
3794rollingMapM :: Monad m => (a -> a -> m b) -> Stream m a -> Stream m b
3795rollingMapM f (Stream step1 state1) = Stream step (RollingMapInit state1)
3797 step gst (RollingMapInit st) = do
3798 r <- step1 (adaptState gst) st
3800 Yield x s -> Skip $ RollingMapGo s x
3801 Skip s -> Skip $ RollingMapInit s
3804 step gst (RollingMapGo s1 x1) = do
3805 r <- step1 (adaptState gst) s1
3809 return $ Yield res $ RollingMapGo s x
3810 Skip s -> return $ Skip $ RollingMapGo s x1
3811 Stop -> return $ Stop
3813{-# INLINE rollingMap #-}
3814rollingMap :: Monad m => (a -> a -> b) -> Stream m a -> Stream m b
3815rollingMap f = rollingMapM (\x y -> return $ f x y)
3817------------------------------------------------------------------------------
3818-- Tapping/Distributing
3819------------------------------------------------------------------------------
3822tap :: Monad m => Fold m a b -> Stream m a -> Stream m a
3823tap (Fold fstep initial extract) (Stream step state) = Stream step' Nothing
3827 step' _ Nothing = do
3829 return $ Skip (Just (r, state))
3831 step' gst (Just (acc, st)) = acc `seq` do
3836 return $ Yield x (Just (acc', s))
3837 Skip s -> return $ Skip (Just (acc, s))
3842{-# INLINE_NORMAL tapOffsetEvery #-}
3843tapOffsetEvery :: Monad m
3844 => Int -> Int -> Fold m a b -> Stream m a -> Stream m a
3845tapOffsetEvery offset n (Fold fstep initial extract) (Stream step state) =
3846 Stream step' Nothing
3850 {-# INLINE_LATE step' #-}
3851 step' _ Nothing = do
3853 return $ Skip (Just (r, state, offset `mod` n))
3855 step' gst (Just (acc, st, count)) | count <= 0 = do
3859 !acc' <- fstep acc x
3860 return $ Yield x (Just (acc', s, n - 1))
3861 Skip s -> return $ Skip (Just (acc, s, count))
3866 step' gst (Just (acc, st, count)) = do
3869 Yield x s -> return $ Yield x (Just (acc, s, count - 1))
3870 Skip s -> return $ Skip (Just (acc, s, count))
3875{-# INLINE_NORMAL pollCounts #-}
3879 -> (Stream m Int -> Stream m Int)
3883pollCounts predicate transf fld (Stream step state) = Stream step' Nothing
3886 {-# INLINE_LATE step' #-}
3887 step' _ Nothing = do
3888 -- As long as we are using an "Int" for counts lockfree reads from
3889 -- Var should work correctly on both 32-bit and 64-bit machines.
3890 -- However, an Int on a 32-bit machine may overflow quickly.
3891 countVar <- liftIO $ Prim.newIORef (0 :: Int)
3893 $ void $ runFold fld
3894 $ transf $ fromPrimIORef countVar
3895 return $ Skip (Just (countVar, tid, state))
3897 step' gst (Just (countVar, tid, st)) = do
3901 when (predicate x) $ liftIO $ Prim.modifyIORef' countVar (+ 1)
3902 return $ Yield x (Just (countVar, tid, s))
3903 Skip s -> return $ Skip (Just (countVar, tid, s))
3905 liftIO $ killThread tid
3908{-# INLINE_NORMAL tapRate #-}
3910 (MonadAsync m, MonadCatch m)
3915tapRate samplingRate action (Stream step state) = Stream step' Nothing
3917 {-# NOINLINE loop #-}
3918 loop countVar prev = do
3921 (do liftIO $ threadDelay (round $ samplingRate * 1000000)
3922 i <- liftIO $ Prim.readIORef countVar
3923 let !diff = i - prev
3926 (\(e :: AsyncException) -> do
3927 i <- liftIO $ Prim.readIORef countVar
3928 let !diff = i - prev
3930 throwM (MC.toException e))
3933 {-# INLINE_LATE step' #-}
3934 step' _ Nothing = do
3935 countVar <- liftIO $ Prim.newIORef 0
3936 tid <- fork $ loop countVar 0
3937 ref <- liftIO $ newIORef ()
3938 _ <- liftIO $ mkWeakIORef ref (killThread tid)
3939 return $ Skip (Just (countVar, tid, state, ref))
3941 step' gst (Just (countVar, tid, st, ref)) = do
3945 liftIO $ Prim.modifyIORef' countVar (+ 1)
3946 return $ Yield x (Just (countVar, tid, s, ref))
3947 Skip s -> return $ Skip (Just (countVar, tid, s, ref))
3949 liftIO $ killThread tid
3953-------------------------------------------------------------------------------
3955-------------------------------------------------------------------------------
3957{-# INLINE_NORMAL takeWhileM #-}
3958takeWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
3959takeWhileM f (Stream step state) = Stream step' state
3961 {-# INLINE_LATE step' #-}
3967 return $ if b then Yield x s else Stop
3968 Skip s -> return $ Skip s
3971{-# INLINE takeWhile #-}
3972takeWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
3973takeWhile f = takeWhileM (return . f)
3975{-# INLINE_NORMAL drop #-}
3976drop :: Monad m => Int -> Stream m a -> Stream m a
3977drop n (Stream step state) = Stream step' (state, Just n)
3979 {-# INLINE_LATE step' #-}
3980 step' gst (st, Just i)
3985 Yield _ s -> Skip (s, Just (i - 1))
3986 Skip s -> Skip (s, Just i)
3988 | otherwise = return $ Skip (st, Nothing)
3990 step' gst (st, Nothing) = do
3994 Yield x s -> Yield x (s, Nothing)
3995 Skip s -> Skip (s, Nothing)
3998data DropWhileState s a
4000 | DropWhileYield a s
4003{-# INLINE_NORMAL dropWhileM #-}
4004dropWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
4005dropWhileM f (Stream step state) = Stream step' (DropWhileDrop state)
4007 {-# INLINE_LATE step' #-}
4008 step' gst (DropWhileDrop st) = do
4014 then return $ Skip (DropWhileDrop s)
4015 else return $ Skip (DropWhileYield x s)
4016 Skip s -> return $ Skip (DropWhileDrop s)
4019 step' gst (DropWhileNext st) = do
4022 Yield x s -> return $ Skip (DropWhileYield x s)
4023 Skip s -> return $ Skip (DropWhileNext s)
4026 step' _ (DropWhileYield x st) = return $ Yield x (DropWhileNext st)
4028{-# INLINE dropWhile #-}
4029dropWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
4030dropWhile f = dropWhileM (return . f)
4032{-# INLINE_NORMAL filterM #-}
4033filterM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
4034filterM f (Stream step state) = Stream step' state
4036 {-# INLINE_LATE step' #-}
4045 Skip s -> return $ Skip s
4048{-# INLINE filter #-}
4049filter :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
4050filter f = filterM (return . f)
4052{-# INLINE_NORMAL uniq #-}
4053uniq :: (Eq a, Monad m) => Stream m a -> Stream m a
4054uniq (Stream step state) = Stream step' (Nothing, state)
4056 {-# INLINE_LATE step' #-}
4057 step' gst (Nothing, st) = do
4060 Yield x s -> return $ Yield x (Just x, s)
4061 Skip s -> return $ Skip (Nothing, s)
4063 step' gst (Just x, st) = do
4066 Yield y s | x == y -> return $ Skip (Just x, s)
4067 | otherwise -> return $ Yield y (Just y, s)
4068 Skip s -> return $ Skip (Just x, s)
4071------------------------------------------------------------------------------
4072-- Transformation by Mapping
4073------------------------------------------------------------------------------
4075{-# INLINE_NORMAL sequence #-}
4076sequence :: Monad m => Stream m (m a) -> Stream m a
4077sequence (Stream step state) = Stream step' state
4079 {-# INLINE_LATE step' #-}
4081 r <- step (adaptState gst) st
4083 Yield x s -> x >>= \a -> return (Yield a s)
4084 Skip s -> return $ Skip s
4087------------------------------------------------------------------------------
4089------------------------------------------------------------------------------
4091data LoopState x s = FirstYield s
4092 | InterspersingYield s
4095{-# INLINE_NORMAL intersperseM #-}
4096intersperseM :: Monad m => m a -> Stream m a -> Stream m a
4097intersperseM m (Stream step state) = Stream step' (FirstYield state)
4099 {-# INLINE_LATE step' #-}
4100 step' gst (FirstYield st) = do
4104 Yield x s -> Skip (YieldAndCarry x s)
4105 Skip s -> Skip (FirstYield s)
4108 step' gst (InterspersingYield st) = do
4113 return $ Yield a (YieldAndCarry x s)
4114 Skip s -> return $ Skip $ InterspersingYield s
4117 step' _ (YieldAndCarry x st) = return $ Yield x (InterspersingYield st)
4119{-# INLINE_NORMAL intersperseM_ #-}
4120intersperseM_ :: Monad m => m b -> Stream m a -> Stream m a
4121intersperseM_ m (Stream step1 state1) = Stream step (Left (pure (), state1))
4123 {-# INLINE_LATE step #-}
4124 step gst (Left (eff, st)) = do
4127 Yield x s -> eff >> return (Yield x (Right s))
4128 Skip s -> return $ Skip (Left (eff, s))
4131 step _ (Right st) = return $ Skip $ Left (void m, st)
4136 | SuffixYield a (SuffixState s a)
4138{-# INLINE_NORMAL intersperseSuffix #-}
4139intersperseSuffix :: forall m a. Monad m => m a -> Stream m a -> Stream m a
4140intersperseSuffix action (Stream step state) = Stream step' (SuffixElem state)
4142 {-# INLINE_LATE step' #-}
4143 step' gst (SuffixElem st) = do
4146 Yield x s -> Skip (SuffixYield x (SuffixSuffix s))
4147 Skip s -> Skip (SuffixElem s)
4150 step' _ (SuffixSuffix st) = do
4151 action >>= \r -> return $ Skip (SuffixYield r (SuffixElem st))
4153 step' _ (SuffixYield x next) = return $ Yield x next
4155{-# INLINE_NORMAL intersperseSuffix_ #-}
4156intersperseSuffix_ :: Monad m => m b -> Stream m a -> Stream m a
4157intersperseSuffix_ m (Stream step1 state1) = Stream step (Left state1)
4159 {-# INLINE_LATE step #-}
4160 step gst (Left st) = do
4163 Yield x s -> return $ Yield x (Right s)
4164 Skip s -> return $ Skip $ Left s
4167 step _ (Right st) = m >> return (Skip (Left st))
4169data SuffixSpanState s a
4170 = SuffixSpanElem s Int
4171 | SuffixSpanSuffix s
4172 | SuffixSpanYield a (SuffixSpanState s a)
4176-- | intersperse after every n items
4177{-# INLINE_NORMAL intersperseSuffixBySpan #-}
4178intersperseSuffixBySpan :: forall m a. Monad m
4179 => Int -> m a -> Stream m a -> Stream m a
4180intersperseSuffixBySpan n action (Stream step state) =
4181 Stream step' (SuffixSpanElem state n)
4183 {-# INLINE_LATE step' #-}
4184 step' gst (SuffixSpanElem st i) | i > 0 = do
4187 Yield x s -> Skip (SuffixSpanYield x (SuffixSpanElem s (i - 1)))
4188 Skip s -> Skip (SuffixSpanElem s i)
4189 Stop -> if i == n then Stop else Skip SuffixSpanLast
4190 step' _ (SuffixSpanElem st _) = return $ Skip (SuffixSpanSuffix st)
4192 step' _ (SuffixSpanSuffix st) = do
4193 action >>= \r -> return $ Skip (SuffixSpanYield r (SuffixSpanElem st n))
4195 step' _ (SuffixSpanLast) = do
4196 action >>= \r -> return $ Skip (SuffixSpanYield r SuffixSpanStop)
4198 step' _ (SuffixSpanYield x next) = return $ Yield x next
4200 step' _ (SuffixSpanStop) = return Stop
4202{-# INLINE intersperse #-}
4203intersperse :: Monad m => a -> Stream m a -> Stream m a
4204intersperse a = intersperseM (return a)
4206{-# INLINE_NORMAL insertBy #-}
4207insertBy :: Monad m => (a -> a -> Ordering) -> a -> Stream m a -> Stream m a
4208insertBy cmp a (Stream step state) = Stream step' (state, False, Nothing)
4210 {-# INLINE_LATE step' #-}
4211 step' gst (st, False, _) = do
4214 Yield x s -> case cmp a x of
4215 GT -> return $ Yield x (s, False, Nothing)
4216 _ -> return $ Yield a (s, True, Just x)
4217 Skip s -> return $ Skip (s, False, Nothing)
4218 Stop -> return $ Yield a (st, True, Nothing)
4220 step' _ (_, True, Nothing) = return Stop
4222 step' gst (st, True, Just prev) = do
4225 Yield x s -> return $ Yield prev (s, True, Just x)
4226 Skip s -> return $ Skip (s, True, Just prev)
4227 Stop -> return $ Yield prev (st, True, Nothing)
4229------------------------------------------------------------------------------
4231------------------------------------------------------------------------------
4233{-# INLINE_NORMAL deleteBy #-}
4234deleteBy :: Monad m => (a -> a -> Bool) -> a -> Stream m a -> Stream m a
4235deleteBy eq x (Stream step state) = Stream step' (state, False)
4237 {-# INLINE_LATE step' #-}
4238 step' gst (st, False) = do
4241 Yield y s -> return $
4242 if eq x y then Skip (s, True) else Yield y (s, False)
4243 Skip s -> return $ Skip (s, False)
4246 step' gst (st, True) = do
4249 Yield y s -> return $ Yield y (s, True)
4250 Skip s -> return $ Skip (s, True)
4253------------------------------------------------------------------------------
4254-- Transformation by Map and Filter
4255------------------------------------------------------------------------------
4257-- XXX Will this always fuse properly?
4258{-# INLINE_NORMAL mapMaybe #-}
4259mapMaybe :: Monad m => (a -> Maybe b) -> Stream m a -> Stream m b
4260mapMaybe f = fmap fromJust . filter isJust . map f
4262{-# INLINE_NORMAL mapMaybeM #-}
4263mapMaybeM :: Monad m => (a -> m (Maybe b)) -> Stream m a -> Stream m b
4264mapMaybeM f = fmap fromJust . filter isJust . mapM f
4266------------------------------------------------------------------------------
4268------------------------------------------------------------------------------
4270{-# INLINE_NORMAL indexed #-}
4271indexed :: Monad m => Stream m a -> Stream m (Int, a)
4272indexed (Stream step state) = Stream step' (state, 0)
4274 {-# INLINE_LATE step' #-}
4275 step' gst (st, i) = i `seq` do
4276 r <- step (adaptState gst) st
4278 Yield x s -> return $ Yield (i, x) (s, i+1)
4279 Skip s -> return $ Skip (s, i)
4282{-# INLINE_NORMAL indexedR #-}
4283indexedR :: Monad m => Int -> Stream m a -> Stream m (Int, a)
4284indexedR m (Stream step state) = Stream step' (state, m)
4286 {-# INLINE_LATE step' #-}
4287 step' gst (st, i) = i `seq` do
4288 r <- step (adaptState gst) st
4290 Yield x s -> let i' = i - 1
4291 in return $ Yield (i, x) (s, i')
4292 Skip s -> return $ Skip (s, i)
4295{-# INLINE_NORMAL zipWithM #-}
4297 => (a -> b -> m c) -> Stream m a -> Stream m b -> Stream m c
4298zipWithM f (Stream stepa ta) (Stream stepb tb) = Stream step (ta, tb, Nothing)
4300 {-# INLINE_LATE step #-}
4301 step gst (sa, sb, Nothing) = do
4302 r <- stepa (adaptState gst) sa
4305 Yield x sa' -> Skip (sa', sb, Just x)
4306 Skip sa' -> Skip (sa', sb, Nothing)
4309 step gst (sa, sb, Just x) = do
4310 r <- stepb (adaptState gst) sb
4314 return $ Yield z (sa, sb', Nothing)
4315 Skip sb' -> return $ Skip (sa, sb', Just x)
4318#if __GLASGOW_HASKELL__ >= 801
4319{-# RULES "zipWithM xs xs"
4320 forall f xs. zipWithM @Identity f xs xs = mapM (\x -> f x x) xs #-}
4323{-# INLINE zipWith #-}
4324zipWith :: Monad m => (a -> b -> c) -> Stream m a -> Stream m b -> Stream m c
4325zipWith f = zipWithM (\a b -> return (f a b))
4327------------------------------------------------------------------------------
4329------------------------------------------------------------------------------
4331{-# INLINE_NORMAL mergeByM #-}
4334 => (a -> a -> m Ordering) -> Stream m a -> Stream m a -> Stream m a
4335mergeByM cmp (Stream stepa ta) (Stream stepb tb) =
4336 Stream step (Just ta, Just tb, Nothing, Nothing)
4338 {-# INLINE_LATE step #-}
4340 -- one of the values is missing, and the corresponding stream is running
4341 step gst (Just sa, sb, Nothing, b) = do
4344 Yield a sa' -> Skip (Just sa', sb, Just a, b)
4345 Skip sa' -> Skip (Just sa', sb, Nothing, b)
4346 Stop -> Skip (Nothing, sb, Nothing, b)
4348 step gst (sa, Just sb, a, Nothing) = do
4351 Yield b sb' -> Skip (sa, Just sb', a, Just b)
4352 Skip sb' -> Skip (sa, Just sb', a, Nothing)
4353 Stop -> Skip (sa, Nothing, a, Nothing)
4355 -- both the values are available
4356 step _ (sa, sb, Just a, Just b) = do
4358 return $ case res of
4359 GT -> Yield b (sa, sb, Just a, Nothing)
4360 _ -> Yield a (sa, sb, Nothing, Just b)
4362 -- one of the values is missing, corresponding stream is done
4363 step _ (Nothing, sb, Nothing, Just b) =
4364 return $ Yield b (Nothing, sb, Nothing, Nothing)
4366 step _ (sa, Nothing, Just a, Nothing) =
4367 return $ Yield a (sa, Nothing, Nothing, Nothing)
4369 step _ (Nothing, Nothing, Nothing, Nothing) = return Stop
4371{-# INLINE mergeBy #-}
4374 => (a -> a -> Ordering) -> Stream m a -> Stream m a -> Stream m a
4375mergeBy cmp = mergeByM (\a b -> return $ cmp a b)
4377------------------------------------------------------------------------------
4378-- Transformation comprehensions
4379------------------------------------------------------------------------------
4381{-# INLINE_NORMAL the #-}
4382the :: (Eq a, Monad m) => Stream m a -> m (Maybe a)
4383the (Stream step state) = go state
4386 r <- step defState st
4388 Yield x s -> go' x s
4390 Stop -> return Nothing
4392 r <- step defState st
4394 Yield x s | x == n -> go' n s
4395 | otherwise -> return Nothing
4397 Stop -> return (Just n)
4399{-# INLINE runFold #-}
4400runFold :: (Monad m) => Fold m a b -> Stream m a -> m b
4401runFold (Fold step begin done) = foldlMx' step begin done
4403-------------------------------------------------------------------------------
4404-- Concurrent application and fold
4405-------------------------------------------------------------------------------
4407-- XXX These functions should be moved to Stream/Parallel.hs
4409-- Using StreamD the worker stream producing code can fuse with the code to
4410-- queue output to the SVar giving some perf boost.
4412-- Note that StreamD can only be used in limited situations, specifically, we
4413-- cannot implement joinStreamVarPar using this.
4415-- XXX make sure that the SVar passed is a Parallel style SVar.
4417-- | Fold the supplied stream to the SVar asynchronously using Parallel
4418-- concurrency style.
4419-- {-# INLINE_NORMAL toSVarParallel #-}
4420{-# INLINE toSVarParallel #-}
4421toSVarParallel :: MonadAsync m
4422 => State t m a -> SVar t m a -> Stream m a -> m ()
4423toSVarParallel st sv xs =
4424 if svarInspectMode sv
4428 case getYieldLimit st of
4429 Nothing -> doFork (work Nothing)
4431 (handleChildException sv)
4432 Just _ -> doFork (workLim Nothing)
4434 (handleChildException sv)
4439 {-# NOINLINE work #-}
4440 work info = (runFold (FL.toParallelSVar sv info) xs)
4442 {-# NOINLINE workLim #-}
4443 workLim info = runFold (FL.toParallelSVarLimited sv info) xs
4445 {-# NOINLINE forkWithDiag #-}
4447 -- We do not use workerCount in case of ParallelVar but still there is
4448 -- no harm in maintaining it correctly.
4449 liftIO $ atomicModifyIORefCAS_ (workerCount sv) $ \n -> n + 1
4451 -- This allocation matters when significant number of workers are being
4452 -- sent. We allocate it only when needed. The overhead increases by 4x.
4454 case yieldRateInfo sv of
4455 Nothing -> return Nothing
4456 Just _ -> liftIO $ do
4457 cntRef <- newIORef 0
4458 t <- getTime Monotonic
4459 lat <- newIORef (0, t)
4460 return $ Just WorkerInfo
4461 { workerYieldMax = 0
4462 , workerYieldCount = cntRef
4463 , workerLatencyStart = lat
4466 case getYieldLimit st of
4467 Nothing -> doFork (work winfo)
4469 (handleChildException sv)
4470 Just _ -> doFork (workLim winfo)
4472 (handleChildException sv)
4475{-# INLINE_NORMAL mkParallelD #-}
4476mkParallelD :: MonadAsync m => Stream m a -> Stream m a
4477mkParallelD m = Stream step Nothing
4480 step gst Nothing = do
4481 sv <- newParallelVar StopNone gst
4482 toSVarParallel gst sv m
4483 -- XXX use unfold instead?
4484 return $ Skip $ Just $ fromSVar sv
4486 step gst (Just (UnStream step1 st)) = do
4489 Yield a s -> Yield a (Just $ Stream step1 s)
4490 Skip s -> Skip (Just $ Stream step1 s)
4493-- Compare with mkAsync. mkAsync uses an Async style SVar whereas this uses a
4494-- parallel style SVar for evaluation. Currently, parallel style cannot use
4495-- rate control whereas Async style can use rate control. In async style SVar
4496-- the worker thread terminates when the buffer is full whereas in Parallel
4499-- | Make the stream producer and consumer run concurrently by introducing a
4500-- buffer between them. The producer thread evaluates the input stream until
4501-- the buffer fills, it blocks if the buffer is full until there is space in
4502-- the buffer. The consumer consumes the stream lazily from the buffer.
4506{-# INLINE_NORMAL mkParallel #-}
4507mkParallel :: (K.IsStream t, MonadAsync m) => t m a -> t m a
4508mkParallel = fromStreamD . mkParallelD . toStreamD
4510-- Note: we can use another API with two callbacks stop and yield if we want
4511-- the callback to be able to indicate end of stream.
4513-- | Generates a callback and a stream pair. The callback returned is used to
4514-- queue values to the stream. The stream is infinite, there is no way for the
4515-- callback to indicate that it is done now.
4519{-# INLINE_NORMAL newCallbackStream #-}
4520newCallbackStream :: (K.IsStream t, MonadAsync m) => m ((a -> m ()), t m a)
4521newCallbackStream = do
4522 sv <- newParallelVar StopNone defState
4524 -- XXX Add our own thread-id to the SVar as we can not know the callback's
4525 -- thread-id and the callback is not run in a managed worker. We need to
4526 -- handle this better.
4527 liftIO myThreadId >>= modifyThread sv
4529 let callback a = liftIO $ void $ send sv (ChildYield a)
4530 -- XXX we can return an SVar and then the consumer can unfold from the
4532 return (callback, fromStreamD (fromSVar sv))
4534-------------------------------------------------------------------------------
4536-------------------------------------------------------------------------------
4538-- | Create an SVar with a fold consumer that will fold any elements sent to it
4539-- using the supplied fold function.
4540{-# INLINE newFoldSVar #-}
4541newFoldSVar :: MonadAsync m => State t m a -> Fold m a b -> m (SVar t m a)
4542newFoldSVar stt f = do
4543 -- Buffer size for the SVar is derived from the current state
4544 sv <- newParallelVar StopAny (adaptState stt)
4545 -- Add the producer thread-id to the SVar.
4546 liftIO myThreadId >>= modifyThread sv
4547 void $ doFork (work sv) (svarMrun sv) (handleFoldException sv)
4552 {-# NOINLINE work #-}
4553 work sv = void $ runFold f $ fromProducer sv
4555data TapState sv st = TapInit | Tapping sv st | TapDone st
4557{-# INLINE_NORMAL tapAsync #-}
4558tapAsync :: MonadAsync m => Fold m a b -> Stream m a -> Stream m a
4559tapAsync f (Stream step1 state1) = Stream step TapInit
4563 -- In general, a Stop event would come equipped with the result
4564 -- of the fold. It is not used here but it would be useful in
4565 -- applicative and distribute.
4566 done <- fromConsumer svr
4567 when (not done) $ do
4568 liftIO $ withDiagMVar svr "teeToSVar: waiting to drain"
4569 $ takeMVar (outputDoorBellFromConsumer svr)
4573 liftIO $ sendStop svr Nothing
4574 -- drain/wait until a stop event arrives from the fold.
4577 {-# INLINE_LATE step #-}
4578 step gst TapInit = do
4579 sv <- newFoldSVar gst f
4580 return $ Skip (Tapping sv state1)
4582 step gst (Tapping sv st) = do
4586 done <- pushToFold sv a
4589 -- XXX we do not need to wait synchronously here
4591 return $ Yield a (TapDone s)
4592 else return $ Yield a (Tapping sv s)
4593 Skip s -> return $ Skip (Tapping sv s)
4598 step gst (TapDone st) = do
4601 Yield a s -> Yield a (TapDone s)
4602 Skip s -> Skip (TapDone s)
4605-- XXX Exported from Array again as this fold is specific to Array
4606-- | Take last 'n' elements from the stream and discard the rest.
4608lastN :: (Storable a, MonadIO m) => Int -> Fold m a (Array a)
4610 | n <= 0 = fmap (const mempty) FL.drain
4611 | otherwise = A.unsafeFreeze <$> Fold step initial done
4613 step (Tuple3' rb rh i) a = do
4614 rh1 <- liftIO $ RB.unsafeInsert rb rh a
4615 return $ Tuple3' rb rh1 (i + 1)
4616 initial = fmap (\(a, b) -> Tuple3' a b (0 :: Int)) $ liftIO $ RB.new n
4617 done (Tuple3' rb rh i) = do
4618 arr <- liftIO $ MA.newArray n
4619 foldFunc i rh snoc' arr rb
4620 snoc' b a = liftIO $ MA.unsafeSnoc b a
4622 | i < n = RB.unsafeFoldRingM
4623 | otherwise = RB.unsafeFoldRingFullM
4625------------------------------------------------------------------------------
4627------------------------------------------------------------------------------
4629-- XXX using getTime in the loop can be pretty expensive especially for
4630-- computations where iterations are lightweight. We have the following
4633-- 1) Run a timeout thread updating a flag asynchronously and check that
4634-- flag here, that way we can have a cheap termination check.
4636-- 2) Use COARSE clock to get time with lower resolution but more efficiently.
4638-- 3) Use rdtscp/rdtsc to get time directly from the processor, compute the
4639-- termination value of rdtsc in the beginning and then in each iteration just
4640-- get rdtsc and check if we should terminate.
4644 | TakeByTimeCheck st s
4645 | TakeByTimeYield st s
4647{-# INLINE_NORMAL takeByTime #-}
4648takeByTime :: (MonadIO m, TimeUnit64 t) => t -> Stream m a -> Stream m a
4649takeByTime duration (Stream step1 state1) = Stream step (TakeByTimeInit state1)
4652 lim = toRelTime64 duration
4654 {-# INLINE_LATE step #-}
4655 step _ (TakeByTimeInit _) | lim == 0 = return Stop
4656 step _ (TakeByTimeInit st) = do
4657 t0 <- liftIO $ getTime Monotonic
4658 return $ Skip (TakeByTimeYield st t0)
4659 step _ (TakeByTimeCheck st t0) = do
4660 t <- liftIO $ getTime Monotonic
4662 if diffAbsTime64 t t0 > lim
4664 else Skip (TakeByTimeYield st t0)
4665 step gst (TakeByTimeYield st t0) = do
4668 Yield x s -> Yield x (TakeByTimeCheck s t0)
4669 Skip s -> Skip (TakeByTimeCheck s t0)
4672data DropByTime st s x
4674 | DropByTimeGen st s
4675 | DropByTimeCheck st s x
4676 | DropByTimeYield st
4678{-# INLINE_NORMAL dropByTime #-}
4679dropByTime :: (MonadIO m, TimeUnit64 t) => t -> Stream m a -> Stream m a
4680dropByTime duration (Stream step1 state1) = Stream step (DropByTimeInit state1)
4683 lim = toRelTime64 duration
4685 {-# INLINE_LATE step #-}
4686 step _ (DropByTimeInit st) = do
4687 t0 <- liftIO $ getTime Monotonic
4688 return $ Skip (DropByTimeGen st t0)
4689 step gst (DropByTimeGen st t0) = do
4692 Yield x s -> Skip (DropByTimeCheck s t0 x)
4693 Skip s -> Skip (DropByTimeGen s t0)
4695 step _ (DropByTimeCheck st t0 x) = do
4696 t <- liftIO $ getTime Monotonic
4697 if diffAbsTime64 t t0 <= lim
4698 then return $ Skip $ DropByTimeGen st t0
4699 else return $ Yield x $ DropByTimeYield st
4700 step gst (DropByTimeYield st) = do
4703 Yield x s -> Yield x (DropByTimeYield s)
4704 Skip s -> Skip (DropByTimeYield s)
4707{-# INLINE updateTimeVar #-}
4708updateTimeVar :: Prim.IORef Int64 -> IO ()
4709updateTimeVar timeVar = do
4710 MicroSecond64 t <- fromAbsTime <$> getTime Monotonic
4711 Prim.modifyIORef' timeVar (const t)
4713{-# INLINE updateWithDelay #-}
4714updateWithDelay :: RealFrac a => a -> Prim.IORef Int64 -> IO ()
4715updateWithDelay precision timeVar = do
4716 threadDelay (delayTime precision)
4717 updateTimeVar timeVar
4721 -- Keep the minimum at least a millisecond to avoid high CPU usage
4722 {-# INLINE delayTime #-}
4724 | g' >= fromIntegral (maxBound :: Int) = maxBound
4726 | otherwise = round g'
4730 g' = g * 10 ^ (6 :: Int)
4732-- XXX we should move this to stream generation section of this file. Also, the
4733-- take/drop combinators above should be moved to filtering section.
4735{-# INLINE_NORMAL times #-}
4736times :: MonadAsync m => Double -> Stream m (AbsTime, RelTime64)
4737times g = Stream step Nothing
4741 {-# INLINE_LATE step #-}
4743 -- XXX note that this is safe only on a 64-bit machine. On a 32-bit
4744 -- machine a 64-bit 'Var' cannot be read consistently without a lock
4745 -- while another thread is writing to it.
4746 timeVar <- liftIO $ Prim.newIORef (0 :: Int64)
4747 liftIO $ updateTimeVar timeVar
4748 tid <- forkManaged $ liftIO $ forever (updateWithDelay g timeVar)
4749 a <- liftIO $ Prim.readIORef timeVar
4750 return $ Skip $ Just (timeVar, tid, a)
4752 step _ s@(Just (timeVar, _, t0)) = do
4753 a <- liftIO $ Prim.readIORef timeVar
4754 -- XXX we can perhaps use an AbsTime64 using a 64 bit Int for
4755 -- efficiency. or maybe we can use a representation using Double for
4756 -- floating precision time
4757 return $ Yield (toAbsTime (MicroSecond64 t0),
4758 (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