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)
276 , intersperseSuffixBySpan
300 -- * Transformation comprehensions
306 , clearIORefFinalizer
318 -- * Concurrent Application
327import Control.Concurrent (killThread, myThreadId, takeMVar, threadDelay)
328import Control.Exception
329 (assert, Exception, SomeException, AsyncException, fromException, mask_)
330import Control.Monad (void, when, forever)
331import Control.Monad.Catch (MonadCatch, MonadThrow, throwM)
332import Control.Monad.IO.Class (MonadIO(..))
333import Control.Monad.Reader (ReaderT)
334import Control.Monad.State.Strict (StateT)
335import Control.Monad.Trans (MonadTrans(lift))
336import Control.Monad.Trans.Control (MonadBaseControl, liftBaseOp_)
337import Data.Bits (shiftR, shiftL, (.|.), (.&.))
338import Data.Functor.Identity (Identity(..))
339import Data.Int (Int64)
340import Data.IORef (newIORef, readIORef, mkWeakIORef, writeIORef, IORef)
341import Data.Maybe (fromJust, isJust, isNothing)
342import Data.Word (Word32)
343import Foreign.Ptr (Ptr)
344import Foreign.Storable (Storable(..))
345import GHC.Types (SPEC(..))
346import System.Mem (performMajorGC)
348 hiding (map, mapM, mapM_, repeat, foldr, last, take, filter,
349 takeWhile, drop, dropWhile, all, any, maximum, minimum, elem,
350 notElem, null, head, tail, zipWith, lookup, foldr1, sequence,
351 (!!), scanl, scanl1, concatMap, replicate, enumFromTo, concat,
352 reverse, iterate, splitAt)
354import qualified Control.Monad.Catch as MC
355import qualified Control.Monad.Reader as Reader
356import qualified Control.Monad.State.Strict as State
357import qualified Prelude
359import Fusion.Plugin.Types (Fuse(..))
360import Streamly.Internal.Mutable.Prim.Var
361 (Prim, Var, readVar, newVar, modifyVar')
362import Streamly.Internal.Data.Time.Units
363 (TimeUnit64, toRelTime64, diffAbsTime64, RelTime64)
365import Streamly.Internal.Data.Atomics (atomicModifyIORefCAS_)
366import Streamly.Internal.Memory.Array.Types (Array(..))
367import Streamly.Internal.Data.Fold.Types (Fold(..))
368import Streamly.Internal.Data.Parser (ParseError(..))
369import Streamly.Internal.Data.Pipe.Types (Pipe(..), PipeState(..))
370import Streamly.Internal.Data.Time.Clock (Clock(Monotonic), getTime)
371import Streamly.Internal.Data.Time.Units
372 (MicroSecond64(..), fromAbsTime, toAbsTime, AbsTime)
373import Streamly.Internal.Data.Unfold.Types (Unfold(..))
374import Streamly.Internal.Data.Strict (Tuple3'(..))
376import Streamly.Internal.Data.Stream.StreamD.Type
377import Streamly.Internal.Data.SVar
378import Streamly.Internal.Data.Stream.SVar (fromConsumer, pushToFold)
380import qualified Streamly.Internal.Data.Pipe.Types as Pipe
381import qualified Streamly.Internal.Memory.Array.Types as A
382import qualified Streamly.Internal.Data.Fold as FL
383import qualified Streamly.Memory.Ring as RB
384import qualified Streamly.Internal.Data.Stream.StreamK as K
385import qualified Streamly.Internal.Data.Parser as PR
386import qualified Streamly.Internal.Data.Parser.ParserD as PRD
388------------------------------------------------------------------------------
390------------------------------------------------------------------------------
392-- | An empty 'Stream'.
393{-# INLINE_NORMAL nil #-}
394nil :: Monad m => Stream m a
395nil = Stream (\_ _ -> return Stop) ()
397-- | An empty 'Stream' with a side effect.
398{-# INLINE_NORMAL nilM #-}
399nilM :: Monad m => m b -> Stream m a
400nilM m = Stream (\_ _ -> m >> return Stop) ()
402{-# INLINE_NORMAL consM #-}
403consM :: Monad m => m a -> Stream m a -> Stream m a
404consM m (Stream step state) = Stream step1 Nothing
406 {-# INLINE_LATE step1 #-}
407 step1 _ Nothing = m >>= \x -> return $ Yield x (Just state)
408 step1 gst (Just st) = do
412 Yield a s -> Yield a (Just s)
413 Skip s -> Skip (Just s)
416-- XXX implement in terms of consM?
417-- cons x = consM (return x)
419-- | Can fuse but has O(n^2) complexity.
420{-# INLINE_NORMAL cons #-}
421cons :: Monad m => a -> Stream m a -> Stream m a
422cons x (Stream step state) = Stream step1 Nothing
424 {-# INLINE_LATE step1 #-}
425 step1 _ Nothing = return $ Yield x (Just state)
426 step1 gst (Just st) = do
430 Yield a s -> Yield a (Just s)
431 Skip s -> Skip (Just s)
434-------------------------------------------------------------------------------
436-------------------------------------------------------------------------------
438-- Does not fuse, has the same performance as the StreamK version.
439{-# INLINE_NORMAL uncons #-}
440uncons :: Monad m => Stream m a -> m (Maybe (a, Stream m a))
441uncons (UnStream step state) = go state
444 r <- step defState st
446 Yield x s -> return $ Just (x, Stream step s)
448 Stop -> return Nothing
450------------------------------------------------------------------------------
451-- Generation by unfold
452------------------------------------------------------------------------------
454{-# INLINE_NORMAL unfoldrM #-}
455unfoldrM :: Monad m => (s -> m (Maybe (a, s))) -> s -> Stream m a
456unfoldrM next state = Stream step state
458 {-# INLINE_LATE step #-}
462 Just (x, s) -> Yield x s
465{-# INLINE_LATE unfoldr #-}
466unfoldr :: Monad m => (s -> Maybe (a, s)) -> s -> Stream m a
467unfoldr f = unfoldrM (return . f)
469-- | Convert an 'Unfold' into a 'Stream' by supplying it a seed.
471{-# INLINE_NORMAL unfold #-}
472unfold :: Monad m => Unfold m a b -> a -> Stream m b
473unfold (Unfold ustep inject) seed = Stream step Nothing
475 {-# INLINE_LATE step #-}
476 step _ Nothing = inject seed >>= return . Skip . Just
477 step _ (Just st) = do
480 Yield x s -> Yield x (Just s)
481 Skip s -> Skip (Just s)
484------------------------------------------------------------------------------
485-- Specialized Generation
486------------------------------------------------------------------------------
488{-# INLINE_NORMAL repeatM #-}
489repeatM :: Monad m => m a -> Stream m a
490repeatM x = Stream (\_ _ -> x >>= \r -> return $ Yield r ()) ()
492{-# INLINE_NORMAL repeat #-}
493repeat :: Monad m => a -> Stream m a
494repeat x = Stream (\_ _ -> return $ Yield x ()) ()
496{-# INLINE_NORMAL iterateM #-}
497iterateM :: Monad m => (a -> m a) -> m a -> Stream m a
498iterateM step = Stream (\_ st -> st >>= \x -> return $ Yield x (step x))
500{-# INLINE_NORMAL iterate #-}
501iterate :: Monad m => (a -> a) -> a -> Stream m a
502iterate step st = iterateM (return . step) (return st)
504{-# INLINE_NORMAL replicateM #-}
505replicateM :: forall m a. Monad m => Int -> m a -> Stream m a
506replicateM n p = Stream step n
508 {-# INLINE_LATE step #-}
510 | i <= 0 = return Stop
513 return $ Yield x (i - 1)
515{-# INLINE_NORMAL replicate #-}
516replicate :: Monad m => Int -> a -> Stream m a
517replicate n x = replicateM n (return x)
519-- This would not work properly for floats, therefore we put an Integral
521-- | Can be used to enumerate unbounded integrals. This does not check for
522-- overflow or underflow for bounded integrals.
523{-# INLINE_NORMAL enumerateFromStepIntegral #-}
524enumerateFromStepIntegral :: (Integral a, Monad m) => a -> a -> Stream m a
525enumerateFromStepIntegral from stride =
526 from `seq` stride `seq` Stream step from
528 {-# INLINE_LATE step #-}
529 step _ !x = return $ Yield x $! (x + stride)
531-- We are assuming that "to" is constrained by the type to be within
533{-# INLINE enumerateFromToIntegral #-}
534enumerateFromToIntegral :: (Monad m, Integral a) => a -> a -> Stream m a
535enumerateFromToIntegral from to =
536 takeWhile (<= to) $ enumerateFromStepIntegral from 1
538{-# INLINE enumerateFromIntegral #-}
539enumerateFromIntegral :: (Monad m, Integral a, Bounded a) => a -> Stream m a
540enumerateFromIntegral from = enumerateFromToIntegral from maxBound
542data EnumState a = EnumInit | EnumYield a a a | EnumStop
544{-# INLINE_NORMAL enumerateFromThenToIntegralUp #-}
545enumerateFromThenToIntegralUp
546 :: (Monad m, Integral a)
547 => a -> a -> a -> Stream m a
548enumerateFromThenToIntegralUp from next to = Stream step EnumInit
550 {-# INLINE_LATE step #-}
556 else Yield from EnumStop
557 else -- from <= next <= to
558 let stride = next - from
559 in Skip $ EnumYield from stride (to - stride)
561 step _ (EnumYield x stride toMinus) =
564 then Yield x EnumStop
565 else Yield x $ EnumYield (x + stride) stride toMinus
567 step _ EnumStop = return Stop
569{-# INLINE_NORMAL enumerateFromThenToIntegralDn #-}
570enumerateFromThenToIntegralDn
571 :: (Monad m, Integral a)
572 => a -> a -> a -> Stream m a
573enumerateFromThenToIntegralDn from next to = Stream step EnumInit
575 {-# INLINE_LATE step #-}
577 return $ if to > next
580 else Yield from EnumStop
581 else -- from >= next >= to
582 let stride = next - from
583 in Skip $ EnumYield from stride (to - stride)
585 step _ (EnumYield x stride toMinus) =
588 then Yield x EnumStop
589 else Yield x $ EnumYield (x + stride) stride toMinus
591 step _ EnumStop = return Stop
593{-# INLINE_NORMAL enumerateFromThenToIntegral #-}
594enumerateFromThenToIntegral
595 :: (Monad m, Integral a)
596 => a -> a -> a -> Stream m a
597enumerateFromThenToIntegral from next to
598 | next >= from = enumerateFromThenToIntegralUp from next to
599 | otherwise = enumerateFromThenToIntegralDn from next to
601{-# INLINE_NORMAL enumerateFromThenIntegral #-}
602enumerateFromThenIntegral
603 :: (Monad m, Integral a, Bounded a)
604 => a -> a -> Stream m a
605enumerateFromThenIntegral from next =
607 then enumerateFromThenToIntegralUp from next maxBound
608 else enumerateFromThenToIntegralDn from next minBound
610-- For floating point numbers if the increment is less than the precision then
611-- it just gets lost. Therefore we cannot always increment it correctly by just
613-- 9007199254740992 + 1 + 1 :: Double => 9.007199254740992e15
614-- 9007199254740992 + 2 :: Double => 9.007199254740994e15
616-- Instead we accumulate the increment counter and compute the increment
617-- every time before adding it to the starting number.
619-- This works for Integrals as well as floating point numbers, but
620-- enumerateFromStepIntegral is faster for integrals.
621{-# INLINE_NORMAL enumerateFromStepNum #-}
622enumerateFromStepNum :: (Monad m, Num a) => a -> a -> Stream m a
623enumerateFromStepNum from stride = Stream step 0
625 {-# INLINE_LATE step #-}
626 step _ !i = return $ (Yield $! (from + i * stride)) $! (i + 1)
628{-# INLINE_NORMAL numFrom #-}
629numFrom :: (Monad m, Num a) => a -> Stream m a
630numFrom from = enumerateFromStepNum from 1
632{-# INLINE_NORMAL numFromThen #-}
633numFromThen :: (Monad m, Num a) => a -> a -> Stream m a
634numFromThen from next = enumerateFromStepNum from (next - from)
636-- We cannot write a general function for Num. The only way to write code
637-- portable between the two is to use a 'Real' constraint and convert between
638-- Fractional and Integral using fromRational which is horribly slow.
639{-# INLINE_NORMAL enumerateFromToFractional #-}
640enumerateFromToFractional
641 :: (Monad m, Fractional a, Ord a)
642 => a -> a -> Stream m a
643enumerateFromToFractional from to =
644 takeWhile (<= to + 1 / 2) $ enumerateFromStepNum from 1
646{-# INLINE_NORMAL enumerateFromThenToFractional #-}
647enumerateFromThenToFractional
648 :: (Monad m, Fractional a, Ord a)
649 => a -> a -> a -> Stream m a
650enumerateFromThenToFractional from next to =
651 takeWhile predicate $ numFromThen from next
653 mid = (next - from) / 2
654 predicate | next >= from = (<= to + mid)
655 | otherwise = (>= to + mid)
657-------------------------------------------------------------------------------
658-- Generation by Conversion
659-------------------------------------------------------------------------------
661{-# INLINE_NORMAL fromIndicesM #-}
662fromIndicesM :: Monad m => (Int -> m a) -> Stream m a
663fromIndicesM gen = Stream step 0
665 {-# INLINE_LATE step #-}
668 return $ Yield x (i + 1)
670{-# INLINE fromIndices #-}
671fromIndices :: Monad m => (Int -> a) -> Stream m a
672fromIndices gen = fromIndicesM (return . gen)
674{-# INLINE_NORMAL generateM #-}
675generateM :: Monad m => Int -> (Int -> m a) -> Stream m a
676generateM n gen = n `seq` Stream step 0
678 {-# INLINE_LATE step #-}
679 step _ i | i < n = do
681 return $ Yield x (i + 1)
682 | otherwise = return Stop
684{-# INLINE generate #-}
685generate :: Monad m => Int -> (Int -> a) -> Stream m a
686generate n gen = generateM n (return . gen)
688-- XXX we need the MonadAsync constraint because of a rewrite rule.
689-- | Convert a list of monadic actions to a 'Stream'
690{-# INLINE_LATE fromListM #-}
691fromListM :: MonadAsync m => [m a] -> Stream m a
692fromListM = Stream step
694 {-# INLINE_LATE step #-}
695 step _ (m:ms) = m >>= \x -> return $ Yield x ms
696 step _ [] = return Stop
698{-# INLINE toStreamD #-}
699toStreamD :: (K.IsStream t, Monad m) => t m a -> Stream m a
700toStreamD = fromStreamK . K.toStream
702{-# INLINE_NORMAL fromPrimVar #-}
703fromPrimVar :: (MonadIO m, Prim a) => Var IO a -> Stream m a
704fromPrimVar var = Stream step ()
706 {-# INLINE_LATE step #-}
707 step _ () = liftIO (readVar var) >>= \x -> return $ Yield x ()
709-------------------------------------------------------------------------------
710-- Generation from SVar
711-------------------------------------------------------------------------------
713data FromSVarState t m a =
715 | FromSVarRead (SVar t m a)
716 | FromSVarLoop (SVar t m a) [ChildEvent a]
717 | FromSVarDone (SVar t m a)
719{-# INLINE_NORMAL fromSVar #-}
720fromSVar :: (MonadAsync m) => SVar t m a -> Stream m a
721fromSVar svar = Stream step FromSVarInit
724 {-# INLINE_LATE step #-}
725 step _ FromSVarInit = do
726 ref <- liftIO $ newIORef ()
727 _ <- liftIO $ mkWeakIORef ref hook
728 -- when this copy of svar gets garbage collected "ref" will get
729 -- garbage collected and our GC hook will be called.
730 let sv = svar{svarRef = Just ref}
731 return $ Skip (FromSVarRead sv)
735 {-# NOINLINE hook #-}
737 when (svarInspectMode svar) $ do
738 r <- liftIO $ readIORef (svarStopTime (svarStats svar))
740 printSVar svar "SVar Garbage Collected"
742 -- If there are any SVars referenced by this SVar a GC will prompt
743 -- them to be cleaned up quickly.
744 when (svarInspectMode svar) performMajorGC
746 step _ (FromSVarRead sv) = do
747 list <- readOutputQ sv
748 -- Reversing the output is important to guarantee that we process the
749 -- outputs in the same order as they were generated by the constituent
751 return $ Skip $ FromSVarLoop sv (Prelude.reverse list)
753 step _ (FromSVarLoop sv []) = do
754 done <- postProcess sv
755 return $ Skip $ if done
756 then (FromSVarDone sv)
757 else (FromSVarRead sv)
759 step _ (FromSVarLoop sv (ev : es)) = do
761 ChildYield a -> return $ Yield a (FromSVarLoop sv es)
762 ChildStop tid e -> do
766 stop <- shouldStop tid
769 liftIO (cleanupSVar sv)
770 return $ Skip (FromSVarDone sv)
771 else return $ Skip (FromSVarLoop sv es)
773 case fromException ex of
775 return $ Skip (FromSVarLoop sv es)
776 Nothing -> liftIO (cleanupSVar sv) >> throwM ex
780 case svarStopStyle sv of
781 StopNone -> return False
782 StopAny -> return True
784 sid <- liftIO $ readIORef (svarStopBy sv)
785 return $ if tid == sid then True else False
787 step _ (FromSVarDone sv) = do
788 when (svarInspectMode sv) $ do
789 t <- liftIO $ getTime Monotonic
790 liftIO $ writeIORef (svarStopTime (svarStats sv)) (Just t)
791 liftIO $ printSVar sv "SVar Done"
794-------------------------------------------------------------------------------
795-- Process events received by a fold consumer from a stream producer
796-------------------------------------------------------------------------------
798{-# INLINE_NORMAL fromProducer #-}
799fromProducer :: (MonadAsync m) => SVar t m a -> Stream m a
800fromProducer svar = Stream step (FromSVarRead svar)
803 {-# INLINE_LATE step #-}
804 step _ (FromSVarRead sv) = do
805 list <- readOutputQ sv
806 -- Reversing the output is important to guarantee that we process the
807 -- outputs in the same order as they were generated by the constituent
809 return $ Skip $ FromSVarLoop sv (Prelude.reverse list)
811 step _ (FromSVarLoop sv []) = return $ Skip $ FromSVarRead sv
812 step _ (FromSVarLoop sv (ev : es)) = do
814 ChildYield a -> return $ Yield a (FromSVarLoop sv es)
815 ChildStop tid e -> do
819 sendStopToProducer sv
820 return $ Skip (FromSVarDone sv)
821 Just _ -> error "Bug: fromProducer: received exception"
823 step _ (FromSVarDone sv) = do
824 when (svarInspectMode sv) $ do
825 t <- liftIO $ getTime Monotonic
826 liftIO $ writeIORef (svarStopTime (svarStats sv)) (Just t)
827 liftIO $ printSVar sv "SVar Done"
830 step _ FromSVarInit = undefined
832-------------------------------------------------------------------------------
833-- Hoisting the inner monad
834-------------------------------------------------------------------------------
836{-# INLINE_NORMAL hoist #-}
837hoist :: Monad n => (forall x. m x -> n x) -> Stream m a -> Stream n a
838hoist f (Stream step state) = (Stream step' state)
840 {-# INLINE_LATE step' #-}
842 r <- f $ step (adaptState gst) st
844 Yield x s -> Yield x s
848{-# INLINE generally #-}
849generally :: Monad m => Stream Identity a -> Stream m a
850generally = hoist (return . runIdentity)
852{-# INLINE_NORMAL liftInner #-}
853liftInner :: (Monad m, MonadTrans t, Monad (t m))
854 => Stream m a -> Stream (t m) a
855liftInner (Stream step state) = Stream step' state
857 {-# INLINE_LATE step' #-}
859 r <- lift $ step (adaptState gst) st
861 Yield x s -> Yield x s
865{-# INLINE_NORMAL runReaderT #-}
866runReaderT :: Monad m => s -> Stream (ReaderT s m) a -> Stream m a
867runReaderT sval (Stream step state) = Stream step' state
869 {-# INLINE_LATE step' #-}
871 r <- Reader.runReaderT (step (adaptState gst) st) sval
873 Yield x s -> Yield x s
877{-# INLINE_NORMAL evalStateT #-}
878evalStateT :: Monad m => s -> Stream (StateT s m) a -> Stream m a
879evalStateT sval (Stream step state) = Stream step' (state, sval)
881 {-# INLINE_LATE step' #-}
882 step' gst (st, sv) = do
883 (r, sv') <- State.runStateT (step (adaptState gst) st) sv
885 Yield x s -> Yield x (s, sv')
886 Skip s -> Skip (s, sv')
889{-# INLINE_NORMAL runStateT #-}
890runStateT :: Monad m => s -> Stream (StateT s m) a -> Stream m (s, a)
891runStateT sval (Stream step state) = Stream step' (state, sval)
893 {-# INLINE_LATE step' #-}
894 step' gst (st, sv) = do
895 (r, sv') <- State.runStateT (step (adaptState gst) st) sv
897 Yield x s -> Yield (sv', x) (s, sv')
898 Skip s -> Skip (s, sv')
901------------------------------------------------------------------------------
902-- Elimination by Folds
903------------------------------------------------------------------------------
905------------------------------------------------------------------------------
907------------------------------------------------------------------------------
909{-# INLINE_NORMAL foldr1 #-}
910foldr1 :: Monad m => (a -> a -> a) -> Stream m a -> m (Maybe a)
914 Nothing -> return Nothing
915 Just (h, t) -> fmap Just (foldr f h t)
917------------------------------------------------------------------------------
919------------------------------------------------------------------------------
921{-# INLINE_NORMAL foldlT #-}
922foldlT :: (Monad m, Monad (s m), MonadTrans s)
923 => (s m b -> a -> s m b) -> s m b -> Stream m a -> s m b
924foldlT fstep begin (Stream step state) = go SPEC begin state
927 r <- lift $ step defState st
929 Yield x s -> go SPEC (fstep acc x) s
930 Skip s -> go SPEC acc s
933-- Note, this is going to have horrible performance, because of the nature of
934-- the stream type (i.e. direct stream vs CPS). Its only for reference, it is
935-- likely be practically unusable.
936{-# INLINE_NORMAL foldlS #-}
938 => (Stream m b -> a -> Stream m b) -> Stream m b -> Stream m a -> Stream m b
939foldlS fstep begin (Stream step state) = Stream step' (Left (state, begin))
941 step' gst (Left (st, acc)) = do
942 r <- step (adaptState gst) st
944 Yield x s -> Skip (Left (s, fstep acc x))
945 Skip s -> Skip (Left (s, acc))
946 Stop -> Skip (Right acc)
948 step' gst (Right (Stream stp stt)) = do
949 r <- stp (adaptState gst) stt
951 Yield x s -> Yield x (Right (Stream stp s))
952 Skip s -> Skip (Right (Stream stp s))
955------------------------------------------------------------------------------
957------------------------------------------------------------------------------
959-- Inlined definition. Without the inline "serially/parser/take" benchmark
960-- degrades and parseMany does not fuse. Even using "inline" at the callsite
962{-# INLINE splitAt #-}
963splitAt :: Int -> [a] -> ([a],[a])
966 | otherwise = splitAt' n ls
968 splitAt' :: Int -> [a] -> ([a], [a])
969 splitAt' _ [] = ([], [])
970 splitAt' 1 (x:xs) = ([x], xs)
971 splitAt' m (x:xs) = (x:xs', xs'')
973 (xs', xs'') = splitAt' (m - 1) xs
975-- | Run a 'Parse' over a stream.
976{-# INLINE_NORMAL parselMx' #-}
979 => (s -> a -> m (PR.Step s b))
984parselMx' pstep initial extract (Stream step state) = do
985 initial >>= go SPEC state []
989 -- XXX currently we are using a dumb list based approach for backtracking
990 -- buffer. This can be replaced by a sliding/ring buffer using Data.Array.
991 -- That will allow us more efficient random back and forth movement.
993 go !_ st buf !pst = do
994 r <- step defState st
999 -- PR.Yield 0 pst1 -> go SPEC s [] pst1
1000 PR.Yield n pst1 -> do
1001 assert (n <= length (x:buf)) (return ())
1002 go SPEC s (Prelude.take n (x:buf)) pst1
1003 PR.YieldB n pst1 -> do
1004 assert (n <= length (x:buf)) (return ())
1005 let src0 = Prelude.take n (x:buf)
1006 src = Prelude.reverse src0
1007 gobuf SPEC s [] src pst1
1008 PR.Skip 0 pst1 -> go SPEC s (x:buf) pst1
1009 PR.Skip n pst1 -> do
1010 assert (n <= length (x:buf)) (return ())
1011 let (src0, buf1) = splitAt n (x:buf)
1012 src = Prelude.reverse src0
1013 gobuf SPEC s buf1 src pst1
1014 PR.Stop _ b -> return b
1015 PR.Error err -> throwM $ ParseError err
1016 Skip s -> go SPEC s buf pst
1019 gobuf !_ s buf [] !pst = go SPEC s buf pst
1020 gobuf !_ s buf (x:xs) !pst = do
1023 -- PR.Yield 0 pst1 -> go SPEC s [] pst1
1024 PR.Yield n pst1 -> do
1025 assert (n <= length (x:buf)) (return ())
1026 gobuf SPEC s (Prelude.take n (x:buf)) xs pst1
1027 PR.YieldB n pst1 -> do
1028 assert (n <= length (x:buf)) (return ())
1029 let src0 = Prelude.take n (x:buf)
1030 src = Prelude.reverse src0 ++ xs
1031 gobuf SPEC s [] src pst1
1032 PR.Skip 0 pst1 -> gobuf SPEC s (x:buf) xs pst1
1033 PR.Skip n pst1 -> do
1034 assert (n <= length (x:buf)) (return ())
1035 let (src0, buf1) = splitAt n (x:buf)
1036 src = Prelude.reverse src0 ++ xs
1037 gobuf SPEC s buf1 src pst1
1038 PR.Stop _ b -> return b
1039 PR.Error err -> throwM $ ParseError err
1041------------------------------------------------------------------------------
1043------------------------------------------------------------------------------
1045{-# ANN type ParseChunksState Fuse #-}
1046data ParseChunksState x inpBuf st pst =
1047 ParseChunksInit inpBuf st
1048 | ParseChunksInitLeftOver inpBuf
1049 | ParseChunksStream st inpBuf pst
1050 | ParseChunksBuf inpBuf st inpBuf pst
1051 | ParseChunksYield x (ParseChunksState x inpBuf st pst)
1053{-# INLINE_NORMAL parseMany #-}
1059parseMany (PRD.Parser pstep initial extract) (Stream step state) =
1060 Stream stepOuter (ParseChunksInit [] state)
1064 {-# INLINE_LATE stepOuter #-}
1065 -- Buffer is empty, go to stream processing loop
1066 stepOuter _ (ParseChunksInit [] st) = do
1067 initial >>= return . Skip . ParseChunksStream st []
1069 -- Buffer is not empty, go to buffered processing loop
1070 stepOuter _ (ParseChunksInit src st) = do
1071 initial >>= return . Skip . ParseChunksBuf src st []
1073 -- XXX we just discard any leftover input at the end
1074 stepOuter _ (ParseChunksInitLeftOver _) = return Stop
1076 -- Buffer is empty process elements from the stream
1077 stepOuter gst (ParseChunksStream st buf pst) = do
1078 r <- step (adaptState gst) st
1083 -- PR.Yield 0 pst1 -> go SPEC s [] pst1
1084 PR.Yield n pst1 -> do
1085 assert (n <= length (x:buf)) (return ())
1086 let buf1 = Prelude.take n (x:buf)
1087 return $ Skip $ ParseChunksStream s buf1 pst1
1088 PR.YieldB n pst1 -> do
1089 assert (n <= length (x:buf)) (return ())
1090 let src0 = Prelude.take n (x:buf)
1091 src = Prelude.reverse src0
1092 return $ Skip $ ParseChunksBuf src s [] pst1
1093 -- PR.Skip 0 pst1 ->
1094 -- return $ Skip $ ParseChunksStream s (x:buf) pst1
1095 PR.Skip n pst1 -> do
1096 assert (n <= length (x:buf)) (return ())
1097 let (src0, buf1) = splitAt n (x:buf)
1098 src = Prelude.reverse src0
1099 return $ Skip $ ParseChunksBuf src s buf1 pst1
1100 -- XXX Specialize for Stop 0 common case?
1102 assert (n <= length (x:buf)) (return ())
1103 let src = Prelude.reverse (Prelude.take n (x:buf))
1105 ParseChunksYield b (ParseChunksInit src s)
1106 PR.Error err -> throwM $ ParseError err
1107 Skip s -> return $ Skip $ ParseChunksStream s buf pst
1110 let src = Prelude.reverse buf
1112 ParseChunksYield b (ParseChunksInitLeftOver src)
1114 -- go back to stream processing mode
1115 stepOuter _ (ParseChunksBuf [] s buf pst) =
1116 return $ Skip $ ParseChunksStream s buf pst
1118 -- buffered processing loop
1119 stepOuter _ (ParseChunksBuf (x:xs) s buf pst) = do
1122 -- PR.Yield 0 pst1 ->
1123 PR.Yield n pst1 -> do
1124 assert (n <= length (x:buf)) (return ())
1125 let buf1 = Prelude.take n (x:buf)
1126 return $ Skip $ ParseChunksBuf xs s buf1 pst1
1127 PR.YieldB n pst1 -> do
1128 assert (n <= length (x:buf)) (return ())
1129 let src0 = Prelude.take n (x:buf)
1130 src = Prelude.reverse src0 ++ xs
1131 return $ Skip $ ParseChunksBuf src s [] pst1
1132 -- PR.Skip 0 pst1 -> return $ Skip $ ParseChunksBuf xs s (x:buf) pst1
1133 PR.Skip n pst1 -> do
1134 assert (n <= length (x:buf)) (return ())
1135 let (src0, buf1) = splitAt n (x:buf)
1136 src = Prelude.reverse src0 ++ xs
1137 return $ Skip $ ParseChunksBuf src s buf1 pst1
1138 -- XXX Specialize for Stop 0 common case?
1140 assert (n <= length (x:buf)) (return ())
1141 let src = Prelude.reverse (Prelude.take n (x:buf)) ++ xs
1142 return $ Skip $ ParseChunksYield b (ParseChunksInit src s)
1143 PR.Error err -> throwM $ ParseError err
1145 stepOuter _ (ParseChunksYield a next) = return $ Yield a next
1147{-# ANN type ConcatParseState Fuse #-}
1148data ConcatParseState x inpBuf st p =
1149 ConcatParseInit inpBuf st p
1150 | ConcatParseInitLeftOver inpBuf
1151 | ConcatParseStream st inpBuf p
1152 | ConcatParseBuf inpBuf st inpBuf p
1153 | ConcatParseYield x (ConcatParseState x inpBuf st p)
1155{-# INLINE_NORMAL parseIterate #-}
1158 => (b -> PRD.Parser m a b)
1162parseIterate func seed (Stream step state) =
1163 Stream stepOuter (ConcatParseInit [] state (func seed))
1167 {-# INLINE_LATE stepOuter #-}
1168 -- Buffer is empty, go to stream processing loop
1169 stepOuter _ (ConcatParseInit [] st (PRD.Parser pstep initial extract)) = do
1170 initial >>= \r -> return $ Skip $ ConcatParseStream st []
1171 (PRD.Parser pstep (return r) extract)
1173 -- Buffer is not empty, go to buffered processing loop
1174 stepOuter _ (ConcatParseInit src st
1175 (PRD.Parser pstep initial extract)) = do
1176 initial >>= \r -> return $ Skip $ ConcatParseBuf src st []
1177 (PRD.Parser pstep (return r) extract)
1179 -- XXX we just discard any leftover input at the end
1180 stepOuter _ (ConcatParseInitLeftOver _) = return Stop
1182 -- Buffer is empty process elements from the stream
1183 stepOuter gst (ConcatParseStream st buf
1184 p@(PRD.Parser pstep initial extract)) = do
1185 r <- step (adaptState gst) st
1191 -- PR.Yield 0 pst1 -> go SPEC s [] pst1
1192 PR.Yield n pst1 -> do
1193 assert (n <= length (x:buf)) (return ())
1194 let buf1 = Prelude.take n (x:buf)
1195 return $ Skip $ ConcatParseStream s buf1
1196 (PRD.Parser pstep (return pst1) extract)
1197 PR.YieldB n pst1 -> do
1198 assert (n <= length (x:buf)) (return ())
1199 let src0 = Prelude.take n (x:buf)
1200 src = Prelude.reverse src0
1201 return $ Skip $ ConcatParseBuf src s []
1202 (PRD.Parser pstep (return pst1) extract)
1203 -- PR.Skip 0 pst1 ->
1204 -- return $ Skip $ ConcatParseStream s (x:buf) pst1
1205 PR.Skip n pst1 -> do
1206 assert (n <= length (x:buf)) (return ())
1207 let (src0, buf1) = splitAt n (x:buf)
1208 src = Prelude.reverse src0
1209 return $ Skip $ ConcatParseBuf src s buf1
1210 (PRD.Parser pstep (return pst1) extract)
1211 -- XXX Specialize for Stop 0 common case?
1213 assert (n <= length (x:buf)) (return ())
1214 let src = Prelude.reverse (Prelude.take n (x:buf))
1216 ConcatParseYield b (ConcatParseInit src s (func b))
1217 PR.Error err -> throwM $ ParseError err
1218 Skip s -> return $ Skip $ ConcatParseStream s buf p
1222 let src = Prelude.reverse buf
1224 ConcatParseYield b (ConcatParseInitLeftOver src)
1226 -- go back to stream processing mode
1227 stepOuter _ (ConcatParseBuf [] s buf p) =
1228 return $ Skip $ ConcatParseStream s buf p
1230 -- buffered processing loop
1231 stepOuter _ (ConcatParseBuf (x:xs) s buf
1232 (PRD.Parser pstep initial extract)) = do
1236 -- PR.Yield 0 pst1 ->
1237 PR.Yield n pst1 -> do
1238 assert (n <= length (x:buf)) (return ())
1239 let buf1 = Prelude.take n (x:buf)
1240 return $ Skip $ ConcatParseBuf xs s buf1
1241 (PRD.Parser pstep (return pst1) extract)
1242 PR.YieldB n pst1 -> do
1243 assert (n <= length (x:buf)) (return ())
1244 let src0 = Prelude.take n (x:buf)
1245 src = Prelude.reverse src0 ++ xs
1246 return $ Skip $ ConcatParseBuf src s []
1247 (PRD.Parser pstep (return pst1) extract)
1248 -- PR.Skip 0 pst1 -> return $ Skip $ ConcatParseBuf xs s (x:buf) pst1
1249 PR.Skip n pst1 -> do
1250 assert (n <= length (x:buf)) (return ())
1251 let (src0, buf1) = splitAt n (x:buf)
1252 src = Prelude.reverse src0 ++ xs
1253 return $ Skip $ ConcatParseBuf src s buf1
1254 (PRD.Parser pstep (return pst1) extract)
1255 -- XXX Specialize for Stop 0 common case?
1257 assert (n <= length (x:buf)) (return ())
1258 let src = Prelude.reverse (Prelude.take n (x:buf)) ++ xs
1259 return $ Skip $ ConcatParseYield b
1260 (ConcatParseInit src s (func b))
1261 PR.Error err -> throwM $ ParseError err
1263 stepOuter _ (ConcatParseYield a next) = return $ Yield a next
1265------------------------------------------------------------------------------
1267------------------------------------------------------------------------------
1269-- | Run a streaming composition, discard the results.
1270{-# INLINE_LATE drain #-}
1271drain :: Monad m => Stream m a -> m ()
1272-- drain = foldrM (\_ xs -> xs) (return ())
1273drain (Stream step state) = go SPEC state
1276 r <- step defState st
1278 Yield _ s -> go SPEC s
1282{-# INLINE_NORMAL null #-}
1283null :: Monad m => Stream m a -> m Bool
1284null m = foldrM (\_ _ -> return False) (return True) m
1286{-# INLINE_NORMAL head #-}
1287head :: Monad m => Stream m a -> m (Maybe a)
1288head m = foldrM (\x _ -> return (Just x)) (return Nothing) m
1290{-# INLINE_NORMAL headElse #-}
1291headElse :: Monad m => a -> Stream m a -> m a
1292headElse a m = foldrM (\x _ -> return x) (return a) m
1294-- Does not fuse, has the same performance as the StreamK version.
1295{-# INLINE_NORMAL tail #-}
1296tail :: Monad m => Stream m a -> m (Maybe (Stream m a))
1297tail (UnStream step state) = go state
1300 r <- step defState st
1302 Yield _ s -> return (Just $ Stream step s)
1304 Stop -> return Nothing
1306-- XXX will it fuse? need custom impl?
1307{-# INLINE_NORMAL last #-}
1308last :: Monad m => Stream m a -> m (Maybe a)
1309last = foldl' (\_ y -> Just y) Nothing
1311{-# INLINE_NORMAL elem #-}
1312elem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
1313-- elem e m = foldrM (\x xs -> if x == e then return True else xs) (return False) m
1314elem e (Stream step state) = go state
1317 r <- step defState st
1320 | x == e -> return True
1323 Stop -> return False
1325{-# INLINE_NORMAL notElem #-}
1326notElem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
1327notElem e s = fmap not (elem e s)
1329{-# INLINE_NORMAL all #-}
1330all :: Monad m => (a -> Bool) -> Stream m a -> m Bool
1331-- all p m = foldrM (\x xs -> if p x then xs else return False) (return True) m
1332all p (Stream step state) = go state
1335 r <- step defState st
1339 | otherwise -> return False
1343{-# INLINE_NORMAL any #-}
1344any :: Monad m => (a -> Bool) -> Stream m a -> m Bool
1345-- any p m = foldrM (\x xs -> if p x then return True else xs) (return False) m
1346any p (Stream step state) = go state
1349 r <- step defState st
1352 | p x -> return True
1355 Stop -> return False
1357{-# INLINE_NORMAL maximum #-}
1358maximum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
1359maximum (Stream step state) = go Nothing state
1362 r <- step defState st
1364 Yield x s -> go (Just x) s
1365 Skip s -> go Nothing s
1366 Stop -> return Nothing
1367 go (Just acc) st = do
1368 r <- step defState st
1371 | acc <= x -> go (Just x) s
1372 | otherwise -> go (Just acc) s
1373 Skip s -> go (Just acc) s
1374 Stop -> return (Just acc)
1376{-# INLINE_NORMAL maximumBy #-}
1377maximumBy :: Monad m => (a -> a -> Ordering) -> Stream m a -> m (Maybe a)
1378maximumBy cmp (Stream step state) = go Nothing state
1381 r <- step defState st
1383 Yield x s -> go (Just x) s
1384 Skip s -> go Nothing s
1385 Stop -> return Nothing
1386 go (Just acc) st = do
1387 r <- step defState st
1389 Yield x s -> case cmp acc x of
1390 GT -> go (Just acc) s
1392 Skip s -> go (Just acc) s
1393 Stop -> return (Just acc)
1395{-# INLINE_NORMAL minimum #-}
1396minimum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
1397minimum (Stream step state) = go Nothing state
1400 r <- step defState st
1402 Yield x s -> go (Just x) s
1403 Skip s -> go Nothing s
1404 Stop -> return Nothing
1405 go (Just acc) st = do
1406 r <- step defState st
1409 | acc <= x -> go (Just acc) s
1410 | otherwise -> go (Just x) s
1411 Skip s -> go (Just acc) s
1412 Stop -> return (Just acc)
1414{-# INLINE_NORMAL minimumBy #-}
1415minimumBy :: Monad m => (a -> a -> Ordering) -> Stream m a -> m (Maybe a)
1416minimumBy cmp (Stream step state) = go Nothing state
1419 r <- step defState st
1421 Yield x s -> go (Just x) s
1422 Skip s -> go Nothing s
1423 Stop -> return Nothing
1424 go (Just acc) st = do
1425 r <- step defState st
1427 Yield x s -> case cmp acc x of
1429 _ -> go (Just acc) s
1430 Skip s -> go (Just acc) s
1431 Stop -> return (Just acc)
1433{-# INLINE_NORMAL (!!) #-}
1434(!!) :: (Monad m) => Stream m a -> Int -> m (Maybe a)
1435(Stream step state) !! i = go i state
1438 r <- step defState st
1440 Yield x s | n < 0 -> return Nothing
1441 | n == 0 -> return $ Just x
1442 | otherwise -> go (n - 1) s
1444 Stop -> return Nothing
1446{-# INLINE_NORMAL lookup #-}
1447lookup :: (Monad m, Eq a) => a -> Stream m (a, b) -> m (Maybe b)
1448lookup e m = foldrM (\(a, b) xs -> if e == a then return (Just b) else xs)
1451{-# INLINE_NORMAL findM #-}
1452findM :: Monad m => (a -> m Bool) -> Stream m a -> m (Maybe a)
1453findM p m = foldrM (\x xs -> p x >>= \r -> if r then return (Just x) else xs)
1457find :: Monad m => (a -> Bool) -> Stream m a -> m (Maybe a)
1458find p = findM (return . p)
1460{-# INLINE_NORMAL findIndices #-}
1461findIndices :: Monad m => (a -> Bool) -> Stream m a -> Stream m Int
1462findIndices p (Stream step state) = Stream step' (state, 0)
1464 {-# INLINE_LATE step' #-}
1465 step' gst (st, i) = i `seq` do
1466 r <- step (adaptState gst) st
1468 Yield x s -> if p x then Yield i (s, i+1) else Skip (s, i+1)
1469 Skip s -> Skip (s, i)
1472{-# INLINE toListRev #-}
1473toListRev :: Monad m => Stream m a -> m [a]
1474toListRev = foldl' (flip (:)) []
1476-- We can implement reverse as:
1478-- > reverse = foldlS (flip cons) nil
1480-- However, this implementation is unusable because of the horrible performance
1481-- of cons. So we just convert it to a list first and then stream from the
1484-- XXX Maybe we can use an Array instead of a list here?
1485{-# INLINE_NORMAL reverse #-}
1486reverse :: Monad m => Stream m a -> Stream m a
1487reverse m = Stream step Nothing
1489 {-# INLINE_LATE step #-}
1492 return $ Skip (Just xs)
1493 step _ (Just (x:xs)) = return $ Yield x (Just xs)
1494 step _ (Just []) = return Stop
1496-- Much faster reverse for Storables
1497{-# INLINE_NORMAL reverse' #-}
1498reverse' :: forall m a. (MonadIO m, Storable a) => Stream m a -> Stream m a
1500-- This commented implementation copies the whole stream into one single array
1501-- and then streams from that array, this is 3-4x faster than the chunked code
1502-- that follows. Though this could be problematic due to unbounded large
1503-- allocations. We need to figure out why the chunked code is slower and if we
1504-- can optimize the chunked code to work as fast as this one. It may be a
1506import Foreign.ForeignPtr (touchForeignPtr)
1507import Foreign.ForeignPtr.Unsafe (unsafeForeignPtrToPtr)
1508import Foreign.Ptr (Ptr, plusPtr)
1509reverse' m = Stream step Nothing
1511 {-# INLINE_LATE step #-}
1513 arr <- A.fromStreamD m
1514 let p = aEnd arr `plusPtr` negate (sizeOf (undefined :: a))
1515 return $ Skip $ Just (aStart arr, p)
1517 step _ (Just (start, p)) | p < unsafeForeignPtrToPtr start = return Stop
1519 step _ (Just (start, p)) = do
1520 let !x = A.unsafeInlineIO $ do
1522 touchForeignPtr start
1524 next = p `plusPtr` negate (sizeOf (undefined :: a))
1525 return $ Yield x (Just (start, next))
1532 $ A.fromStreamDArraysOf A.defaultChunkSize m
1535------------------------------------------------------------------------------
1536-- Grouping/Splitting
1537------------------------------------------------------------------------------
1539{-# INLINE_NORMAL splitSuffixBy' #-}
1540splitSuffixBy' :: Monad m
1541 => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1542splitSuffixBy' predicate f (Stream step state) =
1543 Stream (stepOuter f) (Just state)
1547 {-# INLINE_LATE stepOuter #-}
1548 stepOuter (Fold fstep initial done) gst (Just st) = do
1549 res <- step (adaptState gst) st
1555 then done acc' >>= \val -> return $ Yield val (Just s)
1558 Skip s -> return $ Skip $ Just s
1564 res <- step (adaptState gst) stt
1569 then done acc' >>= \val -> return $ Yield val (Just s)
1571 Skip s -> go SPEC s acc
1572 Stop -> done acc >>= \val -> return $ Yield val Nothing
1574 stepOuter _ _ Nothing = return Stop
1576{-# INLINE_NORMAL groupsBy #-}
1582groupsBy cmp f (Stream step state) = Stream (stepOuter f) (Just state, Nothing)
1586 {-# INLINE_LATE stepOuter #-}
1587 stepOuter (Fold fstep initial done) gst (Just st, Nothing) = do
1588 res <- step (adaptState gst) st
1595 Skip s -> return $ Skip $ (Just s, Nothing)
1600 go !_ prev stt !acc = do
1601 res <- step (adaptState gst) stt
1608 else done acc >>= \r -> return $ Yield r (Just s, Just x)
1609 Skip s -> go SPEC prev s acc
1610 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1612 stepOuter (Fold fstep initial done) gst (Just st, Just prev) = do
1614 acc' <- fstep acc prev
1619 -- XXX code duplicated from the previous equation
1621 res <- step (adaptState gst) stt
1628 else done acc >>= \r -> return $ Yield r (Just s, Just x)
1629 Skip s -> go SPEC s acc
1630 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1632 stepOuter _ _ (Nothing,_) = return Stop
1634{-# INLINE_NORMAL groupsRollingBy #-}
1635groupsRollingBy :: Monad m
1640groupsRollingBy cmp f (Stream step state) =
1641 Stream (stepOuter f) (Just state, Nothing)
1644 {-# INLINE_LATE stepOuter #-}
1645 stepOuter (Fold fstep initial done) gst (Just st, Nothing) = do
1646 res <- step (adaptState gst) st
1653 Skip s -> return $ Skip $ (Just s, Nothing)
1657 go !_ prev stt !acc = do
1658 res <- step (adaptState gst) stt
1666 done acc >>= \r -> return $ Yield r (Just s, Just x)
1667 Skip s -> go SPEC prev s acc
1668 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1670 stepOuter (Fold fstep initial done) gst (Just st, Just prev') = do
1672 acc' <- fstep acc prev'
1673 go SPEC prev' st acc'
1676 go !_ prevv stt !acc = do
1677 res <- step (adaptState gst) stt
1684 else done acc >>= \r -> return $ Yield r (Just s, Just x)
1685 Skip s -> go SPEC prevv s acc
1686 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1688 stepOuter _ _ (Nothing, _) = return Stop
1690{-# INLINE_NORMAL splitBy #-}
1691splitBy :: Monad m => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1692splitBy predicate f (Stream step state) = Stream (step' f) (Just state)
1696 {-# INLINE_LATE step' #-}
1697 step' (Fold fstep initial done) gst (Just st) = initial >>= go SPEC st
1702 res <- step (adaptState gst) stt
1706 then done acc >>= \r -> return $ Yield r (Just s)
1710 Skip s -> go SPEC s acc
1711 Stop -> done acc >>= \r -> return $ Yield r Nothing
1713 step' _ _ Nothing = return Stop
1715-- XXX requires -funfolding-use-threshold=150 in lines-unlines benchmark
1716{-# INLINE_NORMAL splitSuffixBy #-}
1717splitSuffixBy :: Monad m
1718 => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1719splitSuffixBy predicate f (Stream step state) = Stream (step' f) (Just state)
1723 {-# INLINE_LATE step' #-}
1724 step' (Fold fstep initial done) gst (Just st) = do
1725 res <- step (adaptState gst) st
1730 then done acc >>= \val -> return $ Yield val (Just s)
1735 Skip s -> return $ Skip $ Just s
1741 res <- step (adaptState gst) stt
1745 then done acc >>= \r -> return $ Yield r (Just s)
1749 Skip s -> go SPEC s acc
1750 Stop -> done acc >>= \r -> return $ Yield r Nothing
1752 step' _ _ Nothing = return Stop
1754{-# INLINE_NORMAL wordsBy #-}
1755wordsBy :: Monad m => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1756wordsBy predicate f (Stream step state) = Stream (stepOuter f) (Just state)
1760 {-# INLINE_LATE stepOuter #-}
1761 stepOuter (Fold fstep initial done) gst (Just st) = do
1762 res <- step (adaptState gst) st
1766 then return $ Skip (Just s)
1772 Skip s -> return $ Skip $ Just s
1778 res <- step (adaptState gst) stt
1782 then done acc >>= \r -> return $ Yield r (Just s)
1786 Skip s -> go SPEC s acc
1787 Stop -> done acc >>= \r -> return $ Yield r Nothing
1789 stepOuter _ _ Nothing = return Stop
1791-- String search algorithms:
1792-- http://www-igm.univ-mlv.fr/~lecroq/string/index.html
1795-- TODO can we unify the splitting operations using a splitting configuration
1796-- like in the split package.
1798data SplitStyle = Infix | Suffix | Prefix deriving (Eq, Show)
1800data SplitOptions = SplitOptions
1801 { style :: SplitStyle
1802 , withSep :: Bool -- ^ keep the separators in output
1803 -- , compact :: Bool -- ^ treat multiple consecutive separators as one
1804 -- , trimHead :: Bool -- ^ drop blank at head
1805 -- , trimTail :: Bool -- ^ drop blank at tail
1809data SplitOnState s a =
1814 | GO_KARP_RABIN s !(RB.Ring a) !(Ptr a)
1817{-# INLINE_NORMAL splitOn #-}
1819 :: forall m a b. (MonadIO m, Storable a, Enum a, Eq a)
1824splitOn patArr (Fold fstep initial done) (Stream step state) =
1825 Stream stepOuter GO_START
1829 patLen = A.length patArr
1830 maxIndex = patLen - 1
1831 elemBits = sizeOf (undefined :: a) * 8
1833 {-# INLINE_LATE stepOuter #-}
1834 stepOuter _ GO_START =
1836 then return $ Skip $ GO_EMPTY_PAT state
1839 r <- liftIO $ (A.unsafeIndexIO patArr 0)
1840 return $ Skip $ GO_SINGLE_PAT state r
1841 else if sizeOf (undefined :: a) * patLen
1842 <= sizeOf (undefined :: Word)
1843 then return $ Skip $ GO_SHORT_PAT state
1845 (rb, rhead) <- liftIO $ RB.new patLen
1846 return $ Skip $ GO_KARP_RABIN state rb rhead
1848 stepOuter gst (GO_SINGLE_PAT stt pat) = initial >>= go SPEC stt
1853 res <- step (adaptState gst) st
1859 return $ Yield r (GO_SINGLE_PAT s pat)
1860 else fstep acc x >>= go SPEC s
1861 Skip s -> go SPEC s acc
1862 Stop -> done acc >>= \r -> return $ Yield r GO_DONE
1864 stepOuter gst (GO_SHORT_PAT stt) = initial >>= go0 SPEC 0 (0 :: Word) stt
1869 mask = (1 `shiftL` (elemBits * patLen)) - 1
1871 addToWord wrd a = (wrd `shiftL` elemBits) .|. fromIntegral (fromEnum a)
1874 patWord = mask .&. A.foldl' addToWord 0 patArr
1876 go0 !_ !idx wrd st !acc = do
1877 res <- step (adaptState gst) st
1880 let wrd' = addToWord wrd x
1883 if wrd' .&. mask == patWord
1886 return $ Yield r (GO_SHORT_PAT s)
1887 else go1 SPEC wrd' s acc
1888 else go0 SPEC (idx + 1) wrd' s acc
1889 Skip s -> go0 SPEC idx wrd s acc
1892 then go2 wrd idx acc
1894 done acc' >>= \r -> return $ Yield r GO_DONE
1897 go1 !_ wrd st !acc = do
1898 res <- step (adaptState gst) st
1901 let wrd' = addToWord wrd x
1902 old = (mask .&. wrd) `shiftR` (elemBits * (patLen - 1))
1903 acc' <- fstep acc (toEnum $ fromIntegral old)
1904 if wrd' .&. mask == patWord
1905 then done acc' >>= \r -> return $ Yield r (GO_SHORT_PAT s)
1906 else go1 SPEC wrd' s acc'
1907 Skip s -> go1 SPEC wrd s acc
1909 acc' <- go2 wrd patLen acc
1910 done acc' >>= \r -> return $ Yield r GO_DONE
1912 go2 !wrd !n !acc | n > 0 = do
1913 let old = (mask .&. wrd) `shiftR` (elemBits * (n - 1))
1914 fstep acc (toEnum $ fromIntegral old) >>= go2 wrd (n - 1)
1915 go2 _ _ acc = return acc
1917 stepOuter gst (GO_KARP_RABIN stt rb rhead) = do
1918 initial >>= go0 SPEC 0 rhead stt
1922 k = 2891336453 :: Word32
1924 addCksum cksum a = cksum * k + fromIntegral (fromEnum a)
1925 deltaCksum cksum old new =
1926 addCksum cksum new - coeff * fromIntegral (fromEnum old)
1928 -- XXX shall we use a random starting hash or 1 instead of 0?
1929 patHash = A.foldl' addCksum 0 patArr
1932 go0 !_ !idx !rh st !acc = do
1933 res <- step (adaptState gst) st
1936 rh' <- liftIO $ RB.unsafeInsert rb rh x
1939 let fold = RB.unsafeFoldRing (RB.ringBound rb)
1940 let !ringHash = fold addCksum 0 rb
1941 if ringHash == patHash
1942 then go2 SPEC ringHash rh' s acc
1943 else go1 SPEC ringHash rh' s acc
1944 else go0 SPEC (idx + 1) rh' s acc
1945 Skip s -> go0 SPEC idx rh s acc
1947 !acc' <- if idx /= 0
1948 then RB.unsafeFoldRingM rh fstep acc rb
1950 done acc' >>= \r -> return $ Yield r GO_DONE
1952 -- XXX Theoretically this code can do 4 times faster if GHC generates
1953 -- optimal code. If we use just "(cksum' == patHash)" condition it goes
1954 -- 4x faster, as soon as we add the "RB.unsafeEqArray rb v" condition
1955 -- the generated code changes drastically and becomes 4x slower. Need
1956 -- to investigate what is going on with GHC.
1958 go1 !_ !cksum !rh st !acc = do
1959 res <- step (adaptState gst) st
1962 old <- liftIO $ peek rh
1963 let cksum' = deltaCksum cksum old x
1964 acc' <- fstep acc old
1966 if (cksum' == patHash)
1968 rh' <- liftIO (RB.unsafeInsert rb rh x)
1969 go2 SPEC cksum' rh' s acc'
1971 rh' <- liftIO (RB.unsafeInsert rb rh x)
1972 go1 SPEC cksum' rh' s acc'
1973 Skip s -> go1 SPEC cksum rh s acc
1975 acc' <- RB.unsafeFoldRingFullM rh fstep acc rb
1976 done acc' >>= \r -> return $ Yield r GO_DONE
1978 go2 !_ !cksum' !rh' s !acc' = do
1979 if RB.unsafeEqArray rb rh' patArr
1982 return $ Yield r (GO_KARP_RABIN s rb rhead)
1983 else go1 SPEC cksum' rh' s acc'
1985 stepOuter gst (GO_EMPTY_PAT st) = do
1986 res <- step (adaptState gst) st
1991 done acc' >>= \r -> return $ Yield r (GO_EMPTY_PAT s)
1992 Skip s -> return $ Skip (GO_EMPTY_PAT s)
1995 stepOuter _ GO_DONE = return Stop
1997{-# INLINE_NORMAL splitSuffixOn #-}
1999 :: forall m a b. (MonadIO m, Storable a, Enum a, Eq a)
2005splitSuffixOn withSep patArr (Fold fstep initial done)
2006 (Stream step state) =
2007 Stream stepOuter GO_START
2011 patLen = A.length patArr
2012 maxIndex = patLen - 1
2013 elemBits = sizeOf (undefined :: a) * 8
2015 {-# INLINE_LATE stepOuter #-}
2016 stepOuter _ GO_START =
2018 then return $ Skip $ GO_EMPTY_PAT state
2021 r <- liftIO $ (A.unsafeIndexIO patArr 0)
2022 return $ Skip $ GO_SINGLE_PAT state r
2023 else if sizeOf (undefined :: a) * patLen
2024 <= sizeOf (undefined :: Word)
2025 then return $ Skip $ GO_SHORT_PAT state
2027 (rb, rhead) <- liftIO $ RB.new patLen
2028 return $ Skip $ GO_KARP_RABIN state rb rhead
2030 stepOuter gst (GO_SINGLE_PAT stt pat) = do
2031 -- This first part is the only difference between splitOn and
2033 -- If the last element is a separator do not issue a blank segment.
2034 res <- step (adaptState gst) stt
2040 acc' <- if withSep then fstep acc x else return acc
2041 done acc' >>= \r -> return $ Yield r (GO_SINGLE_PAT s pat)
2042 else fstep acc x >>= go SPEC s
2043 Skip s -> return $ Skip $ (GO_SINGLE_PAT s pat)
2048 -- This is identical for splitOn and splitSuffixOn
2050 res <- step (adaptState gst) st
2055 acc' <- if withSep then fstep acc x else return acc
2057 return $ Yield r (GO_SINGLE_PAT s pat)
2058 else fstep acc x >>= go SPEC s
2059 Skip s -> go SPEC s acc
2060 Stop -> done acc >>= \r -> return $ Yield r GO_DONE
2062 stepOuter gst (GO_SHORT_PAT stt) = do
2064 -- Call "initial" only if the stream yields an element, otherwise we
2065 -- may call "initial" but never yield anything. initial may produce a
2066 -- side effect, therefore we will end up doing and discard a side
2071 res <- step (adaptState gst) stt
2075 let wrd' = addToWord wrd x
2076 acc' <- if withSep then fstep acc x else return acc
2079 if wrd' .&. mask == patWord
2080 then done acc' >>= \r -> return $ Yield r (GO_SHORT_PAT s)
2081 else go0 SPEC (idx + 1) wrd' s acc'
2082 else go0 SPEC (idx + 1) wrd' s acc'
2083 Skip s -> return $ Skip (GO_SHORT_PAT s)
2089 mask = (1 `shiftL` (elemBits * patLen)) - 1
2091 addToWord wrd a = (wrd `shiftL` elemBits) .|. fromIntegral (fromEnum a)
2094 patWord = mask .&. A.foldl' addToWord 0 patArr
2096 go0 !_ !idx wrd st !acc = do
2097 res <- step (adaptState gst) st
2100 let wrd' = addToWord wrd x
2101 acc' <- if withSep then fstep acc x else return acc
2104 if wrd' .&. mask == patWord
2107 return $ Yield r (GO_SHORT_PAT s)
2108 else go1 SPEC wrd' s acc'
2109 else go0 SPEC (idx + 1) wrd' s acc'
2110 Skip s -> go0 SPEC idx wrd s acc
2112 if (idx == maxIndex) && (wrd .&. mask == patWord)
2115 acc' <- if idx /= 0 && not withSep
2116 then go2 wrd idx acc
2118 done acc' >>= \r -> return $ Yield r GO_DONE
2121 go1 !_ wrd st !acc = do
2122 res <- step (adaptState gst) st
2125 let wrd' = addToWord wrd x
2126 old = (mask .&. wrd) `shiftR` (elemBits * (patLen - 1))
2129 else fstep acc (toEnum $ fromIntegral old)
2130 if wrd' .&. mask == patWord
2131 then done acc' >>= \r -> return $ Yield r (GO_SHORT_PAT s)
2132 else go1 SPEC wrd' s acc'
2133 Skip s -> go1 SPEC wrd s acc
2135 -- If the last sequence is a separator do not issue a blank
2137 if wrd .&. mask == patWord
2142 else go2 wrd patLen acc
2143 done acc' >>= \r -> return $ Yield r GO_DONE
2145 go2 !wrd !n !acc | n > 0 = do
2146 let old = (mask .&. wrd) `shiftR` (elemBits * (n - 1))
2147 fstep acc (toEnum $ fromIntegral old) >>= go2 wrd (n - 1)
2148 go2 _ _ acc = return acc
2150 stepOuter gst (GO_KARP_RABIN stt rb rhead) = do
2152 res <- step (adaptState gst) stt
2156 acc' <- if withSep then fstep acc x else return acc
2157 rh' <- liftIO (RB.unsafeInsert rb rhead x)
2160 let fold = RB.unsafeFoldRing (RB.ringBound rb)
2161 let !ringHash = fold addCksum 0 rb
2162 if ringHash == patHash
2163 then go2 SPEC ringHash rh' s acc'
2164 else go0 SPEC (idx + 1) rh' s acc'
2165 else go0 SPEC (idx + 1) rh' s acc'
2166 Skip s -> return $ Skip (GO_KARP_RABIN s rb rhead)
2171 k = 2891336453 :: Word32
2173 addCksum cksum a = cksum * k + fromIntegral (fromEnum a)
2174 deltaCksum cksum old new =
2175 addCksum cksum new - coeff * fromIntegral (fromEnum old)
2177 -- XXX shall we use a random starting hash or 1 instead of 0?
2178 patHash = A.foldl' addCksum 0 patArr
2181 go0 !_ !idx !rh st !acc = do
2182 res <- step (adaptState gst) st
2185 acc' <- if withSep then fstep acc x else return acc
2186 rh' <- liftIO (RB.unsafeInsert rb rh x)
2189 let fold = RB.unsafeFoldRing (RB.ringBound rb)
2190 let !ringHash = fold addCksum 0 rb
2191 if ringHash == patHash
2192 then go2 SPEC ringHash rh' s acc'
2193 else go1 SPEC ringHash rh' s acc'
2194 else go0 SPEC (idx + 1) rh' s acc'
2195 Skip s -> go0 SPEC idx rh s acc
2197 -- do not issue a blank segment when we end at pattern
2198 if (idx == maxIndex) && RB.unsafeEqArray rb rh patArr
2201 !acc' <- if idx /= 0 && not withSep
2202 then RB.unsafeFoldRingM rh fstep acc rb
2204 done acc' >>= \r -> return $ Yield r GO_DONE
2206 -- XXX Theoretically this code can do 4 times faster if GHC generates
2207 -- optimal code. If we use just "(cksum' == patHash)" condition it goes
2208 -- 4x faster, as soon as we add the "RB.unsafeEqArray rb v" condition
2209 -- the generated code changes drastically and becomes 4x slower. Need
2210 -- to investigate what is going on with GHC.
2212 go1 !_ !cksum !rh st !acc = do
2213 res <- step (adaptState gst) st
2216 old <- liftIO $ peek rh
2217 let cksum' = deltaCksum cksum old x
2222 if (cksum' == patHash)
2224 rh' <- liftIO (RB.unsafeInsert rb rh x)
2225 go2 SPEC cksum' rh' s acc'
2227 rh' <- liftIO (RB.unsafeInsert rb rh x)
2228 go1 SPEC cksum' rh' s acc'
2229 Skip s -> go1 SPEC cksum rh s acc
2231 if RB.unsafeEqArray rb rh patArr
2236 else RB.unsafeFoldRingFullM rh fstep acc rb
2237 done acc' >>= \r -> return $ Yield r GO_DONE
2239 go2 !_ !cksum' !rh' s !acc' = do
2240 if RB.unsafeEqArray rb rh' patArr
2243 return $ Yield r (GO_KARP_RABIN s rb rhead)
2244 else go1 SPEC cksum' rh' s acc'
2246 stepOuter gst (GO_EMPTY_PAT st) = do
2247 res <- step (adaptState gst) st
2252 done acc' >>= \r -> return $ Yield r (GO_EMPTY_PAT s)
2253 Skip s -> return $ Skip (GO_EMPTY_PAT s)
2256 stepOuter _ GO_DONE = return Stop
2258data SplitState s arr
2260 | SplitBuffering s arr
2261 | SplitSplitting s arr
2262 | SplitYielding arr (SplitState s arr)
2265-- XXX An alternative approach would be to use a partial fold (Fold m a b) to
2266-- split using a splitBy like combinator. The Fold would consume upto the
2267-- separator and return any leftover which can then be fed to the next fold.
2269-- We can revisit this once we have partial folds/parsers.
2271-- | Performs infix separator style splitting.
2272{-# INLINE_NORMAL splitInnerBy #-}
2275 => (f a -> m (f a, Maybe (f a))) -- splitter
2276 -> (f a -> f a -> m (f a)) -- joiner
2279splitInnerBy splitter joiner (Stream step1 state1) =
2280 (Stream step (SplitInitial state1))
2284 {-# INLINE_LATE step #-}
2285 step gst (SplitInitial st) = do
2289 (x1, mx2) <- splitter x
2290 return $ case mx2 of
2291 Nothing -> Skip (SplitBuffering s x1)
2292 Just x2 -> Skip (SplitYielding x1 (SplitSplitting s x2))
2293 Skip s -> return $ Skip (SplitInitial s)
2294 Stop -> return $ Stop
2296 step gst (SplitBuffering st buf) = do
2300 (x1, mx2) <- splitter x
2301 buf' <- joiner buf x1
2302 return $ case mx2 of
2303 Nothing -> Skip (SplitBuffering s buf')
2304 Just x2 -> Skip (SplitYielding buf' (SplitSplitting s x2))
2305 Skip s -> return $ Skip (SplitBuffering s buf)
2306 Stop -> return $ Skip (SplitYielding buf SplitFinishing)
2308 step _ (SplitSplitting st buf) = do
2309 (x1, mx2) <- splitter buf
2310 return $ case mx2 of
2311 Nothing -> Skip $ SplitBuffering st x1
2312 Just x2 -> Skip $ SplitYielding x1 (SplitSplitting st x2)
2314 step _ (SplitYielding x next) = return $ Yield x next
2315 step _ SplitFinishing = return $ Stop
2317-- | Performs infix separator style splitting.
2318{-# INLINE_NORMAL splitInnerBySuffix #-}
2320 :: (Monad m, Eq (f a), Monoid (f a))
2321 => (f a -> m (f a, Maybe (f a))) -- splitter
2322 -> (f a -> f a -> m (f a)) -- joiner
2325splitInnerBySuffix splitter joiner (Stream step1 state1) =
2326 (Stream step (SplitInitial state1))
2330 {-# INLINE_LATE step #-}
2331 step gst (SplitInitial st) = do
2335 (x1, mx2) <- splitter x
2336 return $ case mx2 of
2337 Nothing -> Skip (SplitBuffering s x1)
2338 Just x2 -> Skip (SplitYielding x1 (SplitSplitting s x2))
2339 Skip s -> return $ Skip (SplitInitial s)
2340 Stop -> return $ Stop
2342 step gst (SplitBuffering st buf) = do
2346 (x1, mx2) <- splitter x
2347 buf' <- joiner buf x1
2348 return $ case mx2 of
2349 Nothing -> Skip (SplitBuffering s buf')
2350 Just x2 -> Skip (SplitYielding buf' (SplitSplitting s x2))
2351 Skip s -> return $ Skip (SplitBuffering s buf)
2355 else Skip (SplitYielding buf SplitFinishing)
2357 step _ (SplitSplitting st buf) = do
2358 (x1, mx2) <- splitter buf
2359 return $ case mx2 of
2360 Nothing -> Skip $ SplitBuffering st x1
2361 Just x2 -> Skip $ SplitYielding x1 (SplitSplitting st x2)
2363 step _ (SplitYielding x next) = return $ Yield x next
2364 step _ SplitFinishing = return $ Stop
2366------------------------------------------------------------------------------
2368------------------------------------------------------------------------------
2370{-# INLINE_NORMAL isPrefixOf #-}
2371isPrefixOf :: (Eq a, Monad m) => Stream m a -> Stream m a -> m Bool
2372isPrefixOf (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
2374 go (sa, sb, Nothing) = do
2375 r <- stepa defState sa
2377 Yield x sa' -> go (sa', sb, Just x)
2378 Skip sa' -> go (sa', sb, Nothing)
2381 go (sa, sb, Just x) = do
2382 r <- stepb defState sb
2386 then go (sa, sb', Nothing)
2388 Skip sb' -> go (sa, sb', Just x)
2389 Stop -> return False
2391{-# INLINE_NORMAL isSubsequenceOf #-}
2392isSubsequenceOf :: (Eq a, Monad m) => Stream m a -> Stream m a -> m Bool
2393isSubsequenceOf (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
2395 go (sa, sb, Nothing) = do
2396 r <- stepa defState sa
2398 Yield x sa' -> go (sa', sb, Just x)
2399 Skip sa' -> go (sa', sb, Nothing)
2402 go (sa, sb, Just x) = do
2403 r <- stepb defState sb
2407 then go (sa, sb', Nothing)
2408 else go (sa, sb', Just x)
2409 Skip sb' -> go (sa, sb', Just x)
2410 Stop -> return False
2412{-# INLINE_NORMAL stripPrefix #-}
2415 => Stream m a -> Stream m a -> m (Maybe (Stream m a))
2416stripPrefix (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
2418 go (sa, sb, Nothing) = do
2419 r <- stepa defState sa
2421 Yield x sa' -> go (sa', sb, Just x)
2422 Skip sa' -> go (sa', sb, Nothing)
2423 Stop -> return $ Just (Stream stepb sb)
2425 go (sa, sb, Just x) = do
2426 r <- stepb defState sb
2430 then go (sa, sb', Nothing)
2432 Skip sb' -> go (sa, sb', Just x)
2433 Stop -> return Nothing
2435------------------------------------------------------------------------------
2437------------------------------------------------------------------------------
2439-- | Execute a monadic action for each element of the 'Stream'
2440{-# INLINE_NORMAL mapM_ #-}
2441mapM_ :: Monad m => (a -> m b) -> Stream m a -> m ()
2442mapM_ m = drain . mapM m
2444-------------------------------------------------------------------------------
2445-- Stream transformations using Unfolds
2446-------------------------------------------------------------------------------
2448-- Define a unique structure to use in inspection testing
2449data ConcatMapUState o i =
2451 | ConcatMapUInner o i
2453-- | @concatMapU unfold stream@ uses @unfold@ to map the input stream elements
2454-- to streams and then flattens the generated streams into a single output
2457-- This is like 'concatMap' but uses an unfold with an explicit state to
2458-- generate the stream instead of a 'Stream' type generator. This allows better
2459-- optimization via fusion. This can be many times more efficient than
2462{-# INLINE_NORMAL concatMapU #-}
2463concatMapU :: Monad m => Unfold m a b -> Stream m a -> Stream m b
2464concatMapU (Unfold istep inject) (Stream ostep ost) =
2465 Stream step (ConcatMapUOuter ost)
2467 {-# INLINE_LATE step #-}
2468 step gst (ConcatMapUOuter o) = do
2469 r <- ostep (adaptState gst) o
2473 i `seq` return (Skip (ConcatMapUInner o' i))
2474 Skip o' -> return $ Skip (ConcatMapUOuter o')
2475 Stop -> return $ Stop
2477 step _ (ConcatMapUInner o i) = do
2480 Yield x i' -> Yield x (ConcatMapUInner o i')
2481 Skip i' -> Skip (ConcatMapUInner o i')
2482 Stop -> Skip (ConcatMapUOuter o)
2484data ConcatUnfoldInterleaveState o i =
2485 ConcatUnfoldInterleaveOuter o [i]
2486 | ConcatUnfoldInterleaveInner o [i]
2487 | ConcatUnfoldInterleaveInnerL [i] [i]
2488 | ConcatUnfoldInterleaveInnerR [i] [i]
2490-- XXX use arrays to store state instead of lists.
2491-- XXX In general we can use different scheduling strategies e.g. how to
2492-- schedule the outer vs inner loop or assigning weights to different streams
2493-- or outer and inner loops.
2495-- After a yield, switch to the next stream. Do not switch streams on Skip.
2496-- Yield from outer stream switches to the inner stream.
2498-- There are two choices here, (1) exhaust the outer stream first and then
2499-- start yielding from the inner streams, this is much simpler to implement,
2500-- (2) yield at least one element from an inner stream before going back to
2501-- outer stream and opening the next stream from it.
2503-- Ideally, we need some scheduling bias to inner streams vs outer stream.
2504-- Maybe we can configure the behavior.
2506{-# INLINE_NORMAL concatUnfoldInterleave #-}
2507concatUnfoldInterleave :: Monad m => Unfold m a b -> Stream m a -> Stream m b
2508concatUnfoldInterleave (Unfold istep inject) (Stream ostep ost) =
2509 Stream step (ConcatUnfoldInterleaveOuter ost [])
2511 {-# INLINE_LATE step #-}
2512 step gst (ConcatUnfoldInterleaveOuter o ls) = do
2513 r <- ostep (adaptState gst) o
2517 i `seq` return (Skip (ConcatUnfoldInterleaveInner o' (i : ls)))
2518 Skip o' -> return $ Skip (ConcatUnfoldInterleaveOuter o' ls)
2519 Stop -> return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2521 step _ (ConcatUnfoldInterleaveInner _ []) = undefined
2522 step _ (ConcatUnfoldInterleaveInner o (st:ls)) = do
2525 Yield x s -> Yield x (ConcatUnfoldInterleaveOuter o (s:ls))
2526 Skip s -> Skip (ConcatUnfoldInterleaveInner o (s:ls))
2527 Stop -> Skip (ConcatUnfoldInterleaveOuter o ls)
2529 step _ (ConcatUnfoldInterleaveInnerL [] []) = return Stop
2530 step _ (ConcatUnfoldInterleaveInnerL [] rs) =
2531 return $ Skip (ConcatUnfoldInterleaveInnerR [] rs)
2533 step _ (ConcatUnfoldInterleaveInnerL (st:ls) rs) = do
2536 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerL ls (s:rs))
2537 Skip s -> Skip (ConcatUnfoldInterleaveInnerL (s:ls) rs)
2538 Stop -> Skip (ConcatUnfoldInterleaveInnerL ls rs)
2540 step _ (ConcatUnfoldInterleaveInnerR [] []) = return Stop
2541 step _ (ConcatUnfoldInterleaveInnerR ls []) =
2542 return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2544 step _ (ConcatUnfoldInterleaveInnerR ls (st:rs)) = do
2547 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerR (s:ls) rs)
2548 Skip s -> Skip (ConcatUnfoldInterleaveInnerR ls (s:rs))
2549 Stop -> Skip (ConcatUnfoldInterleaveInnerR ls rs)
2551-- XXX In general we can use different scheduling strategies e.g. how to
2552-- schedule the outer vs inner loop or assigning weights to different streams
2553-- or outer and inner loops.
2555-- This could be inefficient if the tasks are too small.
2557-- Compared to concatUnfoldInterleave this one switches streams on Skips.
2559{-# INLINE_NORMAL concatUnfoldRoundrobin #-}
2560concatUnfoldRoundrobin :: Monad m => Unfold m a b -> Stream m a -> Stream m b
2561concatUnfoldRoundrobin (Unfold istep inject) (Stream ostep ost) =
2562 Stream step (ConcatUnfoldInterleaveOuter ost [])
2564 {-# INLINE_LATE step #-}
2565 step gst (ConcatUnfoldInterleaveOuter o ls) = do
2566 r <- ostep (adaptState gst) o
2570 i `seq` return (Skip (ConcatUnfoldInterleaveInner o' (i : ls)))
2571 Skip o' -> return $ Skip (ConcatUnfoldInterleaveInner o' ls)
2572 Stop -> return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2574 step _ (ConcatUnfoldInterleaveInner o []) =
2575 return $ Skip (ConcatUnfoldInterleaveOuter o [])
2577 step _ (ConcatUnfoldInterleaveInner o (st:ls)) = do
2580 Yield x s -> Yield x (ConcatUnfoldInterleaveOuter o (s:ls))
2581 Skip s -> Skip (ConcatUnfoldInterleaveOuter o (s:ls))
2582 Stop -> Skip (ConcatUnfoldInterleaveOuter o ls)
2584 step _ (ConcatUnfoldInterleaveInnerL [] []) = return Stop
2585 step _ (ConcatUnfoldInterleaveInnerL [] rs) =
2586 return $ Skip (ConcatUnfoldInterleaveInnerR [] rs)
2588 step _ (ConcatUnfoldInterleaveInnerL (st:ls) rs) = do
2591 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerL ls (s:rs))
2592 Skip s -> Skip (ConcatUnfoldInterleaveInnerL ls (s:rs))
2593 Stop -> Skip (ConcatUnfoldInterleaveInnerL ls rs)
2595 step _ (ConcatUnfoldInterleaveInnerR [] []) = return Stop
2596 step _ (ConcatUnfoldInterleaveInnerR ls []) =
2597 return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2599 step _ (ConcatUnfoldInterleaveInnerR ls (st:rs)) = do
2602 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerR (s:ls) rs)
2603 Skip s -> Skip (ConcatUnfoldInterleaveInnerR (s:ls) rs)
2604 Stop -> Skip (ConcatUnfoldInterleaveInnerR ls rs)
2606data AppendState s1 s2 = AppendFirst s1 | AppendSecond s2
2608-- Note that this could be much faster compared to the CPS stream. However, as
2609-- the number of streams being composed increases this may become expensive.
2610-- Need to see where the breaking point is between the two.
2612{-# INLINE_NORMAL append #-}
2613append :: Monad m => Stream m a -> Stream m a -> Stream m a
2614append (Stream step1 state1) (Stream step2 state2) =
2615 Stream step (AppendFirst state1)
2619 {-# INLINE_LATE step #-}
2620 step gst (AppendFirst st) = do
2623 Yield a s -> Yield a (AppendFirst s)
2624 Skip s -> Skip (AppendFirst s)
2625 Stop -> Skip (AppendSecond state2)
2627 step gst (AppendSecond st) = do
2630 Yield a s -> Yield a (AppendSecond s)
2631 Skip s -> Skip (AppendSecond s)
2634data InterleaveState s1 s2 = InterleaveFirst s1 s2 | InterleaveSecond s1 s2
2635 | InterleaveSecondOnly s2 | InterleaveFirstOnly s1
2637{-# INLINE_NORMAL interleave #-}
2638interleave :: Monad m => Stream m a -> Stream m a -> Stream m a
2639interleave (Stream step1 state1) (Stream step2 state2) =
2640 Stream step (InterleaveFirst state1 state2)
2644 {-# INLINE_LATE step #-}
2645 step gst (InterleaveFirst st1 st2) = do
2648 Yield a s -> Yield a (InterleaveSecond s st2)
2649 Skip s -> Skip (InterleaveFirst s st2)
2650 Stop -> Skip (InterleaveSecondOnly st2)
2652 step gst (InterleaveSecond st1 st2) = do
2655 Yield a s -> Yield a (InterleaveFirst st1 s)
2656 Skip s -> Skip (InterleaveSecond st1 s)
2657 Stop -> Skip (InterleaveFirstOnly st1)
2659 step gst (InterleaveFirstOnly st1) = do
2662 Yield a s -> Yield a (InterleaveFirstOnly s)
2663 Skip s -> Skip (InterleaveFirstOnly s)
2666 step gst (InterleaveSecondOnly st2) = do
2669 Yield a s -> Yield a (InterleaveSecondOnly s)
2670 Skip s -> Skip (InterleaveSecondOnly s)
2673{-# INLINE_NORMAL interleaveMin #-}
2674interleaveMin :: Monad m => Stream m a -> Stream m a -> Stream m a
2675interleaveMin (Stream step1 state1) (Stream step2 state2) =
2676 Stream step (InterleaveFirst state1 state2)
2680 {-# INLINE_LATE step #-}
2681 step gst (InterleaveFirst st1 st2) = do
2684 Yield a s -> Yield a (InterleaveSecond s st2)
2685 Skip s -> Skip (InterleaveFirst s st2)
2688 step gst (InterleaveSecond st1 st2) = do
2691 Yield a s -> Yield a (InterleaveFirst st1 s)
2692 Skip s -> Skip (InterleaveSecond st1 s)
2695 step _ (InterleaveFirstOnly _) = undefined
2696 step _ (InterleaveSecondOnly _) = undefined
2698{-# INLINE_NORMAL interleaveSuffix #-}
2699interleaveSuffix :: Monad m => Stream m a -> Stream m a -> Stream m a
2700interleaveSuffix (Stream step1 state1) (Stream step2 state2) =
2701 Stream step (InterleaveFirst state1 state2)
2705 {-# INLINE_LATE step #-}
2706 step gst (InterleaveFirst st1 st2) = do
2709 Yield a s -> Yield a (InterleaveSecond s st2)
2710 Skip s -> Skip (InterleaveFirst s st2)
2713 step gst (InterleaveSecond st1 st2) = do
2716 Yield a s -> Yield a (InterleaveFirst st1 s)
2717 Skip s -> Skip (InterleaveSecond st1 s)
2718 Stop -> Skip (InterleaveFirstOnly st1)
2720 step gst (InterleaveFirstOnly st1) = do
2723 Yield a s -> Yield a (InterleaveFirstOnly s)
2724 Skip s -> Skip (InterleaveFirstOnly s)
2727 step _ (InterleaveSecondOnly _) = undefined
2729data InterleaveInfixState s1 s2 a
2730 = InterleaveInfixFirst s1 s2
2731 | InterleaveInfixSecondBuf s1 s2
2732 | InterleaveInfixSecondYield s1 s2 a
2733 | InterleaveInfixFirstYield s1 s2 a
2734 | InterleaveInfixFirstOnly s1
2736{-# INLINE_NORMAL interleaveInfix #-}
2737interleaveInfix :: Monad m => Stream m a -> Stream m a -> Stream m a
2738interleaveInfix (Stream step1 state1) (Stream step2 state2) =
2739 Stream step (InterleaveInfixFirst state1 state2)
2743 {-# INLINE_LATE step #-}
2744 step gst (InterleaveInfixFirst st1 st2) = do
2747 Yield a s -> Yield a (InterleaveInfixSecondBuf s st2)
2748 Skip s -> Skip (InterleaveInfixFirst s st2)
2751 step gst (InterleaveInfixSecondBuf st1 st2) = do
2754 Yield a s -> Skip (InterleaveInfixSecondYield st1 s a)
2755 Skip s -> Skip (InterleaveInfixSecondBuf st1 s)
2756 Stop -> Skip (InterleaveInfixFirstOnly st1)
2758 step gst (InterleaveInfixSecondYield st1 st2 x) = do
2761 Yield a s -> Yield x (InterleaveInfixFirstYield s st2 a)
2762 Skip s -> Skip (InterleaveInfixSecondYield s st2 x)
2765 step _ (InterleaveInfixFirstYield st1 st2 x) = do
2766 return $ Yield x (InterleaveInfixSecondBuf st1 st2)
2768 step gst (InterleaveInfixFirstOnly st1) = do
2771 Yield a s -> Yield a (InterleaveInfixFirstOnly s)
2772 Skip s -> Skip (InterleaveInfixFirstOnly s)
2775{-# INLINE_NORMAL roundRobin #-}
2776roundRobin :: Monad m => Stream m a -> Stream m a -> Stream m a
2777roundRobin (Stream step1 state1) (Stream step2 state2) =
2778 Stream step (InterleaveFirst state1 state2)
2782 {-# INLINE_LATE step #-}
2783 step gst (InterleaveFirst st1 st2) = do
2786 Yield a s -> Yield a (InterleaveSecond s st2)
2787 Skip s -> Skip (InterleaveSecond s st2)
2788 Stop -> Skip (InterleaveSecondOnly st2)
2790 step gst (InterleaveSecond st1 st2) = do
2793 Yield a s -> Yield a (InterleaveFirst st1 s)
2794 Skip s -> Skip (InterleaveFirst st1 s)
2795 Stop -> Skip (InterleaveFirstOnly st1)
2797 step gst (InterleaveSecondOnly st2) = do
2800 Yield a s -> Yield a (InterleaveSecondOnly s)
2801 Skip s -> Skip (InterleaveSecondOnly s)
2804 step gst (InterleaveFirstOnly st1) = do
2807 Yield a s -> Yield a (InterleaveFirstOnly s)
2808 Skip s -> Skip (InterleaveFirstOnly s)
2811data ICUState s1 s2 i1 i2 =
2816 | ICUFirstInner s1 s2 i1
2817 | ICUSecondInner s1 s2 i2
2818 | ICUFirstOnlyInner s1 i1
2819 | ICUSecondOnlyInner s2 i2
2821-- | Interleave streams (full streams, not the elements) unfolded from two
2822-- input streams and concat. Stop when the first stream stops. If the second
2823-- stream ends before the first one then first stream still keeps running alone
2824-- without any interleaving with the second stream.
2826-- [a1, a2, ... an] [b1, b2 ...]
2827-- => [streamA1, streamA2, ... streamAn] [streamB1, streamB2, ...]
2828-- => [streamA1, streamB1, streamA2...StreamAn, streamBn]
2829-- => [a11, a12, ...a1j, b11, b12, ...b1k, a21, a22, ...]
2831{-# INLINE_NORMAL gintercalateSuffix #-}
2834 => Unfold m a c -> Stream m a -> Unfold m b c -> Stream m b -> Stream m c
2836 (Unfold istep1 inject1) (Stream step1 state1)
2837 (Unfold istep2 inject2) (Stream step2 state2) =
2838 Stream step (ICUFirst state1 state2)
2842 {-# INLINE_LATE step #-}
2843 step gst (ICUFirst s1 s2) = do
2844 r <- step1 (adaptState gst) s1
2848 i `seq` return (Skip (ICUFirstInner s s2 i))
2849 Skip s -> return $ Skip (ICUFirst s s2)
2852 step gst (ICUFirstOnly s1) = do
2853 r <- step1 (adaptState gst) s1
2857 i `seq` return (Skip (ICUFirstOnlyInner s i))
2858 Skip s -> return $ Skip (ICUFirstOnly s)
2861 step _ (ICUFirstInner s1 s2 i1) = do
2864 Yield x i' -> Yield x (ICUFirstInner s1 s2 i')
2865 Skip i' -> Skip (ICUFirstInner s1 s2 i')
2866 Stop -> Skip (ICUSecond s1 s2)
2868 step _ (ICUFirstOnlyInner s1 i1) = do
2871 Yield x i' -> Yield x (ICUFirstOnlyInner s1 i')
2872 Skip i' -> Skip (ICUFirstOnlyInner s1 i')
2873 Stop -> Skip (ICUFirstOnly s1)
2875 step gst (ICUSecond s1 s2) = do
2876 r <- step2 (adaptState gst) s2
2880 i `seq` return (Skip (ICUSecondInner s1 s i))
2881 Skip s -> return $ Skip (ICUSecond s1 s)
2882 Stop -> return $ Skip (ICUFirstOnly s1)
2884 step _ (ICUSecondInner s1 s2 i2) = do
2887 Yield x i' -> Yield x (ICUSecondInner s1 s2 i')
2888 Skip i' -> Skip (ICUSecondInner s1 s2 i')
2889 Stop -> Skip (ICUFirst s1 s2)
2891 step _ (ICUSecondOnly _s2) = undefined
2892 step _ (ICUSecondOnlyInner _s2 _i2) = undefined
2894data InterposeSuffixState s1 i1 =
2895 InterposeSuffixFirst s1
2896 -- | InterposeSuffixFirstYield s1 i1
2897 | InterposeSuffixFirstInner s1 i1
2898 | InterposeSuffixSecond s1
2900-- Note that if an unfolded layer turns out to be nil we still emit the
2901-- separator effect. An alternate behavior could be to emit the separator
2902-- effect only if at least one element has been yielded by the unfolding.
2903-- However, that becomes a bit complicated, so we have chosen the former
2905{-# INLINE_NORMAL interposeSuffix #-}
2908 => m c -> Unfold m b c -> Stream m b -> Stream m c
2911 (Unfold istep1 inject1) (Stream step1 state1) =
2912 Stream step (InterposeSuffixFirst state1)
2916 {-# INLINE_LATE step #-}
2917 step gst (InterposeSuffixFirst s1) = do
2918 r <- step1 (adaptState gst) s1
2922 i `seq` return (Skip (InterposeSuffixFirstInner s i))
2923 -- i `seq` return (Skip (InterposeSuffixFirstYield s i))
2924 Skip s -> return $ Skip (InterposeSuffixFirst s)
2928 step _ (InterposeSuffixFirstYield s1 i1) = do
2931 Yield x i' -> Yield x (InterposeSuffixFirstInner s1 i')
2932 Skip i' -> Skip (InterposeSuffixFirstYield s1 i')
2933 Stop -> Skip (InterposeSuffixFirst s1)
2936 step _ (InterposeSuffixFirstInner s1 i1) = do
2939 Yield x i' -> Yield x (InterposeSuffixFirstInner s1 i')
2940 Skip i' -> Skip (InterposeSuffixFirstInner s1 i')
2941 Stop -> Skip (InterposeSuffixSecond s1)
2943 step _ (InterposeSuffixSecond s1) = do
2945 return $ Yield r (InterposeSuffixFirst s1)
2947data ICALState s1 s2 i1 i2 a =
2949 -- | ICALFirstYield s1 s2 i1
2950 | ICALFirstInner s1 s2 i1
2952 | ICALFirstOnlyInner s1 i1
2953 | ICALSecondInject s1 s2
2954 | ICALFirstInject s1 s2 i2
2955 -- | ICALFirstBuf s1 s2 i1 i2
2956 | ICALSecondInner s1 s2 i1 i2
2957 -- -- | ICALSecondInner s1 s2 i1 i2 a
2958 -- -- | ICALFirstResume s1 s2 i1 i2 a
2960-- | Interleave streams (full streams, not the elements) unfolded from two
2961-- input streams and concat. Stop when the first stream stops. If the second
2962-- stream ends before the first one then first stream still keeps running alone
2963-- without any interleaving with the second stream.
2965-- [a1, a2, ... an] [b1, b2 ...]
2966-- => [streamA1, streamA2, ... streamAn] [streamB1, streamB2, ...]
2967-- => [streamA1, streamB1, streamA2...StreamAn, streamBn]
2968-- => [a11, a12, ...a1j, b11, b12, ...b1k, a21, a22, ...]
2970{-# INLINE_NORMAL gintercalate #-}
2973 => Unfold m a c -> Stream m a -> Unfold m b c -> Stream m b -> Stream m c
2975 (Unfold istep1 inject1) (Stream step1 state1)
2976 (Unfold istep2 inject2) (Stream step2 state2) =
2977 Stream step (ICALFirst state1 state2)
2981 {-# INLINE_LATE step #-}
2982 step gst (ICALFirst s1 s2) = do
2983 r <- step1 (adaptState gst) s1
2987 i `seq` return (Skip (ICALFirstInner s s2 i))
2988 -- i `seq` return (Skip (ICALFirstYield s s2 i))
2989 Skip s -> return $ Skip (ICALFirst s s2)
2993 step _ (ICALFirstYield s1 s2 i1) = do
2996 Yield x i' -> Yield x (ICALFirstInner s1 s2 i')
2997 Skip i' -> Skip (ICALFirstYield s1 s2 i')
2998 Stop -> Skip (ICALFirst s1 s2)
3001 step _ (ICALFirstInner s1 s2 i1) = do
3004 Yield x i' -> Yield x (ICALFirstInner s1 s2 i')
3005 Skip i' -> Skip (ICALFirstInner s1 s2 i')
3006 Stop -> Skip (ICALSecondInject s1 s2)
3008 step gst (ICALFirstOnly s1) = do
3009 r <- step1 (adaptState gst) s1
3013 i `seq` return (Skip (ICALFirstOnlyInner s i))
3014 Skip s -> return $ Skip (ICALFirstOnly s)
3017 step _ (ICALFirstOnlyInner s1 i1) = do
3020 Yield x i' -> Yield x (ICALFirstOnlyInner s1 i')
3021 Skip i' -> Skip (ICALFirstOnlyInner s1 i')
3022 Stop -> Skip (ICALFirstOnly s1)
3024 -- We inject the second stream even before checking if the first stream
3025 -- would yield any more elements. There is no clear choice whether we
3026 -- should do this before or after that. Doing it after may make the state
3027 -- machine a bit simpler though.
3028 step gst (ICALSecondInject s1 s2) = do
3029 r <- step2 (adaptState gst) s2
3033 i `seq` return (Skip (ICALFirstInject s1 s i))
3034 Skip s -> return $ Skip (ICALSecondInject s1 s)
3035 Stop -> return $ Skip (ICALFirstOnly s1)
3037 step gst (ICALFirstInject s1 s2 i2) = do
3038 r <- step1 (adaptState gst) s1
3042 i `seq` return (Skip (ICALSecondInner s s2 i i2))
3043 -- i `seq` return (Skip (ICALFirstBuf s s2 i i2))
3044 Skip s -> return $ Skip (ICALFirstInject s s2 i2)
3048 step _ (ICALFirstBuf s1 s2 i1 i2) = do
3051 Yield x i' -> Skip (ICALSecondInner s1 s2 i' i2 x)
3052 Skip i' -> Skip (ICALFirstBuf s1 s2 i' i2)
3055 step _ (ICALSecondInner s1 s2 i1 i2 v) = do
3058 Yield x i' -> Yield x (ICALSecondInner s1 s2 i1 i' v)
3059 Skip i' -> Skip (ICALSecondInner s1 s2 i1 i' v)
3060 Stop -> Skip (ICALFirstResume s1 s2 i1 i2 v)
3063 step _ (ICALSecondInner s1 s2 i1 i2) = do
3066 Yield x i' -> Yield x (ICALSecondInner s1 s2 i1 i')
3067 Skip i' -> Skip (ICALSecondInner s1 s2 i1 i')
3068 Stop -> Skip (ICALFirstInner s1 s2 i1)
3069 -- Stop -> Skip (ICALFirstResume s1 s2 i1 i2)
3072 step _ (ICALFirstResume s1 s2 i1 i2 x) = do
3073 return $ Yield x (ICALFirstInner s1 s2 i1 i2)
3076data InterposeState s1 i1 a =
3078 -- | InterposeFirstYield s1 i1
3079 | InterposeFirstInner s1 i1
3080 | InterposeFirstInject s1
3081 -- | InterposeFirstBuf s1 i1
3082 | InterposeSecondYield s1 i1
3083 -- -- | InterposeSecondYield s1 i1 a
3084 -- -- | InterposeFirstResume s1 i1 a
3086-- Note that this only interposes the pure values, we may run many effects to
3087-- generate those values as some effects may not generate anything (Skip).
3088{-# INLINE_NORMAL interpose #-}
3089interpose :: Monad m => m c -> Unfold m b c -> Stream m b -> Stream m c
3092 (Unfold istep1 inject1) (Stream step1 state1) =
3093 Stream step (InterposeFirst state1)
3097 {-# INLINE_LATE step #-}
3098 step gst (InterposeFirst s1) = do
3099 r <- step1 (adaptState gst) s1
3103 i `seq` return (Skip (InterposeFirstInner s i))
3104 -- i `seq` return (Skip (InterposeFirstYield s i))
3105 Skip s -> return $ Skip (InterposeFirst s)
3109 step _ (InterposeFirstYield s1 i1) = do
3112 Yield x i' -> Yield x (InterposeFirstInner s1 i')
3113 Skip i' -> Skip (InterposeFirstYield s1 i')
3114 Stop -> Skip (InterposeFirst s1)
3117 step _ (InterposeFirstInner s1 i1) = do
3120 Yield x i' -> Yield x (InterposeFirstInner s1 i')
3121 Skip i' -> Skip (InterposeFirstInner s1 i')
3122 Stop -> Skip (InterposeFirstInject s1)
3124 step gst (InterposeFirstInject s1) = do
3125 r <- step1 (adaptState gst) s1
3129 -- i `seq` return (Skip (InterposeFirstBuf s i))
3130 i `seq` return (Skip (InterposeSecondYield s i))
3131 Skip s -> return $ Skip (InterposeFirstInject s)
3135 step _ (InterposeFirstBuf s1 i1) = do
3138 Yield x i' -> Skip (InterposeSecondYield s1 i' x)
3139 Skip i' -> Skip (InterposeFirstBuf s1 i')
3144 step _ (InterposeSecondYield s1 i1 v) = do
3146 return $ Yield r (InterposeFirstResume s1 i1 v)
3148 step _ (InterposeSecondYield s1 i1) = do
3150 return $ Yield r (InterposeFirstInner s1 i1)
3153 step _ (InterposeFirstResume s1 i1 v) = do
3154 return $ Yield v (InterposeFirstInner s1 i1)
3157------------------------------------------------------------------------------
3159------------------------------------------------------------------------------
3161data GbracketState s1 s2 v
3163 | GBracketNormal s1 v
3164 | GBracketException s2
3166-- | The most general bracketing and exception combinator. All other
3167-- combinators can be expressed in terms of this combinator. This can also be
3168-- used for cases which are not covered by the standard combinators.
3172{-# INLINE_NORMAL gbracket #-}
3176 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
3177 -> (c -> m d) -- ^ after, on normal stop
3178 -> (c -> e -> Stream m b) -- ^ on exception
3179 -> (c -> Stream m b) -- ^ stream generator
3181gbracket bef exc aft fexc fnormal =
3182 Stream step GBracketInit
3186 {-# INLINE_LATE step #-}
3187 step _ GBracketInit = do
3189 return $ Skip $ GBracketNormal (fnormal r) r
3191 step gst (GBracketNormal (UnStream step1 st) v) = do
3192 res <- exc $ step1 gst st
3194 Right r -> case r of
3196 return $ Yield x (GBracketNormal (Stream step1 s) v)
3197 Skip s -> return $ Skip (GBracketNormal (Stream step1 s) v)
3198 Stop -> aft v >> return Stop
3199 Left e -> return $ Skip (GBracketException (fexc v e))
3200 step gst (GBracketException (UnStream step1 st)) = do
3203 Yield x s -> return $ Yield x (GBracketException (Stream step1 s))
3204 Skip s -> return $ Skip (GBracketException (Stream step1 s))
3207-- | Create an IORef holding a finalizer that is called automatically when the
3208-- IORef is garbage collected. The IORef can be written to with a 'Nothing'
3209-- value to deactivate the finalizer.
3210newFinalizedIORef :: (MonadIO m, MonadBaseControl IO m)
3211 => m a -> m (IORef (Maybe (IO ())))
3212newFinalizedIORef finalizer = do
3213 mrun <- captureMonadState
3214 ref <- liftIO $ newIORef $ Just $ liftIO $ void $ do
3215 _ <- runInIO mrun finalizer
3218 res <- readIORef ref
3220 Nothing -> return ()
3222 _ <- liftIO $ mkWeakIORef ref finalizer1
3225-- | Run the finalizer stored in an IORef and deactivate it so that it is run
3228runIORefFinalizer :: MonadIO m => IORef (Maybe (IO ())) -> m ()
3229runIORefFinalizer ref = liftIO $ do
3230 res <- readIORef ref
3232 Nothing -> return ()
3233 Just f -> writeIORef ref Nothing >> f
3235-- | Deactivate the finalizer stored in an IORef without running it.
3237clearIORefFinalizer :: MonadIO m => IORef (Maybe (IO ())) -> m ()
3238clearIORefFinalizer ref = liftIO $ writeIORef ref Nothing
3240data GbracketIOState s1 s2 v wref
3242 | GBracketIONormal s1 v wref
3243 | GBracketIOException s2
3245-- | Like gbracket but also uses a finalizer to make sure when the stream is
3246-- garbage collected we run the finalizing action. This requires a MonadIO and
3247-- MonadBaseControl IO constraint.
3249-- | The most general bracketing and exception combinator. All other
3250-- combinators can be expressed in terms of this combinator. This can also be
3251-- used for cases which are not covered by the standard combinators.
3255{-# INLINE_NORMAL gbracketIO #-}
3257 :: (MonadIO m, MonadBaseControl IO m)
3259 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
3260 -> (c -> m d) -- ^ after, on normal stop or GC
3261 -> (c -> e -> Stream m b) -- ^ on exception
3262 -> (c -> Stream m b) -- ^ stream generator
3264gbracketIO bef exc aft fexc fnormal =
3265 Stream step GBracketIOInit
3269 -- If the stream is never evaluated the "aft" action will never be
3270 -- called. For that to occur we will need the user of this API to pass a
3271 -- weak pointer to us.
3272 {-# INLINE_LATE step #-}
3273 step _ GBracketIOInit = do
3274 -- We mask asynchronous exceptions to make the execution
3275 -- of 'bef' and the registration of 'aft' atomic.
3276 -- A similar thing is done in the resourcet package: https://git.io/JvKV3
3277 -- Tutorial: https://markkarpov.com/tutorial/exceptions.html
3278 (r, ref) <- liftBaseOp_ mask_ $ do
3280 ref <- newFinalizedIORef (aft r)
3282 return $ Skip $ GBracketIONormal (fnormal r) r ref
3284 step gst (GBracketIONormal (UnStream step1 st) v ref) = do
3285 res <- exc $ step1 gst st
3287 Right r -> case r of
3289 return $ Yield x (GBracketIONormal (Stream step1 s) v ref)
3291 return $ Skip (GBracketIONormal (Stream step1 s) v ref)
3293 runIORefFinalizer ref
3296 clearIORefFinalizer ref
3297 return $ Skip (GBracketIOException (fexc v e))
3298 step gst (GBracketIOException (UnStream step1 st)) = do
3302 return $ Yield x (GBracketIOException (Stream step1 s))
3303 Skip s -> return $ Skip (GBracketIOException (Stream step1 s))
3306-- | Run a side effect before the stream yields its first element.
3307{-# INLINE_NORMAL before #-}
3308before :: Monad m => m b -> Stream m a -> Stream m a
3309before action (Stream step state) = Stream step' Nothing
3313 {-# INLINE_LATE step' #-}
3314 step' _ Nothing = action >> return (Skip (Just state))
3316 step' gst (Just st) = do
3319 Yield x s -> return $ Yield x (Just s)
3320 Skip s -> return $ Skip (Just s)
3323-- | Run a side effect whenever the stream stops normally.
3324{-# INLINE_NORMAL after #-}
3325after :: Monad m => m b -> Stream m a -> Stream m a
3326after action (Stream step state) = Stream step' state
3330 {-# INLINE_LATE step' #-}
3334 Yield x s -> return $ Yield x s
3335 Skip s -> return $ Skip s
3336 Stop -> action >> return Stop
3338{-# INLINE_NORMAL afterIO #-}
3339afterIO :: (MonadIO m, MonadBaseControl IO m)
3340 => m b -> Stream m a -> Stream m a
3341afterIO action (Stream step state) = Stream step' Nothing
3345 {-# INLINE_LATE step' #-}
3346 step' _ Nothing = do
3347 ref <- newFinalizedIORef action
3348 return $ Skip $ Just (state, ref)
3349 step' gst (Just (st, ref)) = do
3352 Yield x s -> return $ Yield x (Just (s, ref))
3353 Skip s -> return $ Skip (Just (s, ref))
3355 runIORefFinalizer ref
3358-- XXX These combinators are expensive due to the call to
3359-- onException/handle/try on each step. Therefore, when possible, they should
3360-- be called in an outer loop where we perform less iterations. For example, we
3361-- cannot call them on each iteration in a char stream, instead we can call
3362-- them when doing an IO on an array.
3364-- XXX For high performance error checks in busy streams we may need another
3365-- Error constructor in step.
3367-- | Run a side effect whenever the stream aborts due to an exception. The
3368-- exception is not caught, simply rethrown.
3369{-# INLINE_NORMAL onException #-}
3370onException :: MonadCatch m => m b -> Stream m a -> Stream m a
3371onException action str =
3372 gbracket (return ()) MC.try return
3373 (\_ (e :: MC.SomeException) -> nilM (action >> MC.throwM e))
3376{-# INLINE_NORMAL _onException #-}
3377_onException :: MonadCatch m => m b -> Stream m a -> Stream m a
3378_onException action (Stream step state) = Stream step' state
3382 {-# INLINE_LATE step' #-}
3384 res <- step gst st `MC.onException` action
3386 Yield x s -> return $ Yield x s
3387 Skip s -> return $ Skip s
3390-- XXX bracket is like concatMap, it generates a stream and then flattens it.
3391-- Like concatMap it has 10x worse performance compared to linear fused
3394-- | Run the first action before the stream starts and remember its output,
3395-- generate a stream using the output, run the second action providing the
3396-- remembered value as an argument whenever the stream ends normally or due to
3398{-# INLINE_NORMAL bracket #-}
3399bracket :: MonadCatch m => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
3400bracket bef aft bet =
3401 gbracket bef MC.try aft
3402 (\a (e :: SomeException) -> nilM (aft a >> MC.throwM e)) bet
3404{-# INLINE_NORMAL bracketIO #-}
3405bracketIO :: (MonadAsync m, MonadCatch m)
3406 => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
3407bracketIO bef aft bet =
3408 gbracketIO bef MC.try aft
3409 (\a (e :: SomeException) -> nilM (aft a >> MC.throwM e)) bet
3411data BracketState s v = BracketInit | BracketRun s v
3413{-# INLINE_NORMAL _bracket #-}
3414_bracket :: MonadCatch m => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
3415_bracket bef aft bet = Stream step' BracketInit
3419 {-# INLINE_LATE step' #-}
3420 step' _ BracketInit = bef >>= \x -> return (Skip (BracketRun (bet x) x))
3422 -- NOTE: It is important to use UnStream instead of the Stream pattern
3423 -- here, otherwise we get huge perf degradation, see note in concatMap.
3424 step' gst (BracketRun (UnStream step state) v) = do
3425 -- res <- step gst state `MC.onException` aft v
3426 res <- MC.try $ step gst state
3428 Left (e :: SomeException) -> aft v >> MC.throwM e >> return Stop
3429 Right r -> case r of
3430 Yield x s -> return $ Yield x (BracketRun (Stream step s) v)
3431 Skip s -> return $ Skip (BracketRun (Stream step s) v)
3432 Stop -> aft v >> return Stop
3434-- | Run a side effect whenever the stream stops normally or aborts due to an
3436{-# INLINE finally #-}
3437finally :: MonadCatch m => m b -> Stream m a -> Stream m a
3438-- finally action xs = after action $ onException action xs
3439finally action xs = bracket (return ()) (\_ -> action) (const xs)
3441{-# INLINE finallyIO #-}
3442finallyIO :: (MonadAsync m, MonadCatch m) => m b -> Stream m a -> Stream m a
3443finallyIO action xs = bracketIO (return ()) (\_ -> action) (const xs)
3445-- | When evaluating a stream if an exception occurs, stream evaluation aborts
3446-- and the specified exception handler is run with the exception as argument.
3447{-# INLINE_NORMAL handle #-}
3448handle :: (MonadCatch m, Exception e)
3449 => (e -> Stream m a) -> Stream m a -> Stream m a
3451 gbracket (return ()) MC.try return (\_ e -> f e) (\_ -> str)
3453{-# INLINE_NORMAL _handle #-}
3454_handle :: (MonadCatch m, Exception e)
3455 => (e -> Stream m a) -> Stream m a -> Stream m a
3456_handle f (Stream step state) = Stream step' (Left state)
3460 {-# INLINE_LATE step' #-}
3461 step' gst (Left st) = do
3462 res <- MC.try $ step gst st
3464 Left e -> return $ Skip $ Right (f e)
3465 Right r -> case r of
3466 Yield x s -> return $ Yield x (Left s)
3467 Skip s -> return $ Skip (Left s)
3470 step' gst (Right (UnStream step1 st)) = do
3473 Yield x s -> return $ Yield x (Right (Stream step1 s))
3474 Skip s -> return $ Skip (Right (Stream step1 s))
3477-------------------------------------------------------------------------------
3478-- General transformation
3479-------------------------------------------------------------------------------
3481{-# INLINE_NORMAL transform #-}
3482transform :: Monad m => Pipe m a b -> Stream m a -> Stream m b
3483transform (Pipe pstep1 pstep2 pstate) (Stream step state) =
3484 Stream step' (Consume pstate, state)
3488 {-# INLINE_LATE step' #-}
3490 step' gst (Consume pst, st) = pst `seq` do
3491 r <- step (adaptState gst) st
3496 Pipe.Yield b pst' -> return $ Yield b (pst', s)
3497 Pipe.Continue pst' -> return $ Skip (pst', s)
3498 Skip s -> return $ Skip (Consume pst, s)
3501 step' _ (Produce pst, st) = pst `seq` do
3504 Pipe.Yield b pst' -> return $ Yield b (pst', st)
3505 Pipe.Continue pst' -> return $ Skip (pst', st)
3507------------------------------------------------------------------------------
3508-- Transformation by Folding (Scans)
3509------------------------------------------------------------------------------
3511------------------------------------------------------------------------------
3513------------------------------------------------------------------------------
3515-- XXX Is a prescan useful, discarding the last step does not sound useful? I
3516-- am not sure about the utility of this function, so this is implemented but
3517-- not exposed. We can expose it if someone provides good reasons why this is
3520-- XXX We have to execute the stream one step ahead to know that we are at the
3521-- last step. The vector implementation of prescan executes the last fold step
3522-- but does not yield the result. This means we have executed the effect but
3523-- discarded value. This does not sound right. In this implementation we are
3524-- not executing the last fold step.
3525{-# INLINE_NORMAL prescanlM' #-}
3526prescanlM' :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
3527prescanlM' f mz (Stream step state) = Stream step' (state, mz)
3529 {-# INLINE_LATE step' #-}
3530 step' gst (st, prev) = do
3531 r <- step (adaptState gst) st
3535 return $ Yield acc (s, f acc x)
3536 Skip s -> return $ Skip (s, prev)
3539{-# INLINE prescanl' #-}
3540prescanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
3541prescanl' f z = prescanlM' (\a b -> return (f a b)) (return z)
3543------------------------------------------------------------------------------
3544-- Monolithic postscans (postscan followed by a map)
3545------------------------------------------------------------------------------
3547-- The performance of a modular postscan followed by a map seems to be
3548-- equivalent to this monolithic scan followed by map therefore we may not need
3549-- this implementation. We just have it for performance comparison and in case
3550-- modular version does not perform well in some situation.
3552{-# INLINE_NORMAL postscanlMx' #-}
3553postscanlMx' :: Monad m
3554 => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> Stream m b
3555postscanlMx' fstep begin done (Stream step state) = do
3556 Stream step' (state, begin)
3558 {-# INLINE_LATE step' #-}
3559 step' gst (st, acc) = do
3560 r <- step (adaptState gst) st
3566 v `seq` y `seq` return (Yield v (s, return y))
3567 Skip s -> return $ Skip (s, acc)
3570{-# INLINE_NORMAL postscanlx' #-}
3571postscanlx' :: Monad m
3572 => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> Stream m b
3573postscanlx' fstep begin done s =
3574 postscanlMx' (\b a -> return (fstep b a)) (return begin) (return . done) s
3576-- XXX do we need consM strict to evaluate the begin value?
3577{-# INLINE scanlMx' #-}
3579 => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> Stream m b
3580scanlMx' fstep begin done s =
3581 (begin >>= \x -> x `seq` done x) `consM` postscanlMx' fstep begin done s
3583{-# INLINE scanlx' #-}
3585 => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> Stream m b
3586scanlx' fstep begin done s =
3587 scanlMx' (\b a -> return (fstep b a)) (return begin) (return . done) s
3589------------------------------------------------------------------------------
3591------------------------------------------------------------------------------
3593{-# INLINE_NORMAL postscanlM' #-}
3594postscanlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
3595postscanlM' fstep begin (Stream step state) =
3596 begin `seq` Stream step' (state, begin)
3598 {-# INLINE_LATE step' #-}
3599 step' gst (st, acc) = acc `seq` do
3600 r <- step (adaptState gst) st
3604 y `seq` return (Yield y (s, y))
3605 Skip s -> return $ Skip (s, acc)
3608{-# INLINE_NORMAL postscanl' #-}
3609postscanl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
3610postscanl' f = postscanlM' (\a b -> return (f a b))
3612{-# INLINE_NORMAL postscanlM #-}
3613postscanlM :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
3614postscanlM fstep begin (Stream step state) = Stream step' (state, begin)
3616 {-# INLINE_LATE step' #-}
3617 step' gst (st, acc) = do
3618 r <- step (adaptState gst) st
3622 return (Yield y (s, y))
3623 Skip s -> return $ Skip (s, acc)
3626{-# INLINE_NORMAL postscanl #-}
3627postscanl :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
3628postscanl f = postscanlM (\a b -> return (f a b))
3630{-# INLINE_NORMAL scanlM' #-}
3631scanlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
3632scanlM' fstep begin s = begin `seq` (begin `cons` postscanlM' fstep begin s)
3634{-# INLINE scanl' #-}
3635scanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
3636scanl' f = scanlM' (\a b -> return (f a b))
3638{-# INLINE_NORMAL scanlM #-}
3639scanlM :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
3640scanlM fstep begin s = begin `cons` postscanlM fstep begin s
3643scanl :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
3644scanl f = scanlM (\a b -> return (f a b))
3646{-# INLINE_NORMAL scanl1M #-}
3647scanl1M :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
3648scanl1M fstep (Stream step state) = Stream step' (state, Nothing)
3650 {-# INLINE_LATE step' #-}
3651 step' gst (st, Nothing) = do
3654 Yield x s -> return $ Yield x (s, Just x)
3655 Skip s -> return $ Skip (s, Nothing)
3658 step' gst (st, Just acc) = do
3663 return $ Yield z (s, Just z)
3664 Skip s -> return $ Skip (s, Just acc)
3667{-# INLINE scanl1 #-}
3668scanl1 :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
3669scanl1 f = scanl1M (\x y -> return (f x y))
3671{-# INLINE_NORMAL scanl1M' #-}
3672scanl1M' :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
3673scanl1M' fstep (Stream step state) = Stream step' (state, Nothing)
3675 {-# INLINE_LATE step' #-}
3676 step' gst (st, Nothing) = do
3679 Yield x s -> x `seq` return $ Yield x (s, Just x)
3680 Skip s -> return $ Skip (s, Nothing)
3683 step' gst (st, Just acc) = acc `seq` do
3688 z `seq` return $ Yield z (s, Just z)
3689 Skip s -> return $ Skip (s, Just acc)
3692{-# INLINE scanl1' #-}
3693scanl1' :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
3694scanl1' f = scanl1M' (\x y -> return (f x y))
3696------------------------------------------------------------------------------
3698------------------------------------------------------------------------------
3700data RollingMapState s a = RollingMapInit s | RollingMapGo s a
3702{-# INLINE rollingMapM #-}
3703rollingMapM :: Monad m => (a -> a -> m b) -> Stream m a -> Stream m b
3704rollingMapM f (Stream step1 state1) = Stream step (RollingMapInit state1)
3706 step gst (RollingMapInit st) = do
3707 r <- step1 (adaptState gst) st
3709 Yield x s -> Skip $ RollingMapGo s x
3710 Skip s -> Skip $ RollingMapInit s
3713 step gst (RollingMapGo s1 x1) = do
3714 r <- step1 (adaptState gst) s1
3718 return $ Yield res $ RollingMapGo s x
3719 Skip s -> return $ Skip $ RollingMapGo s x1
3720 Stop -> return $ Stop
3722{-# INLINE rollingMap #-}
3723rollingMap :: Monad m => (a -> a -> b) -> Stream m a -> Stream m b
3724rollingMap f = rollingMapM (\x y -> return $ f x y)
3726------------------------------------------------------------------------------
3727-- Tapping/Distributing
3728------------------------------------------------------------------------------
3731tap :: Monad m => Fold m a b -> Stream m a -> Stream m a
3732tap (Fold fstep initial extract) (Stream step state) = Stream step' Nothing
3736 step' _ Nothing = do
3738 return $ Skip (Just (r, state))
3740 step' gst (Just (acc, st)) = acc `seq` do
3745 return $ Yield x (Just (acc', s))
3746 Skip s -> return $ Skip (Just (acc, s))
3751{-# INLINE_NORMAL tapOffsetEvery #-}
3752tapOffsetEvery :: Monad m
3753 => Int -> Int -> Fold m a b -> Stream m a -> Stream m a
3754tapOffsetEvery offset n (Fold fstep initial extract) (Stream step state) =
3755 Stream step' Nothing
3759 {-# INLINE_LATE step' #-}
3760 step' _ Nothing = do
3762 return $ Skip (Just (r, state, offset `mod` n))
3764 step' gst (Just (acc, st, count)) | count <= 0 = do
3768 !acc' <- fstep acc x
3769 return $ Yield x (Just (acc', s, n - 1))
3770 Skip s -> return $ Skip (Just (acc, s, count))
3775 step' gst (Just (acc, st, count)) = do
3778 Yield x s -> return $ Yield x (Just (acc, s, count - 1))
3779 Skip s -> return $ Skip (Just (acc, s, count))
3784{-# INLINE_NORMAL pollCounts #-}
3788 -> (Stream m Int -> Stream m Int)
3792pollCounts predicate transf fld (Stream step state) = Stream step' Nothing
3795 {-# INLINE_LATE step' #-}
3796 step' _ Nothing = do
3797 -- As long as we are using an "Int" for counts lockfree reads from
3798 -- Var should work correctly on both 32-bit and 64-bit machines.
3799 -- However, an Int on a 32-bit machine may overflow quickly.
3800 countVar <- liftIO $ newVar (0 :: Int)
3802 $ void $ runFold fld
3803 $ transf $ fromPrimVar countVar
3804 return $ Skip (Just (countVar, tid, state))
3806 step' gst (Just (countVar, tid, st)) = do
3810 when (predicate x) $ liftIO $ modifyVar' countVar (+ 1)
3811 return $ Yield x (Just (countVar, tid, s))
3812 Skip s -> return $ Skip (Just (countVar, tid, s))
3814 liftIO $ killThread tid
3817{-# INLINE_NORMAL tapRate #-}
3819 (MonadAsync m, MonadCatch m)
3824tapRate samplingRate action (Stream step state) = Stream step' Nothing
3826 {-# NOINLINE loop #-}
3827 loop countVar prev = do
3830 (do liftIO $ threadDelay (round $ samplingRate * 1000000)
3831 i <- liftIO $ readVar countVar
3832 let !diff = i - prev
3835 (\(e :: AsyncException) -> do
3836 i <- liftIO $ readVar countVar
3837 let !diff = i - prev
3839 throwM (MC.toException e))
3842 {-# INLINE_LATE step' #-}
3843 step' _ Nothing = do
3844 countVar <- liftIO $ newVar 0
3845 tid <- fork $ loop countVar 0
3846 ref <- liftIO $ newIORef ()
3847 _ <- liftIO $ mkWeakIORef ref (killThread tid)
3848 return $ Skip (Just (countVar, tid, state, ref))
3850 step' gst (Just (countVar, tid, st, ref)) = do
3854 liftIO $ modifyVar' countVar (+ 1)
3855 return $ Yield x (Just (countVar, tid, s, ref))
3856 Skip s -> return $ Skip (Just (countVar, tid, s, ref))
3858 liftIO $ killThread tid
3862-------------------------------------------------------------------------------
3864-------------------------------------------------------------------------------
3866{-# INLINE_NORMAL takeWhileM #-}
3867takeWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
3868takeWhileM f (Stream step state) = Stream step' state
3870 {-# INLINE_LATE step' #-}
3876 return $ if b then Yield x s else Stop
3877 Skip s -> return $ Skip s
3880{-# INLINE takeWhile #-}
3881takeWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
3882takeWhile f = takeWhileM (return . f)
3884{-# INLINE_NORMAL drop #-}
3885drop :: Monad m => Int -> Stream m a -> Stream m a
3886drop n (Stream step state) = Stream step' (state, Just n)
3888 {-# INLINE_LATE step' #-}
3889 step' gst (st, Just i)
3894 Yield _ s -> Skip (s, Just (i - 1))
3895 Skip s -> Skip (s, Just i)
3897 | otherwise = return $ Skip (st, Nothing)
3899 step' gst (st, Nothing) = do
3903 Yield x s -> Yield x (s, Nothing)
3904 Skip s -> Skip (s, Nothing)
3907data DropWhileState s a
3909 | DropWhileYield a s
3912{-# INLINE_NORMAL dropWhileM #-}
3913dropWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
3914dropWhileM f (Stream step state) = Stream step' (DropWhileDrop state)
3916 {-# INLINE_LATE step' #-}
3917 step' gst (DropWhileDrop st) = do
3923 then return $ Skip (DropWhileDrop s)
3924 else return $ Skip (DropWhileYield x s)
3925 Skip s -> return $ Skip (DropWhileDrop s)
3928 step' gst (DropWhileNext st) = do
3931 Yield x s -> return $ Skip (DropWhileYield x s)
3932 Skip s -> return $ Skip (DropWhileNext s)
3935 step' _ (DropWhileYield x st) = return $ Yield x (DropWhileNext st)
3937{-# INLINE dropWhile #-}
3938dropWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
3939dropWhile f = dropWhileM (return . f)
3941{-# INLINE_NORMAL filterM #-}
3942filterM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
3943filterM f (Stream step state) = Stream step' state
3945 {-# INLINE_LATE step' #-}
3954 Skip s -> return $ Skip s
3957{-# INLINE filter #-}
3958filter :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
3959filter f = filterM (return . f)
3961{-# INLINE_NORMAL uniq #-}
3962uniq :: (Eq a, Monad m) => Stream m a -> Stream m a
3963uniq (Stream step state) = Stream step' (Nothing, state)
3965 {-# INLINE_LATE step' #-}
3966 step' gst (Nothing, st) = do
3969 Yield x s -> return $ Yield x (Just x, s)
3970 Skip s -> return $ Skip (Nothing, s)
3972 step' gst (Just x, st) = do
3975 Yield y s | x == y -> return $ Skip (Just x, s)
3976 | otherwise -> return $ Yield y (Just y, s)
3977 Skip s -> return $ Skip (Just x, s)
3980------------------------------------------------------------------------------
3981-- Transformation by Mapping
3982------------------------------------------------------------------------------
3984{-# INLINE_NORMAL sequence #-}
3985sequence :: Monad m => Stream m (m a) -> Stream m a
3986sequence (Stream step state) = Stream step' state
3988 {-# INLINE_LATE step' #-}
3990 r <- step (adaptState gst) st
3992 Yield x s -> x >>= \a -> return (Yield a s)
3993 Skip s -> return $ Skip s
3996------------------------------------------------------------------------------
3998------------------------------------------------------------------------------
4000data LoopState x s = FirstYield s
4001 | InterspersingYield s
4004{-# INLINE_NORMAL intersperseM #-}
4005intersperseM :: Monad m => m a -> Stream m a -> Stream m a
4006intersperseM m (Stream step state) = Stream step' (FirstYield state)
4008 {-# INLINE_LATE step' #-}
4009 step' gst (FirstYield st) = do
4013 Yield x s -> Skip (YieldAndCarry x s)
4014 Skip s -> Skip (FirstYield s)
4017 step' gst (InterspersingYield st) = do
4022 return $ Yield a (YieldAndCarry x s)
4023 Skip s -> return $ Skip $ InterspersingYield s
4026 step' _ (YieldAndCarry x st) = return $ Yield x (InterspersingYield st)
4028{-# INLINE_NORMAL intersperseM_ #-}
4029intersperseM_ :: Monad m => m b -> Stream m a -> Stream m a
4030intersperseM_ m (Stream step1 state1) = Stream step (Left (pure (), state1))
4032 {-# INLINE_LATE step #-}
4033 step gst (Left (eff, st)) = do
4036 Yield x s -> eff >> return (Yield x (Right s))
4037 Skip s -> return $ Skip (Left (eff, s))
4040 step _ (Right st) = return $ Skip $ Left (void m, st)
4045 | SuffixYield a (SuffixState s a)
4047{-# INLINE_NORMAL intersperseSuffix #-}
4048intersperseSuffix :: forall m a. Monad m => m a -> Stream m a -> Stream m a
4049intersperseSuffix action (Stream step state) = Stream step' (SuffixElem state)
4051 {-# INLINE_LATE step' #-}
4052 step' gst (SuffixElem st) = do
4055 Yield x s -> Skip (SuffixYield x (SuffixSuffix s))
4056 Skip s -> Skip (SuffixElem s)
4059 step' _ (SuffixSuffix st) = do
4060 action >>= \r -> return $ Skip (SuffixYield r (SuffixElem st))
4062 step' _ (SuffixYield x next) = return $ Yield x next
4064{-# INLINE_NORMAL intersperseSuffix_ #-}
4065intersperseSuffix_ :: Monad m => m b -> Stream m a -> Stream m a
4066intersperseSuffix_ m (Stream step1 state1) = Stream step (Left state1)
4068 {-# INLINE_LATE step #-}
4069 step gst (Left st) = do
4072 Yield x s -> return $ Yield x (Right s)
4073 Skip s -> return $ Skip $ Left s
4076 step _ (Right st) = m >> return (Skip (Left st))
4078data SuffixSpanState s a
4079 = SuffixSpanElem s Int
4080 | SuffixSpanSuffix s
4081 | SuffixSpanYield a (SuffixSpanState s a)
4085-- | intersperse after every n items
4086{-# INLINE_NORMAL intersperseSuffixBySpan #-}
4087intersperseSuffixBySpan :: forall m a. Monad m
4088 => Int -> m a -> Stream m a -> Stream m a
4089intersperseSuffixBySpan n action (Stream step state) =
4090 Stream step' (SuffixSpanElem state n)
4092 {-# INLINE_LATE step' #-}
4093 step' gst (SuffixSpanElem st i) | i > 0 = do
4096 Yield x s -> Skip (SuffixSpanYield x (SuffixSpanElem s (i - 1)))
4097 Skip s -> Skip (SuffixSpanElem s i)
4098 Stop -> if i == n then Stop else Skip SuffixSpanLast
4099 step' _ (SuffixSpanElem st _) = return $ Skip (SuffixSpanSuffix st)
4101 step' _ (SuffixSpanSuffix st) = do
4102 action >>= \r -> return $ Skip (SuffixSpanYield r (SuffixSpanElem st n))
4104 step' _ (SuffixSpanLast) = do
4105 action >>= \r -> return $ Skip (SuffixSpanYield r SuffixSpanStop)
4107 step' _ (SuffixSpanYield x next) = return $ Yield x next
4109 step' _ (SuffixSpanStop) = return Stop
4111{-# INLINE intersperse #-}
4112intersperse :: Monad m => a -> Stream m a -> Stream m a
4113intersperse a = intersperseM (return a)
4115{-# INLINE_NORMAL insertBy #-}
4116insertBy :: Monad m => (a -> a -> Ordering) -> a -> Stream m a -> Stream m a
4117insertBy cmp a (Stream step state) = Stream step' (state, False, Nothing)
4119 {-# INLINE_LATE step' #-}
4120 step' gst (st, False, _) = do
4123 Yield x s -> case cmp a x of
4124 GT -> return $ Yield x (s, False, Nothing)
4125 _ -> return $ Yield a (s, True, Just x)
4126 Skip s -> return $ Skip (s, False, Nothing)
4127 Stop -> return $ Yield a (st, True, Nothing)
4129 step' _ (_, True, Nothing) = return Stop
4131 step' gst (st, True, Just prev) = do
4134 Yield x s -> return $ Yield prev (s, True, Just x)
4135 Skip s -> return $ Skip (s, True, Just prev)
4136 Stop -> return $ Yield prev (st, True, Nothing)
4138------------------------------------------------------------------------------
4140------------------------------------------------------------------------------
4142{-# INLINE_NORMAL deleteBy #-}
4143deleteBy :: Monad m => (a -> a -> Bool) -> a -> Stream m a -> Stream m a
4144deleteBy eq x (Stream step state) = Stream step' (state, False)
4146 {-# INLINE_LATE step' #-}
4147 step' gst (st, False) = do
4150 Yield y s -> return $
4151 if eq x y then Skip (s, True) else Yield y (s, False)
4152 Skip s -> return $ Skip (s, False)
4155 step' gst (st, True) = do
4158 Yield y s -> return $ Yield y (s, True)
4159 Skip s -> return $ Skip (s, True)
4162------------------------------------------------------------------------------
4163-- Transformation by Map and Filter
4164------------------------------------------------------------------------------
4166-- XXX Will this always fuse properly?
4167{-# INLINE_NORMAL mapMaybe #-}
4168mapMaybe :: Monad m => (a -> Maybe b) -> Stream m a -> Stream m b
4169mapMaybe f = fmap fromJust . filter isJust . map f
4171{-# INLINE_NORMAL mapMaybeM #-}
4172mapMaybeM :: Monad m => (a -> m (Maybe b)) -> Stream m a -> Stream m b
4173mapMaybeM f = fmap fromJust . filter isJust . mapM f
4175------------------------------------------------------------------------------
4177------------------------------------------------------------------------------
4179{-# INLINE_NORMAL indexed #-}
4180indexed :: Monad m => Stream m a -> Stream m (Int, a)
4181indexed (Stream step state) = Stream step' (state, 0)
4183 {-# INLINE_LATE step' #-}
4184 step' gst (st, i) = i `seq` do
4185 r <- step (adaptState gst) st
4187 Yield x s -> return $ Yield (i, x) (s, i+1)
4188 Skip s -> return $ Skip (s, i)
4191{-# INLINE_NORMAL indexedR #-}
4192indexedR :: Monad m => Int -> Stream m a -> Stream m (Int, a)
4193indexedR m (Stream step state) = Stream step' (state, m)
4195 {-# INLINE_LATE step' #-}
4196 step' gst (st, i) = i `seq` do
4197 r <- step (adaptState gst) st
4199 Yield x s -> let i' = i - 1
4200 in return $ Yield (i, x) (s, i')
4201 Skip s -> return $ Skip (s, i)
4204{-# INLINE_NORMAL zipWithM #-}
4206 => (a -> b -> m c) -> Stream m a -> Stream m b -> Stream m c
4207zipWithM f (Stream stepa ta) (Stream stepb tb) = Stream step (ta, tb, Nothing)
4209 {-# INLINE_LATE step #-}
4210 step gst (sa, sb, Nothing) = do
4211 r <- stepa (adaptState gst) sa
4214 Yield x sa' -> Skip (sa', sb, Just x)
4215 Skip sa' -> Skip (sa', sb, Nothing)
4218 step gst (sa, sb, Just x) = do
4219 r <- stepb (adaptState gst) sb
4223 return $ Yield z (sa, sb', Nothing)
4224 Skip sb' -> return $ Skip (sa, sb', Just x)
4227#if __GLASGOW_HASKELL__ >= 801
4228{-# RULES "zipWithM xs xs"
4229 forall f xs. zipWithM @Identity f xs xs = mapM (\x -> f x x) xs #-}
4232{-# INLINE zipWith #-}
4233zipWith :: Monad m => (a -> b -> c) -> Stream m a -> Stream m b -> Stream m c
4234zipWith f = zipWithM (\a b -> return (f a b))
4236------------------------------------------------------------------------------
4238------------------------------------------------------------------------------
4240{-# INLINE_NORMAL mergeByM #-}
4243 => (a -> a -> m Ordering) -> Stream m a -> Stream m a -> Stream m a
4244mergeByM cmp (Stream stepa ta) (Stream stepb tb) =
4245 Stream step (Just ta, Just tb, Nothing, Nothing)
4247 {-# INLINE_LATE step #-}
4249 -- one of the values is missing, and the corresponding stream is running
4250 step gst (Just sa, sb, Nothing, b) = do
4253 Yield a sa' -> Skip (Just sa', sb, Just a, b)
4254 Skip sa' -> Skip (Just sa', sb, Nothing, b)
4255 Stop -> Skip (Nothing, sb, Nothing, b)
4257 step gst (sa, Just sb, a, Nothing) = do
4260 Yield b sb' -> Skip (sa, Just sb', a, Just b)
4261 Skip sb' -> Skip (sa, Just sb', a, Nothing)
4262 Stop -> Skip (sa, Nothing, a, Nothing)
4264 -- both the values are available
4265 step _ (sa, sb, Just a, Just b) = do
4267 return $ case res of
4268 GT -> Yield b (sa, sb, Just a, Nothing)
4269 _ -> Yield a (sa, sb, Nothing, Just b)
4271 -- one of the values is missing, corresponding stream is done
4272 step _ (Nothing, sb, Nothing, Just b) =
4273 return $ Yield b (Nothing, sb, Nothing, Nothing)
4275 step _ (sa, Nothing, Just a, Nothing) =
4276 return $ Yield a (sa, Nothing, Nothing, Nothing)
4278 step _ (Nothing, Nothing, Nothing, Nothing) = return Stop
4280{-# INLINE mergeBy #-}
4283 => (a -> a -> Ordering) -> Stream m a -> Stream m a -> Stream m a
4284mergeBy cmp = mergeByM (\a b -> return $ cmp a b)
4286------------------------------------------------------------------------------
4287-- Transformation comprehensions
4288------------------------------------------------------------------------------
4290{-# INLINE_NORMAL the #-}
4291the :: (Eq a, Monad m) => Stream m a -> m (Maybe a)
4292the (Stream step state) = go state
4295 r <- step defState st
4297 Yield x s -> go' x s
4299 Stop -> return Nothing
4301 r <- step defState st
4303 Yield x s | x == n -> go' n s
4304 | otherwise -> return Nothing
4306 Stop -> return (Just n)
4308{-# INLINE runFold #-}
4309runFold :: (Monad m) => Fold m a b -> Stream m a -> m b
4310runFold (Fold step begin done) = foldlMx' step begin done
4312-------------------------------------------------------------------------------
4313-- Concurrent application and fold
4314-------------------------------------------------------------------------------
4316-- XXX These functions should be moved to Stream/Parallel.hs
4318-- Using StreamD the worker stream producing code can fuse with the code to
4319-- queue output to the SVar giving some perf boost.
4321-- Note that StreamD can only be used in limited situations, specifically, we
4322-- cannot implement joinStreamVarPar using this.
4324-- XXX make sure that the SVar passed is a Parallel style SVar.
4326-- | Fold the supplied stream to the SVar asynchronously using Parallel
4327-- concurrency style.
4328-- {-# INLINE_NORMAL toSVarParallel #-}
4329{-# INLINE toSVarParallel #-}
4330toSVarParallel :: MonadAsync m
4331 => State t m a -> SVar t m a -> Stream m a -> m ()
4332toSVarParallel st sv xs =
4333 if svarInspectMode sv
4337 case getYieldLimit st of
4338 Nothing -> doFork (work Nothing)
4340 (handleChildException sv)
4341 Just _ -> doFork (workLim Nothing)
4343 (handleChildException sv)
4348 {-# NOINLINE work #-}
4349 work info = (runFold (FL.toParallelSVar sv info) xs)
4351 {-# NOINLINE workLim #-}
4352 workLim info = runFold (FL.toParallelSVarLimited sv info) xs
4354 {-# NOINLINE forkWithDiag #-}
4356 -- We do not use workerCount in case of ParallelVar but still there is
4357 -- no harm in maintaining it correctly.
4358 liftIO $ atomicModifyIORefCAS_ (workerCount sv) $ \n -> n + 1
4360 -- This allocation matters when significant number of workers are being
4361 -- sent. We allocate it only when needed. The overhead increases by 4x.
4363 case yieldRateInfo sv of
4364 Nothing -> return Nothing
4365 Just _ -> liftIO $ do
4366 cntRef <- newIORef 0
4367 t <- getTime Monotonic
4368 lat <- newIORef (0, t)
4369 return $ Just WorkerInfo
4370 { workerYieldMax = 0
4371 , workerYieldCount = cntRef
4372 , workerLatencyStart = lat
4375 case getYieldLimit st of
4376 Nothing -> doFork (work winfo)
4378 (handleChildException sv)
4379 Just _ -> doFork (workLim winfo)
4381 (handleChildException sv)
4384{-# INLINE_NORMAL mkParallelD #-}
4385mkParallelD :: MonadAsync m => Stream m a -> Stream m a
4386mkParallelD m = Stream step Nothing
4389 step gst Nothing = do
4390 sv <- newParallelVar StopNone gst
4391 toSVarParallel gst sv m
4392 -- XXX use unfold instead?
4393 return $ Skip $ Just $ fromSVar sv
4395 step gst (Just (UnStream step1 st)) = do
4398 Yield a s -> Yield a (Just $ Stream step1 s)
4399 Skip s -> Skip (Just $ Stream step1 s)
4402-- Compare with mkAsync. mkAsync uses an Async style SVar whereas this uses a
4403-- parallel style SVar for evaluation. Currently, parallel style cannot use
4404-- rate control whereas Async style can use rate control. In async style SVar
4405-- the worker thread terminates when the buffer is full whereas in Parallel
4408-- | Make the stream producer and consumer run concurrently by introducing a
4409-- buffer between them. The producer thread evaluates the input stream until
4410-- the buffer fills, it blocks if the buffer is full until there is space in
4411-- the buffer. The consumer consumes the stream lazily from the buffer.
4415{-# INLINE_NORMAL mkParallel #-}
4416mkParallel :: (K.IsStream t, MonadAsync m) => t m a -> t m a
4417mkParallel = fromStreamD . mkParallelD . toStreamD
4419-- Note: we can use another API with two callbacks stop and yield if we want
4420-- the callback to be able to indicate end of stream.
4422-- | Generates a callback and a stream pair. The callback returned is used to
4423-- queue values to the stream. The stream is infinite, there is no way for the
4424-- callback to indicate that it is done now.
4428{-# INLINE_NORMAL newCallbackStream #-}
4429newCallbackStream :: (K.IsStream t, MonadAsync m) => m ((a -> m ()), t m a)
4430newCallbackStream = do
4431 sv <- newParallelVar StopNone defState
4433 -- XXX Add our own thread-id to the SVar as we can not know the callback's
4434 -- thread-id and the callback is not run in a managed worker. We need to
4435 -- handle this better.
4436 liftIO myThreadId >>= modifyThread sv
4438 let callback a = liftIO $ void $ send sv (ChildYield a)
4439 -- XXX we can return an SVar and then the consumer can unfold from the
4441 return (callback, fromStreamD (fromSVar sv))
4443-------------------------------------------------------------------------------
4445-------------------------------------------------------------------------------
4447-- | Create an SVar with a fold consumer that will fold any elements sent to it
4448-- using the supplied fold function.
4449{-# INLINE newFoldSVar #-}
4450newFoldSVar :: MonadAsync m => State t m a -> Fold m a b -> m (SVar t m a)
4451newFoldSVar stt f = do
4452 -- Buffer size for the SVar is derived from the current state
4453 sv <- newParallelVar StopAny (adaptState stt)
4454 -- Add the producer thread-id to the SVar.
4455 liftIO myThreadId >>= modifyThread sv
4456 void $ doFork (work sv) (svarMrun sv) (handleFoldException sv)
4461 {-# NOINLINE work #-}
4462 work sv = void $ runFold f $ fromProducer sv
4464data TapState sv st = TapInit | Tapping sv st | TapDone st
4466{-# INLINE_NORMAL tapAsync #-}
4467tapAsync :: MonadAsync m => Fold m a b -> Stream m a -> Stream m a
4468tapAsync f (Stream step1 state1) = Stream step TapInit
4472 -- In general, a Stop event would come equipped with the result
4473 -- of the fold. It is not used here but it would be useful in
4474 -- applicative and distribute.
4475 done <- fromConsumer svr
4476 when (not done) $ do
4477 liftIO $ withDiagMVar svr "teeToSVar: waiting to drain"
4478 $ takeMVar (outputDoorBellFromConsumer svr)
4482 liftIO $ sendStop svr Nothing
4483 -- drain/wait until a stop event arrives from the fold.
4486 {-# INLINE_LATE step #-}
4487 step gst TapInit = do
4488 sv <- newFoldSVar gst f
4489 return $ Skip (Tapping sv state1)
4491 step gst (Tapping sv st) = do
4495 done <- pushToFold sv a
4498 -- XXX we do not need to wait synchronously here
4500 return $ Yield a (TapDone s)
4501 else return $ Yield a (Tapping sv s)
4502 Skip s -> return $ Skip (Tapping sv s)
4507 step gst (TapDone st) = do
4510 Yield a s -> Yield a (TapDone s)
4511 Skip s -> Skip (TapDone s)
4514-- XXX Exported from Array again as this fold is specific to Array
4515-- | Take last 'n' elements from the stream and discard the rest.
4517lastN :: (Storable a, MonadIO m) => Int -> Fold m a (Array a)
4519 | n <= 0 = fmap (const mempty) FL.drain
4520 | otherwise = Fold step initial done
4522 step (Tuple3' rb rh i) a = do
4523 rh1 <- liftIO $ RB.unsafeInsert rb rh a
4524 return $ Tuple3' rb rh1 (i + 1)
4525 initial = fmap (\(a, b) -> Tuple3' a b (0 :: Int)) $ liftIO $ RB.new n
4526 done (Tuple3' rb rh i) = do
4527 arr <- liftIO $ A.newArray n
4528 foldFunc i rh snoc' arr rb
4529 snoc' b a = liftIO $ A.unsafeSnoc b a
4531 | i < n = RB.unsafeFoldRingM
4532 | otherwise = RB.unsafeFoldRingFullM
4534------------------------------------------------------------------------------
4536------------------------------------------------------------------------------
4538-- XXX using getTime in the loop can be pretty expensive especially for
4539-- computations where iterations are lightweight. We have the following
4542-- 1) Run a timeout thread updating a flag asynchronously and check that
4543-- flag here, that way we can have a cheap termination check.
4545-- 2) Use COARSE clock to get time with lower resolution but more efficiently.
4547-- 3) Use rdtscp/rdtsc to get time directly from the processor, compute the
4548-- termination value of rdtsc in the beginning and then in each iteration just
4549-- get rdtsc and check if we should terminate.
4553 | TakeByTimeCheck st s
4554 | TakeByTimeYield st s
4556{-# INLINE_NORMAL takeByTime #-}
4557takeByTime :: (MonadIO m, TimeUnit64 t) => t -> Stream m a -> Stream m a
4558takeByTime duration (Stream step1 state1) = Stream step (TakeByTimeInit state1)
4561 lim = toRelTime64 duration
4563 {-# INLINE_LATE step #-}
4564 step _ (TakeByTimeInit _) | lim == 0 = return Stop
4565 step _ (TakeByTimeInit st) = do
4566 t0 <- liftIO $ getTime Monotonic
4567 return $ Skip (TakeByTimeYield st t0)
4568 step _ (TakeByTimeCheck st t0) = do
4569 t <- liftIO $ getTime Monotonic
4571 if diffAbsTime64 t t0 > lim
4573 else Skip (TakeByTimeYield st t0)
4574 step gst (TakeByTimeYield st t0) = do
4577 Yield x s -> Yield x (TakeByTimeCheck s t0)
4578 Skip s -> Skip (TakeByTimeCheck s t0)
4581data DropByTime st s x
4583 | DropByTimeGen st s
4584 | DropByTimeCheck st s x
4585 | DropByTimeYield st
4587{-# INLINE_NORMAL dropByTime #-}
4588dropByTime :: (MonadIO m, TimeUnit64 t) => t -> Stream m a -> Stream m a
4589dropByTime duration (Stream step1 state1) = Stream step (DropByTimeInit state1)
4592 lim = toRelTime64 duration
4594 {-# INLINE_LATE step #-}
4595 step _ (DropByTimeInit st) = do
4596 t0 <- liftIO $ getTime Monotonic
4597 return $ Skip (DropByTimeGen st t0)
4598 step gst (DropByTimeGen st t0) = do
4601 Yield x s -> Skip (DropByTimeCheck s t0 x)
4602 Skip s -> Skip (DropByTimeGen s t0)
4604 step _ (DropByTimeCheck st t0 x) = do
4605 t <- liftIO $ getTime Monotonic
4606 if diffAbsTime64 t t0 <= lim
4607 then return $ Skip $ DropByTimeGen st t0
4608 else return $ Yield x $ DropByTimeYield st
4609 step gst (DropByTimeYield st) = do
4612 Yield x s -> Yield x (DropByTimeYield s)
4613 Skip s -> Skip (DropByTimeYield s)
4616{-# INLINE updateTimeVar #-}
4617updateTimeVar :: Var IO Int64 -> IO ()
4618updateTimeVar timeVar = do
4619 MicroSecond64 t <- fromAbsTime <$> getTime Monotonic
4620 modifyVar' timeVar (const t)
4622{-# INLINE updateWithDelay #-}
4623updateWithDelay :: RealFrac a => a -> Var IO Int64 -> IO ()
4624updateWithDelay precision timeVar = do
4625 threadDelay (delayTime precision)
4626 updateTimeVar timeVar
4630 -- Keep the minimum at least a millisecond to avoid high CPU usage
4631 {-# INLINE delayTime #-}
4633 | g' >= fromIntegral (maxBound :: Int) = maxBound
4635 | otherwise = round g'
4639 g' = g * 10 ^ (6 :: Int)
4641-- XXX we should move this to stream generation section of this file. Also, the
4642-- take/drop combinators above should be moved to filtering section.
4644{-# INLINE_NORMAL times #-}
4645times :: MonadAsync m => Double -> Stream m (AbsTime, RelTime64)
4646times g = Stream step Nothing
4650 {-# INLINE_LATE step #-}
4652 -- XXX note that this is safe only on a 64-bit machine. On a 32-bit
4653 -- machine a 64-bit 'Var' cannot be read consistently without a lock
4654 -- while another thread is writing to it.
4655 timeVar <- liftIO $ newVar (0 :: Int64)
4656 liftIO $ updateTimeVar timeVar
4657 tid <- forkManaged $ liftIO $ forever (updateWithDelay g timeVar)
4658 a <- liftIO $ readVar timeVar
4659 return $ Skip $ Just (timeVar, tid, a)
4661 step _ s@(Just (timeVar, _, t0)) = do
4662 a <- liftIO $ readVar timeVar
4663 -- XXX we can perhaps use an AbsTime64 using a 64 bit Int for
4664 -- efficiency. or maybe we can use a representation using Double for
4665 -- floating precision time
4666 return $ Yield (toAbsTime (MicroSecond64 t0),
4667 (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