1{-# LANGUAGE BangPatterns #-}
3{-# LANGUAGE ConstraintKinds #-}
4{-# LANGUAGE ExistentialQuantification #-}
5{-# LANGUAGE FlexibleContexts #-}
6{-# LANGUAGE FlexibleInstances #-}
7{-# LANGUAGE MultiParamTypeClasses #-}
8{-# LANGUAGE PatternSynonyms #-}
9{-# LANGUAGE RecordWildCards #-}
10{-# LANGUAGE ScopedTypeVariables #-}
11{-# LANGUAGE ViewPatterns #-}
12{-# LANGUAGE RankNTypes #-}
13{-# LANGUAGE MagicHash #-}
15#if __GLASGOW_HASKELL__ >= 801
16{-# LANGUAGE TypeApplications #-}
22-- Module : Streamly.Internal.Data.Stream.StreamD
23-- Copyright : (c) 2018 Harendra Kumar
24-- (c) Roman Leshchinskiy 2008-2010
25-- (c) The University of Glasgow, 2009
28-- Maintainer : streamly@composewell.com
29-- Stability : experimental
32-- Direct style re-implementation of CPS style stream in StreamK module. The
33-- symbol or suffix 'D' in this module denotes the "Direct" style. GHC is able
34-- to INLINE and fuse direct style better, providing better performance than
38-- import qualified Streamly.Internal.Data.Stream.StreamD as D
41-- Some of the functions in this file have been adapted from the vector
42-- library, https://hackage.haskell.org/package/vector.
44module Streamly.Internal.Data.Stream.StreamD
49#if __GLASGOW_HASKELL__ >= 800
50 , Stream (Stream, UnStream)
70 -- ** Specialized Generation
71 -- | Generate a monadic stream from a seed.
84 , enumerateFromStepIntegral
85 , enumerateFromIntegral
86 , enumerateFromThenIntegral
87 , enumerateFromToIntegral
88 , enumerateFromThenToIntegral
90 , enumerateFromStepNum
93 , enumerateFromToFractional
94 , enumerateFromThenToFractional
100 -- | Transform an input structure into a stream.
101 -- | Direct style stream does not support @fromFoldable@.
135 -- ** Specialized Folds
162 -- ** Flattening nested streams
165 , ConcatMapUState (..)
167 , ConcatUnfoldInterleaveState (..)
168 , concatUnfoldInterleave
169 , concatUnfoldRoundrobin
172 , InterleaveState(..)
177 , roundRobin -- interleaveFair?/ParallelFair
210 -- | Transform a stream into another type.
227 -- ** By folding (scans)
278 , intersperseSuffixBySpan
302 -- * Transformation comprehensions
308 , clearIORefFinalizer
320 -- * Concurrent Application
329import Control.Concurrent (killThread, myThreadId, takeMVar, threadDelay)
330import Control.Exception
331 (assert, Exception, SomeException, AsyncException, fromException, mask_)
332import Control.Monad (void, when, forever)
333import Control.Monad.Catch (MonadCatch, MonadThrow, throwM)
334import Control.Monad.IO.Class (MonadIO(..))
335import Control.Monad.Trans.Reader (ReaderT)
336import Control.Monad.Trans.State.Strict (StateT)
337import Control.Monad.Trans.Class (MonadTrans(lift))
338import Control.Monad.Trans.Control (MonadBaseControl, liftBaseOp_)
339import Data.Bits (shiftR, shiftL, (.|.), (.&.))
340import Data.Functor.Identity (Identity(..))
341import Data.Int (Int64)
342import Data.IORef (newIORef, readIORef, mkWeakIORef, writeIORef, IORef)
343import Data.Maybe (fromJust, isJust, isNothing)
344import Data.Word (Word32)
345import Foreign.Ptr (Ptr)
346import Foreign.Storable (Storable(..))
347import GHC.Types (SPEC(..))
348import System.Mem (performMajorGC)
350 hiding (map, mapM, mapM_, repeat, foldr, last, take, filter,
351 takeWhile, drop, dropWhile, all, any, maximum, minimum, elem,
352 notElem, null, head, tail, zipWith, lookup, foldr1, sequence,
353 (!!), scanl, scanl1, concatMap, replicate, enumFromTo, concat,
354 reverse, iterate, splitAt)
356import qualified Control.Monad.Catch as MC
357import qualified Control.Monad.Trans.Reader as Reader
358import qualified Control.Monad.Trans.State.Strict as State
359import qualified Prelude
361import Fusion.Plugin.Types (Fuse(..))
362import Streamly.Internal.Mutable.Prim.Var
363 (Prim, Var, readVar, newVar, modifyVar')
364import Streamly.Internal.Data.Time.Units
365 (TimeUnit64, toRelTime64, diffAbsTime64, RelTime64)
367import Streamly.Internal.Data.Atomics (atomicModifyIORefCAS_)
368import Streamly.Internal.Memory.Array.Types (Array(..))
369import Streamly.Internal.Data.Fold.Types (Fold(..))
370import Streamly.Internal.Data.Parser (ParseError(..))
371import Streamly.Internal.Data.Pipe.Types (Pipe(..), PipeState(..))
372import Streamly.Internal.Data.Time.Clock (Clock(Monotonic), getTime)
373import Streamly.Internal.Data.Time.Units
374 (MicroSecond64(..), fromAbsTime, toAbsTime, AbsTime)
375import Streamly.Internal.Data.Unfold.Types (Unfold(..))
376import Streamly.Internal.Data.Strict (Tuple3'(..))
378import Streamly.Internal.Data.Stream.StreamD.Type
379import Streamly.Internal.Data.SVar
380import Streamly.Internal.Data.Stream.SVar (fromConsumer, pushToFold)
382import qualified Streamly.Internal.Data.Pipe.Types as Pipe
383import qualified Streamly.Internal.Memory.Array.Types as A
384import qualified Streamly.Internal.Data.Fold as FL
385import qualified Streamly.Memory.Ring as RB
386import qualified Streamly.Internal.Data.Stream.StreamK as K
387import qualified Streamly.Internal.Data.Parser as PR
388import qualified Streamly.Internal.Data.Parser.ParserD as PRD
390------------------------------------------------------------------------------
392------------------------------------------------------------------------------
394-- | An empty 'Stream'.
395{-# INLINE_NORMAL nil #-}
396nil :: Monad m => Stream m a
397nil = Stream (\_ _ -> return Stop) ()
399-- | An empty 'Stream' with a side effect.
400{-# INLINE_NORMAL nilM #-}
401nilM :: Monad m => m b -> Stream m a
402nilM m = Stream (\_ _ -> m >> return Stop) ()
404{-# INLINE_NORMAL consM #-}
405consM :: Monad m => m a -> Stream m a -> Stream m a
406consM m (Stream step state) = Stream step1 Nothing
408 {-# INLINE_LATE step1 #-}
409 step1 _ Nothing = m >>= \x -> return $ Yield x (Just state)
410 step1 gst (Just st) = do
414 Yield a s -> Yield a (Just s)
415 Skip s -> Skip (Just s)
418-- XXX implement in terms of consM?
419-- cons x = consM (return x)
421-- | Can fuse but has O(n^2) complexity.
422{-# INLINE_NORMAL cons #-}
423cons :: Monad m => a -> Stream m a -> Stream m a
424cons x (Stream step state) = Stream step1 Nothing
426 {-# INLINE_LATE step1 #-}
427 step1 _ Nothing = return $ Yield x (Just state)
428 step1 gst (Just st) = do
432 Yield a s -> Yield a (Just s)
433 Skip s -> Skip (Just s)
436-------------------------------------------------------------------------------
438-------------------------------------------------------------------------------
440-- Does not fuse, has the same performance as the StreamK version.
441{-# INLINE_NORMAL uncons #-}
442uncons :: Monad m => Stream m a -> m (Maybe (a, Stream m a))
443uncons (UnStream step state) = go state
446 r <- step defState st
448 Yield x s -> return $ Just (x, Stream step s)
450 Stop -> return Nothing
452------------------------------------------------------------------------------
453-- Generation by unfold
454------------------------------------------------------------------------------
456{-# INLINE_NORMAL unfoldrM #-}
457unfoldrM :: Monad m => (s -> m (Maybe (a, s))) -> s -> Stream m a
458unfoldrM next state = Stream step state
460 {-# INLINE_LATE step #-}
464 Just (x, s) -> Yield x s
467{-# INLINE_LATE unfoldr #-}
468unfoldr :: Monad m => (s -> Maybe (a, s)) -> s -> Stream m a
469unfoldr f = unfoldrM (return . f)
471-- | Convert an 'Unfold' into a 'Stream' by supplying it a seed.
473{-# INLINE_NORMAL unfold #-}
474unfold :: Monad m => Unfold m a b -> a -> Stream m b
475unfold (Unfold ustep inject) seed = Stream step Nothing
477 {-# INLINE_LATE step #-}
478 step _ Nothing = inject seed >>= return . Skip . Just
479 step _ (Just st) = do
482 Yield x s -> Yield x (Just s)
483 Skip s -> Skip (Just s)
486------------------------------------------------------------------------------
487-- Specialized Generation
488------------------------------------------------------------------------------
490{-# INLINE_NORMAL repeatM #-}
491repeatM :: Monad m => m a -> Stream m a
492repeatM x = Stream (\_ _ -> x >>= \r -> return $ Yield r ()) ()
494{-# INLINE_NORMAL repeat #-}
495repeat :: Monad m => a -> Stream m a
496repeat x = Stream (\_ _ -> return $ Yield x ()) ()
498{-# INLINE_NORMAL iterateM #-}
499iterateM :: Monad m => (a -> m a) -> m a -> Stream m a
500iterateM step = Stream (\_ st -> st >>= \x -> return $ Yield x (step x))
502{-# INLINE_NORMAL iterate #-}
503iterate :: Monad m => (a -> a) -> a -> Stream m a
504iterate step st = iterateM (return . step) (return st)
506{-# INLINE_NORMAL replicateM #-}
507replicateM :: forall m a. Monad m => Int -> m a -> Stream m a
508replicateM n p = Stream step n
510 {-# INLINE_LATE step #-}
512 | i <= 0 = return Stop
515 return $ Yield x (i - 1)
517{-# INLINE_NORMAL replicate #-}
518replicate :: Monad m => Int -> a -> Stream m a
519replicate n x = replicateM n (return x)
521-- This would not work properly for floats, therefore we put an Integral
523-- | Can be used to enumerate unbounded integrals. This does not check for
524-- overflow or underflow for bounded integrals.
525{-# INLINE_NORMAL enumerateFromStepIntegral #-}
526enumerateFromStepIntegral :: (Integral a, Monad m) => a -> a -> Stream m a
527enumerateFromStepIntegral from stride =
528 from `seq` stride `seq` Stream step from
530 {-# INLINE_LATE step #-}
531 step _ !x = return $ Yield x $! (x + stride)
533-- We are assuming that "to" is constrained by the type to be within
535{-# INLINE enumerateFromToIntegral #-}
536enumerateFromToIntegral :: (Monad m, Integral a) => a -> a -> Stream m a
537enumerateFromToIntegral from to =
538 takeWhile (<= to) $ enumerateFromStepIntegral from 1
540{-# INLINE enumerateFromIntegral #-}
541enumerateFromIntegral :: (Monad m, Integral a, Bounded a) => a -> Stream m a
542enumerateFromIntegral from = enumerateFromToIntegral from maxBound
544data EnumState a = EnumInit | EnumYield a a a | EnumStop
546{-# INLINE_NORMAL enumerateFromThenToIntegralUp #-}
547enumerateFromThenToIntegralUp
548 :: (Monad m, Integral a)
549 => a -> a -> a -> Stream m a
550enumerateFromThenToIntegralUp from next to = Stream step EnumInit
552 {-# INLINE_LATE step #-}
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 enumerateFromThenToIntegralDn #-}
572enumerateFromThenToIntegralDn
573 :: (Monad m, Integral a)
574 => a -> a -> a -> Stream m a
575enumerateFromThenToIntegralDn from next to = Stream step EnumInit
577 {-# INLINE_LATE step #-}
579 return $ if to > next
582 else Yield from EnumStop
583 else -- from >= next >= to
584 let stride = next - from
585 in Skip $ EnumYield from stride (to - stride)
587 step _ (EnumYield x stride toMinus) =
590 then Yield x EnumStop
591 else Yield x $ EnumYield (x + stride) stride toMinus
593 step _ EnumStop = return Stop
595{-# INLINE_NORMAL enumerateFromThenToIntegral #-}
596enumerateFromThenToIntegral
597 :: (Monad m, Integral a)
598 => a -> a -> a -> Stream m a
599enumerateFromThenToIntegral from next to
600 | next >= from = enumerateFromThenToIntegralUp from next to
601 | otherwise = enumerateFromThenToIntegralDn from next to
603{-# INLINE_NORMAL enumerateFromThenIntegral #-}
604enumerateFromThenIntegral
605 :: (Monad m, Integral a, Bounded a)
606 => a -> a -> Stream m a
607enumerateFromThenIntegral from next =
609 then enumerateFromThenToIntegralUp from next maxBound
610 else enumerateFromThenToIntegralDn from next minBound
612-- For floating point numbers if the increment is less than the precision then
613-- it just gets lost. Therefore we cannot always increment it correctly by just
615-- 9007199254740992 + 1 + 1 :: Double => 9.007199254740992e15
616-- 9007199254740992 + 2 :: Double => 9.007199254740994e15
618-- Instead we accumulate the increment counter and compute the increment
619-- every time before adding it to the starting number.
621-- This works for Integrals as well as floating point numbers, but
622-- enumerateFromStepIntegral is faster for integrals.
623{-# INLINE_NORMAL enumerateFromStepNum #-}
624enumerateFromStepNum :: (Monad m, Num a) => a -> a -> Stream m a
625enumerateFromStepNum from stride = Stream step 0
627 {-# INLINE_LATE step #-}
628 step _ !i = return $ (Yield $! (from + i * stride)) $! (i + 1)
630{-# INLINE_NORMAL numFrom #-}
631numFrom :: (Monad m, Num a) => a -> Stream m a
632numFrom from = enumerateFromStepNum from 1
634{-# INLINE_NORMAL numFromThen #-}
635numFromThen :: (Monad m, Num a) => a -> a -> Stream m a
636numFromThen from next = enumerateFromStepNum from (next - from)
638-- We cannot write a general function for Num. The only way to write code
639-- portable between the two is to use a 'Real' constraint and convert between
640-- Fractional and Integral using fromRational which is horribly slow.
641{-# INLINE_NORMAL enumerateFromToFractional #-}
642enumerateFromToFractional
643 :: (Monad m, Fractional a, Ord a)
644 => a -> a -> Stream m a
645enumerateFromToFractional from to =
646 takeWhile (<= to + 1 / 2) $ enumerateFromStepNum from 1
648{-# INLINE_NORMAL enumerateFromThenToFractional #-}
649enumerateFromThenToFractional
650 :: (Monad m, Fractional a, Ord a)
651 => a -> a -> a -> Stream m a
652enumerateFromThenToFractional from next to =
653 takeWhile predicate $ numFromThen from next
655 mid = (next - from) / 2
656 predicate | next >= from = (<= to + mid)
657 | otherwise = (>= to + mid)
659-------------------------------------------------------------------------------
660-- Generation by Conversion
661-------------------------------------------------------------------------------
663{-# INLINE_NORMAL fromIndicesM #-}
664fromIndicesM :: Monad m => (Int -> m a) -> Stream m a
665fromIndicesM gen = Stream step 0
667 {-# INLINE_LATE step #-}
670 return $ Yield x (i + 1)
672{-# INLINE fromIndices #-}
673fromIndices :: Monad m => (Int -> a) -> Stream m a
674fromIndices gen = fromIndicesM (return . gen)
676{-# INLINE_NORMAL generateM #-}
677generateM :: Monad m => Int -> (Int -> m a) -> Stream m a
678generateM n gen = n `seq` Stream step 0
680 {-# INLINE_LATE step #-}
681 step _ i | i < n = do
683 return $ Yield x (i + 1)
684 | otherwise = return Stop
686{-# INLINE generate #-}
687generate :: Monad m => Int -> (Int -> a) -> Stream m a
688generate n gen = generateM n (return . gen)
690-- XXX we need the MonadAsync constraint because of a rewrite rule.
691-- | Convert a list of monadic actions to a 'Stream'
692{-# INLINE_LATE fromListM #-}
693fromListM :: MonadAsync m => [m a] -> Stream m a
694fromListM = Stream step
696 {-# INLINE_LATE step #-}
697 step _ (m:ms) = m >>= \x -> return $ Yield x ms
698 step _ [] = return Stop
700{-# INLINE toStreamD #-}
701toStreamD :: (K.IsStream t, Monad m) => t m a -> Stream m a
702toStreamD = fromStreamK . K.toStream
704{-# INLINE_NORMAL fromPrimVar #-}
705fromPrimVar :: (MonadIO m, Prim a) => Var IO a -> Stream m a
706fromPrimVar var = Stream step ()
708 {-# INLINE_LATE step #-}
709 step _ () = liftIO (readVar var) >>= \x -> return $ Yield x ()
711-------------------------------------------------------------------------------
712-- Generation from SVar
713-------------------------------------------------------------------------------
715data FromSVarState t m a =
717 | FromSVarRead (SVar t m a)
718 | FromSVarLoop (SVar t m a) [ChildEvent a]
719 | FromSVarDone (SVar t m a)
721{-# INLINE_NORMAL fromSVar #-}
722fromSVar :: (MonadAsync m) => SVar t m a -> Stream m a
723fromSVar svar = Stream step FromSVarInit
726 {-# INLINE_LATE step #-}
727 step _ FromSVarInit = do
728 ref <- liftIO $ newIORef ()
729 _ <- liftIO $ mkWeakIORef ref hook
730 -- when this copy of svar gets garbage collected "ref" will get
731 -- garbage collected and our GC hook will be called.
732 let sv = svar{svarRef = Just ref}
733 return $ Skip (FromSVarRead sv)
737 {-# NOINLINE hook #-}
739 when (svarInspectMode svar) $ do
740 r <- liftIO $ readIORef (svarStopTime (svarStats svar))
742 printSVar svar "SVar Garbage Collected"
744 -- If there are any SVars referenced by this SVar a GC will prompt
745 -- them to be cleaned up quickly.
746 when (svarInspectMode svar) performMajorGC
748 step _ (FromSVarRead sv) = do
749 list <- readOutputQ sv
750 -- Reversing the output is important to guarantee that we process the
751 -- outputs in the same order as they were generated by the constituent
753 return $ Skip $ FromSVarLoop sv (Prelude.reverse list)
755 step _ (FromSVarLoop sv []) = do
756 done <- postProcess sv
757 return $ Skip $ if done
758 then (FromSVarDone sv)
759 else (FromSVarRead sv)
761 step _ (FromSVarLoop sv (ev : es)) = do
763 ChildYield a -> return $ Yield a (FromSVarLoop sv es)
764 ChildStop tid e -> do
768 stop <- shouldStop tid
771 liftIO (cleanupSVar sv)
772 return $ Skip (FromSVarDone sv)
773 else return $ Skip (FromSVarLoop sv es)
775 case fromException ex of
777 return $ Skip (FromSVarLoop sv es)
778 Nothing -> liftIO (cleanupSVar sv) >> throwM ex
782 case svarStopStyle sv of
783 StopNone -> return False
784 StopAny -> return True
786 sid <- liftIO $ readIORef (svarStopBy sv)
787 return $ if tid == sid then True else False
789 step _ (FromSVarDone sv) = do
790 when (svarInspectMode sv) $ do
791 t <- liftIO $ getTime Monotonic
792 liftIO $ writeIORef (svarStopTime (svarStats sv)) (Just t)
793 liftIO $ printSVar sv "SVar Done"
796-------------------------------------------------------------------------------
797-- Process events received by a fold consumer from a stream producer
798-------------------------------------------------------------------------------
800{-# INLINE_NORMAL fromProducer #-}
801fromProducer :: (MonadAsync m) => SVar t m a -> Stream m a
802fromProducer svar = Stream step (FromSVarRead svar)
805 {-# INLINE_LATE step #-}
806 step _ (FromSVarRead sv) = do
807 list <- readOutputQ sv
808 -- Reversing the output is important to guarantee that we process the
809 -- outputs in the same order as they were generated by the constituent
811 return $ Skip $ FromSVarLoop sv (Prelude.reverse list)
813 step _ (FromSVarLoop sv []) = return $ Skip $ FromSVarRead sv
814 step _ (FromSVarLoop sv (ev : es)) = do
816 ChildYield a -> return $ Yield a (FromSVarLoop sv es)
817 ChildStop tid e -> do
821 sendStopToProducer sv
822 return $ Skip (FromSVarDone sv)
823 Just _ -> error "Bug: fromProducer: received exception"
825 step _ (FromSVarDone sv) = do
826 when (svarInspectMode sv) $ do
827 t <- liftIO $ getTime Monotonic
828 liftIO $ writeIORef (svarStopTime (svarStats sv)) (Just t)
829 liftIO $ printSVar sv "SVar Done"
832 step _ FromSVarInit = undefined
834-------------------------------------------------------------------------------
835-- Hoisting the inner monad
836-------------------------------------------------------------------------------
838{-# INLINE_NORMAL hoist #-}
839hoist :: Monad n => (forall x. m x -> n x) -> Stream m a -> Stream n a
840hoist f (Stream step state) = (Stream step' state)
842 {-# INLINE_LATE step' #-}
844 r <- f $ step (adaptState gst) st
846 Yield x s -> Yield x s
850{-# INLINE generally #-}
851generally :: Monad m => Stream Identity a -> Stream m a
852generally = hoist (return . runIdentity)
854{-# INLINE_NORMAL liftInner #-}
855liftInner :: (Monad m, MonadTrans t, Monad (t m))
856 => Stream m a -> Stream (t m) a
857liftInner (Stream step state) = Stream step' state
859 {-# INLINE_LATE step' #-}
861 r <- lift $ step (adaptState gst) st
863 Yield x s -> Yield x s
867{-# INLINE_NORMAL runReaderT #-}
868runReaderT :: Monad m => s -> Stream (ReaderT s m) a -> Stream m a
869runReaderT sval (Stream step state) = Stream step' state
871 {-# INLINE_LATE step' #-}
873 r <- Reader.runReaderT (step (adaptState gst) st) sval
875 Yield x s -> Yield x s
879{-# INLINE_NORMAL evalStateT #-}
880evalStateT :: Monad m => s -> Stream (StateT s m) a -> Stream m a
881evalStateT sval (Stream step state) = Stream step' (state, sval)
883 {-# INLINE_LATE step' #-}
884 step' gst (st, sv) = do
885 (r, sv') <- State.runStateT (step (adaptState gst) st) sv
887 Yield x s -> Yield x (s, sv')
888 Skip s -> Skip (s, sv')
891{-# INLINE_NORMAL runStateT #-}
892runStateT :: Monad m => s -> Stream (StateT s m) a -> Stream m (s, a)
893runStateT sval (Stream step state) = Stream step' (state, sval)
895 {-# INLINE_LATE step' #-}
896 step' gst (st, sv) = do
897 (r, sv') <- State.runStateT (step (adaptState gst) st) sv
899 Yield x s -> Yield (sv', x) (s, sv')
900 Skip s -> Skip (s, sv')
903------------------------------------------------------------------------------
904-- Elimination by Folds
905------------------------------------------------------------------------------
907------------------------------------------------------------------------------
909------------------------------------------------------------------------------
911{-# INLINE_NORMAL foldr1 #-}
912foldr1 :: Monad m => (a -> a -> a) -> Stream m a -> m (Maybe a)
916 Nothing -> return Nothing
917 Just (h, t) -> fmap Just (foldr f h t)
919------------------------------------------------------------------------------
921------------------------------------------------------------------------------
923{-# INLINE_NORMAL foldlT #-}
924foldlT :: (Monad m, Monad (s m), MonadTrans s)
925 => (s m b -> a -> s m b) -> s m b -> Stream m a -> s m b
926foldlT fstep begin (Stream step state) = go SPEC begin state
929 r <- lift $ step defState st
931 Yield x s -> go SPEC (fstep acc x) s
932 Skip s -> go SPEC acc s
935-- Note, this is going to have horrible performance, because of the nature of
936-- the stream type (i.e. direct stream vs CPS). Its only for reference, it is
937-- likely be practically unusable.
938{-# INLINE_NORMAL foldlS #-}
940 => (Stream m b -> a -> Stream m b) -> Stream m b -> Stream m a -> Stream m b
941foldlS fstep begin (Stream step state) = Stream step' (Left (state, begin))
943 step' gst (Left (st, acc)) = do
944 r <- step (adaptState gst) st
946 Yield x s -> Skip (Left (s, fstep acc x))
947 Skip s -> Skip (Left (s, acc))
948 Stop -> Skip (Right acc)
950 step' gst (Right (Stream stp stt)) = do
951 r <- stp (adaptState gst) stt
953 Yield x s -> Yield x (Right (Stream stp s))
954 Skip s -> Skip (Right (Stream stp s))
957------------------------------------------------------------------------------
959------------------------------------------------------------------------------
961-- Inlined definition. Without the inline "serially/parser/take" benchmark
962-- degrades and parseMany does not fuse. Even using "inline" at the callsite
964{-# INLINE splitAt #-}
965splitAt :: Int -> [a] -> ([a],[a])
968 | otherwise = splitAt' n ls
970 splitAt' :: Int -> [a] -> ([a], [a])
971 splitAt' _ [] = ([], [])
972 splitAt' 1 (x:xs) = ([x], xs)
973 splitAt' m (x:xs) = (x:xs', xs'')
975 (xs', xs'') = splitAt' (m - 1) xs
977-- | Run a 'Parse' over a stream.
978{-# INLINE_NORMAL parselMx' #-}
981 => (s -> a -> m (PR.Step s b))
986parselMx' pstep initial extract (Stream step state) = do
987 initial >>= go SPEC state []
991 -- XXX currently we are using a dumb list based approach for backtracking
992 -- buffer. This can be replaced by a sliding/ring buffer using Data.Array.
993 -- That will allow us more efficient random back and forth movement.
995 go !_ st buf !pst = do
996 r <- step defState st
1001 PR.Partial 0 pst1 -> go SPEC s [] pst1
1002 PR.Partial n pst1 -> do
1003 assert (n <= length (x:buf)) (return ())
1004 let src0 = Prelude.take n (x:buf)
1005 src = Prelude.reverse src0
1006 gobuf SPEC s [] src pst1
1007 PR.Continue 0 pst1 -> go SPEC s (x:buf) pst1
1008 PR.Continue n pst1 -> do
1009 assert (n <= length (x:buf)) (return ())
1010 let (src0, buf1) = splitAt n (x:buf)
1011 src = Prelude.reverse src0
1012 gobuf SPEC s buf1 src pst1
1013 PR.Done _ b -> return b
1014 PR.Error err -> throwM $ ParseError err
1015 Skip s -> go SPEC s buf pst
1018 gobuf !_ s buf [] !pst = go SPEC s buf pst
1019 gobuf !_ s buf (x:xs) !pst = do
1022 PR.Partial 0 pst1 ->
1023 gobuf SPEC s [] xs pst1
1024 PR.Partial n pst1 -> do
1025 assert (n <= length (x:buf)) (return ())
1026 let src0 = Prelude.take n (x:buf)
1027 src = Prelude.reverse src0 ++ xs
1028 gobuf SPEC s [] src pst1
1029 PR.Continue 0 pst1 -> gobuf SPEC s (x:buf) xs pst1
1030 PR.Continue n pst1 -> do
1031 assert (n <= length (x:buf)) (return ())
1032 let (src0, buf1) = splitAt n (x:buf)
1033 src = Prelude.reverse src0 ++ xs
1034 gobuf SPEC s buf1 src pst1
1035 PR.Done _ b -> return b
1036 PR.Error err -> throwM $ ParseError err
1038------------------------------------------------------------------------------
1040------------------------------------------------------------------------------
1042{-# ANN type ParseChunksState Fuse #-}
1043data ParseChunksState x inpBuf st pst =
1044 ParseChunksInit inpBuf st
1045 | ParseChunksInitLeftOver inpBuf
1046 | ParseChunksStream st inpBuf pst
1047 | ParseChunksBuf inpBuf st inpBuf pst
1048 | ParseChunksYield x (ParseChunksState x inpBuf st pst)
1050{-# INLINE_NORMAL parseMany #-}
1056parseMany (PRD.Parser pstep initial extract) (Stream step state) =
1057 Stream stepOuter (ParseChunksInit [] state)
1061 {-# INLINE_LATE stepOuter #-}
1062 -- Buffer is empty, go to stream processing loop
1063 stepOuter _ (ParseChunksInit [] st) = do
1064 initial >>= return . Skip . ParseChunksStream st []
1066 -- Buffer is not empty, go to buffered processing loop
1067 stepOuter _ (ParseChunksInit src st) = do
1068 initial >>= return . Skip . ParseChunksBuf src st []
1070 -- XXX we just discard any leftover input at the end
1071 stepOuter _ (ParseChunksInitLeftOver _) = return Stop
1073 -- Buffer is empty process elements from the stream
1074 stepOuter gst (ParseChunksStream st buf pst) = do
1075 r <- step (adaptState gst) st
1080 PR.Partial 0 pst1 ->
1081 return $ Skip $ ParseChunksStream s [] pst1
1082 PR.Partial n pst1 -> do
1083 assert (n <= length (x:buf)) (return ())
1084 let src0 = Prelude.take n (x:buf)
1085 src = Prelude.reverse src0
1086 return $ Skip $ ParseChunksBuf src s [] pst1
1087 -- PR.Continue 0 pst1 ->
1088 -- return $ Skip $ ParseChunksStream s (x:buf) pst1
1089 PR.Continue n pst1 -> do
1090 assert (n <= length (x:buf)) (return ())
1091 let (src0, buf1) = splitAt n (x:buf)
1092 src = Prelude.reverse src0
1093 return $ Skip $ ParseChunksBuf src s buf1 pst1
1094 -- XXX Specialize for Stop 0 common case?
1096 assert (n <= length (x:buf)) (return ())
1097 let src = Prelude.reverse (Prelude.take n (x:buf))
1099 ParseChunksYield b (ParseChunksInit src s)
1100 PR.Error err -> throwM $ ParseError err
1101 Skip s -> return $ Skip $ ParseChunksStream s buf pst
1104 let src = Prelude.reverse buf
1106 ParseChunksYield b (ParseChunksInitLeftOver src)
1108 -- go back to stream processing mode
1109 stepOuter _ (ParseChunksBuf [] s buf pst) =
1110 return $ Skip $ ParseChunksStream s buf pst
1112 -- buffered processing loop
1113 stepOuter _ (ParseChunksBuf (x:xs) s buf pst) = do
1116 PR.Partial 0 pst1 ->
1117 return $ Skip $ ParseChunksBuf xs s [] pst1
1118 PR.Partial n pst1 -> do
1119 assert (n <= length (x:buf)) (return ())
1120 let src0 = Prelude.take n (x:buf)
1121 src = Prelude.reverse src0 ++ xs
1122 return $ Skip $ ParseChunksBuf src s [] pst1
1123 -- PR.Continue 0 pst1 -> return $ Skip $ ParseChunksBuf xs s (x:buf) pst1
1124 PR.Continue n pst1 -> do
1125 assert (n <= length (x:buf)) (return ())
1126 let (src0, buf1) = splitAt n (x:buf)
1127 src = Prelude.reverse src0 ++ xs
1128 return $ Skip $ ParseChunksBuf src s buf1 pst1
1129 -- XXX Specialize for Stop 0 common case?
1131 assert (n <= length (x:buf)) (return ())
1132 let src = Prelude.reverse (Prelude.take n (x:buf)) ++ xs
1133 return $ Skip $ ParseChunksYield b (ParseChunksInit src s)
1134 PR.Error err -> throwM $ ParseError err
1136 stepOuter _ (ParseChunksYield a next) = return $ Yield a next
1138{-# ANN type ConcatParseState Fuse #-}
1139data ConcatParseState x inpBuf st p =
1140 ConcatParseInit inpBuf st p
1141 | ConcatParseInitLeftOver inpBuf
1142 | ConcatParseStream st inpBuf p
1143 | ConcatParseBuf inpBuf st inpBuf p
1144 | ConcatParseYield x (ConcatParseState x inpBuf st p)
1146{-# INLINE_NORMAL parseIterate #-}
1149 => (b -> PRD.Parser m a b)
1153parseIterate func seed (Stream step state) =
1154 Stream stepOuter (ConcatParseInit [] state (func seed))
1158 {-# INLINE_LATE stepOuter #-}
1159 -- Buffer is empty, go to stream processing loop
1160 stepOuter _ (ConcatParseInit [] st (PRD.Parser pstep initial extract)) = do
1161 initial >>= \r -> return $ Skip $ ConcatParseStream st []
1162 (PRD.Parser pstep (return r) extract)
1164 -- Buffer is not empty, go to buffered processing loop
1165 stepOuter _ (ConcatParseInit src st
1166 (PRD.Parser pstep initial extract)) = do
1167 initial >>= \r -> return $ Skip $ ConcatParseBuf src st []
1168 (PRD.Parser pstep (return r) extract)
1170 -- XXX we just discard any leftover input at the end
1171 stepOuter _ (ConcatParseInitLeftOver _) = return Stop
1173 -- Buffer is empty process elements from the stream
1174 stepOuter gst (ConcatParseStream st buf
1175 p@(PRD.Parser pstep initial extract)) = do
1176 r <- step (adaptState gst) st
1182 PR.Partial 0 pst1 ->
1183 return $ Skip $ ConcatParseStream s []
1184 (PRD.Parser pstep (return pst1) extract)
1185 PR.Partial n pst1 -> do
1186 assert (n <= length (x:buf)) (return ())
1187 let src0 = Prelude.take n (x:buf)
1188 src = Prelude.reverse src0
1189 return $ Skip $ ConcatParseBuf src s []
1190 (PRD.Parser pstep (return pst1) extract)
1191 -- PR.Continue 0 pst1 ->
1192 -- return $ Skip $ ConcatParseStream s (x:buf) pst1
1193 PR.Continue n pst1 -> do
1194 assert (n <= length (x:buf)) (return ())
1195 let (src0, buf1) = splitAt n (x:buf)
1196 src = Prelude.reverse src0
1197 return $ Skip $ ConcatParseBuf src s buf1
1198 (PRD.Parser pstep (return pst1) extract)
1199 -- XXX Specialize for Stop 0 common case?
1201 assert (n <= length (x:buf)) (return ())
1202 let src = Prelude.reverse (Prelude.take n (x:buf))
1204 ConcatParseYield b (ConcatParseInit src s (func b))
1205 PR.Error err -> throwM $ ParseError err
1206 Skip s -> return $ Skip $ ConcatParseStream s buf p
1210 let src = Prelude.reverse buf
1212 ConcatParseYield b (ConcatParseInitLeftOver src)
1214 -- go back to stream processing mode
1215 stepOuter _ (ConcatParseBuf [] s buf p) =
1216 return $ Skip $ ConcatParseStream s buf p
1218 -- buffered processing loop
1219 stepOuter _ (ConcatParseBuf (x:xs) s buf
1220 (PRD.Parser pstep initial extract)) = do
1224 PR.Partial 0 pst1 ->
1225 return $ Skip $ ConcatParseBuf xs s []
1226 (PRD.Parser pstep (return pst1) extract)
1227 PR.Partial n pst1 -> do
1228 assert (n <= length (x:buf)) (return ())
1229 let src0 = Prelude.take n (x:buf)
1230 src = Prelude.reverse src0 ++ xs
1231 return $ Skip $ ConcatParseBuf src s []
1232 (PRD.Parser pstep (return pst1) extract)
1233 -- PR.Continue 0 pst1 -> return $ Skip $ ConcatParseBuf xs s (x:buf) pst1
1234 PR.Continue n pst1 -> do
1235 assert (n <= length (x:buf)) (return ())
1236 let (src0, buf1) = splitAt n (x:buf)
1237 src = Prelude.reverse src0 ++ xs
1238 return $ Skip $ ConcatParseBuf src s buf1
1239 (PRD.Parser pstep (return pst1) extract)
1240 -- XXX Specialize for Stop 0 common case?
1242 assert (n <= length (x:buf)) (return ())
1243 let src = Prelude.reverse (Prelude.take n (x:buf)) ++ xs
1244 return $ Skip $ ConcatParseYield b
1245 (ConcatParseInit src s (func b))
1246 PR.Error err -> throwM $ ParseError err
1248 stepOuter _ (ConcatParseYield a next) = return $ Yield a next
1250------------------------------------------------------------------------------
1252------------------------------------------------------------------------------
1254-- | Run a streaming composition, discard the results.
1255{-# INLINE_LATE drain #-}
1256drain :: Monad m => Stream m a -> m ()
1257-- drain = foldrM (\_ xs -> xs) (return ())
1258drain (Stream step state) = go SPEC state
1261 r <- step defState st
1263 Yield _ s -> go SPEC s
1267{-# INLINE_NORMAL null #-}
1268null :: Monad m => Stream m a -> m Bool
1269null m = foldrM (\_ _ -> return False) (return True) m
1271{-# INLINE_NORMAL head #-}
1272head :: Monad m => Stream m a -> m (Maybe a)
1273head m = foldrM (\x _ -> return (Just x)) (return Nothing) m
1275{-# INLINE_NORMAL headElse #-}
1276headElse :: Monad m => a -> Stream m a -> m a
1277headElse a m = foldrM (\x _ -> return x) (return a) m
1279-- Does not fuse, has the same performance as the StreamK version.
1280{-# INLINE_NORMAL tail #-}
1281tail :: Monad m => Stream m a -> m (Maybe (Stream m a))
1282tail (UnStream step state) = go state
1285 r <- step defState st
1287 Yield _ s -> return (Just $ Stream step s)
1289 Stop -> return Nothing
1291-- XXX will it fuse? need custom impl?
1292{-# INLINE_NORMAL last #-}
1293last :: Monad m => Stream m a -> m (Maybe a)
1294last = foldl' (\_ y -> Just y) Nothing
1296{-# INLINE_NORMAL elem #-}
1297elem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
1298-- elem e m = foldrM (\x xs -> if x == e then return True else xs) (return False) m
1299elem e (Stream step state) = go state
1302 r <- step defState st
1305 | x == e -> return True
1308 Stop -> return False
1310{-# INLINE_NORMAL notElem #-}
1311notElem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
1312notElem e s = fmap not (elem e s)
1314{-# INLINE_NORMAL all #-}
1315all :: Monad m => (a -> Bool) -> Stream m a -> m Bool
1316-- all p m = foldrM (\x xs -> if p x then xs else return False) (return True) m
1317all p (Stream step state) = go state
1320 r <- step defState st
1324 | otherwise -> return False
1328{-# INLINE_NORMAL any #-}
1329any :: Monad m => (a -> Bool) -> Stream m a -> m Bool
1330-- any p m = foldrM (\x xs -> if p x then return True else xs) (return False) m
1331any p (Stream step state) = go state
1334 r <- step defState st
1337 | p x -> return True
1340 Stop -> return False
1342{-# INLINE_NORMAL maximum #-}
1343maximum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
1344maximum (Stream step state) = go Nothing state
1347 r <- step defState st
1349 Yield x s -> go (Just x) s
1350 Skip s -> go Nothing s
1351 Stop -> return Nothing
1352 go (Just acc) st = do
1353 r <- step defState st
1356 | acc <= x -> go (Just x) s
1357 | otherwise -> go (Just acc) s
1358 Skip s -> go (Just acc) s
1359 Stop -> return (Just acc)
1361{-# INLINE_NORMAL maximumBy #-}
1362maximumBy :: Monad m => (a -> a -> Ordering) -> Stream m a -> m (Maybe a)
1363maximumBy cmp (Stream step state) = go Nothing state
1366 r <- step defState st
1368 Yield x s -> go (Just x) s
1369 Skip s -> go Nothing s
1370 Stop -> return Nothing
1371 go (Just acc) st = do
1372 r <- step defState st
1374 Yield x s -> case cmp acc x of
1375 GT -> go (Just acc) s
1377 Skip s -> go (Just acc) s
1378 Stop -> return (Just acc)
1380{-# INLINE_NORMAL minimum #-}
1381minimum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
1382minimum (Stream step state) = go Nothing state
1385 r <- step defState st
1387 Yield x s -> go (Just x) s
1388 Skip s -> go Nothing s
1389 Stop -> return Nothing
1390 go (Just acc) st = do
1391 r <- step defState st
1394 | acc <= x -> go (Just acc) s
1395 | otherwise -> go (Just x) s
1396 Skip s -> go (Just acc) s
1397 Stop -> return (Just acc)
1399{-# INLINE_NORMAL minimumBy #-}
1400minimumBy :: Monad m => (a -> a -> Ordering) -> Stream m a -> m (Maybe a)
1401minimumBy cmp (Stream step state) = go Nothing state
1404 r <- step defState st
1406 Yield x s -> go (Just x) s
1407 Skip s -> go Nothing s
1408 Stop -> return Nothing
1409 go (Just acc) st = do
1410 r <- step defState st
1412 Yield x s -> case cmp acc x of
1414 _ -> go (Just acc) s
1415 Skip s -> go (Just acc) s
1416 Stop -> return (Just acc)
1418{-# INLINE_NORMAL (!!) #-}
1419(!!) :: (Monad m) => Stream m a -> Int -> m (Maybe a)
1420(Stream step state) !! i = go i state
1423 r <- step defState st
1425 Yield x s | n < 0 -> return Nothing
1426 | n == 0 -> return $ Just x
1427 | otherwise -> go (n - 1) s
1429 Stop -> return Nothing
1431{-# INLINE_NORMAL lookup #-}
1432lookup :: (Monad m, Eq a) => a -> Stream m (a, b) -> m (Maybe b)
1433lookup e m = foldrM (\(a, b) xs -> if e == a then return (Just b) else xs)
1436{-# INLINE_NORMAL findM #-}
1437findM :: Monad m => (a -> m Bool) -> Stream m a -> m (Maybe a)
1438findM p m = foldrM (\x xs -> p x >>= \r -> if r then return (Just x) else xs)
1442find :: Monad m => (a -> Bool) -> Stream m a -> m (Maybe a)
1443find p = findM (return . p)
1445{-# INLINE_NORMAL findIndices #-}
1446findIndices :: Monad m => (a -> Bool) -> Stream m a -> Stream m Int
1447findIndices p (Stream step state) = Stream step' (state, 0)
1449 {-# INLINE_LATE step' #-}
1450 step' gst (st, i) = i `seq` do
1451 r <- step (adaptState gst) st
1453 Yield x s -> if p x then Yield i (s, i+1) else Skip (s, i+1)
1454 Skip s -> Skip (s, i)
1457{-# INLINE toListRev #-}
1458toListRev :: Monad m => Stream m a -> m [a]
1459toListRev = foldl' (flip (:)) []
1461-- We can implement reverse as:
1463-- > reverse = foldlS (flip cons) nil
1465-- However, this implementation is unusable because of the horrible performance
1466-- of cons. So we just convert it to a list first and then stream from the
1469-- XXX Maybe we can use an Array instead of a list here?
1470{-# INLINE_NORMAL reverse #-}
1471reverse :: Monad m => Stream m a -> Stream m a
1472reverse m = Stream step Nothing
1474 {-# INLINE_LATE step #-}
1477 return $ Skip (Just xs)
1478 step _ (Just (x:xs)) = return $ Yield x (Just xs)
1479 step _ (Just []) = return Stop
1481-- Much faster reverse for Storables
1482{-# INLINE_NORMAL reverse' #-}
1483reverse' :: forall m a. (MonadIO m, Storable a) => Stream m a -> Stream m a
1485-- This commented implementation copies the whole stream into one single array
1486-- and then streams from that array, this is 3-4x faster than the chunked code
1487-- that follows. Though this could be problematic due to unbounded large
1488-- allocations. We need to figure out why the chunked code is slower and if we
1489-- can optimize the chunked code to work as fast as this one. It may be a
1491import Foreign.ForeignPtr (touchForeignPtr)
1492import Foreign.ForeignPtr.Unsafe (unsafeForeignPtrToPtr)
1493import Foreign.Ptr (Ptr, plusPtr)
1494reverse' m = Stream step Nothing
1496 {-# INLINE_LATE step #-}
1498 arr <- A.fromStreamD m
1499 let p = aEnd arr `plusPtr` negate (sizeOf (undefined :: a))
1500 return $ Skip $ Just (aStart arr, p)
1502 step _ (Just (start, p)) | p < unsafeForeignPtrToPtr start = return Stop
1504 step _ (Just (start, p)) = do
1505 let !x = A.unsafeInlineIO $ do
1507 touchForeignPtr start
1509 next = p `plusPtr` negate (sizeOf (undefined :: a))
1510 return $ Yield x (Just (start, next))
1517 $ A.fromStreamDArraysOf A.defaultChunkSize m
1520------------------------------------------------------------------------------
1521-- Grouping/Splitting
1522------------------------------------------------------------------------------
1524{-# INLINE_NORMAL splitSuffixBy' #-}
1525splitSuffixBy' :: Monad m
1526 => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1527splitSuffixBy' predicate f (Stream step state) =
1528 Stream (stepOuter f) (Just state)
1532 {-# INLINE_LATE stepOuter #-}
1533 stepOuter (Fold fstep initial done) gst (Just st) = do
1534 res <- step (adaptState gst) st
1540 then done acc' >>= \val -> return $ Yield val (Just s)
1543 Skip s -> return $ Skip $ Just s
1549 res <- step (adaptState gst) stt
1554 then done acc' >>= \val -> return $ Yield val (Just s)
1556 Skip s -> go SPEC s acc
1557 Stop -> done acc >>= \val -> return $ Yield val Nothing
1559 stepOuter _ _ Nothing = return Stop
1561{-# INLINE_NORMAL groupsBy #-}
1567groupsBy cmp f (Stream step state) = Stream (stepOuter f) (Just state, Nothing)
1571 {-# INLINE_LATE stepOuter #-}
1572 stepOuter (Fold fstep initial done) gst (Just st, Nothing) = do
1573 res <- step (adaptState gst) st
1580 Skip s -> return $ Skip $ (Just s, Nothing)
1585 go !_ prev stt !acc = do
1586 res <- step (adaptState gst) stt
1593 else done acc >>= \r -> return $ Yield r (Just s, Just x)
1594 Skip s -> go SPEC prev s acc
1595 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1597 stepOuter (Fold fstep initial done) gst (Just st, Just prev) = do
1599 acc' <- fstep acc prev
1604 -- XXX code duplicated from the previous equation
1606 res <- step (adaptState gst) stt
1613 else done acc >>= \r -> return $ Yield r (Just s, Just x)
1614 Skip s -> go SPEC s acc
1615 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1617 stepOuter _ _ (Nothing,_) = return Stop
1619{-# INLINE_NORMAL groupsRollingBy #-}
1620groupsRollingBy :: Monad m
1625groupsRollingBy cmp f (Stream step state) =
1626 Stream (stepOuter f) (Just state, Nothing)
1629 {-# INLINE_LATE stepOuter #-}
1630 stepOuter (Fold fstep initial done) gst (Just st, Nothing) = do
1631 res <- step (adaptState gst) st
1638 Skip s -> return $ Skip $ (Just s, Nothing)
1642 go !_ prev stt !acc = do
1643 res <- step (adaptState gst) stt
1651 done acc >>= \r -> return $ Yield r (Just s, Just x)
1652 Skip s -> go SPEC prev s acc
1653 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1655 stepOuter (Fold fstep initial done) gst (Just st, Just prev') = do
1657 acc' <- fstep acc prev'
1658 go SPEC prev' st acc'
1661 go !_ prevv stt !acc = do
1662 res <- step (adaptState gst) stt
1669 else done acc >>= \r -> return $ Yield r (Just s, Just x)
1670 Skip s -> go SPEC prevv s acc
1671 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1673 stepOuter _ _ (Nothing, _) = return Stop
1675{-# INLINE_NORMAL splitBy #-}
1676splitBy :: Monad m => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1677splitBy predicate f (Stream step state) = Stream (step' f) (Just state)
1681 {-# INLINE_LATE step' #-}
1682 step' (Fold fstep initial done) gst (Just st) = initial >>= go SPEC st
1687 res <- step (adaptState gst) stt
1691 then done acc >>= \r -> return $ Yield r (Just s)
1695 Skip s -> go SPEC s acc
1696 Stop -> done acc >>= \r -> return $ Yield r Nothing
1698 step' _ _ Nothing = return Stop
1700-- XXX requires -funfolding-use-threshold=150 in lines-unlines benchmark
1701{-# INLINE_NORMAL splitSuffixBy #-}
1702splitSuffixBy :: Monad m
1703 => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1704splitSuffixBy predicate f (Stream step state) = Stream (step' f) (Just state)
1708 {-# INLINE_LATE step' #-}
1709 step' (Fold fstep initial done) gst (Just st) = do
1710 res <- step (adaptState gst) st
1715 then done acc >>= \val -> return $ Yield val (Just s)
1720 Skip s -> return $ Skip $ Just s
1726 res <- step (adaptState gst) stt
1730 then done acc >>= \r -> return $ Yield r (Just s)
1734 Skip s -> go SPEC s acc
1735 Stop -> done acc >>= \r -> return $ Yield r Nothing
1737 step' _ _ Nothing = return Stop
1739{-# INLINE_NORMAL wordsBy #-}
1740wordsBy :: Monad m => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1741wordsBy predicate f (Stream step state) = Stream (stepOuter f) (Just state)
1745 {-# INLINE_LATE stepOuter #-}
1746 stepOuter (Fold fstep initial done) gst (Just st) = do
1747 res <- step (adaptState gst) st
1751 then return $ Skip (Just s)
1757 Skip s -> return $ Skip $ Just s
1763 res <- step (adaptState gst) stt
1767 then done acc >>= \r -> return $ Yield r (Just s)
1771 Skip s -> go SPEC s acc
1772 Stop -> done acc >>= \r -> return $ Yield r Nothing
1774 stepOuter _ _ Nothing = return Stop
1776-- String search algorithms:
1777-- http://www-igm.univ-mlv.fr/~lecroq/string/index.html
1780-- TODO can we unify the splitting operations using a splitting configuration
1781-- like in the split package.
1783data SplitStyle = Infix | Suffix | Prefix deriving (Eq, Show)
1785data SplitOptions = SplitOptions
1786 { style :: SplitStyle
1787 , withSep :: Bool -- ^ keep the separators in output
1788 -- , compact :: Bool -- ^ treat multiple consecutive separators as one
1789 -- , trimHead :: Bool -- ^ drop blank at head
1790 -- , trimTail :: Bool -- ^ drop blank at tail
1794data SplitOnState s a =
1799 | GO_KARP_RABIN s !(RB.Ring a) !(Ptr a)
1802{-# INLINE_NORMAL splitOn #-}
1804 :: forall m a b. (MonadIO m, Storable a, Enum a, Eq a)
1809splitOn patArr (Fold fstep initial done) (Stream step state) =
1810 Stream stepOuter GO_START
1814 patLen = A.length patArr
1815 maxIndex = patLen - 1
1816 elemBits = sizeOf (undefined :: a) * 8
1818 {-# INLINE_LATE stepOuter #-}
1819 stepOuter _ GO_START =
1821 then return $ Skip $ GO_EMPTY_PAT state
1824 r <- liftIO $ (A.unsafeIndexIO patArr 0)
1825 return $ Skip $ GO_SINGLE_PAT state r
1826 else if sizeOf (undefined :: a) * patLen
1827 <= sizeOf (undefined :: Word)
1828 then return $ Skip $ GO_SHORT_PAT state
1830 (rb, rhead) <- liftIO $ RB.new patLen
1831 return $ Skip $ GO_KARP_RABIN state rb rhead
1833 stepOuter gst (GO_SINGLE_PAT stt pat) = initial >>= go SPEC stt
1838 res <- step (adaptState gst) st
1844 return $ Yield r (GO_SINGLE_PAT s pat)
1845 else fstep acc x >>= go SPEC s
1846 Skip s -> go SPEC s acc
1847 Stop -> done acc >>= \r -> return $ Yield r GO_DONE
1849 stepOuter gst (GO_SHORT_PAT stt) = initial >>= go0 SPEC 0 (0 :: Word) stt
1854 mask = (1 `shiftL` (elemBits * patLen)) - 1
1856 addToWord wrd a = (wrd `shiftL` elemBits) .|. fromIntegral (fromEnum a)
1859 patWord = mask .&. A.foldl' addToWord 0 patArr
1861 go0 !_ !idx wrd st !acc = do
1862 res <- step (adaptState gst) st
1865 let wrd' = addToWord wrd x
1868 if wrd' .&. mask == patWord
1871 return $ Yield r (GO_SHORT_PAT s)
1872 else go1 SPEC wrd' s acc
1873 else go0 SPEC (idx + 1) wrd' s acc
1874 Skip s -> go0 SPEC idx wrd s acc
1877 then go2 wrd idx acc
1879 done acc' >>= \r -> return $ Yield r GO_DONE
1882 go1 !_ wrd st !acc = do
1883 res <- step (adaptState gst) st
1886 let wrd' = addToWord wrd x
1887 old = (mask .&. wrd) `shiftR` (elemBits * (patLen - 1))
1888 acc' <- fstep acc (toEnum $ fromIntegral old)
1889 if wrd' .&. mask == patWord
1890 then done acc' >>= \r -> return $ Yield r (GO_SHORT_PAT s)
1891 else go1 SPEC wrd' s acc'
1892 Skip s -> go1 SPEC wrd s acc
1894 acc' <- go2 wrd patLen acc
1895 done acc' >>= \r -> return $ Yield r GO_DONE
1897 go2 !wrd !n !acc | n > 0 = do
1898 let old = (mask .&. wrd) `shiftR` (elemBits * (n - 1))
1899 fstep acc (toEnum $ fromIntegral old) >>= go2 wrd (n - 1)
1900 go2 _ _ acc = return acc
1902 stepOuter gst (GO_KARP_RABIN stt rb rhead) = do
1903 initial >>= go0 SPEC 0 rhead stt
1907 k = 2891336453 :: Word32
1909 addCksum cksum a = cksum * k + fromIntegral (fromEnum a)
1910 deltaCksum cksum old new =
1911 addCksum cksum new - coeff * fromIntegral (fromEnum old)
1913 -- XXX shall we use a random starting hash or 1 instead of 0?
1914 patHash = A.foldl' addCksum 0 patArr
1917 go0 !_ !idx !rh st !acc = do
1918 res <- step (adaptState gst) st
1921 rh' <- liftIO $ RB.unsafeInsert rb rh x
1924 let fold = RB.unsafeFoldRing (RB.ringBound rb)
1925 let !ringHash = fold addCksum 0 rb
1926 if ringHash == patHash
1927 then go2 SPEC ringHash rh' s acc
1928 else go1 SPEC ringHash rh' s acc
1929 else go0 SPEC (idx + 1) rh' s acc
1930 Skip s -> go0 SPEC idx rh s acc
1932 !acc' <- if idx /= 0
1933 then RB.unsafeFoldRingM rh fstep acc rb
1935 done acc' >>= \r -> return $ Yield r GO_DONE
1937 -- XXX Theoretically this code can do 4 times faster if GHC generates
1938 -- optimal code. If we use just "(cksum' == patHash)" condition it goes
1939 -- 4x faster, as soon as we add the "RB.unsafeEqArray rb v" condition
1940 -- the generated code changes drastically and becomes 4x slower. Need
1941 -- to investigate what is going on with GHC.
1943 go1 !_ !cksum !rh st !acc = do
1944 res <- step (adaptState gst) st
1947 old <- liftIO $ peek rh
1948 let cksum' = deltaCksum cksum old x
1949 acc' <- fstep acc old
1951 if (cksum' == patHash)
1953 rh' <- liftIO (RB.unsafeInsert rb rh x)
1954 go2 SPEC cksum' rh' s acc'
1956 rh' <- liftIO (RB.unsafeInsert rb rh x)
1957 go1 SPEC cksum' rh' s acc'
1958 Skip s -> go1 SPEC cksum rh s acc
1960 acc' <- RB.unsafeFoldRingFullM rh fstep acc rb
1961 done acc' >>= \r -> return $ Yield r GO_DONE
1963 go2 !_ !cksum' !rh' s !acc' = do
1964 if RB.unsafeEqArray rb rh' patArr
1967 return $ Yield r (GO_KARP_RABIN s rb rhead)
1968 else go1 SPEC cksum' rh' s acc'
1970 stepOuter gst (GO_EMPTY_PAT st) = do
1971 res <- step (adaptState gst) st
1976 done acc' >>= \r -> return $ Yield r (GO_EMPTY_PAT s)
1977 Skip s -> return $ Skip (GO_EMPTY_PAT s)
1980 stepOuter _ GO_DONE = return Stop
1982{-# INLINE_NORMAL splitSuffixOn #-}
1984 :: forall m a b. (MonadIO m, Storable a, Enum a, Eq a)
1990splitSuffixOn withSep patArr (Fold fstep initial done)
1991 (Stream step state) =
1992 Stream stepOuter GO_START
1996 patLen = A.length patArr
1997 maxIndex = patLen - 1
1998 elemBits = sizeOf (undefined :: a) * 8
2000 {-# INLINE_LATE stepOuter #-}
2001 stepOuter _ GO_START =
2003 then return $ Skip $ GO_EMPTY_PAT state
2006 r <- liftIO $ (A.unsafeIndexIO patArr 0)
2007 return $ Skip $ GO_SINGLE_PAT state r
2008 else if sizeOf (undefined :: a) * patLen
2009 <= sizeOf (undefined :: Word)
2010 then return $ Skip $ GO_SHORT_PAT state
2012 (rb, rhead) <- liftIO $ RB.new patLen
2013 return $ Skip $ GO_KARP_RABIN state rb rhead
2015 stepOuter gst (GO_SINGLE_PAT stt pat) = do
2016 -- This first part is the only difference between splitOn and
2018 -- If the last element is a separator do not issue a blank segment.
2019 res <- step (adaptState gst) stt
2025 acc' <- if withSep then fstep acc x else return acc
2026 done acc' >>= \r -> return $ Yield r (GO_SINGLE_PAT s pat)
2027 else fstep acc x >>= go SPEC s
2028 Skip s -> return $ Skip $ (GO_SINGLE_PAT s pat)
2033 -- This is identical for splitOn and splitSuffixOn
2035 res <- step (adaptState gst) st
2040 acc' <- if withSep then fstep acc x else return acc
2042 return $ Yield r (GO_SINGLE_PAT s pat)
2043 else fstep acc x >>= go SPEC s
2044 Skip s -> go SPEC s acc
2045 Stop -> done acc >>= \r -> return $ Yield r GO_DONE
2047 stepOuter gst (GO_SHORT_PAT stt) = do
2049 -- Call "initial" only if the stream yields an element, otherwise we
2050 -- may call "initial" but never yield anything. initial may produce a
2051 -- side effect, therefore we will end up doing and discard a side
2056 res <- step (adaptState gst) stt
2060 let wrd' = addToWord wrd x
2061 acc' <- if withSep then fstep acc x else return acc
2064 if wrd' .&. mask == patWord
2065 then done acc' >>= \r -> return $ Yield r (GO_SHORT_PAT s)
2066 else go0 SPEC (idx + 1) wrd' s acc'
2067 else go0 SPEC (idx + 1) wrd' s acc'
2068 Skip s -> return $ Skip (GO_SHORT_PAT s)
2074 mask = (1 `shiftL` (elemBits * patLen)) - 1
2076 addToWord wrd a = (wrd `shiftL` elemBits) .|. fromIntegral (fromEnum a)
2079 patWord = mask .&. A.foldl' addToWord 0 patArr
2081 go0 !_ !idx wrd st !acc = do
2082 res <- step (adaptState gst) st
2085 let wrd' = addToWord wrd x
2086 acc' <- if withSep then fstep acc x else return acc
2089 if wrd' .&. mask == patWord
2092 return $ Yield r (GO_SHORT_PAT s)
2093 else go1 SPEC wrd' s acc'
2094 else go0 SPEC (idx + 1) wrd' s acc'
2095 Skip s -> go0 SPEC idx wrd s acc
2097 if (idx == maxIndex) && (wrd .&. mask == patWord)
2100 acc' <- if idx /= 0 && not withSep
2101 then go2 wrd idx acc
2103 done acc' >>= \r -> return $ Yield r GO_DONE
2106 go1 !_ wrd st !acc = do
2107 res <- step (adaptState gst) st
2110 let wrd' = addToWord wrd x
2111 old = (mask .&. wrd) `shiftR` (elemBits * (patLen - 1))
2114 else fstep acc (toEnum $ fromIntegral old)
2115 if wrd' .&. mask == patWord
2116 then done acc' >>= \r -> return $ Yield r (GO_SHORT_PAT s)
2117 else go1 SPEC wrd' s acc'
2118 Skip s -> go1 SPEC wrd s acc
2120 -- If the last sequence is a separator do not issue a blank
2122 if wrd .&. mask == patWord
2127 else go2 wrd patLen acc
2128 done acc' >>= \r -> return $ Yield r GO_DONE
2130 go2 !wrd !n !acc | n > 0 = do
2131 let old = (mask .&. wrd) `shiftR` (elemBits * (n - 1))
2132 fstep acc (toEnum $ fromIntegral old) >>= go2 wrd (n - 1)
2133 go2 _ _ acc = return acc
2135 stepOuter gst (GO_KARP_RABIN stt rb rhead) = do
2137 res <- step (adaptState gst) stt
2141 acc' <- if withSep then fstep acc x else return acc
2142 rh' <- liftIO (RB.unsafeInsert rb rhead x)
2145 let fold = RB.unsafeFoldRing (RB.ringBound rb)
2146 let !ringHash = fold addCksum 0 rb
2147 if ringHash == patHash
2148 then go2 SPEC ringHash rh' s acc'
2149 else go0 SPEC (idx + 1) rh' s acc'
2150 else go0 SPEC (idx + 1) rh' s acc'
2151 Skip s -> return $ Skip (GO_KARP_RABIN s rb rhead)
2156 k = 2891336453 :: Word32
2158 addCksum cksum a = cksum * k + fromIntegral (fromEnum a)
2159 deltaCksum cksum old new =
2160 addCksum cksum new - coeff * fromIntegral (fromEnum old)
2162 -- XXX shall we use a random starting hash or 1 instead of 0?
2163 patHash = A.foldl' addCksum 0 patArr
2166 go0 !_ !idx !rh st !acc = do
2167 res <- step (adaptState gst) st
2170 acc' <- if withSep then fstep acc x else return acc
2171 rh' <- liftIO (RB.unsafeInsert rb rh x)
2174 let fold = RB.unsafeFoldRing (RB.ringBound rb)
2175 let !ringHash = fold addCksum 0 rb
2176 if ringHash == patHash
2177 then go2 SPEC ringHash rh' s acc'
2178 else go1 SPEC ringHash rh' s acc'
2179 else go0 SPEC (idx + 1) rh' s acc'
2180 Skip s -> go0 SPEC idx rh s acc
2182 -- do not issue a blank segment when we end at pattern
2183 if (idx == maxIndex) && RB.unsafeEqArray rb rh patArr
2186 !acc' <- if idx /= 0 && not withSep
2187 then RB.unsafeFoldRingM rh fstep acc rb
2189 done acc' >>= \r -> return $ Yield r GO_DONE
2191 -- XXX Theoretically this code can do 4 times faster if GHC generates
2192 -- optimal code. If we use just "(cksum' == patHash)" condition it goes
2193 -- 4x faster, as soon as we add the "RB.unsafeEqArray rb v" condition
2194 -- the generated code changes drastically and becomes 4x slower. Need
2195 -- to investigate what is going on with GHC.
2197 go1 !_ !cksum !rh st !acc = do
2198 res <- step (adaptState gst) st
2201 old <- liftIO $ peek rh
2202 let cksum' = deltaCksum cksum old x
2207 if (cksum' == patHash)
2209 rh' <- liftIO (RB.unsafeInsert rb rh x)
2210 go2 SPEC cksum' rh' s acc'
2212 rh' <- liftIO (RB.unsafeInsert rb rh x)
2213 go1 SPEC cksum' rh' s acc'
2214 Skip s -> go1 SPEC cksum rh s acc
2216 if RB.unsafeEqArray rb rh patArr
2221 else RB.unsafeFoldRingFullM rh fstep acc rb
2222 done acc' >>= \r -> return $ Yield r GO_DONE
2224 go2 !_ !cksum' !rh' s !acc' = do
2225 if RB.unsafeEqArray rb rh' patArr
2228 return $ Yield r (GO_KARP_RABIN s rb rhead)
2229 else go1 SPEC cksum' rh' s acc'
2231 stepOuter gst (GO_EMPTY_PAT st) = do
2232 res <- step (adaptState gst) st
2237 done acc' >>= \r -> return $ Yield r (GO_EMPTY_PAT s)
2238 Skip s -> return $ Skip (GO_EMPTY_PAT s)
2241 stepOuter _ GO_DONE = return Stop
2243data SplitState s arr
2245 | SplitBuffering s arr
2246 | SplitSplitting s arr
2247 | SplitYielding arr (SplitState s arr)
2250-- XXX An alternative approach would be to use a partial fold (Fold m a b) to
2251-- split using a splitBy like combinator. The Fold would consume upto the
2252-- separator and return any leftover which can then be fed to the next fold.
2254-- We can revisit this once we have partial folds/parsers.
2256-- | Performs infix separator style splitting.
2257{-# INLINE_NORMAL splitInnerBy #-}
2260 => (f a -> m (f a, Maybe (f a))) -- splitter
2261 -> (f a -> f a -> m (f a)) -- joiner
2264splitInnerBy splitter joiner (Stream step1 state1) =
2265 (Stream step (SplitInitial state1))
2269 {-# INLINE_LATE step #-}
2270 step gst (SplitInitial st) = do
2274 (x1, mx2) <- splitter x
2275 return $ case mx2 of
2276 Nothing -> Skip (SplitBuffering s x1)
2277 Just x2 -> Skip (SplitYielding x1 (SplitSplitting s x2))
2278 Skip s -> return $ Skip (SplitInitial s)
2279 Stop -> return $ Stop
2281 step gst (SplitBuffering st buf) = do
2285 (x1, mx2) <- splitter x
2286 buf' <- joiner buf x1
2287 return $ case mx2 of
2288 Nothing -> Skip (SplitBuffering s buf')
2289 Just x2 -> Skip (SplitYielding buf' (SplitSplitting s x2))
2290 Skip s -> return $ Skip (SplitBuffering s buf)
2291 Stop -> return $ Skip (SplitYielding buf SplitFinishing)
2293 step _ (SplitSplitting st buf) = do
2294 (x1, mx2) <- splitter buf
2295 return $ case mx2 of
2296 Nothing -> Skip $ SplitBuffering st x1
2297 Just x2 -> Skip $ SplitYielding x1 (SplitSplitting st x2)
2299 step _ (SplitYielding x next) = return $ Yield x next
2300 step _ SplitFinishing = return $ Stop
2302-- | Performs infix separator style splitting.
2303{-# INLINE_NORMAL splitInnerBySuffix #-}
2305 :: (Monad m, Eq (f a), Monoid (f a))
2306 => (f a -> m (f a, Maybe (f a))) -- splitter
2307 -> (f a -> f a -> m (f a)) -- joiner
2310splitInnerBySuffix splitter joiner (Stream step1 state1) =
2311 (Stream step (SplitInitial state1))
2315 {-# INLINE_LATE step #-}
2316 step gst (SplitInitial st) = do
2320 (x1, mx2) <- splitter x
2321 return $ case mx2 of
2322 Nothing -> Skip (SplitBuffering s x1)
2323 Just x2 -> Skip (SplitYielding x1 (SplitSplitting s x2))
2324 Skip s -> return $ Skip (SplitInitial s)
2325 Stop -> return $ Stop
2327 step gst (SplitBuffering st buf) = do
2331 (x1, mx2) <- splitter x
2332 buf' <- joiner buf x1
2333 return $ case mx2 of
2334 Nothing -> Skip (SplitBuffering s buf')
2335 Just x2 -> Skip (SplitYielding buf' (SplitSplitting s x2))
2336 Skip s -> return $ Skip (SplitBuffering s buf)
2340 else Skip (SplitYielding buf SplitFinishing)
2342 step _ (SplitSplitting st buf) = do
2343 (x1, mx2) <- splitter buf
2344 return $ case mx2 of
2345 Nothing -> Skip $ SplitBuffering st x1
2346 Just x2 -> Skip $ SplitYielding x1 (SplitSplitting st x2)
2348 step _ (SplitYielding x next) = return $ Yield x next
2349 step _ SplitFinishing = return $ Stop
2351------------------------------------------------------------------------------
2353------------------------------------------------------------------------------
2355{-# INLINE_NORMAL isPrefixOf #-}
2356isPrefixOf :: (Eq a, Monad m) => Stream m a -> Stream m a -> m Bool
2357isPrefixOf (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
2359 go (sa, sb, Nothing) = do
2360 r <- stepa defState sa
2362 Yield x sa' -> go (sa', sb, Just x)
2363 Skip sa' -> go (sa', sb, Nothing)
2366 go (sa, sb, Just x) = do
2367 r <- stepb defState sb
2371 then go (sa, sb', Nothing)
2373 Skip sb' -> go (sa, sb', Just x)
2374 Stop -> return False
2376{-# INLINE_NORMAL isSubsequenceOf #-}
2377isSubsequenceOf :: (Eq a, Monad m) => Stream m a -> Stream m a -> m Bool
2378isSubsequenceOf (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
2380 go (sa, sb, Nothing) = do
2381 r <- stepa defState sa
2383 Yield x sa' -> go (sa', sb, Just x)
2384 Skip sa' -> go (sa', sb, Nothing)
2387 go (sa, sb, Just x) = do
2388 r <- stepb defState sb
2392 then go (sa, sb', Nothing)
2393 else go (sa, sb', Just x)
2394 Skip sb' -> go (sa, sb', Just x)
2395 Stop -> return False
2397{-# INLINE_NORMAL stripPrefix #-}
2400 => Stream m a -> Stream m a -> m (Maybe (Stream m a))
2401stripPrefix (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
2403 go (sa, sb, Nothing) = do
2404 r <- stepa defState sa
2406 Yield x sa' -> go (sa', sb, Just x)
2407 Skip sa' -> go (sa', sb, Nothing)
2408 Stop -> return $ Just (Stream stepb sb)
2410 go (sa, sb, Just x) = do
2411 r <- stepb defState sb
2415 then go (sa, sb', Nothing)
2417 Skip sb' -> go (sa, sb', Just x)
2418 Stop -> return Nothing
2420------------------------------------------------------------------------------
2422------------------------------------------------------------------------------
2424-- | Execute a monadic action for each element of the 'Stream'
2425{-# INLINE_NORMAL mapM_ #-}
2426mapM_ :: Monad m => (a -> m b) -> Stream m a -> m ()
2427mapM_ m = drain . mapM m
2429-------------------------------------------------------------------------------
2430-- Stream transformations using Unfolds
2431-------------------------------------------------------------------------------
2433-- Define a unique structure to use in inspection testing
2434data ConcatMapUState o i =
2436 | ConcatMapUInner o i
2438-- | @concatMapU unfold stream@ uses @unfold@ to map the input stream elements
2439-- to streams and then flattens the generated streams into a single output
2442-- This is like 'concatMap' but uses an unfold with an explicit state to
2443-- generate the stream instead of a 'Stream' type generator. This allows better
2444-- optimization via fusion. This can be many times more efficient than
2447{-# INLINE_NORMAL concatMapU #-}
2448concatMapU :: Monad m => Unfold m a b -> Stream m a -> Stream m b
2449concatMapU (Unfold istep inject) (Stream ostep ost) =
2450 Stream step (ConcatMapUOuter ost)
2452 {-# INLINE_LATE step #-}
2453 step gst (ConcatMapUOuter o) = do
2454 r <- ostep (adaptState gst) o
2458 i `seq` return (Skip (ConcatMapUInner o' i))
2459 Skip o' -> return $ Skip (ConcatMapUOuter o')
2460 Stop -> return $ Stop
2462 step _ (ConcatMapUInner o i) = do
2465 Yield x i' -> Yield x (ConcatMapUInner o i')
2466 Skip i' -> Skip (ConcatMapUInner o i')
2467 Stop -> Skip (ConcatMapUOuter o)
2469data ConcatUnfoldInterleaveState o i =
2470 ConcatUnfoldInterleaveOuter o [i]
2471 | ConcatUnfoldInterleaveInner o [i]
2472 | ConcatUnfoldInterleaveInnerL [i] [i]
2473 | ConcatUnfoldInterleaveInnerR [i] [i]
2475-- XXX use arrays to store state instead of lists.
2476-- XXX In general we can use different scheduling strategies e.g. how to
2477-- schedule the outer vs inner loop or assigning weights to different streams
2478-- or outer and inner loops.
2480-- After a yield, switch to the next stream. Do not switch streams on Skip.
2481-- Yield from outer stream switches to the inner stream.
2483-- There are two choices here, (1) exhaust the outer stream first and then
2484-- start yielding from the inner streams, this is much simpler to implement,
2485-- (2) yield at least one element from an inner stream before going back to
2486-- outer stream and opening the next stream from it.
2488-- Ideally, we need some scheduling bias to inner streams vs outer stream.
2489-- Maybe we can configure the behavior.
2491{-# INLINE_NORMAL concatUnfoldInterleave #-}
2492concatUnfoldInterleave :: Monad m => Unfold m a b -> Stream m a -> Stream m b
2493concatUnfoldInterleave (Unfold istep inject) (Stream ostep ost) =
2494 Stream step (ConcatUnfoldInterleaveOuter ost [])
2496 {-# INLINE_LATE step #-}
2497 step gst (ConcatUnfoldInterleaveOuter o ls) = do
2498 r <- ostep (adaptState gst) o
2502 i `seq` return (Skip (ConcatUnfoldInterleaveInner o' (i : ls)))
2503 Skip o' -> return $ Skip (ConcatUnfoldInterleaveOuter o' ls)
2504 Stop -> return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2506 step _ (ConcatUnfoldInterleaveInner _ []) = undefined
2507 step _ (ConcatUnfoldInterleaveInner o (st:ls)) = do
2510 Yield x s -> Yield x (ConcatUnfoldInterleaveOuter o (s:ls))
2511 Skip s -> Skip (ConcatUnfoldInterleaveInner o (s:ls))
2512 Stop -> Skip (ConcatUnfoldInterleaveOuter o ls)
2514 step _ (ConcatUnfoldInterleaveInnerL [] []) = return Stop
2515 step _ (ConcatUnfoldInterleaveInnerL [] rs) =
2516 return $ Skip (ConcatUnfoldInterleaveInnerR [] rs)
2518 step _ (ConcatUnfoldInterleaveInnerL (st:ls) rs) = do
2521 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerL ls (s:rs))
2522 Skip s -> Skip (ConcatUnfoldInterleaveInnerL (s:ls) rs)
2523 Stop -> Skip (ConcatUnfoldInterleaveInnerL ls rs)
2525 step _ (ConcatUnfoldInterleaveInnerR [] []) = return Stop
2526 step _ (ConcatUnfoldInterleaveInnerR ls []) =
2527 return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2529 step _ (ConcatUnfoldInterleaveInnerR ls (st:rs)) = do
2532 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerR (s:ls) rs)
2533 Skip s -> Skip (ConcatUnfoldInterleaveInnerR ls (s:rs))
2534 Stop -> Skip (ConcatUnfoldInterleaveInnerR ls rs)
2536-- XXX In general we can use different scheduling strategies e.g. how to
2537-- schedule the outer vs inner loop or assigning weights to different streams
2538-- or outer and inner loops.
2540-- This could be inefficient if the tasks are too small.
2542-- Compared to concatUnfoldInterleave this one switches streams on Skips.
2544{-# INLINE_NORMAL concatUnfoldRoundrobin #-}
2545concatUnfoldRoundrobin :: Monad m => Unfold m a b -> Stream m a -> Stream m b
2546concatUnfoldRoundrobin (Unfold istep inject) (Stream ostep ost) =
2547 Stream step (ConcatUnfoldInterleaveOuter ost [])
2549 {-# INLINE_LATE step #-}
2550 step gst (ConcatUnfoldInterleaveOuter o ls) = do
2551 r <- ostep (adaptState gst) o
2555 i `seq` return (Skip (ConcatUnfoldInterleaveInner o' (i : ls)))
2556 Skip o' -> return $ Skip (ConcatUnfoldInterleaveInner o' ls)
2557 Stop -> return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2559 step _ (ConcatUnfoldInterleaveInner o []) =
2560 return $ Skip (ConcatUnfoldInterleaveOuter o [])
2562 step _ (ConcatUnfoldInterleaveInner o (st:ls)) = do
2565 Yield x s -> Yield x (ConcatUnfoldInterleaveOuter o (s:ls))
2566 Skip s -> Skip (ConcatUnfoldInterleaveOuter o (s:ls))
2567 Stop -> Skip (ConcatUnfoldInterleaveOuter o ls)
2569 step _ (ConcatUnfoldInterleaveInnerL [] []) = return Stop
2570 step _ (ConcatUnfoldInterleaveInnerL [] rs) =
2571 return $ Skip (ConcatUnfoldInterleaveInnerR [] rs)
2573 step _ (ConcatUnfoldInterleaveInnerL (st:ls) rs) = do
2576 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerL ls (s:rs))
2577 Skip s -> Skip (ConcatUnfoldInterleaveInnerL ls (s:rs))
2578 Stop -> Skip (ConcatUnfoldInterleaveInnerL ls rs)
2580 step _ (ConcatUnfoldInterleaveInnerR [] []) = return Stop
2581 step _ (ConcatUnfoldInterleaveInnerR ls []) =
2582 return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2584 step _ (ConcatUnfoldInterleaveInnerR ls (st:rs)) = do
2587 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerR (s:ls) rs)
2588 Skip s -> Skip (ConcatUnfoldInterleaveInnerR (s:ls) rs)
2589 Stop -> Skip (ConcatUnfoldInterleaveInnerR ls rs)
2591data AppendState s1 s2 = AppendFirst s1 | AppendSecond s2
2593-- Note that this could be much faster compared to the CPS stream. However, as
2594-- the number of streams being composed increases this may become expensive.
2595-- Need to see where the breaking point is between the two.
2597{-# INLINE_NORMAL append #-}
2598append :: Monad m => Stream m a -> Stream m a -> Stream m a
2599append (Stream step1 state1) (Stream step2 state2) =
2600 Stream step (AppendFirst state1)
2604 {-# INLINE_LATE step #-}
2605 step gst (AppendFirst st) = do
2608 Yield a s -> Yield a (AppendFirst s)
2609 Skip s -> Skip (AppendFirst s)
2610 Stop -> Skip (AppendSecond state2)
2612 step gst (AppendSecond st) = do
2615 Yield a s -> Yield a (AppendSecond s)
2616 Skip s -> Skip (AppendSecond s)
2619data InterleaveState s1 s2 = InterleaveFirst s1 s2 | InterleaveSecond s1 s2
2620 | InterleaveSecondOnly s2 | InterleaveFirstOnly s1
2622{-# INLINE_NORMAL interleave #-}
2623interleave :: Monad m => Stream m a -> Stream m a -> Stream m a
2624interleave (Stream step1 state1) (Stream step2 state2) =
2625 Stream step (InterleaveFirst state1 state2)
2629 {-# INLINE_LATE step #-}
2630 step gst (InterleaveFirst st1 st2) = do
2633 Yield a s -> Yield a (InterleaveSecond s st2)
2634 Skip s -> Skip (InterleaveFirst s st2)
2635 Stop -> Skip (InterleaveSecondOnly st2)
2637 step gst (InterleaveSecond st1 st2) = do
2640 Yield a s -> Yield a (InterleaveFirst st1 s)
2641 Skip s -> Skip (InterleaveSecond st1 s)
2642 Stop -> Skip (InterleaveFirstOnly st1)
2644 step gst (InterleaveFirstOnly st1) = do
2647 Yield a s -> Yield a (InterleaveFirstOnly s)
2648 Skip s -> Skip (InterleaveFirstOnly s)
2651 step gst (InterleaveSecondOnly st2) = do
2654 Yield a s -> Yield a (InterleaveSecondOnly s)
2655 Skip s -> Skip (InterleaveSecondOnly s)
2658{-# INLINE_NORMAL interleaveMin #-}
2659interleaveMin :: Monad m => Stream m a -> Stream m a -> Stream m a
2660interleaveMin (Stream step1 state1) (Stream step2 state2) =
2661 Stream step (InterleaveFirst state1 state2)
2665 {-# INLINE_LATE step #-}
2666 step gst (InterleaveFirst st1 st2) = do
2669 Yield a s -> Yield a (InterleaveSecond s st2)
2670 Skip s -> Skip (InterleaveFirst s st2)
2673 step gst (InterleaveSecond st1 st2) = do
2676 Yield a s -> Yield a (InterleaveFirst st1 s)
2677 Skip s -> Skip (InterleaveSecond st1 s)
2680 step _ (InterleaveFirstOnly _) = undefined
2681 step _ (InterleaveSecondOnly _) = undefined
2683{-# INLINE_NORMAL interleaveSuffix #-}
2684interleaveSuffix :: Monad m => Stream m a -> Stream m a -> Stream m a
2685interleaveSuffix (Stream step1 state1) (Stream step2 state2) =
2686 Stream step (InterleaveFirst state1 state2)
2690 {-# INLINE_LATE step #-}
2691 step gst (InterleaveFirst st1 st2) = do
2694 Yield a s -> Yield a (InterleaveSecond s st2)
2695 Skip s -> Skip (InterleaveFirst s st2)
2698 step gst (InterleaveSecond st1 st2) = do
2701 Yield a s -> Yield a (InterleaveFirst st1 s)
2702 Skip s -> Skip (InterleaveSecond st1 s)
2703 Stop -> Skip (InterleaveFirstOnly st1)
2705 step gst (InterleaveFirstOnly st1) = do
2708 Yield a s -> Yield a (InterleaveFirstOnly s)
2709 Skip s -> Skip (InterleaveFirstOnly s)
2712 step _ (InterleaveSecondOnly _) = undefined
2714data InterleaveInfixState s1 s2 a
2715 = InterleaveInfixFirst s1 s2
2716 | InterleaveInfixSecondBuf s1 s2
2717 | InterleaveInfixSecondYield s1 s2 a
2718 | InterleaveInfixFirstYield s1 s2 a
2719 | InterleaveInfixFirstOnly s1
2721{-# INLINE_NORMAL interleaveInfix #-}
2722interleaveInfix :: Monad m => Stream m a -> Stream m a -> Stream m a
2723interleaveInfix (Stream step1 state1) (Stream step2 state2) =
2724 Stream step (InterleaveInfixFirst state1 state2)
2728 {-# INLINE_LATE step #-}
2729 step gst (InterleaveInfixFirst st1 st2) = do
2732 Yield a s -> Yield a (InterleaveInfixSecondBuf s st2)
2733 Skip s -> Skip (InterleaveInfixFirst s st2)
2736 step gst (InterleaveInfixSecondBuf st1 st2) = do
2739 Yield a s -> Skip (InterleaveInfixSecondYield st1 s a)
2740 Skip s -> Skip (InterleaveInfixSecondBuf st1 s)
2741 Stop -> Skip (InterleaveInfixFirstOnly st1)
2743 step gst (InterleaveInfixSecondYield st1 st2 x) = do
2746 Yield a s -> Yield x (InterleaveInfixFirstYield s st2 a)
2747 Skip s -> Skip (InterleaveInfixSecondYield s st2 x)
2750 step _ (InterleaveInfixFirstYield st1 st2 x) = do
2751 return $ Yield x (InterleaveInfixSecondBuf st1 st2)
2753 step gst (InterleaveInfixFirstOnly st1) = do
2756 Yield a s -> Yield a (InterleaveInfixFirstOnly s)
2757 Skip s -> Skip (InterleaveInfixFirstOnly s)
2760{-# INLINE_NORMAL roundRobin #-}
2761roundRobin :: Monad m => Stream m a -> Stream m a -> Stream m a
2762roundRobin (Stream step1 state1) (Stream step2 state2) =
2763 Stream step (InterleaveFirst state1 state2)
2767 {-# INLINE_LATE step #-}
2768 step gst (InterleaveFirst st1 st2) = do
2771 Yield a s -> Yield a (InterleaveSecond s st2)
2772 Skip s -> Skip (InterleaveSecond s st2)
2773 Stop -> Skip (InterleaveSecondOnly st2)
2775 step gst (InterleaveSecond st1 st2) = do
2778 Yield a s -> Yield a (InterleaveFirst st1 s)
2779 Skip s -> Skip (InterleaveFirst st1 s)
2780 Stop -> Skip (InterleaveFirstOnly st1)
2782 step gst (InterleaveSecondOnly st2) = do
2785 Yield a s -> Yield a (InterleaveSecondOnly s)
2786 Skip s -> Skip (InterleaveSecondOnly s)
2789 step gst (InterleaveFirstOnly st1) = do
2792 Yield a s -> Yield a (InterleaveFirstOnly s)
2793 Skip s -> Skip (InterleaveFirstOnly s)
2796data ICUState s1 s2 i1 i2 =
2801 | ICUFirstInner s1 s2 i1
2802 | ICUSecondInner s1 s2 i2
2803 | ICUFirstOnlyInner s1 i1
2804 | ICUSecondOnlyInner s2 i2
2806-- | Interleave streams (full streams, not the elements) unfolded from two
2807-- input streams and concat. Stop when the first stream stops. If the second
2808-- stream ends before the first one then first stream still keeps running alone
2809-- without any interleaving with the second stream.
2811-- [a1, a2, ... an] [b1, b2 ...]
2812-- => [streamA1, streamA2, ... streamAn] [streamB1, streamB2, ...]
2813-- => [streamA1, streamB1, streamA2...StreamAn, streamBn]
2814-- => [a11, a12, ...a1j, b11, b12, ...b1k, a21, a22, ...]
2816{-# INLINE_NORMAL gintercalateSuffix #-}
2819 => Unfold m a c -> Stream m a -> Unfold m b c -> Stream m b -> Stream m c
2821 (Unfold istep1 inject1) (Stream step1 state1)
2822 (Unfold istep2 inject2) (Stream step2 state2) =
2823 Stream step (ICUFirst state1 state2)
2827 {-# INLINE_LATE step #-}
2828 step gst (ICUFirst s1 s2) = do
2829 r <- step1 (adaptState gst) s1
2833 i `seq` return (Skip (ICUFirstInner s s2 i))
2834 Skip s -> return $ Skip (ICUFirst s s2)
2837 step gst (ICUFirstOnly s1) = do
2838 r <- step1 (adaptState gst) s1
2842 i `seq` return (Skip (ICUFirstOnlyInner s i))
2843 Skip s -> return $ Skip (ICUFirstOnly s)
2846 step _ (ICUFirstInner s1 s2 i1) = do
2849 Yield x i' -> Yield x (ICUFirstInner s1 s2 i')
2850 Skip i' -> Skip (ICUFirstInner s1 s2 i')
2851 Stop -> Skip (ICUSecond s1 s2)
2853 step _ (ICUFirstOnlyInner s1 i1) = do
2856 Yield x i' -> Yield x (ICUFirstOnlyInner s1 i')
2857 Skip i' -> Skip (ICUFirstOnlyInner s1 i')
2858 Stop -> Skip (ICUFirstOnly s1)
2860 step gst (ICUSecond s1 s2) = do
2861 r <- step2 (adaptState gst) s2
2865 i `seq` return (Skip (ICUSecondInner s1 s i))
2866 Skip s -> return $ Skip (ICUSecond s1 s)
2867 Stop -> return $ Skip (ICUFirstOnly s1)
2869 step _ (ICUSecondInner s1 s2 i2) = do
2872 Yield x i' -> Yield x (ICUSecondInner s1 s2 i')
2873 Skip i' -> Skip (ICUSecondInner s1 s2 i')
2874 Stop -> Skip (ICUFirst s1 s2)
2876 step _ (ICUSecondOnly _s2) = undefined
2877 step _ (ICUSecondOnlyInner _s2 _i2) = undefined
2879data InterposeSuffixState s1 i1 =
2880 InterposeSuffixFirst s1
2881 -- | InterposeSuffixFirstYield s1 i1
2882 | InterposeSuffixFirstInner s1 i1
2883 | InterposeSuffixSecond s1
2885-- Note that if an unfolded layer turns out to be nil we still emit the
2886-- separator effect. An alternate behavior could be to emit the separator
2887-- effect only if at least one element has been yielded by the unfolding.
2888-- However, that becomes a bit complicated, so we have chosen the former
2890{-# INLINE_NORMAL interposeSuffix #-}
2893 => m c -> Unfold m b c -> Stream m b -> Stream m c
2896 (Unfold istep1 inject1) (Stream step1 state1) =
2897 Stream step (InterposeSuffixFirst state1)
2901 {-# INLINE_LATE step #-}
2902 step gst (InterposeSuffixFirst s1) = do
2903 r <- step1 (adaptState gst) s1
2907 i `seq` return (Skip (InterposeSuffixFirstInner s i))
2908 -- i `seq` return (Skip (InterposeSuffixFirstYield s i))
2909 Skip s -> return $ Skip (InterposeSuffixFirst s)
2913 step _ (InterposeSuffixFirstYield s1 i1) = do
2916 Yield x i' -> Yield x (InterposeSuffixFirstInner s1 i')
2917 Skip i' -> Skip (InterposeSuffixFirstYield s1 i')
2918 Stop -> Skip (InterposeSuffixFirst s1)
2921 step _ (InterposeSuffixFirstInner s1 i1) = do
2924 Yield x i' -> Yield x (InterposeSuffixFirstInner s1 i')
2925 Skip i' -> Skip (InterposeSuffixFirstInner s1 i')
2926 Stop -> Skip (InterposeSuffixSecond s1)
2928 step _ (InterposeSuffixSecond s1) = do
2930 return $ Yield r (InterposeSuffixFirst s1)
2932data ICALState s1 s2 i1 i2 a =
2934 -- | ICALFirstYield s1 s2 i1
2935 | ICALFirstInner s1 s2 i1
2937 | ICALFirstOnlyInner s1 i1
2938 | ICALSecondInject s1 s2
2939 | ICALFirstInject s1 s2 i2
2940 -- | ICALFirstBuf s1 s2 i1 i2
2941 | ICALSecondInner s1 s2 i1 i2
2942 -- -- | ICALSecondInner s1 s2 i1 i2 a
2943 -- -- | ICALFirstResume s1 s2 i1 i2 a
2945-- | Interleave streams (full streams, not the elements) unfolded from two
2946-- input streams and concat. Stop when the first stream stops. If the second
2947-- stream ends before the first one then first stream still keeps running alone
2948-- without any interleaving with the second stream.
2950-- [a1, a2, ... an] [b1, b2 ...]
2951-- => [streamA1, streamA2, ... streamAn] [streamB1, streamB2, ...]
2952-- => [streamA1, streamB1, streamA2...StreamAn, streamBn]
2953-- => [a11, a12, ...a1j, b11, b12, ...b1k, a21, a22, ...]
2955{-# INLINE_NORMAL gintercalate #-}
2958 => Unfold m a c -> Stream m a -> Unfold m b c -> Stream m b -> Stream m c
2960 (Unfold istep1 inject1) (Stream step1 state1)
2961 (Unfold istep2 inject2) (Stream step2 state2) =
2962 Stream step (ICALFirst state1 state2)
2966 {-# INLINE_LATE step #-}
2967 step gst (ICALFirst s1 s2) = do
2968 r <- step1 (adaptState gst) s1
2972 i `seq` return (Skip (ICALFirstInner s s2 i))
2973 -- i `seq` return (Skip (ICALFirstYield s s2 i))
2974 Skip s -> return $ Skip (ICALFirst s s2)
2978 step _ (ICALFirstYield s1 s2 i1) = do
2981 Yield x i' -> Yield x (ICALFirstInner s1 s2 i')
2982 Skip i' -> Skip (ICALFirstYield s1 s2 i')
2983 Stop -> Skip (ICALFirst s1 s2)
2986 step _ (ICALFirstInner s1 s2 i1) = do
2989 Yield x i' -> Yield x (ICALFirstInner s1 s2 i')
2990 Skip i' -> Skip (ICALFirstInner s1 s2 i')
2991 Stop -> Skip (ICALSecondInject s1 s2)
2993 step gst (ICALFirstOnly s1) = do
2994 r <- step1 (adaptState gst) s1
2998 i `seq` return (Skip (ICALFirstOnlyInner s i))
2999 Skip s -> return $ Skip (ICALFirstOnly s)
3002 step _ (ICALFirstOnlyInner s1 i1) = do
3005 Yield x i' -> Yield x (ICALFirstOnlyInner s1 i')
3006 Skip i' -> Skip (ICALFirstOnlyInner s1 i')
3007 Stop -> Skip (ICALFirstOnly s1)
3009 -- We inject the second stream even before checking if the first stream
3010 -- would yield any more elements. There is no clear choice whether we
3011 -- should do this before or after that. Doing it after may make the state
3012 -- machine a bit simpler though.
3013 step gst (ICALSecondInject s1 s2) = do
3014 r <- step2 (adaptState gst) s2
3018 i `seq` return (Skip (ICALFirstInject s1 s i))
3019 Skip s -> return $ Skip (ICALSecondInject s1 s)
3020 Stop -> return $ Skip (ICALFirstOnly s1)
3022 step gst (ICALFirstInject s1 s2 i2) = do
3023 r <- step1 (adaptState gst) s1
3027 i `seq` return (Skip (ICALSecondInner s s2 i i2))
3028 -- i `seq` return (Skip (ICALFirstBuf s s2 i i2))
3029 Skip s -> return $ Skip (ICALFirstInject s s2 i2)
3033 step _ (ICALFirstBuf s1 s2 i1 i2) = do
3036 Yield x i' -> Skip (ICALSecondInner s1 s2 i' i2 x)
3037 Skip i' -> Skip (ICALFirstBuf s1 s2 i' i2)
3040 step _ (ICALSecondInner s1 s2 i1 i2 v) = do
3043 Yield x i' -> Yield x (ICALSecondInner s1 s2 i1 i' v)
3044 Skip i' -> Skip (ICALSecondInner s1 s2 i1 i' v)
3045 Stop -> Skip (ICALFirstResume s1 s2 i1 i2 v)
3048 step _ (ICALSecondInner s1 s2 i1 i2) = do
3051 Yield x i' -> Yield x (ICALSecondInner s1 s2 i1 i')
3052 Skip i' -> Skip (ICALSecondInner s1 s2 i1 i')
3053 Stop -> Skip (ICALFirstInner s1 s2 i1)
3054 -- Stop -> Skip (ICALFirstResume s1 s2 i1 i2)
3057 step _ (ICALFirstResume s1 s2 i1 i2 x) = do
3058 return $ Yield x (ICALFirstInner s1 s2 i1 i2)
3061data InterposeState s1 i1 a =
3063 -- | InterposeFirstYield s1 i1
3064 | InterposeFirstInner s1 i1
3065 | InterposeFirstInject s1
3066 -- | InterposeFirstBuf s1 i1
3067 | InterposeSecondYield s1 i1
3068 -- -- | InterposeSecondYield s1 i1 a
3069 -- -- | InterposeFirstResume s1 i1 a
3071-- Note that this only interposes the pure values, we may run many effects to
3072-- generate those values as some effects may not generate anything (Skip).
3073{-# INLINE_NORMAL interpose #-}
3074interpose :: Monad m => m c -> Unfold m b c -> Stream m b -> Stream m c
3077 (Unfold istep1 inject1) (Stream step1 state1) =
3078 Stream step (InterposeFirst state1)
3082 {-# INLINE_LATE step #-}
3083 step gst (InterposeFirst s1) = do
3084 r <- step1 (adaptState gst) s1
3088 i `seq` return (Skip (InterposeFirstInner s i))
3089 -- i `seq` return (Skip (InterposeFirstYield s i))
3090 Skip s -> return $ Skip (InterposeFirst s)
3094 step _ (InterposeFirstYield s1 i1) = do
3097 Yield x i' -> Yield x (InterposeFirstInner s1 i')
3098 Skip i' -> Skip (InterposeFirstYield s1 i')
3099 Stop -> Skip (InterposeFirst s1)
3102 step _ (InterposeFirstInner s1 i1) = do
3105 Yield x i' -> Yield x (InterposeFirstInner s1 i')
3106 Skip i' -> Skip (InterposeFirstInner s1 i')
3107 Stop -> Skip (InterposeFirstInject s1)
3109 step gst (InterposeFirstInject s1) = do
3110 r <- step1 (adaptState gst) s1
3114 -- i `seq` return (Skip (InterposeFirstBuf s i))
3115 i `seq` return (Skip (InterposeSecondYield s i))
3116 Skip s -> return $ Skip (InterposeFirstInject s)
3120 step _ (InterposeFirstBuf s1 i1) = do
3123 Yield x i' -> Skip (InterposeSecondYield s1 i' x)
3124 Skip i' -> Skip (InterposeFirstBuf s1 i')
3129 step _ (InterposeSecondYield s1 i1 v) = do
3131 return $ Yield r (InterposeFirstResume s1 i1 v)
3133 step _ (InterposeSecondYield s1 i1) = do
3135 return $ Yield r (InterposeFirstInner s1 i1)
3138 step _ (InterposeFirstResume s1 i1 v) = do
3139 return $ Yield v (InterposeFirstInner s1 i1)
3142------------------------------------------------------------------------------
3144------------------------------------------------------------------------------
3146data GbracketState s1 s2 v
3148 | GBracketNormal s1 v
3149 | GBracketException s2
3151-- | The most general bracketing and exception combinator. All other
3152-- combinators can be expressed in terms of this combinator. This can also be
3153-- used for cases which are not covered by the standard combinators.
3157{-# INLINE_NORMAL gbracket #-}
3161 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
3162 -> (c -> m d) -- ^ after, on normal stop
3163 -> (c -> e -> Stream m b) -- ^ on exception
3164 -> (c -> Stream m b) -- ^ stream generator
3166gbracket bef exc aft fexc fnormal =
3167 Stream step GBracketInit
3171 {-# INLINE_LATE step #-}
3172 step _ GBracketInit = do
3174 return $ Skip $ GBracketNormal (fnormal r) r
3176 step gst (GBracketNormal (UnStream step1 st) v) = do
3177 res <- exc $ step1 gst st
3179 Right r -> case r of
3181 return $ Yield x (GBracketNormal (Stream step1 s) v)
3182 Skip s -> return $ Skip (GBracketNormal (Stream step1 s) v)
3183 Stop -> aft v >> return Stop
3184 Left e -> return $ Skip (GBracketException (fexc v e))
3185 step gst (GBracketException (UnStream step1 st)) = do
3188 Yield x s -> return $ Yield x (GBracketException (Stream step1 s))
3189 Skip s -> return $ Skip (GBracketException (Stream step1 s))
3192-- | Create an IORef holding a finalizer that is called automatically when the
3193-- IORef is garbage collected. The IORef can be written to with a 'Nothing'
3194-- value to deactivate the finalizer.
3195newFinalizedIORef :: (MonadIO m, MonadBaseControl IO m)
3196 => m a -> m (IORef (Maybe (IO ())))
3197newFinalizedIORef finalizer = do
3198 mrun <- captureMonadState
3199 ref <- liftIO $ newIORef $ Just $ liftIO $ void $ do
3200 _ <- runInIO mrun finalizer
3203 res <- readIORef ref
3205 Nothing -> return ()
3207 _ <- liftIO $ mkWeakIORef ref finalizer1
3210-- | Run the finalizer stored in an IORef and deactivate it so that it is run
3213runIORefFinalizer :: MonadIO m => IORef (Maybe (IO ())) -> m ()
3214runIORefFinalizer ref = liftIO $ do
3215 res <- readIORef ref
3217 Nothing -> return ()
3218 Just f -> writeIORef ref Nothing >> f
3220-- | Deactivate the finalizer stored in an IORef without running it.
3222clearIORefFinalizer :: MonadIO m => IORef (Maybe (IO ())) -> m ()
3223clearIORefFinalizer ref = liftIO $ writeIORef ref Nothing
3225data GbracketIOState s1 s2 v wref
3227 | GBracketIONormal s1 v wref
3228 | GBracketIOException s2
3230-- | Like gbracket but also uses a finalizer to make sure when the stream is
3231-- garbage collected we run the finalizing action. This requires a MonadIO and
3232-- MonadBaseControl IO constraint.
3234-- | The most general bracketing and exception combinator. All other
3235-- combinators can be expressed in terms of this combinator. This can also be
3236-- used for cases which are not covered by the standard combinators.
3240{-# INLINE_NORMAL gbracketIO #-}
3242 :: (MonadIO m, MonadBaseControl IO m)
3244 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
3245 -> (c -> m d) -- ^ after, on normal stop or GC
3246 -> (c -> e -> Stream m b) -- ^ on exception
3247 -> (c -> Stream m b) -- ^ stream generator
3249gbracketIO bef exc aft fexc fnormal =
3250 Stream step GBracketIOInit
3254 -- If the stream is never evaluated the "aft" action will never be
3255 -- called. For that to occur we will need the user of this API to pass a
3256 -- weak pointer to us.
3257 {-# INLINE_LATE step #-}
3258 step _ GBracketIOInit = do
3259 -- We mask asynchronous exceptions to make the execution
3260 -- of 'bef' and the registration of 'aft' atomic.
3261 -- A similar thing is done in the resourcet package: https://git.io/JvKV3
3262 -- Tutorial: https://markkarpov.com/tutorial/exceptions.html
3263 (r, ref) <- liftBaseOp_ mask_ $ do
3265 ref <- newFinalizedIORef (aft r)
3267 return $ Skip $ GBracketIONormal (fnormal r) r ref
3269 step gst (GBracketIONormal (UnStream step1 st) v ref) = do
3270 res <- exc $ step1 gst st
3272 Right r -> case r of
3274 return $ Yield x (GBracketIONormal (Stream step1 s) v ref)
3276 return $ Skip (GBracketIONormal (Stream step1 s) v ref)
3278 runIORefFinalizer ref
3281 clearIORefFinalizer ref
3282 return $ Skip (GBracketIOException (fexc v e))
3283 step gst (GBracketIOException (UnStream step1 st)) = do
3287 return $ Yield x (GBracketIOException (Stream step1 s))
3288 Skip s -> return $ Skip (GBracketIOException (Stream step1 s))
3291-- | Run a side effect before the stream yields its first element.
3292{-# INLINE_NORMAL before #-}
3293before :: Monad m => m b -> Stream m a -> Stream m a
3294before action (Stream step state) = Stream step' Nothing
3298 {-# INLINE_LATE step' #-}
3299 step' _ Nothing = action >> return (Skip (Just state))
3301 step' gst (Just st) = do
3304 Yield x s -> return $ Yield x (Just s)
3305 Skip s -> return $ Skip (Just s)
3308-- | Run a side effect whenever the stream stops normally.
3309{-# INLINE_NORMAL after #-}
3310after :: Monad m => m b -> Stream m a -> Stream m a
3311after action (Stream step state) = Stream step' state
3315 {-# INLINE_LATE step' #-}
3319 Yield x s -> return $ Yield x s
3320 Skip s -> return $ Skip s
3321 Stop -> action >> return Stop
3323{-# INLINE_NORMAL afterIO #-}
3324afterIO :: (MonadIO m, MonadBaseControl IO m)
3325 => m b -> Stream m a -> Stream m a
3326afterIO action (Stream step state) = Stream step' Nothing
3330 {-# INLINE_LATE step' #-}
3331 step' _ Nothing = do
3332 ref <- newFinalizedIORef action
3333 return $ Skip $ Just (state, ref)
3334 step' gst (Just (st, ref)) = do
3337 Yield x s -> return $ Yield x (Just (s, ref))
3338 Skip s -> return $ Skip (Just (s, ref))
3340 runIORefFinalizer ref
3343-- XXX These combinators are expensive due to the call to
3344-- onException/handle/try on each step. Therefore, when possible, they should
3345-- be called in an outer loop where we perform less iterations. For example, we
3346-- cannot call them on each iteration in a char stream, instead we can call
3347-- them when doing an IO on an array.
3349-- XXX For high performance error checks in busy streams we may need another
3350-- Error constructor in step.
3352-- | Run a side effect whenever the stream aborts due to an exception. The
3353-- exception is not caught, simply rethrown.
3354{-# INLINE_NORMAL onException #-}
3355onException :: MonadCatch m => m b -> Stream m a -> Stream m a
3356onException action str =
3357 gbracket (return ()) MC.try return
3358 (\_ (e :: MC.SomeException) -> nilM (action >> MC.throwM e))
3361{-# INLINE_NORMAL _onException #-}
3362_onException :: MonadCatch m => m b -> Stream m a -> Stream m a
3363_onException action (Stream step state) = Stream step' state
3367 {-# INLINE_LATE step' #-}
3369 res <- step gst st `MC.onException` action
3371 Yield x s -> return $ Yield x s
3372 Skip s -> return $ Skip s
3375-- XXX bracket is like concatMap, it generates a stream and then flattens it.
3376-- Like concatMap it has 10x worse performance compared to linear fused
3379-- | Run the first action before the stream starts and remember its output,
3380-- generate a stream using the output, run the second action providing the
3381-- remembered value as an argument whenever the stream ends normally or due to
3383{-# INLINE_NORMAL bracket #-}
3384bracket :: MonadCatch m => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
3385bracket bef aft bet =
3386 gbracket bef MC.try aft
3387 (\a (e :: SomeException) -> nilM (aft a >> MC.throwM e)) bet
3389{-# INLINE_NORMAL bracketIO #-}
3390bracketIO :: (MonadAsync m, MonadCatch m)
3391 => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
3392bracketIO bef aft bet =
3393 gbracketIO bef MC.try aft
3394 (\a (e :: SomeException) -> nilM (aft a >> MC.throwM e)) bet
3396data BracketState s v = BracketInit | BracketRun s v
3398{-# INLINE_NORMAL _bracket #-}
3399_bracket :: MonadCatch m => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
3400_bracket bef aft bet = Stream step' BracketInit
3404 {-# INLINE_LATE step' #-}
3405 step' _ BracketInit = bef >>= \x -> return (Skip (BracketRun (bet x) x))
3407 -- NOTE: It is important to use UnStream instead of the Stream pattern
3408 -- here, otherwise we get huge perf degradation, see note in concatMap.
3409 step' gst (BracketRun (UnStream step state) v) = do
3410 -- res <- step gst state `MC.onException` aft v
3411 res <- MC.try $ step gst state
3413 Left (e :: SomeException) -> aft v >> MC.throwM e >> return Stop
3414 Right r -> case r of
3415 Yield x s -> return $ Yield x (BracketRun (Stream step s) v)
3416 Skip s -> return $ Skip (BracketRun (Stream step s) v)
3417 Stop -> aft v >> return Stop
3419-- | Run a side effect whenever the stream stops normally or aborts due to an
3421{-# INLINE finally #-}
3422finally :: MonadCatch m => m b -> Stream m a -> Stream m a
3423-- finally action xs = after action $ onException action xs
3424finally action xs = bracket (return ()) (\_ -> action) (const xs)
3426{-# INLINE finallyIO #-}
3427finallyIO :: (MonadAsync m, MonadCatch m) => m b -> Stream m a -> Stream m a
3428finallyIO action xs = bracketIO (return ()) (\_ -> action) (const xs)
3430-- | When evaluating a stream if an exception occurs, stream evaluation aborts
3431-- and the specified exception handler is run with the exception as argument.
3432{-# INLINE_NORMAL handle #-}
3433handle :: (MonadCatch m, Exception e)
3434 => (e -> Stream m a) -> Stream m a -> Stream m a
3436 gbracket (return ()) MC.try return (\_ e -> f e) (\_ -> str)
3438{-# INLINE_NORMAL _handle #-}
3439_handle :: (MonadCatch m, Exception e)
3440 => (e -> Stream m a) -> Stream m a -> Stream m a
3441_handle f (Stream step state) = Stream step' (Left state)
3445 {-# INLINE_LATE step' #-}
3446 step' gst (Left st) = do
3447 res <- MC.try $ step gst st
3449 Left e -> return $ Skip $ Right (f e)
3450 Right r -> case r of
3451 Yield x s -> return $ Yield x (Left s)
3452 Skip s -> return $ Skip (Left s)
3455 step' gst (Right (UnStream step1 st)) = do
3458 Yield x s -> return $ Yield x (Right (Stream step1 s))
3459 Skip s -> return $ Skip (Right (Stream step1 s))
3462-------------------------------------------------------------------------------
3463-- General transformation
3464-------------------------------------------------------------------------------
3466{-# INLINE_NORMAL transform #-}
3467transform :: Monad m => Pipe m a b -> Stream m a -> Stream m b
3468transform (Pipe pstep1 pstep2 pstate) (Stream step state) =
3469 Stream step' (Consume pstate, state)
3473 {-# INLINE_LATE step' #-}
3475 step' gst (Consume pst, st) = pst `seq` do
3476 r <- step (adaptState gst) st
3481 Pipe.Yield b pst' -> return $ Yield b (pst', s)
3482 Pipe.Continue pst' -> return $ Skip (pst', s)
3483 Skip s -> return $ Skip (Consume pst, s)
3486 step' _ (Produce pst, st) = pst `seq` do
3489 Pipe.Yield b pst' -> return $ Yield b (pst', st)
3490 Pipe.Continue pst' -> return $ Skip (pst', st)
3492------------------------------------------------------------------------------
3493-- Transformation by Folding (Scans)
3494------------------------------------------------------------------------------
3496------------------------------------------------------------------------------
3498------------------------------------------------------------------------------
3500-- XXX Is a prescan useful, discarding the last step does not sound useful? I
3501-- am not sure about the utility of this function, so this is implemented but
3502-- not exposed. We can expose it if someone provides good reasons why this is
3505-- XXX We have to execute the stream one step ahead to know that we are at the
3506-- last step. The vector implementation of prescan executes the last fold step
3507-- but does not yield the result. This means we have executed the effect but
3508-- discarded value. This does not sound right. In this implementation we are
3509-- not executing the last fold step.
3510{-# INLINE_NORMAL prescanlM' #-}
3511prescanlM' :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
3512prescanlM' f mz (Stream step state) = Stream step' (state, mz)
3514 {-# INLINE_LATE step' #-}
3515 step' gst (st, prev) = do
3516 r <- step (adaptState gst) st
3520 return $ Yield acc (s, f acc x)
3521 Skip s -> return $ Skip (s, prev)
3524{-# INLINE prescanl' #-}
3525prescanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
3526prescanl' f z = prescanlM' (\a b -> return (f a b)) (return z)
3528------------------------------------------------------------------------------
3529-- Monolithic postscans (postscan followed by a map)
3530------------------------------------------------------------------------------
3532-- The performance of a modular postscan followed by a map seems to be
3533-- equivalent to this monolithic scan followed by map therefore we may not need
3534-- this implementation. We just have it for performance comparison and in case
3535-- modular version does not perform well in some situation.
3537{-# INLINE_NORMAL postscanlMx' #-}
3538postscanlMx' :: Monad m
3539 => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> Stream m b
3540postscanlMx' fstep begin done (Stream step state) = do
3541 Stream step' (state, begin)
3543 {-# INLINE_LATE step' #-}
3544 step' gst (st, acc) = do
3545 r <- step (adaptState gst) st
3551 v `seq` y `seq` return (Yield v (s, return y))
3552 Skip s -> return $ Skip (s, acc)
3555{-# INLINE_NORMAL postscanlx' #-}
3556postscanlx' :: Monad m
3557 => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> Stream m b
3558postscanlx' fstep begin done s =
3559 postscanlMx' (\b a -> return (fstep b a)) (return begin) (return . done) s
3561-- XXX do we need consM strict to evaluate the begin value?
3562{-# INLINE scanlMx' #-}
3564 => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> Stream m b
3565scanlMx' fstep begin done s =
3566 (begin >>= \x -> x `seq` done x) `consM` postscanlMx' fstep begin done s
3568{-# INLINE scanlx' #-}
3570 => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> Stream m b
3571scanlx' fstep begin done s =
3572 scanlMx' (\b a -> return (fstep b a)) (return begin) (return . done) s
3574------------------------------------------------------------------------------
3576------------------------------------------------------------------------------
3578{-# INLINE_NORMAL postscanlM' #-}
3579postscanlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
3580postscanlM' fstep begin (Stream step state) =
3581 begin `seq` Stream step' (state, begin)
3583 {-# INLINE_LATE step' #-}
3584 step' gst (st, acc) = acc `seq` do
3585 r <- step (adaptState gst) st
3589 y `seq` return (Yield y (s, y))
3590 Skip s -> return $ Skip (s, acc)
3593{-# INLINE_NORMAL postscanl' #-}
3594postscanl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
3595postscanl' f = postscanlM' (\a b -> return (f a b))
3597-- We can possibly have the "done" function as a Maybe to provide an option to
3598-- emit or not emit the accumulator when the stream stops.
3600-- TBD: use a single Yield point
3602{-# INLINE_NORMAL postscanlMAfter' #-}
3603postscanlMAfter' :: Monad m
3604 => (b -> a -> m b) -> m b -> (b -> m b) -> Stream m a -> Stream m b
3605postscanlMAfter' fstep initial done (Stream step1 state1) = do
3606 Stream step (Just (state1, initial))
3610 {-# INLINE_LATE step #-}
3611 step gst (Just (st, acc)) = do
3612 r <- step1 (adaptState gst) st
3617 y `seq` return (Yield y (Just (s, return y)))
3618 Skip s -> return $ Skip $ Just (s, acc)
3619 Stop -> acc >>= done >>= \res -> return (Yield res Nothing)
3620 step _ Nothing = return Stop
3622{-# INLINE_NORMAL postscanlM #-}
3623postscanlM :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
3624postscanlM fstep begin (Stream step state) = Stream step' (state, begin)
3626 {-# INLINE_LATE step' #-}
3627 step' gst (st, acc) = do
3628 r <- step (adaptState gst) st
3632 return (Yield y (s, y))
3633 Skip s -> return $ Skip (s, acc)
3636{-# INLINE_NORMAL postscanl #-}
3637postscanl :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
3638postscanl f = postscanlM (\a b -> return (f a b))
3640{-# INLINE_NORMAL scanlM' #-}
3641scanlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
3642scanlM' fstep begin s = begin `seq` (begin `cons` postscanlM' fstep begin s)
3644{-# INLINE scanlMAfter' #-}
3645scanlMAfter' :: Monad m
3646 => (b -> a -> m b) -> m b -> (b -> m b) -> Stream m a -> Stream m b
3647scanlMAfter' fstep initial done s =
3648 (initial >>= \x -> x `seq` return x) `consM`
3649 postscanlMAfter' fstep initial done s
3651{-# INLINE scanl' #-}
3652scanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
3653scanl' f = scanlM' (\a b -> return (f a b))
3655{-# INLINE_NORMAL scanlM #-}
3656scanlM :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
3657scanlM fstep begin s = begin `cons` postscanlM fstep begin s
3660scanl :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
3661scanl f = scanlM (\a b -> return (f a b))
3663{-# INLINE_NORMAL scanl1M #-}
3664scanl1M :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
3665scanl1M fstep (Stream step state) = Stream step' (state, Nothing)
3667 {-# INLINE_LATE step' #-}
3668 step' gst (st, Nothing) = do
3671 Yield x s -> return $ Yield x (s, Just x)
3672 Skip s -> return $ Skip (s, Nothing)
3675 step' gst (st, Just acc) = do
3680 return $ Yield z (s, Just z)
3681 Skip s -> return $ Skip (s, Just acc)
3684{-# INLINE scanl1 #-}
3685scanl1 :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
3686scanl1 f = scanl1M (\x y -> return (f x y))
3688{-# INLINE_NORMAL scanl1M' #-}
3689scanl1M' :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
3690scanl1M' fstep (Stream step state) = Stream step' (state, Nothing)
3692 {-# INLINE_LATE step' #-}
3693 step' gst (st, Nothing) = do
3696 Yield x s -> x `seq` return $ Yield x (s, Just x)
3697 Skip s -> return $ Skip (s, Nothing)
3700 step' gst (st, Just acc) = acc `seq` do
3705 z `seq` return $ Yield z (s, Just z)
3706 Skip s -> return $ Skip (s, Just acc)
3709{-# INLINE scanl1' #-}
3710scanl1' :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
3711scanl1' f = scanl1M' (\x y -> return (f x y))
3713------------------------------------------------------------------------------
3715------------------------------------------------------------------------------
3717data RollingMapState s a = RollingMapInit s | RollingMapGo s a
3719{-# INLINE rollingMapM #-}
3720rollingMapM :: Monad m => (a -> a -> m b) -> Stream m a -> Stream m b
3721rollingMapM f (Stream step1 state1) = Stream step (RollingMapInit state1)
3723 step gst (RollingMapInit st) = do
3724 r <- step1 (adaptState gst) st
3726 Yield x s -> Skip $ RollingMapGo s x
3727 Skip s -> Skip $ RollingMapInit s
3730 step gst (RollingMapGo s1 x1) = do
3731 r <- step1 (adaptState gst) s1
3735 return $ Yield res $ RollingMapGo s x
3736 Skip s -> return $ Skip $ RollingMapGo s x1
3737 Stop -> return $ Stop
3739{-# INLINE rollingMap #-}
3740rollingMap :: Monad m => (a -> a -> b) -> Stream m a -> Stream m b
3741rollingMap f = rollingMapM (\x y -> return $ f x y)
3743------------------------------------------------------------------------------
3744-- Tapping/Distributing
3745------------------------------------------------------------------------------
3748tap :: Monad m => Fold m a b -> Stream m a -> Stream m a
3749tap (Fold fstep initial extract) (Stream step state) = Stream step' Nothing
3753 step' _ Nothing = do
3755 return $ Skip (Just (r, state))
3757 step' gst (Just (acc, st)) = acc `seq` do
3762 return $ Yield x (Just (acc', s))
3763 Skip s -> return $ Skip (Just (acc, s))
3768{-# INLINE_NORMAL tapOffsetEvery #-}
3769tapOffsetEvery :: Monad m
3770 => Int -> Int -> Fold m a b -> Stream m a -> Stream m a
3771tapOffsetEvery offset n (Fold fstep initial extract) (Stream step state) =
3772 Stream step' Nothing
3776 {-# INLINE_LATE step' #-}
3777 step' _ Nothing = do
3779 return $ Skip (Just (r, state, offset `mod` n))
3781 step' gst (Just (acc, st, count)) | count <= 0 = do
3785 !acc' <- fstep acc x
3786 return $ Yield x (Just (acc', s, n - 1))
3787 Skip s -> return $ Skip (Just (acc, s, count))
3792 step' gst (Just (acc, st, count)) = do
3795 Yield x s -> return $ Yield x (Just (acc, s, count - 1))
3796 Skip s -> return $ Skip (Just (acc, s, count))
3801{-# INLINE_NORMAL pollCounts #-}
3805 -> (Stream m Int -> Stream m Int)
3809pollCounts predicate transf fld (Stream step state) = Stream step' Nothing
3812 {-# INLINE_LATE step' #-}
3813 step' _ Nothing = do
3814 -- As long as we are using an "Int" for counts lockfree reads from
3815 -- Var should work correctly on both 32-bit and 64-bit machines.
3816 -- However, an Int on a 32-bit machine may overflow quickly.
3817 countVar <- liftIO $ newVar (0 :: Int)
3819 $ void $ runFold fld
3820 $ transf $ fromPrimVar countVar
3821 return $ Skip (Just (countVar, tid, state))
3823 step' gst (Just (countVar, tid, st)) = do
3827 when (predicate x) $ liftIO $ modifyVar' countVar (+ 1)
3828 return $ Yield x (Just (countVar, tid, s))
3829 Skip s -> return $ Skip (Just (countVar, tid, s))
3831 liftIO $ killThread tid
3834{-# INLINE_NORMAL tapRate #-}
3836 (MonadAsync m, MonadCatch m)
3841tapRate samplingRate action (Stream step state) = Stream step' Nothing
3843 {-# NOINLINE loop #-}
3844 loop countVar prev = do
3847 (do liftIO $ threadDelay (round $ samplingRate * 1000000)
3848 i <- liftIO $ readVar countVar
3849 let !diff = i - prev
3852 (\(e :: AsyncException) -> do
3853 i <- liftIO $ readVar countVar
3854 let !diff = i - prev
3856 throwM (MC.toException e))
3859 {-# INLINE_LATE step' #-}
3860 step' _ Nothing = do
3861 countVar <- liftIO $ newVar 0
3862 tid <- fork $ loop countVar 0
3863 ref <- liftIO $ newIORef ()
3864 _ <- liftIO $ mkWeakIORef ref (killThread tid)
3865 return $ Skip (Just (countVar, tid, state, ref))
3867 step' gst (Just (countVar, tid, st, ref)) = do
3871 liftIO $ modifyVar' countVar (+ 1)
3872 return $ Yield x (Just (countVar, tid, s, ref))
3873 Skip s -> return $ Skip (Just (countVar, tid, s, ref))
3875 liftIO $ killThread tid
3879-------------------------------------------------------------------------------
3881-------------------------------------------------------------------------------
3883{-# INLINE_NORMAL takeWhileM #-}
3884takeWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
3885takeWhileM f (Stream step state) = Stream step' state
3887 {-# INLINE_LATE step' #-}
3893 return $ if b then Yield x s else Stop
3894 Skip s -> return $ Skip s
3897{-# INLINE takeWhile #-}
3898takeWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
3899takeWhile f = takeWhileM (return . f)
3901{-# INLINE_NORMAL drop #-}
3902drop :: Monad m => Int -> Stream m a -> Stream m a
3903drop n (Stream step state) = Stream step' (state, Just n)
3905 {-# INLINE_LATE step' #-}
3906 step' gst (st, Just i)
3911 Yield _ s -> Skip (s, Just (i - 1))
3912 Skip s -> Skip (s, Just i)
3914 | otherwise = return $ Skip (st, Nothing)
3916 step' gst (st, Nothing) = do
3920 Yield x s -> Yield x (s, Nothing)
3921 Skip s -> Skip (s, Nothing)
3924data DropWhileState s a
3926 | DropWhileYield a s
3929{-# INLINE_NORMAL dropWhileM #-}
3930dropWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
3931dropWhileM f (Stream step state) = Stream step' (DropWhileDrop state)
3933 {-# INLINE_LATE step' #-}
3934 step' gst (DropWhileDrop st) = do
3940 then return $ Skip (DropWhileDrop s)
3941 else return $ Skip (DropWhileYield x s)
3942 Skip s -> return $ Skip (DropWhileDrop s)
3945 step' gst (DropWhileNext st) = do
3948 Yield x s -> return $ Skip (DropWhileYield x s)
3949 Skip s -> return $ Skip (DropWhileNext s)
3952 step' _ (DropWhileYield x st) = return $ Yield x (DropWhileNext st)
3954{-# INLINE dropWhile #-}
3955dropWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
3956dropWhile f = dropWhileM (return . f)
3958{-# INLINE_NORMAL filterM #-}
3959filterM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
3960filterM f (Stream step state) = Stream step' state
3962 {-# INLINE_LATE step' #-}
3971 Skip s -> return $ Skip s
3974{-# INLINE filter #-}
3975filter :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
3976filter f = filterM (return . f)
3978{-# INLINE_NORMAL uniq #-}
3979uniq :: (Eq a, Monad m) => Stream m a -> Stream m a
3980uniq (Stream step state) = Stream step' (Nothing, state)
3982 {-# INLINE_LATE step' #-}
3983 step' gst (Nothing, st) = do
3986 Yield x s -> return $ Yield x (Just x, s)
3987 Skip s -> return $ Skip (Nothing, s)
3989 step' gst (Just x, st) = do
3992 Yield y s | x == y -> return $ Skip (Just x, s)
3993 | otherwise -> return $ Yield y (Just y, s)
3994 Skip s -> return $ Skip (Just x, s)
3997------------------------------------------------------------------------------
3998-- Transformation by Mapping
3999------------------------------------------------------------------------------
4001{-# INLINE_NORMAL sequence #-}
4002sequence :: Monad m => Stream m (m a) -> Stream m a
4003sequence (Stream step state) = Stream step' state
4005 {-# INLINE_LATE step' #-}
4007 r <- step (adaptState gst) st
4009 Yield x s -> x >>= \a -> return (Yield a s)
4010 Skip s -> return $ Skip s
4013------------------------------------------------------------------------------
4015------------------------------------------------------------------------------
4017data LoopState x s = FirstYield s
4018 | InterspersingYield s
4021{-# INLINE_NORMAL intersperseM #-}
4022intersperseM :: Monad m => m a -> Stream m a -> Stream m a
4023intersperseM m (Stream step state) = Stream step' (FirstYield state)
4025 {-# INLINE_LATE step' #-}
4026 step' gst (FirstYield st) = do
4030 Yield x s -> Skip (YieldAndCarry x s)
4031 Skip s -> Skip (FirstYield s)
4034 step' gst (InterspersingYield st) = do
4039 return $ Yield a (YieldAndCarry x s)
4040 Skip s -> return $ Skip $ InterspersingYield s
4043 step' _ (YieldAndCarry x st) = return $ Yield x (InterspersingYield st)
4045{-# INLINE_NORMAL intersperseM_ #-}
4046intersperseM_ :: Monad m => m b -> Stream m a -> Stream m a
4047intersperseM_ m (Stream step1 state1) = Stream step (Left (pure (), state1))
4049 {-# INLINE_LATE step #-}
4050 step gst (Left (eff, st)) = do
4053 Yield x s -> eff >> return (Yield x (Right s))
4054 Skip s -> return $ Skip (Left (eff, s))
4057 step _ (Right st) = return $ Skip $ Left (void m, st)
4062 | SuffixYield a (SuffixState s a)
4064{-# INLINE_NORMAL intersperseSuffix #-}
4065intersperseSuffix :: forall m a. Monad m => m a -> Stream m a -> Stream m a
4066intersperseSuffix action (Stream step state) = Stream step' (SuffixElem state)
4068 {-# INLINE_LATE step' #-}
4069 step' gst (SuffixElem st) = do
4072 Yield x s -> Skip (SuffixYield x (SuffixSuffix s))
4073 Skip s -> Skip (SuffixElem s)
4076 step' _ (SuffixSuffix st) = do
4077 action >>= \r -> return $ Skip (SuffixYield r (SuffixElem st))
4079 step' _ (SuffixYield x next) = return $ Yield x next
4081{-# INLINE_NORMAL intersperseSuffix_ #-}
4082intersperseSuffix_ :: Monad m => m b -> Stream m a -> Stream m a
4083intersperseSuffix_ m (Stream step1 state1) = Stream step (Left state1)
4085 {-# INLINE_LATE step #-}
4086 step gst (Left st) = do
4089 Yield x s -> return $ Yield x (Right s)
4090 Skip s -> return $ Skip $ Left s
4093 step _ (Right st) = m >> return (Skip (Left st))
4095data SuffixSpanState s a
4096 = SuffixSpanElem s Int
4097 | SuffixSpanSuffix s
4098 | SuffixSpanYield a (SuffixSpanState s a)
4102-- | intersperse after every n items
4103{-# INLINE_NORMAL intersperseSuffixBySpan #-}
4104intersperseSuffixBySpan :: forall m a. Monad m
4105 => Int -> m a -> Stream m a -> Stream m a
4106intersperseSuffixBySpan n action (Stream step state) =
4107 Stream step' (SuffixSpanElem state n)
4109 {-# INLINE_LATE step' #-}
4110 step' gst (SuffixSpanElem st i) | i > 0 = do
4113 Yield x s -> Skip (SuffixSpanYield x (SuffixSpanElem s (i - 1)))
4114 Skip s -> Skip (SuffixSpanElem s i)
4115 Stop -> if i == n then Stop else Skip SuffixSpanLast
4116 step' _ (SuffixSpanElem st _) = return $ Skip (SuffixSpanSuffix st)
4118 step' _ (SuffixSpanSuffix st) = do
4119 action >>= \r -> return $ Skip (SuffixSpanYield r (SuffixSpanElem st n))
4121 step' _ (SuffixSpanLast) = do
4122 action >>= \r -> return $ Skip (SuffixSpanYield r SuffixSpanStop)
4124 step' _ (SuffixSpanYield x next) = return $ Yield x next
4126 step' _ (SuffixSpanStop) = return Stop
4128{-# INLINE intersperse #-}
4129intersperse :: Monad m => a -> Stream m a -> Stream m a
4130intersperse a = intersperseM (return a)
4132{-# INLINE_NORMAL insertBy #-}
4133insertBy :: Monad m => (a -> a -> Ordering) -> a -> Stream m a -> Stream m a
4134insertBy cmp a (Stream step state) = Stream step' (state, False, Nothing)
4136 {-# INLINE_LATE step' #-}
4137 step' gst (st, False, _) = do
4140 Yield x s -> case cmp a x of
4141 GT -> return $ Yield x (s, False, Nothing)
4142 _ -> return $ Yield a (s, True, Just x)
4143 Skip s -> return $ Skip (s, False, Nothing)
4144 Stop -> return $ Yield a (st, True, Nothing)
4146 step' _ (_, True, Nothing) = return Stop
4148 step' gst (st, True, Just prev) = do
4151 Yield x s -> return $ Yield prev (s, True, Just x)
4152 Skip s -> return $ Skip (s, True, Just prev)
4153 Stop -> return $ Yield prev (st, True, Nothing)
4155------------------------------------------------------------------------------
4157------------------------------------------------------------------------------
4159{-# INLINE_NORMAL deleteBy #-}
4160deleteBy :: Monad m => (a -> a -> Bool) -> a -> Stream m a -> Stream m a
4161deleteBy eq x (Stream step state) = Stream step' (state, False)
4163 {-# INLINE_LATE step' #-}
4164 step' gst (st, False) = do
4167 Yield y s -> return $
4168 if eq x y then Skip (s, True) else Yield y (s, False)
4169 Skip s -> return $ Skip (s, False)
4172 step' gst (st, True) = do
4175 Yield y s -> return $ Yield y (s, True)
4176 Skip s -> return $ Skip (s, True)
4179------------------------------------------------------------------------------
4180-- Transformation by Map and Filter
4181------------------------------------------------------------------------------
4183-- XXX Will this always fuse properly?
4184{-# INLINE_NORMAL mapMaybe #-}
4185mapMaybe :: Monad m => (a -> Maybe b) -> Stream m a -> Stream m b
4186mapMaybe f = fmap fromJust . filter isJust . map f
4188{-# INLINE_NORMAL mapMaybeM #-}
4189mapMaybeM :: Monad m => (a -> m (Maybe b)) -> Stream m a -> Stream m b
4190mapMaybeM f = fmap fromJust . filter isJust . mapM f
4192------------------------------------------------------------------------------
4194------------------------------------------------------------------------------
4196{-# INLINE_NORMAL indexed #-}
4197indexed :: Monad m => Stream m a -> Stream m (Int, a)
4198indexed (Stream step state) = Stream step' (state, 0)
4200 {-# INLINE_LATE step' #-}
4201 step' gst (st, i) = i `seq` do
4202 r <- step (adaptState gst) st
4204 Yield x s -> return $ Yield (i, x) (s, i+1)
4205 Skip s -> return $ Skip (s, i)
4208{-# INLINE_NORMAL indexedR #-}
4209indexedR :: Monad m => Int -> Stream m a -> Stream m (Int, a)
4210indexedR m (Stream step state) = Stream step' (state, m)
4212 {-# INLINE_LATE step' #-}
4213 step' gst (st, i) = i `seq` do
4214 r <- step (adaptState gst) st
4216 Yield x s -> let i' = i - 1
4217 in return $ Yield (i, x) (s, i')
4218 Skip s -> return $ Skip (s, i)
4221{-# INLINE_NORMAL zipWithM #-}
4223 => (a -> b -> m c) -> Stream m a -> Stream m b -> Stream m c
4224zipWithM f (Stream stepa ta) (Stream stepb tb) = Stream step (ta, tb, Nothing)
4226 {-# INLINE_LATE step #-}
4227 step gst (sa, sb, Nothing) = do
4228 r <- stepa (adaptState gst) sa
4231 Yield x sa' -> Skip (sa', sb, Just x)
4232 Skip sa' -> Skip (sa', sb, Nothing)
4235 step gst (sa, sb, Just x) = do
4236 r <- stepb (adaptState gst) sb
4240 return $ Yield z (sa, sb', Nothing)
4241 Skip sb' -> return $ Skip (sa, sb', Just x)
4244#if __GLASGOW_HASKELL__ >= 801
4245{-# RULES "zipWithM xs xs"
4246 forall f xs. zipWithM @Identity f xs xs = mapM (\x -> f x x) xs #-}
4249{-# INLINE zipWith #-}
4250zipWith :: Monad m => (a -> b -> c) -> Stream m a -> Stream m b -> Stream m c
4251zipWith f = zipWithM (\a b -> return (f a b))
4253------------------------------------------------------------------------------
4255------------------------------------------------------------------------------
4257{-# INLINE_NORMAL mergeByM #-}
4260 => (a -> a -> m Ordering) -> Stream m a -> Stream m a -> Stream m a
4261mergeByM cmp (Stream stepa ta) (Stream stepb tb) =
4262 Stream step (Just ta, Just tb, Nothing, Nothing)
4264 {-# INLINE_LATE step #-}
4266 -- one of the values is missing, and the corresponding stream is running
4267 step gst (Just sa, sb, Nothing, b) = do
4270 Yield a sa' -> Skip (Just sa', sb, Just a, b)
4271 Skip sa' -> Skip (Just sa', sb, Nothing, b)
4272 Stop -> Skip (Nothing, sb, Nothing, b)
4274 step gst (sa, Just sb, a, Nothing) = do
4277 Yield b sb' -> Skip (sa, Just sb', a, Just b)
4278 Skip sb' -> Skip (sa, Just sb', a, Nothing)
4279 Stop -> Skip (sa, Nothing, a, Nothing)
4281 -- both the values are available
4282 step _ (sa, sb, Just a, Just b) = do
4284 return $ case res of
4285 GT -> Yield b (sa, sb, Just a, Nothing)
4286 _ -> Yield a (sa, sb, Nothing, Just b)
4288 -- one of the values is missing, corresponding stream is done
4289 step _ (Nothing, sb, Nothing, Just b) =
4290 return $ Yield b (Nothing, sb, Nothing, Nothing)
4292 step _ (sa, Nothing, Just a, Nothing) =
4293 return $ Yield a (sa, Nothing, Nothing, Nothing)
4295 step _ (Nothing, Nothing, Nothing, Nothing) = return Stop
4297{-# INLINE mergeBy #-}
4300 => (a -> a -> Ordering) -> Stream m a -> Stream m a -> Stream m a
4301mergeBy cmp = mergeByM (\a b -> return $ cmp a b)
4303------------------------------------------------------------------------------
4304-- Transformation comprehensions
4305------------------------------------------------------------------------------
4307{-# INLINE_NORMAL the #-}
4308the :: (Eq a, Monad m) => Stream m a -> m (Maybe a)
4309the (Stream step state) = go state
4312 r <- step defState st
4314 Yield x s -> go' x s
4316 Stop -> return Nothing
4318 r <- step defState st
4320 Yield x s | x == n -> go' n s
4321 | otherwise -> return Nothing
4323 Stop -> return (Just n)
4325{-# INLINE runFold #-}
4326runFold :: (Monad m) => Fold m a b -> Stream m a -> m b
4327runFold (Fold step begin done) = foldlMx' step begin done
4329-------------------------------------------------------------------------------
4330-- Concurrent application and fold
4331-------------------------------------------------------------------------------
4333-- XXX These functions should be moved to Stream/Parallel.hs
4335-- Using StreamD the worker stream producing code can fuse with the code to
4336-- queue output to the SVar giving some perf boost.
4338-- Note that StreamD can only be used in limited situations, specifically, we
4339-- cannot implement joinStreamVarPar using this.
4341-- XXX make sure that the SVar passed is a Parallel style SVar.
4343-- | Fold the supplied stream to the SVar asynchronously using Parallel
4344-- concurrency style.
4345-- {-# INLINE_NORMAL toSVarParallel #-}
4346{-# INLINE toSVarParallel #-}
4347toSVarParallel :: MonadAsync m
4348 => State t m a -> SVar t m a -> Stream m a -> m ()
4349toSVarParallel st sv xs =
4350 if svarInspectMode sv
4354 case getYieldLimit st of
4355 Nothing -> doFork (work Nothing)
4357 (handleChildException sv)
4358 Just _ -> doFork (workLim Nothing)
4360 (handleChildException sv)
4365 {-# NOINLINE work #-}
4366 work info = (runFold (FL.toParallelSVar sv info) xs)
4368 {-# NOINLINE workLim #-}
4369 workLim info = runFold (FL.toParallelSVarLimited sv info) xs
4371 {-# NOINLINE forkWithDiag #-}
4373 -- We do not use workerCount in case of ParallelVar but still there is
4374 -- no harm in maintaining it correctly.
4375 liftIO $ atomicModifyIORefCAS_ (workerCount sv) $ \n -> n + 1
4377 -- This allocation matters when significant number of workers are being
4378 -- sent. We allocate it only when needed. The overhead increases by 4x.
4380 case yieldRateInfo sv of
4381 Nothing -> return Nothing
4382 Just _ -> liftIO $ do
4383 cntRef <- newIORef 0
4384 t <- getTime Monotonic
4385 lat <- newIORef (0, t)
4386 return $ Just WorkerInfo
4387 { workerYieldMax = 0
4388 , workerYieldCount = cntRef
4389 , workerLatencyStart = lat
4392 case getYieldLimit st of
4393 Nothing -> doFork (work winfo)
4395 (handleChildException sv)
4396 Just _ -> doFork (workLim winfo)
4398 (handleChildException sv)
4401{-# INLINE_NORMAL mkParallelD #-}
4402mkParallelD :: MonadAsync m => Stream m a -> Stream m a
4403mkParallelD m = Stream step Nothing
4406 step gst Nothing = do
4407 sv <- newParallelVar StopNone gst
4408 toSVarParallel gst sv m
4409 -- XXX use unfold instead?
4410 return $ Skip $ Just $ fromSVar sv
4412 step gst (Just (UnStream step1 st)) = do
4415 Yield a s -> Yield a (Just $ Stream step1 s)
4416 Skip s -> Skip (Just $ Stream step1 s)
4419-- Compare with mkAsync. mkAsync uses an Async style SVar whereas this uses a
4420-- parallel style SVar for evaluation. Currently, parallel style cannot use
4421-- rate control whereas Async style can use rate control. In async style SVar
4422-- the worker thread terminates when the buffer is full whereas in Parallel
4425-- | Make the stream producer and consumer run concurrently by introducing a
4426-- buffer between them. The producer thread evaluates the input stream until
4427-- the buffer fills, it blocks if the buffer is full until there is space in
4428-- the buffer. The consumer consumes the stream lazily from the buffer.
4432{-# INLINE_NORMAL mkParallel #-}
4433mkParallel :: (K.IsStream t, MonadAsync m) => t m a -> t m a
4434mkParallel = fromStreamD . mkParallelD . toStreamD
4436-- Note: we can use another API with two callbacks stop and yield if we want
4437-- the callback to be able to indicate end of stream.
4439-- | Generates a callback and a stream pair. The callback returned is used to
4440-- queue values to the stream. The stream is infinite, there is no way for the
4441-- callback to indicate that it is done now.
4445{-# INLINE_NORMAL newCallbackStream #-}
4446newCallbackStream :: (K.IsStream t, MonadAsync m) => m ((a -> m ()), t m a)
4447newCallbackStream = do
4448 sv <- newParallelVar StopNone defState
4450 -- XXX Add our own thread-id to the SVar as we can not know the callback's
4451 -- thread-id and the callback is not run in a managed worker. We need to
4452 -- handle this better.
4453 liftIO myThreadId >>= modifyThread sv
4455 let callback a = liftIO $ void $ send sv (ChildYield a)
4456 -- XXX we can return an SVar and then the consumer can unfold from the
4458 return (callback, fromStreamD (fromSVar sv))
4460-------------------------------------------------------------------------------
4462-------------------------------------------------------------------------------
4464-- | Create an SVar with a fold consumer that will fold any elements sent to it
4465-- using the supplied fold function.
4466{-# INLINE newFoldSVar #-}
4467newFoldSVar :: MonadAsync m => State t m a -> Fold m a b -> m (SVar t m a)
4468newFoldSVar stt f = do
4469 -- Buffer size for the SVar is derived from the current state
4470 sv <- newParallelVar StopAny (adaptState stt)
4471 -- Add the producer thread-id to the SVar.
4472 liftIO myThreadId >>= modifyThread sv
4473 void $ doFork (work sv) (svarMrun sv) (handleFoldException sv)
4478 {-# NOINLINE work #-}
4479 work sv = void $ runFold f $ fromProducer sv
4481data TapState sv st = TapInit | Tapping sv st | TapDone st
4483{-# INLINE_NORMAL tapAsync #-}
4484tapAsync :: MonadAsync m => Fold m a b -> Stream m a -> Stream m a
4485tapAsync f (Stream step1 state1) = Stream step TapInit
4489 -- In general, a Stop event would come equipped with the result
4490 -- of the fold. It is not used here but it would be useful in
4491 -- applicative and distribute.
4492 done <- fromConsumer svr
4493 when (not done) $ do
4494 liftIO $ withDiagMVar svr "teeToSVar: waiting to drain"
4495 $ takeMVar (outputDoorBellFromConsumer svr)
4499 liftIO $ sendStop svr Nothing
4500 -- drain/wait until a stop event arrives from the fold.
4503 {-# INLINE_LATE step #-}
4504 step gst TapInit = do
4505 sv <- newFoldSVar gst f
4506 return $ Skip (Tapping sv state1)
4508 step gst (Tapping sv st) = do
4512 done <- pushToFold sv a
4515 -- XXX we do not need to wait synchronously here
4517 return $ Yield a (TapDone s)
4518 else return $ Yield a (Tapping sv s)
4519 Skip s -> return $ Skip (Tapping sv s)
4524 step gst (TapDone st) = do
4527 Yield a s -> Yield a (TapDone s)
4528 Skip s -> Skip (TapDone s)
4531-- XXX Exported from Array again as this fold is specific to Array
4532-- | Take last 'n' elements from the stream and discard the rest.
4534lastN :: (Storable a, MonadIO m) => Int -> Fold m a (Array a)
4536 | n <= 0 = fmap (const mempty) FL.drain
4537 | otherwise = Fold step initial done
4539 step (Tuple3' rb rh i) a = do
4540 rh1 <- liftIO $ RB.unsafeInsert rb rh a
4541 return $ Tuple3' rb rh1 (i + 1)
4542 initial = fmap (\(a, b) -> Tuple3' a b (0 :: Int)) $ liftIO $ RB.new n
4543 done (Tuple3' rb rh i) = do
4544 arr <- liftIO $ A.newArray n
4545 foldFunc i rh snoc' arr rb
4546 snoc' b a = liftIO $ A.unsafeSnoc b a
4548 | i < n = RB.unsafeFoldRingM
4549 | otherwise = RB.unsafeFoldRingFullM
4551------------------------------------------------------------------------------
4553------------------------------------------------------------------------------
4555-- XXX using getTime in the loop can be pretty expensive especially for
4556-- computations where iterations are lightweight. We have the following
4559-- 1) Run a timeout thread updating a flag asynchronously and check that
4560-- flag here, that way we can have a cheap termination check.
4562-- 2) Use COARSE clock to get time with lower resolution but more efficiently.
4564-- 3) Use rdtscp/rdtsc to get time directly from the processor, compute the
4565-- termination value of rdtsc in the beginning and then in each iteration just
4566-- get rdtsc and check if we should terminate.
4570 | TakeByTimeCheck st s
4571 | TakeByTimeYield st s
4573{-# INLINE_NORMAL takeByTime #-}
4574takeByTime :: (MonadIO m, TimeUnit64 t) => t -> Stream m a -> Stream m a
4575takeByTime duration (Stream step1 state1) = Stream step (TakeByTimeInit state1)
4578 lim = toRelTime64 duration
4580 {-# INLINE_LATE step #-}
4581 step _ (TakeByTimeInit _) | lim == 0 = return Stop
4582 step _ (TakeByTimeInit st) = do
4583 t0 <- liftIO $ getTime Monotonic
4584 return $ Skip (TakeByTimeYield st t0)
4585 step _ (TakeByTimeCheck st t0) = do
4586 t <- liftIO $ getTime Monotonic
4588 if diffAbsTime64 t t0 > lim
4590 else Skip (TakeByTimeYield st t0)
4591 step gst (TakeByTimeYield st t0) = do
4594 Yield x s -> Yield x (TakeByTimeCheck s t0)
4595 Skip s -> Skip (TakeByTimeCheck s t0)
4598data DropByTime st s x
4600 | DropByTimeGen st s
4601 | DropByTimeCheck st s x
4602 | DropByTimeYield st
4604{-# INLINE_NORMAL dropByTime #-}
4605dropByTime :: (MonadIO m, TimeUnit64 t) => t -> Stream m a -> Stream m a
4606dropByTime duration (Stream step1 state1) = Stream step (DropByTimeInit state1)
4609 lim = toRelTime64 duration
4611 {-# INLINE_LATE step #-}
4612 step _ (DropByTimeInit st) = do
4613 t0 <- liftIO $ getTime Monotonic
4614 return $ Skip (DropByTimeGen st t0)
4615 step gst (DropByTimeGen st t0) = do
4618 Yield x s -> Skip (DropByTimeCheck s t0 x)
4619 Skip s -> Skip (DropByTimeGen s t0)
4621 step _ (DropByTimeCheck st t0 x) = do
4622 t <- liftIO $ getTime Monotonic
4623 if diffAbsTime64 t t0 <= lim
4624 then return $ Skip $ DropByTimeGen st t0
4625 else return $ Yield x $ DropByTimeYield st
4626 step gst (DropByTimeYield st) = do
4629 Yield x s -> Yield x (DropByTimeYield s)
4630 Skip s -> Skip (DropByTimeYield s)
4633{-# INLINE updateTimeVar #-}
4634updateTimeVar :: Var IO Int64 -> IO ()
4635updateTimeVar timeVar = do
4636 MicroSecond64 t <- fromAbsTime <$> getTime Monotonic
4637 modifyVar' timeVar (const t)
4639{-# INLINE updateWithDelay #-}
4640updateWithDelay :: RealFrac a => a -> Var IO Int64 -> IO ()
4641updateWithDelay precision timeVar = do
4642 threadDelay (delayTime precision)
4643 updateTimeVar timeVar
4647 -- Keep the minimum at least a millisecond to avoid high CPU usage
4648 {-# INLINE delayTime #-}
4650 | g' >= fromIntegral (maxBound :: Int) = maxBound
4652 | otherwise = round g'
4656 g' = g * 10 ^ (6 :: Int)
4658-- XXX we should move this to stream generation section of this file. Also, the
4659-- take/drop combinators above should be moved to filtering section.
4661{-# INLINE_NORMAL times #-}
4662times :: MonadAsync m => Double -> Stream m (AbsTime, RelTime64)
4663times g = Stream step Nothing
4667 {-# INLINE_LATE step #-}
4669 -- XXX note that this is safe only on a 64-bit machine. On a 32-bit
4670 -- machine a 64-bit 'Var' cannot be read consistently without a lock
4671 -- while another thread is writing to it.
4672 timeVar <- liftIO $ newVar (0 :: Int64)
4673 liftIO $ updateTimeVar timeVar
4674 tid <- forkManaged $ liftIO $ forever (updateWithDelay g timeVar)
4675 a <- liftIO $ readVar timeVar
4676 return $ Skip $ Just (timeVar, tid, a)
4678 step _ s@(Just (timeVar, _, t0)) = do
4679 a <- liftIO $ readVar timeVar
4680 -- XXX we can perhaps use an AbsTime64 using a 64 bit Int for
4681 -- efficiency. or maybe we can use a representation using Double for
4682 -- floating precision time
4683 return $ Yield (toAbsTime (MicroSecond64 t0),
4684 (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