1{-# LANGUAGE BangPatterns #-}
3{-# LANGUAGE ConstraintKinds #-}
4{-# LANGUAGE ExistentialQuantification #-}
5{-# LANGUAGE FlexibleContexts #-}
6{-# LANGUAGE FlexibleInstances #-}
7{-# LANGUAGE MultiParamTypeClasses #-}
8{-# LANGUAGE PatternSynonyms #-}
9{-# LANGUAGE RecordWildCards #-}
10{-# LANGUAGE ScopedTypeVariables #-}
11{-# LANGUAGE ViewPatterns #-}
12{-# LANGUAGE RankNTypes #-}
13{-# LANGUAGE MagicHash #-}
15#if __GLASGOW_HASKELL__ >= 801
16{-# LANGUAGE TypeApplications #-}
22-- Module : Streamly.Internal.Data.Stream.StreamD
23-- Copyright : (c) 2018 Harendra Kumar
24-- (c) Roman Leshchinskiy 2008-2010
25-- (c) The University of Glasgow, 2009
28-- Maintainer : streamly@composewell.com
29-- Stability : experimental
32-- Direct style re-implementation of CPS style stream in StreamK module. The
33-- symbol or suffix 'D' in this module denotes the "Direct" style. GHC is able
34-- to INLINE and fuse direct style better, providing better performance than
38-- import qualified Streamly.Internal.Data.Stream.StreamD as D
41-- Some of the functions in this file have been adapted from the vector
42-- library, https://hackage.haskell.org/package/vector.
44module Streamly.Internal.Data.Stream.StreamD
49#if __GLASGOW_HASKELL__ >= 800
50 , Stream (Stream, UnStream)
70 -- ** Specialized Generation
71 -- | Generate a monadic stream from a seed.
84 , enumerateFromStepIntegral
85 , enumerateFromIntegral
86 , enumerateFromThenIntegral
87 , enumerateFromToIntegral
88 , enumerateFromThenToIntegral
90 , enumerateFromStepNum
93 , enumerateFromToFractional
94 , enumerateFromThenToFractional
100 -- | Transform an input structure into a stream.
101 -- | Direct style stream does not support @fromFoldable@.
135 -- ** Specialized Folds
162 -- ** Flattening nested streams
165 , ConcatMapUState (..)
167 , ConcatUnfoldInterleaveState (..)
168 , concatUnfoldInterleave
169 , concatUnfoldRoundrobin
172 , InterleaveState(..)
177 , roundRobin -- interleaveFair?/ParallelFair
210 -- | Transform a stream into another type.
227 -- ** By folding (scans)
278 , intersperseSuffixBySpan
302 -- * Transformation comprehensions
308 , clearIORefFinalizer
320 -- * Concurrent Application
329import Control.Concurrent (killThread, myThreadId, takeMVar, threadDelay)
330import Control.Exception
331 (assert, Exception, SomeException, AsyncException, fromException, mask_)
332import Control.Monad (void, when, forever)
333import Control.Monad.Catch (MonadCatch, MonadThrow, throwM)
334import Control.Monad.IO.Class (MonadIO(..))
335import Control.Monad.Reader (ReaderT)
336import Control.Monad.State.Strict (StateT)
337import Control.Monad.Trans (MonadTrans(lift))
338import Control.Monad.Trans.Control (MonadBaseControl, liftBaseOp_)
339import Data.Bits (shiftR, shiftL, (.|.), (.&.))
340import Data.Functor.Identity (Identity(..))
341import Data.Int (Int64)
342import Data.IORef (newIORef, readIORef, mkWeakIORef, writeIORef, IORef)
343import Data.Maybe (fromJust, isJust, isNothing)
344import Data.Word (Word32)
345import Foreign.Ptr (Ptr)
346import Foreign.Storable (Storable(..))
347import GHC.Types (SPEC(..))
348import System.Mem (performMajorGC)
350 hiding (map, mapM, mapM_, repeat, foldr, last, take, filter,
351 takeWhile, drop, dropWhile, all, any, maximum, minimum, elem,
352 notElem, null, head, tail, zipWith, lookup, foldr1, sequence,
353 (!!), scanl, scanl1, concatMap, replicate, enumFromTo, concat,
354 reverse, iterate, splitAt)
356import qualified Control.Monad.Catch as MC
357import qualified Control.Monad.Reader as Reader
358import qualified Control.Monad.State.Strict as State
359import qualified Prelude
361import Fusion.Plugin.Types (Fuse(..))
362import Streamly.Internal.Mutable.Prim.Var
363 (Prim, Var, readVar, newVar, modifyVar')
364import Streamly.Internal.Data.Time.Units
365 (TimeUnit64, toRelTime64, diffAbsTime64, RelTime64)
367import Streamly.Internal.Data.Atomics (atomicModifyIORefCAS_)
368import Streamly.Internal.Memory.Array.Types (Array(..))
369import Streamly.Internal.Data.Fold.Types (Fold(..))
370import Streamly.Internal.Data.Parser (ParseError(..))
371import Streamly.Internal.Data.Pipe.Types (Pipe(..), PipeState(..))
372import Streamly.Internal.Data.Time.Clock (Clock(Monotonic), getTime)
373import Streamly.Internal.Data.Time.Units
374 (MicroSecond64(..), fromAbsTime, toAbsTime, AbsTime)
375import Streamly.Internal.Data.Unfold.Types (Unfold(..))
376import Streamly.Internal.Data.Strict (Tuple3'(..))
378import Streamly.Internal.Data.Stream.StreamD.Type
379import Streamly.Internal.Data.SVar
380import Streamly.Internal.Data.Stream.SVar (fromConsumer, pushToFold)
382import qualified Streamly.Internal.Data.Pipe.Types as Pipe
383import qualified Streamly.Internal.Memory.Array.Types as A
384import qualified Streamly.Internal.Data.Fold as FL
385import qualified Streamly.Memory.Ring as RB
386import qualified Streamly.Internal.Data.Stream.StreamK as K
387import qualified Streamly.Internal.Data.Parser as PR
388import qualified Streamly.Internal.Data.Parser.ParserD as PRD
390------------------------------------------------------------------------------
392------------------------------------------------------------------------------
394-- | An empty 'Stream'.
395{-# INLINE_NORMAL nil #-}
396nil :: Monad m => Stream m a
397nil = Stream (\_ _ -> return Stop) ()
399-- | An empty 'Stream' with a side effect.
400{-# INLINE_NORMAL nilM #-}
401nilM :: Monad m => m b -> Stream m a
402nilM m = Stream (\_ _ -> m >> return Stop) ()
404{-# INLINE_NORMAL consM #-}
405consM :: Monad m => m a -> Stream m a -> Stream m a
406consM m (Stream step state) = Stream step1 Nothing
408 {-# INLINE_LATE step1 #-}
409 step1 _ Nothing = m >>= \x -> return $ Yield x (Just state)
410 step1 gst (Just st) = do
414 Yield a s -> Yield a (Just s)
415 Skip s -> Skip (Just s)
418-- XXX implement in terms of consM?
419-- cons x = consM (return x)
421-- | Can fuse but has O(n^2) complexity.
422{-# INLINE_NORMAL cons #-}
423cons :: Monad m => a -> Stream m a -> Stream m a
424cons x (Stream step state) = Stream step1 Nothing
426 {-# INLINE_LATE step1 #-}
427 step1 _ Nothing = return $ Yield x (Just state)
428 step1 gst (Just st) = do
432 Yield a s -> Yield a (Just s)
433 Skip s -> Skip (Just s)
436-------------------------------------------------------------------------------
438-------------------------------------------------------------------------------
440-- Does not fuse, has the same performance as the StreamK version.
441{-# INLINE_NORMAL uncons #-}
442uncons :: Monad m => Stream m a -> m (Maybe (a, Stream m a))
443uncons (UnStream step state) = go state
446 r <- step defState st
448 Yield x s -> return $ Just (x, Stream step s)
450 Stop -> return Nothing
452------------------------------------------------------------------------------
453-- Generation by unfold
454------------------------------------------------------------------------------
456{-# INLINE_NORMAL unfoldrM #-}
457unfoldrM :: Monad m => (s -> m (Maybe (a, s))) -> s -> Stream m a
458unfoldrM next state = Stream step state
460 {-# INLINE_LATE step #-}
464 Just (x, s) -> Yield x s
467{-# INLINE_LATE unfoldr #-}
468unfoldr :: Monad m => (s -> Maybe (a, s)) -> s -> Stream m a
469unfoldr f = unfoldrM (return . f)
471-- | Convert an 'Unfold' into a 'Stream' by supplying it a seed.
473{-# INLINE_NORMAL unfold #-}
474unfold :: Monad m => Unfold m a b -> a -> Stream m b
475unfold (Unfold ustep inject) seed = Stream step Nothing
477 {-# INLINE_LATE step #-}
478 step _ Nothing = inject seed >>= return . Skip . Just
479 step _ (Just st) = do
482 Yield x s -> Yield x (Just s)
483 Skip s -> Skip (Just s)
486------------------------------------------------------------------------------
487-- Specialized Generation
488------------------------------------------------------------------------------
490{-# INLINE_NORMAL repeatM #-}
491repeatM :: Monad m => m a -> Stream m a
492repeatM x = Stream (\_ _ -> x >>= \r -> return $ Yield r ()) ()
494{-# INLINE_NORMAL repeat #-}
495repeat :: Monad m => a -> Stream m a
496repeat x = Stream (\_ _ -> return $ Yield x ()) ()
498{-# INLINE_NORMAL iterateM #-}
499iterateM :: Monad m => (a -> m a) -> m a -> Stream m a
500iterateM step = Stream (\_ st -> st >>= \x -> return $ Yield x (step x))
502{-# INLINE_NORMAL iterate #-}
503iterate :: Monad m => (a -> a) -> a -> Stream m a
504iterate step st = iterateM (return . step) (return st)
506{-# INLINE_NORMAL replicateM #-}
507replicateM :: forall m a. Monad m => Int -> m a -> Stream m a
508replicateM n p = Stream step n
510 {-# INLINE_LATE step #-}
512 | i <= 0 = return Stop
515 return $ Yield x (i - 1)
517{-# INLINE_NORMAL replicate #-}
518replicate :: Monad m => Int -> a -> Stream m a
519replicate n x = replicateM n (return x)
521-- This would not work properly for floats, therefore we put an Integral
523-- | Can be used to enumerate unbounded integrals. This does not check for
524-- overflow or underflow for bounded integrals.
525{-# INLINE_NORMAL enumerateFromStepIntegral #-}
526enumerateFromStepIntegral :: (Integral a, Monad m) => a -> a -> Stream m a
527enumerateFromStepIntegral from stride =
528 from `seq` stride `seq` Stream step from
530 {-# INLINE_LATE step #-}
531 step _ !x = return $ Yield x $! (x + stride)
533-- We are assuming that "to" is constrained by the type to be within
535{-# INLINE enumerateFromToIntegral #-}
536enumerateFromToIntegral :: (Monad m, Integral a) => a -> a -> Stream m a
537enumerateFromToIntegral from to =
538 takeWhile (<= to) $ enumerateFromStepIntegral from 1
540{-# INLINE enumerateFromIntegral #-}
541enumerateFromIntegral :: (Monad m, Integral a, Bounded a) => a -> Stream m a
542enumerateFromIntegral from = enumerateFromToIntegral from maxBound
544data EnumState a = EnumInit | EnumYield a a a | EnumStop
546{-# INLINE_NORMAL enumerateFromThenToIntegralUp #-}
547enumerateFromThenToIntegralUp
548 :: (Monad m, Integral a)
549 => a -> a -> a -> Stream m a
550enumerateFromThenToIntegralUp from next to = Stream step EnumInit
552 {-# INLINE_LATE step #-}
558 else Yield from EnumStop
559 else -- from <= next <= to
560 let stride = next - from
561 in Skip $ EnumYield from stride (to - stride)
563 step _ (EnumYield x stride toMinus) =
566 then Yield x EnumStop
567 else Yield x $ EnumYield (x + stride) stride toMinus
569 step _ EnumStop = return Stop
571{-# INLINE_NORMAL enumerateFromThenToIntegralDn #-}
572enumerateFromThenToIntegralDn
573 :: (Monad m, Integral a)
574 => a -> a -> a -> Stream m a
575enumerateFromThenToIntegralDn from next to = Stream step EnumInit
577 {-# INLINE_LATE step #-}
579 return $ if to > next
582 else Yield from EnumStop
583 else -- from >= next >= to
584 let stride = next - from
585 in Skip $ EnumYield from stride (to - stride)
587 step _ (EnumYield x stride toMinus) =
590 then Yield x EnumStop
591 else Yield x $ EnumYield (x + stride) stride toMinus
593 step _ EnumStop = return Stop
595{-# INLINE_NORMAL enumerateFromThenToIntegral #-}
596enumerateFromThenToIntegral
597 :: (Monad m, Integral a)
598 => a -> a -> a -> Stream m a
599enumerateFromThenToIntegral from next to
600 | next >= from = enumerateFromThenToIntegralUp from next to
601 | otherwise = enumerateFromThenToIntegralDn from next to
603{-# INLINE_NORMAL enumerateFromThenIntegral #-}
604enumerateFromThenIntegral
605 :: (Monad m, Integral a, Bounded a)
606 => a -> a -> Stream m a
607enumerateFromThenIntegral from next =
609 then enumerateFromThenToIntegralUp from next maxBound
610 else enumerateFromThenToIntegralDn from next minBound
612-- For floating point numbers if the increment is less than the precision then
613-- it just gets lost. Therefore we cannot always increment it correctly by just
615-- 9007199254740992 + 1 + 1 :: Double => 9.007199254740992e15
616-- 9007199254740992 + 2 :: Double => 9.007199254740994e15
618-- Instead we accumulate the increment counter and compute the increment
619-- every time before adding it to the starting number.
621-- This works for Integrals as well as floating point numbers, but
622-- enumerateFromStepIntegral is faster for integrals.
623{-# INLINE_NORMAL enumerateFromStepNum #-}
624enumerateFromStepNum :: (Monad m, Num a) => a -> a -> Stream m a
625enumerateFromStepNum from stride = Stream step 0
627 {-# INLINE_LATE step #-}
628 step _ !i = return $ (Yield $! (from + i * stride)) $! (i + 1)
630{-# INLINE_NORMAL numFrom #-}
631numFrom :: (Monad m, Num a) => a -> Stream m a
632numFrom from = enumerateFromStepNum from 1
634{-# INLINE_NORMAL numFromThen #-}
635numFromThen :: (Monad m, Num a) => a -> a -> Stream m a
636numFromThen from next = enumerateFromStepNum from (next - from)
638-- We cannot write a general function for Num. The only way to write code
639-- portable between the two is to use a 'Real' constraint and convert between
640-- Fractional and Integral using fromRational which is horribly slow.
641{-# INLINE_NORMAL enumerateFromToFractional #-}
642enumerateFromToFractional
643 :: (Monad m, Fractional a, Ord a)
644 => a -> a -> Stream m a
645enumerateFromToFractional from to =
646 takeWhile (<= to + 1 / 2) $ enumerateFromStepNum from 1
648{-# INLINE_NORMAL enumerateFromThenToFractional #-}
649enumerateFromThenToFractional
650 :: (Monad m, Fractional a, Ord a)
651 => a -> a -> a -> Stream m a
652enumerateFromThenToFractional from next to =
653 takeWhile predicate $ numFromThen from next
655 mid = (next - from) / 2
656 predicate | next >= from = (<= to + mid)
657 | otherwise = (>= to + mid)
659-------------------------------------------------------------------------------
660-- Generation by Conversion
661-------------------------------------------------------------------------------
663{-# INLINE_NORMAL fromIndicesM #-}
664fromIndicesM :: Monad m => (Int -> m a) -> Stream m a
665fromIndicesM gen = Stream step 0
667 {-# INLINE_LATE step #-}
670 return $ Yield x (i + 1)
672{-# INLINE fromIndices #-}
673fromIndices :: Monad m => (Int -> a) -> Stream m a
674fromIndices gen = fromIndicesM (return . gen)
676{-# INLINE_NORMAL generateM #-}
677generateM :: Monad m => Int -> (Int -> m a) -> Stream m a
678generateM n gen = n `seq` Stream step 0
680 {-# INLINE_LATE step #-}
681 step _ i | i < n = do
683 return $ Yield x (i + 1)
684 | otherwise = return Stop
686{-# INLINE generate #-}
687generate :: Monad m => Int -> (Int -> a) -> Stream m a
688generate n gen = generateM n (return . gen)
690-- XXX we need the MonadAsync constraint because of a rewrite rule.
691-- | Convert a list of monadic actions to a 'Stream'
692{-# INLINE_LATE fromListM #-}
693fromListM :: MonadAsync m => [m a] -> Stream m a
694fromListM = Stream step
696 {-# INLINE_LATE step #-}
697 step _ (m:ms) = m >>= \x -> return $ Yield x ms
698 step _ [] = return Stop
700{-# INLINE toStreamD #-}
701toStreamD :: (K.IsStream t, Monad m) => t m a -> Stream m a
702toStreamD = fromStreamK . K.toStream
704{-# INLINE_NORMAL fromPrimVar #-}
705fromPrimVar :: (MonadIO m, Prim a) => Var IO a -> Stream m a
706fromPrimVar var = Stream step ()
708 {-# INLINE_LATE step #-}
709 step _ () = liftIO (readVar var) >>= \x -> return $ Yield x ()
711-------------------------------------------------------------------------------
712-- Generation from SVar
713-------------------------------------------------------------------------------
715data FromSVarState t m a =
717 | FromSVarRead (SVar t m a)
718 | FromSVarLoop (SVar t m a) [ChildEvent a]
719 | FromSVarDone (SVar t m a)
721{-# INLINE_NORMAL fromSVar #-}
722fromSVar :: (MonadAsync m) => SVar t m a -> Stream m a
723fromSVar svar = Stream step FromSVarInit
726 {-# INLINE_LATE step #-}
727 step _ FromSVarInit = do
728 ref <- liftIO $ newIORef ()
729 _ <- liftIO $ mkWeakIORef ref hook
730 -- when this copy of svar gets garbage collected "ref" will get
731 -- garbage collected and our GC hook will be called.
732 let sv = svar{svarRef = Just ref}
733 return $ Skip (FromSVarRead sv)
737 {-# NOINLINE hook #-}
739 when (svarInspectMode svar) $ do
740 r <- liftIO $ readIORef (svarStopTime (svarStats svar))
742 printSVar svar "SVar Garbage Collected"
744 -- If there are any SVars referenced by this SVar a GC will prompt
745 -- them to be cleaned up quickly.
746 when (svarInspectMode svar) performMajorGC
748 step _ (FromSVarRead sv) = do
749 list <- readOutputQ sv
750 -- Reversing the output is important to guarantee that we process the
751 -- outputs in the same order as they were generated by the constituent
753 return $ Skip $ FromSVarLoop sv (Prelude.reverse list)
755 step _ (FromSVarLoop sv []) = do
756 done <- postProcess sv
757 return $ Skip $ if done
758 then (FromSVarDone sv)
759 else (FromSVarRead sv)
761 step _ (FromSVarLoop sv (ev : es)) = do
763 ChildYield a -> return $ Yield a (FromSVarLoop sv es)
764 ChildStop tid e -> do
768 stop <- shouldStop tid
771 liftIO (cleanupSVar sv)
772 return $ Skip (FromSVarDone sv)
773 else return $ Skip (FromSVarLoop sv es)
775 case fromException ex of
777 return $ Skip (FromSVarLoop sv es)
778 Nothing -> liftIO (cleanupSVar sv) >> throwM ex
782 case svarStopStyle sv of
783 StopNone -> return False
784 StopAny -> return True
786 sid <- liftIO $ readIORef (svarStopBy sv)
787 return $ if tid == sid then True else False
789 step _ (FromSVarDone sv) = do
790 when (svarInspectMode sv) $ do
791 t <- liftIO $ getTime Monotonic
792 liftIO $ writeIORef (svarStopTime (svarStats sv)) (Just t)
793 liftIO $ printSVar sv "SVar Done"
796-------------------------------------------------------------------------------
797-- Process events received by a fold consumer from a stream producer
798-------------------------------------------------------------------------------
800{-# INLINE_NORMAL fromProducer #-}
801fromProducer :: (MonadAsync m) => SVar t m a -> Stream m a
802fromProducer svar = Stream step (FromSVarRead svar)
805 {-# INLINE_LATE step #-}
806 step _ (FromSVarRead sv) = do
807 list <- readOutputQ sv
808 -- Reversing the output is important to guarantee that we process the
809 -- outputs in the same order as they were generated by the constituent
811 return $ Skip $ FromSVarLoop sv (Prelude.reverse list)
813 step _ (FromSVarLoop sv []) = return $ Skip $ FromSVarRead sv
814 step _ (FromSVarLoop sv (ev : es)) = do
816 ChildYield a -> return $ Yield a (FromSVarLoop sv es)
817 ChildStop tid e -> do
821 sendStopToProducer sv
822 return $ Skip (FromSVarDone sv)
823 Just _ -> error "Bug: fromProducer: received exception"
825 step _ (FromSVarDone sv) = do
826 when (svarInspectMode sv) $ do
827 t <- liftIO $ getTime Monotonic
828 liftIO $ writeIORef (svarStopTime (svarStats sv)) (Just t)
829 liftIO $ printSVar sv "SVar Done"
832 step _ FromSVarInit = undefined
834-------------------------------------------------------------------------------
835-- Hoisting the inner monad
836-------------------------------------------------------------------------------
838{-# INLINE_NORMAL hoist #-}
839hoist :: Monad n => (forall x. m x -> n x) -> Stream m a -> Stream n a
840hoist f (Stream step state) = (Stream step' state)
842 {-# INLINE_LATE step' #-}
844 r <- f $ step (adaptState gst) st
846 Yield x s -> Yield x s
850{-# INLINE generally #-}
851generally :: Monad m => Stream Identity a -> Stream m a
852generally = hoist (return . runIdentity)
854{-# INLINE_NORMAL liftInner #-}
855liftInner :: (Monad m, MonadTrans t, Monad (t m))
856 => Stream m a -> Stream (t m) a
857liftInner (Stream step state) = Stream step' state
859 {-# INLINE_LATE step' #-}
861 r <- lift $ step (adaptState gst) st
863 Yield x s -> Yield x s
867{-# INLINE_NORMAL runReaderT #-}
868runReaderT :: Monad m => s -> Stream (ReaderT s m) a -> Stream m a
869runReaderT sval (Stream step state) = Stream step' state
871 {-# INLINE_LATE step' #-}
873 r <- Reader.runReaderT (step (adaptState gst) st) sval
875 Yield x s -> Yield x s
879{-# INLINE_NORMAL evalStateT #-}
880evalStateT :: Monad m => s -> Stream (StateT s m) a -> Stream m a
881evalStateT sval (Stream step state) = Stream step' (state, sval)
883 {-# INLINE_LATE step' #-}
884 step' gst (st, sv) = do
885 (r, sv') <- State.runStateT (step (adaptState gst) st) sv
887 Yield x s -> Yield x (s, sv')
888 Skip s -> Skip (s, sv')
891{-# INLINE_NORMAL runStateT #-}
892runStateT :: Monad m => s -> Stream (StateT s m) a -> Stream m (s, a)
893runStateT sval (Stream step state) = Stream step' (state, sval)
895 {-# INLINE_LATE step' #-}
896 step' gst (st, sv) = do
897 (r, sv') <- State.runStateT (step (adaptState gst) st) sv
899 Yield x s -> Yield (sv', x) (s, sv')
900 Skip s -> Skip (s, sv')
903------------------------------------------------------------------------------
904-- Elimination by Folds
905------------------------------------------------------------------------------
907------------------------------------------------------------------------------
909------------------------------------------------------------------------------
911{-# INLINE_NORMAL foldr1 #-}
912foldr1 :: Monad m => (a -> a -> a) -> Stream m a -> m (Maybe a)
916 Nothing -> return Nothing
917 Just (h, t) -> fmap Just (foldr f h t)
919------------------------------------------------------------------------------
921------------------------------------------------------------------------------
923{-# INLINE_NORMAL foldlT #-}
924foldlT :: (Monad m, Monad (s m), MonadTrans s)
925 => (s m b -> a -> s m b) -> s m b -> Stream m a -> s m b
926foldlT fstep begin (Stream step state) = go SPEC begin state
929 r <- lift $ step defState st
931 Yield x s -> go SPEC (fstep acc x) s
932 Skip s -> go SPEC acc s
935-- Note, this is going to have horrible performance, because of the nature of
936-- the stream type (i.e. direct stream vs CPS). Its only for reference, it is
937-- likely be practically unusable.
938{-# INLINE_NORMAL foldlS #-}
940 => (Stream m b -> a -> Stream m b) -> Stream m b -> Stream m a -> Stream m b
941foldlS fstep begin (Stream step state) = Stream step' (Left (state, begin))
943 step' gst (Left (st, acc)) = do
944 r <- step (adaptState gst) st
946 Yield x s -> Skip (Left (s, fstep acc x))
947 Skip s -> Skip (Left (s, acc))
948 Stop -> Skip (Right acc)
950 step' gst (Right (Stream stp stt)) = do
951 r <- stp (adaptState gst) stt
953 Yield x s -> Yield x (Right (Stream stp s))
954 Skip s -> Skip (Right (Stream stp s))
957------------------------------------------------------------------------------
959------------------------------------------------------------------------------
961-- Inlined definition. Without the inline "serially/parser/take" benchmark
962-- degrades and parseMany does not fuse. Even using "inline" at the callsite
964{-# INLINE splitAt #-}
965splitAt :: Int -> [a] -> ([a],[a])
968 | otherwise = splitAt' n ls
970 splitAt' :: Int -> [a] -> ([a], [a])
971 splitAt' _ [] = ([], [])
972 splitAt' 1 (x:xs) = ([x], xs)
973 splitAt' m (x:xs) = (x:xs', xs'')
975 (xs', xs'') = splitAt' (m - 1) xs
977-- | Run a 'Parse' over a stream.
978{-# INLINE_NORMAL parselMx' #-}
981 => (s -> a -> m (PR.Step s b))
986parselMx' pstep initial extract (Stream step state) = do
987 initial >>= go SPEC state []
991 -- XXX currently we are using a dumb list based approach for backtracking
992 -- buffer. This can be replaced by a sliding/ring buffer using Data.Array.
993 -- That will allow us more efficient random back and forth movement.
995 go !_ st buf !pst = do
996 r <- step defState st
1001 -- PR.Yield 0 pst1 -> go SPEC s [] pst1
1002 PR.Yield n pst1 -> do
1003 assert (n <= length (x:buf)) (return ())
1004 go SPEC s (Prelude.take n (x:buf)) pst1
1005 PR.YieldB n pst1 -> do
1006 assert (n <= length (x:buf)) (return ())
1007 let src0 = Prelude.take n (x:buf)
1008 src = Prelude.reverse src0
1009 gobuf SPEC s [] src pst1
1010 PR.Skip 0 pst1 -> go SPEC s (x:buf) pst1
1011 PR.Skip n pst1 -> do
1012 assert (n <= length (x:buf)) (return ())
1013 let (src0, buf1) = splitAt n (x:buf)
1014 src = Prelude.reverse src0
1015 gobuf SPEC s buf1 src pst1
1016 PR.Stop _ b -> return b
1017 PR.Error err -> throwM $ ParseError err
1018 Skip s -> go SPEC s buf pst
1021 gobuf !_ s buf [] !pst = go SPEC s buf pst
1022 gobuf !_ s buf (x:xs) !pst = do
1025 -- PR.Yield 0 pst1 -> go SPEC s [] pst1
1026 PR.Yield n pst1 -> do
1027 assert (n <= length (x:buf)) (return ())
1028 gobuf SPEC s (Prelude.take n (x:buf)) xs pst1
1029 PR.YieldB n pst1 -> do
1030 assert (n <= length (x:buf)) (return ())
1031 let src0 = Prelude.take n (x:buf)
1032 src = Prelude.reverse src0 ++ xs
1033 gobuf SPEC s [] src pst1
1034 PR.Skip 0 pst1 -> gobuf SPEC s (x:buf) xs pst1
1035 PR.Skip n pst1 -> do
1036 assert (n <= length (x:buf)) (return ())
1037 let (src0, buf1) = splitAt n (x:buf)
1038 src = Prelude.reverse src0 ++ xs
1039 gobuf SPEC s buf1 src pst1
1040 PR.Stop _ b -> return b
1041 PR.Error err -> throwM $ ParseError err
1043------------------------------------------------------------------------------
1045------------------------------------------------------------------------------
1047{-# ANN type ParseChunksState Fuse #-}
1048data ParseChunksState x inpBuf st pst =
1049 ParseChunksInit inpBuf st
1050 | ParseChunksInitLeftOver inpBuf
1051 | ParseChunksStream st inpBuf pst
1052 | ParseChunksBuf inpBuf st inpBuf pst
1053 | ParseChunksYield x (ParseChunksState x inpBuf st pst)
1055{-# INLINE_NORMAL parseMany #-}
1061parseMany (PRD.Parser pstep initial extract) (Stream step state) =
1062 Stream stepOuter (ParseChunksInit [] state)
1066 {-# INLINE_LATE stepOuter #-}
1067 -- Buffer is empty, go to stream processing loop
1068 stepOuter _ (ParseChunksInit [] st) = do
1069 initial >>= return . Skip . ParseChunksStream st []
1071 -- Buffer is not empty, go to buffered processing loop
1072 stepOuter _ (ParseChunksInit src st) = do
1073 initial >>= return . Skip . ParseChunksBuf src st []
1075 -- XXX we just discard any leftover input at the end
1076 stepOuter _ (ParseChunksInitLeftOver _) = return Stop
1078 -- Buffer is empty process elements from the stream
1079 stepOuter gst (ParseChunksStream st buf pst) = do
1080 r <- step (adaptState gst) st
1085 -- PR.Yield 0 pst1 -> go SPEC s [] pst1
1086 PR.Yield n pst1 -> do
1087 assert (n <= length (x:buf)) (return ())
1088 let buf1 = Prelude.take n (x:buf)
1089 return $ Skip $ ParseChunksStream s buf1 pst1
1090 PR.YieldB n pst1 -> do
1091 assert (n <= length (x:buf)) (return ())
1092 let src0 = Prelude.take n (x:buf)
1093 src = Prelude.reverse src0
1094 return $ Skip $ ParseChunksBuf src s [] pst1
1095 -- PR.Skip 0 pst1 ->
1096 -- return $ Skip $ ParseChunksStream s (x:buf) pst1
1097 PR.Skip n pst1 -> do
1098 assert (n <= length (x:buf)) (return ())
1099 let (src0, buf1) = splitAt n (x:buf)
1100 src = Prelude.reverse src0
1101 return $ Skip $ ParseChunksBuf src s buf1 pst1
1102 -- XXX Specialize for Stop 0 common case?
1104 assert (n <= length (x:buf)) (return ())
1105 let src = Prelude.reverse (Prelude.take n (x:buf))
1107 ParseChunksYield b (ParseChunksInit src s)
1108 PR.Error err -> throwM $ ParseError err
1109 Skip s -> return $ Skip $ ParseChunksStream s buf pst
1112 let src = Prelude.reverse buf
1114 ParseChunksYield b (ParseChunksInitLeftOver src)
1116 -- go back to stream processing mode
1117 stepOuter _ (ParseChunksBuf [] s buf pst) =
1118 return $ Skip $ ParseChunksStream s buf pst
1120 -- buffered processing loop
1121 stepOuter _ (ParseChunksBuf (x:xs) s buf pst) = do
1124 -- PR.Yield 0 pst1 ->
1125 PR.Yield n pst1 -> do
1126 assert (n <= length (x:buf)) (return ())
1127 let buf1 = Prelude.take n (x:buf)
1128 return $ Skip $ ParseChunksBuf xs s buf1 pst1
1129 PR.YieldB n pst1 -> do
1130 assert (n <= length (x:buf)) (return ())
1131 let src0 = Prelude.take n (x:buf)
1132 src = Prelude.reverse src0 ++ xs
1133 return $ Skip $ ParseChunksBuf src s [] pst1
1134 -- PR.Skip 0 pst1 -> return $ Skip $ ParseChunksBuf xs s (x:buf) pst1
1135 PR.Skip n pst1 -> do
1136 assert (n <= length (x:buf)) (return ())
1137 let (src0, buf1) = splitAt n (x:buf)
1138 src = Prelude.reverse src0 ++ xs
1139 return $ Skip $ ParseChunksBuf src s buf1 pst1
1140 -- XXX Specialize for Stop 0 common case?
1142 assert (n <= length (x:buf)) (return ())
1143 let src = Prelude.reverse (Prelude.take n (x:buf)) ++ xs
1144 return $ Skip $ ParseChunksYield b (ParseChunksInit src s)
1145 PR.Error err -> throwM $ ParseError err
1147 stepOuter _ (ParseChunksYield a next) = return $ Yield a next
1149{-# ANN type ConcatParseState Fuse #-}
1150data ConcatParseState x inpBuf st p =
1151 ConcatParseInit inpBuf st p
1152 | ConcatParseInitLeftOver inpBuf
1153 | ConcatParseStream st inpBuf p
1154 | ConcatParseBuf inpBuf st inpBuf p
1155 | ConcatParseYield x (ConcatParseState x inpBuf st p)
1157{-# INLINE_NORMAL parseIterate #-}
1160 => (b -> PRD.Parser m a b)
1164parseIterate func seed (Stream step state) =
1165 Stream stepOuter (ConcatParseInit [] state (func seed))
1169 {-# INLINE_LATE stepOuter #-}
1170 -- Buffer is empty, go to stream processing loop
1171 stepOuter _ (ConcatParseInit [] st (PRD.Parser pstep initial extract)) = do
1172 initial >>= \r -> return $ Skip $ ConcatParseStream st []
1173 (PRD.Parser pstep (return r) extract)
1175 -- Buffer is not empty, go to buffered processing loop
1176 stepOuter _ (ConcatParseInit src st
1177 (PRD.Parser pstep initial extract)) = do
1178 initial >>= \r -> return $ Skip $ ConcatParseBuf src st []
1179 (PRD.Parser pstep (return r) extract)
1181 -- XXX we just discard any leftover input at the end
1182 stepOuter _ (ConcatParseInitLeftOver _) = return Stop
1184 -- Buffer is empty process elements from the stream
1185 stepOuter gst (ConcatParseStream st buf
1186 p@(PRD.Parser pstep initial extract)) = do
1187 r <- step (adaptState gst) st
1193 -- PR.Yield 0 pst1 -> go SPEC s [] pst1
1194 PR.Yield n pst1 -> do
1195 assert (n <= length (x:buf)) (return ())
1196 let buf1 = Prelude.take n (x:buf)
1197 return $ Skip $ ConcatParseStream s buf1
1198 (PRD.Parser pstep (return pst1) extract)
1199 PR.YieldB n pst1 -> do
1200 assert (n <= length (x:buf)) (return ())
1201 let src0 = Prelude.take n (x:buf)
1202 src = Prelude.reverse src0
1203 return $ Skip $ ConcatParseBuf src s []
1204 (PRD.Parser pstep (return pst1) extract)
1205 -- PR.Skip 0 pst1 ->
1206 -- return $ Skip $ ConcatParseStream s (x:buf) pst1
1207 PR.Skip n pst1 -> do
1208 assert (n <= length (x:buf)) (return ())
1209 let (src0, buf1) = splitAt n (x:buf)
1210 src = Prelude.reverse src0
1211 return $ Skip $ ConcatParseBuf src s buf1
1212 (PRD.Parser pstep (return pst1) extract)
1213 -- XXX Specialize for Stop 0 common case?
1215 assert (n <= length (x:buf)) (return ())
1216 let src = Prelude.reverse (Prelude.take n (x:buf))
1218 ConcatParseYield b (ConcatParseInit src s (func b))
1219 PR.Error err -> throwM $ ParseError err
1220 Skip s -> return $ Skip $ ConcatParseStream s buf p
1224 let src = Prelude.reverse buf
1226 ConcatParseYield b (ConcatParseInitLeftOver src)
1228 -- go back to stream processing mode
1229 stepOuter _ (ConcatParseBuf [] s buf p) =
1230 return $ Skip $ ConcatParseStream s buf p
1232 -- buffered processing loop
1233 stepOuter _ (ConcatParseBuf (x:xs) s buf
1234 (PRD.Parser pstep initial extract)) = do
1238 -- PR.Yield 0 pst1 ->
1239 PR.Yield n pst1 -> do
1240 assert (n <= length (x:buf)) (return ())
1241 let buf1 = Prelude.take n (x:buf)
1242 return $ Skip $ ConcatParseBuf xs s buf1
1243 (PRD.Parser pstep (return pst1) extract)
1244 PR.YieldB n pst1 -> do
1245 assert (n <= length (x:buf)) (return ())
1246 let src0 = Prelude.take n (x:buf)
1247 src = Prelude.reverse src0 ++ xs
1248 return $ Skip $ ConcatParseBuf src s []
1249 (PRD.Parser pstep (return pst1) extract)
1250 -- PR.Skip 0 pst1 -> return $ Skip $ ConcatParseBuf xs s (x:buf) pst1
1251 PR.Skip n pst1 -> do
1252 assert (n <= length (x:buf)) (return ())
1253 let (src0, buf1) = splitAt n (x:buf)
1254 src = Prelude.reverse src0 ++ xs
1255 return $ Skip $ ConcatParseBuf src s buf1
1256 (PRD.Parser pstep (return pst1) extract)
1257 -- XXX Specialize for Stop 0 common case?
1259 assert (n <= length (x:buf)) (return ())
1260 let src = Prelude.reverse (Prelude.take n (x:buf)) ++ xs
1261 return $ Skip $ ConcatParseYield b
1262 (ConcatParseInit src s (func b))
1263 PR.Error err -> throwM $ ParseError err
1265 stepOuter _ (ConcatParseYield a next) = return $ Yield a next
1267------------------------------------------------------------------------------
1269------------------------------------------------------------------------------
1271-- | Run a streaming composition, discard the results.
1272{-# INLINE_LATE drain #-}
1273drain :: Monad m => Stream m a -> m ()
1274-- drain = foldrM (\_ xs -> xs) (return ())
1275drain (Stream step state) = go SPEC state
1278 r <- step defState st
1280 Yield _ s -> go SPEC s
1284{-# INLINE_NORMAL null #-}
1285null :: Monad m => Stream m a -> m Bool
1286null m = foldrM (\_ _ -> return False) (return True) m
1288{-# INLINE_NORMAL head #-}
1289head :: Monad m => Stream m a -> m (Maybe a)
1290head m = foldrM (\x _ -> return (Just x)) (return Nothing) m
1292{-# INLINE_NORMAL headElse #-}
1293headElse :: Monad m => a -> Stream m a -> m a
1294headElse a m = foldrM (\x _ -> return x) (return a) m
1296-- Does not fuse, has the same performance as the StreamK version.
1297{-# INLINE_NORMAL tail #-}
1298tail :: Monad m => Stream m a -> m (Maybe (Stream m a))
1299tail (UnStream step state) = go state
1302 r <- step defState st
1304 Yield _ s -> return (Just $ Stream step s)
1306 Stop -> return Nothing
1308-- XXX will it fuse? need custom impl?
1309{-# INLINE_NORMAL last #-}
1310last :: Monad m => Stream m a -> m (Maybe a)
1311last = foldl' (\_ y -> Just y) Nothing
1313{-# INLINE_NORMAL elem #-}
1314elem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
1315-- elem e m = foldrM (\x xs -> if x == e then return True else xs) (return False) m
1316elem e (Stream step state) = go state
1319 r <- step defState st
1322 | x == e -> return True
1325 Stop -> return False
1327{-# INLINE_NORMAL notElem #-}
1328notElem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
1329notElem e s = fmap not (elem e s)
1331{-# INLINE_NORMAL all #-}
1332all :: Monad m => (a -> Bool) -> Stream m a -> m Bool
1333-- all p m = foldrM (\x xs -> if p x then xs else return False) (return True) m
1334all p (Stream step state) = go state
1337 r <- step defState st
1341 | otherwise -> return False
1345{-# INLINE_NORMAL any #-}
1346any :: Monad m => (a -> Bool) -> Stream m a -> m Bool
1347-- any p m = foldrM (\x xs -> if p x then return True else xs) (return False) m
1348any p (Stream step state) = go state
1351 r <- step defState st
1354 | p x -> return True
1357 Stop -> return False
1359{-# INLINE_NORMAL maximum #-}
1360maximum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
1361maximum (Stream step state) = go Nothing state
1364 r <- step defState st
1366 Yield x s -> go (Just x) s
1367 Skip s -> go Nothing s
1368 Stop -> return Nothing
1369 go (Just acc) st = do
1370 r <- step defState st
1373 | acc <= x -> go (Just x) s
1374 | otherwise -> go (Just acc) s
1375 Skip s -> go (Just acc) s
1376 Stop -> return (Just acc)
1378{-# INLINE_NORMAL maximumBy #-}
1379maximumBy :: Monad m => (a -> a -> Ordering) -> Stream m a -> m (Maybe a)
1380maximumBy cmp (Stream step state) = go Nothing state
1383 r <- step defState st
1385 Yield x s -> go (Just x) s
1386 Skip s -> go Nothing s
1387 Stop -> return Nothing
1388 go (Just acc) st = do
1389 r <- step defState st
1391 Yield x s -> case cmp acc x of
1392 GT -> go (Just acc) s
1394 Skip s -> go (Just acc) s
1395 Stop -> return (Just acc)
1397{-# INLINE_NORMAL minimum #-}
1398minimum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
1399minimum (Stream step state) = go Nothing state
1402 r <- step defState st
1404 Yield x s -> go (Just x) s
1405 Skip s -> go Nothing s
1406 Stop -> return Nothing
1407 go (Just acc) st = do
1408 r <- step defState st
1411 | acc <= x -> go (Just acc) s
1412 | otherwise -> go (Just x) s
1413 Skip s -> go (Just acc) s
1414 Stop -> return (Just acc)
1416{-# INLINE_NORMAL minimumBy #-}
1417minimumBy :: Monad m => (a -> a -> Ordering) -> Stream m a -> m (Maybe a)
1418minimumBy cmp (Stream step state) = go Nothing state
1421 r <- step defState st
1423 Yield x s -> go (Just x) s
1424 Skip s -> go Nothing s
1425 Stop -> return Nothing
1426 go (Just acc) st = do
1427 r <- step defState st
1429 Yield x s -> case cmp acc x of
1431 _ -> go (Just acc) s
1432 Skip s -> go (Just acc) s
1433 Stop -> return (Just acc)
1435{-# INLINE_NORMAL (!!) #-}
1436(!!) :: (Monad m) => Stream m a -> Int -> m (Maybe a)
1437(Stream step state) !! i = go i state
1440 r <- step defState st
1442 Yield x s | n < 0 -> return Nothing
1443 | n == 0 -> return $ Just x
1444 | otherwise -> go (n - 1) s
1446 Stop -> return Nothing
1448{-# INLINE_NORMAL lookup #-}
1449lookup :: (Monad m, Eq a) => a -> Stream m (a, b) -> m (Maybe b)
1450lookup e m = foldrM (\(a, b) xs -> if e == a then return (Just b) else xs)
1453{-# INLINE_NORMAL findM #-}
1454findM :: Monad m => (a -> m Bool) -> Stream m a -> m (Maybe a)
1455findM p m = foldrM (\x xs -> p x >>= \r -> if r then return (Just x) else xs)
1459find :: Monad m => (a -> Bool) -> Stream m a -> m (Maybe a)
1460find p = findM (return . p)
1462{-# INLINE_NORMAL findIndices #-}
1463findIndices :: Monad m => (a -> Bool) -> Stream m a -> Stream m Int
1464findIndices p (Stream step state) = Stream step' (state, 0)
1466 {-# INLINE_LATE step' #-}
1467 step' gst (st, i) = i `seq` do
1468 r <- step (adaptState gst) st
1470 Yield x s -> if p x then Yield i (s, i+1) else Skip (s, i+1)
1471 Skip s -> Skip (s, i)
1474{-# INLINE toListRev #-}
1475toListRev :: Monad m => Stream m a -> m [a]
1476toListRev = foldl' (flip (:)) []
1478-- We can implement reverse as:
1480-- > reverse = foldlS (flip cons) nil
1482-- However, this implementation is unusable because of the horrible performance
1483-- of cons. So we just convert it to a list first and then stream from the
1486-- XXX Maybe we can use an Array instead of a list here?
1487{-# INLINE_NORMAL reverse #-}
1488reverse :: Monad m => Stream m a -> Stream m a
1489reverse m = Stream step Nothing
1491 {-# INLINE_LATE step #-}
1494 return $ Skip (Just xs)
1495 step _ (Just (x:xs)) = return $ Yield x (Just xs)
1496 step _ (Just []) = return Stop
1498-- Much faster reverse for Storables
1499{-# INLINE_NORMAL reverse' #-}
1500reverse' :: forall m a. (MonadIO m, Storable a) => Stream m a -> Stream m a
1502-- This commented implementation copies the whole stream into one single array
1503-- and then streams from that array, this is 3-4x faster than the chunked code
1504-- that follows. Though this could be problematic due to unbounded large
1505-- allocations. We need to figure out why the chunked code is slower and if we
1506-- can optimize the chunked code to work as fast as this one. It may be a
1508import Foreign.ForeignPtr (touchForeignPtr)
1509import Foreign.ForeignPtr.Unsafe (unsafeForeignPtrToPtr)
1510import Foreign.Ptr (Ptr, plusPtr)
1511reverse' m = Stream step Nothing
1513 {-# INLINE_LATE step #-}
1515 arr <- A.fromStreamD m
1516 let p = aEnd arr `plusPtr` negate (sizeOf (undefined :: a))
1517 return $ Skip $ Just (aStart arr, p)
1519 step _ (Just (start, p)) | p < unsafeForeignPtrToPtr start = return Stop
1521 step _ (Just (start, p)) = do
1522 let !x = A.unsafeInlineIO $ do
1524 touchForeignPtr start
1526 next = p `plusPtr` negate (sizeOf (undefined :: a))
1527 return $ Yield x (Just (start, next))
1534 $ A.fromStreamDArraysOf A.defaultChunkSize m
1537------------------------------------------------------------------------------
1538-- Grouping/Splitting
1539------------------------------------------------------------------------------
1541{-# INLINE_NORMAL splitSuffixBy' #-}
1542splitSuffixBy' :: Monad m
1543 => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1544splitSuffixBy' predicate f (Stream step state) =
1545 Stream (stepOuter f) (Just state)
1549 {-# INLINE_LATE stepOuter #-}
1550 stepOuter (Fold fstep initial done) gst (Just st) = do
1551 res <- step (adaptState gst) st
1557 then done acc' >>= \val -> return $ Yield val (Just s)
1560 Skip s -> return $ Skip $ Just s
1566 res <- step (adaptState gst) stt
1571 then done acc' >>= \val -> return $ Yield val (Just s)
1573 Skip s -> go SPEC s acc
1574 Stop -> done acc >>= \val -> return $ Yield val Nothing
1576 stepOuter _ _ Nothing = return Stop
1578{-# INLINE_NORMAL groupsBy #-}
1584groupsBy cmp f (Stream step state) = Stream (stepOuter f) (Just state, Nothing)
1588 {-# INLINE_LATE stepOuter #-}
1589 stepOuter (Fold fstep initial done) gst (Just st, Nothing) = do
1590 res <- step (adaptState gst) st
1597 Skip s -> return $ Skip $ (Just s, Nothing)
1602 go !_ prev stt !acc = do
1603 res <- step (adaptState gst) stt
1610 else done acc >>= \r -> return $ Yield r (Just s, Just x)
1611 Skip s -> go SPEC prev s acc
1612 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1614 stepOuter (Fold fstep initial done) gst (Just st, Just prev) = do
1616 acc' <- fstep acc prev
1621 -- XXX code duplicated from the previous equation
1623 res <- step (adaptState gst) stt
1630 else done acc >>= \r -> return $ Yield r (Just s, Just x)
1631 Skip s -> go SPEC s acc
1632 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1634 stepOuter _ _ (Nothing,_) = return Stop
1636{-# INLINE_NORMAL groupsRollingBy #-}
1637groupsRollingBy :: Monad m
1642groupsRollingBy cmp f (Stream step state) =
1643 Stream (stepOuter f) (Just state, Nothing)
1646 {-# INLINE_LATE stepOuter #-}
1647 stepOuter (Fold fstep initial done) gst (Just st, Nothing) = do
1648 res <- step (adaptState gst) st
1655 Skip s -> return $ Skip $ (Just s, Nothing)
1659 go !_ prev stt !acc = do
1660 res <- step (adaptState gst) stt
1668 done acc >>= \r -> return $ Yield r (Just s, Just x)
1669 Skip s -> go SPEC prev s acc
1670 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1672 stepOuter (Fold fstep initial done) gst (Just st, Just prev') = do
1674 acc' <- fstep acc prev'
1675 go SPEC prev' st acc'
1678 go !_ prevv stt !acc = do
1679 res <- step (adaptState gst) stt
1686 else done acc >>= \r -> return $ Yield r (Just s, Just x)
1687 Skip s -> go SPEC prevv s acc
1688 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1690 stepOuter _ _ (Nothing, _) = return Stop
1692{-# INLINE_NORMAL splitBy #-}
1693splitBy :: Monad m => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1694splitBy predicate f (Stream step state) = Stream (step' f) (Just state)
1698 {-# INLINE_LATE step' #-}
1699 step' (Fold fstep initial done) gst (Just st) = initial >>= go SPEC st
1704 res <- step (adaptState gst) stt
1708 then done acc >>= \r -> return $ Yield r (Just s)
1712 Skip s -> go SPEC s acc
1713 Stop -> done acc >>= \r -> return $ Yield r Nothing
1715 step' _ _ Nothing = return Stop
1717-- XXX requires -funfolding-use-threshold=150 in lines-unlines benchmark
1718{-# INLINE_NORMAL splitSuffixBy #-}
1719splitSuffixBy :: Monad m
1720 => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1721splitSuffixBy predicate f (Stream step state) = Stream (step' f) (Just state)
1725 {-# INLINE_LATE step' #-}
1726 step' (Fold fstep initial done) gst (Just st) = do
1727 res <- step (adaptState gst) st
1732 then done acc >>= \val -> return $ Yield val (Just s)
1737 Skip s -> return $ Skip $ Just s
1743 res <- step (adaptState gst) stt
1747 then done acc >>= \r -> return $ Yield r (Just s)
1751 Skip s -> go SPEC s acc
1752 Stop -> done acc >>= \r -> return $ Yield r Nothing
1754 step' _ _ Nothing = return Stop
1756{-# INLINE_NORMAL wordsBy #-}
1757wordsBy :: Monad m => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1758wordsBy predicate f (Stream step state) = Stream (stepOuter f) (Just state)
1762 {-# INLINE_LATE stepOuter #-}
1763 stepOuter (Fold fstep initial done) gst (Just st) = do
1764 res <- step (adaptState gst) st
1768 then return $ Skip (Just s)
1774 Skip s -> return $ Skip $ Just s
1780 res <- step (adaptState gst) stt
1784 then done acc >>= \r -> return $ Yield r (Just s)
1788 Skip s -> go SPEC s acc
1789 Stop -> done acc >>= \r -> return $ Yield r Nothing
1791 stepOuter _ _ Nothing = return Stop
1793-- String search algorithms:
1794-- http://www-igm.univ-mlv.fr/~lecroq/string/index.html
1797-- TODO can we unify the splitting operations using a splitting configuration
1798-- like in the split package.
1800data SplitStyle = Infix | Suffix | Prefix deriving (Eq, Show)
1802data SplitOptions = SplitOptions
1803 { style :: SplitStyle
1804 , withSep :: Bool -- ^ keep the separators in output
1805 -- , compact :: Bool -- ^ treat multiple consecutive separators as one
1806 -- , trimHead :: Bool -- ^ drop blank at head
1807 -- , trimTail :: Bool -- ^ drop blank at tail
1811data SplitOnState s a =
1816 | GO_KARP_RABIN s !(RB.Ring a) !(Ptr a)
1819{-# INLINE_NORMAL splitOn #-}
1821 :: forall m a b. (MonadIO m, Storable a, Enum a, Eq a)
1826splitOn patArr (Fold fstep initial done) (Stream step state) =
1827 Stream stepOuter GO_START
1831 patLen = A.length patArr
1832 maxIndex = patLen - 1
1833 elemBits = sizeOf (undefined :: a) * 8
1835 {-# INLINE_LATE stepOuter #-}
1836 stepOuter _ GO_START =
1838 then return $ Skip $ GO_EMPTY_PAT state
1841 r <- liftIO $ (A.unsafeIndexIO patArr 0)
1842 return $ Skip $ GO_SINGLE_PAT state r
1843 else if sizeOf (undefined :: a) * patLen
1844 <= sizeOf (undefined :: Word)
1845 then return $ Skip $ GO_SHORT_PAT state
1847 (rb, rhead) <- liftIO $ RB.new patLen
1848 return $ Skip $ GO_KARP_RABIN state rb rhead
1850 stepOuter gst (GO_SINGLE_PAT stt pat) = initial >>= go SPEC stt
1855 res <- step (adaptState gst) st
1861 return $ Yield r (GO_SINGLE_PAT s pat)
1862 else fstep acc x >>= go SPEC s
1863 Skip s -> go SPEC s acc
1864 Stop -> done acc >>= \r -> return $ Yield r GO_DONE
1866 stepOuter gst (GO_SHORT_PAT stt) = initial >>= go0 SPEC 0 (0 :: Word) stt
1871 mask = (1 `shiftL` (elemBits * patLen)) - 1
1873 addToWord wrd a = (wrd `shiftL` elemBits) .|. fromIntegral (fromEnum a)
1876 patWord = mask .&. A.foldl' addToWord 0 patArr
1878 go0 !_ !idx wrd st !acc = do
1879 res <- step (adaptState gst) st
1882 let wrd' = addToWord wrd x
1885 if wrd' .&. mask == patWord
1888 return $ Yield r (GO_SHORT_PAT s)
1889 else go1 SPEC wrd' s acc
1890 else go0 SPEC (idx + 1) wrd' s acc
1891 Skip s -> go0 SPEC idx wrd s acc
1894 then go2 wrd idx acc
1896 done acc' >>= \r -> return $ Yield r GO_DONE
1899 go1 !_ wrd st !acc = do
1900 res <- step (adaptState gst) st
1903 let wrd' = addToWord wrd x
1904 old = (mask .&. wrd) `shiftR` (elemBits * (patLen - 1))
1905 acc' <- fstep acc (toEnum $ fromIntegral old)
1906 if wrd' .&. mask == patWord
1907 then done acc' >>= \r -> return $ Yield r (GO_SHORT_PAT s)
1908 else go1 SPEC wrd' s acc'
1909 Skip s -> go1 SPEC wrd s acc
1911 acc' <- go2 wrd patLen acc
1912 done acc' >>= \r -> return $ Yield r GO_DONE
1914 go2 !wrd !n !acc | n > 0 = do
1915 let old = (mask .&. wrd) `shiftR` (elemBits * (n - 1))
1916 fstep acc (toEnum $ fromIntegral old) >>= go2 wrd (n - 1)
1917 go2 _ _ acc = return acc
1919 stepOuter gst (GO_KARP_RABIN stt rb rhead) = do
1920 initial >>= go0 SPEC 0 rhead stt
1924 k = 2891336453 :: Word32
1926 addCksum cksum a = cksum * k + fromIntegral (fromEnum a)
1927 deltaCksum cksum old new =
1928 addCksum cksum new - coeff * fromIntegral (fromEnum old)
1930 -- XXX shall we use a random starting hash or 1 instead of 0?
1931 patHash = A.foldl' addCksum 0 patArr
1934 go0 !_ !idx !rh st !acc = do
1935 res <- step (adaptState gst) st
1938 rh' <- liftIO $ RB.unsafeInsert rb rh x
1941 let fold = RB.unsafeFoldRing (RB.ringBound rb)
1942 let !ringHash = fold addCksum 0 rb
1943 if ringHash == patHash
1944 then go2 SPEC ringHash rh' s acc
1945 else go1 SPEC ringHash rh' s acc
1946 else go0 SPEC (idx + 1) rh' s acc
1947 Skip s -> go0 SPEC idx rh s acc
1949 !acc' <- if idx /= 0
1950 then RB.unsafeFoldRingM rh fstep acc rb
1952 done acc' >>= \r -> return $ Yield r GO_DONE
1954 -- XXX Theoretically this code can do 4 times faster if GHC generates
1955 -- optimal code. If we use just "(cksum' == patHash)" condition it goes
1956 -- 4x faster, as soon as we add the "RB.unsafeEqArray rb v" condition
1957 -- the generated code changes drastically and becomes 4x slower. Need
1958 -- to investigate what is going on with GHC.
1960 go1 !_ !cksum !rh st !acc = do
1961 res <- step (adaptState gst) st
1964 old <- liftIO $ peek rh
1965 let cksum' = deltaCksum cksum old x
1966 acc' <- fstep acc old
1968 if (cksum' == patHash)
1970 rh' <- liftIO (RB.unsafeInsert rb rh x)
1971 go2 SPEC cksum' rh' s acc'
1973 rh' <- liftIO (RB.unsafeInsert rb rh x)
1974 go1 SPEC cksum' rh' s acc'
1975 Skip s -> go1 SPEC cksum rh s acc
1977 acc' <- RB.unsafeFoldRingFullM rh fstep acc rb
1978 done acc' >>= \r -> return $ Yield r GO_DONE
1980 go2 !_ !cksum' !rh' s !acc' = do
1981 if RB.unsafeEqArray rb rh' patArr
1984 return $ Yield r (GO_KARP_RABIN s rb rhead)
1985 else go1 SPEC cksum' rh' s acc'
1987 stepOuter gst (GO_EMPTY_PAT st) = do
1988 res <- step (adaptState gst) st
1993 done acc' >>= \r -> return $ Yield r (GO_EMPTY_PAT s)
1994 Skip s -> return $ Skip (GO_EMPTY_PAT s)
1997 stepOuter _ GO_DONE = return Stop
1999{-# INLINE_NORMAL splitSuffixOn #-}
2001 :: forall m a b. (MonadIO m, Storable a, Enum a, Eq a)
2007splitSuffixOn withSep patArr (Fold fstep initial done)
2008 (Stream step state) =
2009 Stream stepOuter GO_START
2013 patLen = A.length patArr
2014 maxIndex = patLen - 1
2015 elemBits = sizeOf (undefined :: a) * 8
2017 {-# INLINE_LATE stepOuter #-}
2018 stepOuter _ GO_START =
2020 then return $ Skip $ GO_EMPTY_PAT state
2023 r <- liftIO $ (A.unsafeIndexIO patArr 0)
2024 return $ Skip $ GO_SINGLE_PAT state r
2025 else if sizeOf (undefined :: a) * patLen
2026 <= sizeOf (undefined :: Word)
2027 then return $ Skip $ GO_SHORT_PAT state
2029 (rb, rhead) <- liftIO $ RB.new patLen
2030 return $ Skip $ GO_KARP_RABIN state rb rhead
2032 stepOuter gst (GO_SINGLE_PAT stt pat) = do
2033 -- This first part is the only difference between splitOn and
2035 -- If the last element is a separator do not issue a blank segment.
2036 res <- step (adaptState gst) stt
2042 acc' <- if withSep then fstep acc x else return acc
2043 done acc' >>= \r -> return $ Yield r (GO_SINGLE_PAT s pat)
2044 else fstep acc x >>= go SPEC s
2045 Skip s -> return $ Skip $ (GO_SINGLE_PAT s pat)
2050 -- This is identical for splitOn and splitSuffixOn
2052 res <- step (adaptState gst) st
2057 acc' <- if withSep then fstep acc x else return acc
2059 return $ Yield r (GO_SINGLE_PAT s pat)
2060 else fstep acc x >>= go SPEC s
2061 Skip s -> go SPEC s acc
2062 Stop -> done acc >>= \r -> return $ Yield r GO_DONE
2064 stepOuter gst (GO_SHORT_PAT stt) = do
2066 -- Call "initial" only if the stream yields an element, otherwise we
2067 -- may call "initial" but never yield anything. initial may produce a
2068 -- side effect, therefore we will end up doing and discard a side
2073 res <- step (adaptState gst) stt
2077 let wrd' = addToWord wrd x
2078 acc' <- if withSep then fstep acc x else return acc
2081 if wrd' .&. mask == patWord
2082 then done acc' >>= \r -> return $ Yield r (GO_SHORT_PAT s)
2083 else go0 SPEC (idx + 1) wrd' s acc'
2084 else go0 SPEC (idx + 1) wrd' s acc'
2085 Skip s -> return $ Skip (GO_SHORT_PAT s)
2091 mask = (1 `shiftL` (elemBits * patLen)) - 1
2093 addToWord wrd a = (wrd `shiftL` elemBits) .|. fromIntegral (fromEnum a)
2096 patWord = mask .&. A.foldl' addToWord 0 patArr
2098 go0 !_ !idx wrd st !acc = do
2099 res <- step (adaptState gst) st
2102 let wrd' = addToWord wrd x
2103 acc' <- if withSep then fstep acc x else return acc
2106 if wrd' .&. mask == patWord
2109 return $ Yield r (GO_SHORT_PAT s)
2110 else go1 SPEC wrd' s acc'
2111 else go0 SPEC (idx + 1) wrd' s acc'
2112 Skip s -> go0 SPEC idx wrd s acc
2114 if (idx == maxIndex) && (wrd .&. mask == patWord)
2117 acc' <- if idx /= 0 && not withSep
2118 then go2 wrd idx acc
2120 done acc' >>= \r -> return $ Yield r GO_DONE
2123 go1 !_ wrd st !acc = do
2124 res <- step (adaptState gst) st
2127 let wrd' = addToWord wrd x
2128 old = (mask .&. wrd) `shiftR` (elemBits * (patLen - 1))
2131 else fstep acc (toEnum $ fromIntegral old)
2132 if wrd' .&. mask == patWord
2133 then done acc' >>= \r -> return $ Yield r (GO_SHORT_PAT s)
2134 else go1 SPEC wrd' s acc'
2135 Skip s -> go1 SPEC wrd s acc
2137 -- If the last sequence is a separator do not issue a blank
2139 if wrd .&. mask == patWord
2144 else go2 wrd patLen acc
2145 done acc' >>= \r -> return $ Yield r GO_DONE
2147 go2 !wrd !n !acc | n > 0 = do
2148 let old = (mask .&. wrd) `shiftR` (elemBits * (n - 1))
2149 fstep acc (toEnum $ fromIntegral old) >>= go2 wrd (n - 1)
2150 go2 _ _ acc = return acc
2152 stepOuter gst (GO_KARP_RABIN stt rb rhead) = do
2154 res <- step (adaptState gst) stt
2158 acc' <- if withSep then fstep acc x else return acc
2159 rh' <- liftIO (RB.unsafeInsert rb rhead x)
2162 let fold = RB.unsafeFoldRing (RB.ringBound rb)
2163 let !ringHash = fold addCksum 0 rb
2164 if ringHash == patHash
2165 then go2 SPEC ringHash rh' s acc'
2166 else go0 SPEC (idx + 1) rh' s acc'
2167 else go0 SPEC (idx + 1) rh' s acc'
2168 Skip s -> return $ Skip (GO_KARP_RABIN s rb rhead)
2173 k = 2891336453 :: Word32
2175 addCksum cksum a = cksum * k + fromIntegral (fromEnum a)
2176 deltaCksum cksum old new =
2177 addCksum cksum new - coeff * fromIntegral (fromEnum old)
2179 -- XXX shall we use a random starting hash or 1 instead of 0?
2180 patHash = A.foldl' addCksum 0 patArr
2183 go0 !_ !idx !rh st !acc = do
2184 res <- step (adaptState gst) st
2187 acc' <- if withSep then fstep acc x else return acc
2188 rh' <- liftIO (RB.unsafeInsert rb rh x)
2191 let fold = RB.unsafeFoldRing (RB.ringBound rb)
2192 let !ringHash = fold addCksum 0 rb
2193 if ringHash == patHash
2194 then go2 SPEC ringHash rh' s acc'
2195 else go1 SPEC ringHash rh' s acc'
2196 else go0 SPEC (idx + 1) rh' s acc'
2197 Skip s -> go0 SPEC idx rh s acc
2199 -- do not issue a blank segment when we end at pattern
2200 if (idx == maxIndex) && RB.unsafeEqArray rb rh patArr
2203 !acc' <- if idx /= 0 && not withSep
2204 then RB.unsafeFoldRingM rh fstep acc rb
2206 done acc' >>= \r -> return $ Yield r GO_DONE
2208 -- XXX Theoretically this code can do 4 times faster if GHC generates
2209 -- optimal code. If we use just "(cksum' == patHash)" condition it goes
2210 -- 4x faster, as soon as we add the "RB.unsafeEqArray rb v" condition
2211 -- the generated code changes drastically and becomes 4x slower. Need
2212 -- to investigate what is going on with GHC.
2214 go1 !_ !cksum !rh st !acc = do
2215 res <- step (adaptState gst) st
2218 old <- liftIO $ peek rh
2219 let cksum' = deltaCksum cksum old x
2224 if (cksum' == patHash)
2226 rh' <- liftIO (RB.unsafeInsert rb rh x)
2227 go2 SPEC cksum' rh' s acc'
2229 rh' <- liftIO (RB.unsafeInsert rb rh x)
2230 go1 SPEC cksum' rh' s acc'
2231 Skip s -> go1 SPEC cksum rh s acc
2233 if RB.unsafeEqArray rb rh patArr
2238 else RB.unsafeFoldRingFullM rh fstep acc rb
2239 done acc' >>= \r -> return $ Yield r GO_DONE
2241 go2 !_ !cksum' !rh' s !acc' = do
2242 if RB.unsafeEqArray rb rh' patArr
2245 return $ Yield r (GO_KARP_RABIN s rb rhead)
2246 else go1 SPEC cksum' rh' s acc'
2248 stepOuter gst (GO_EMPTY_PAT st) = do
2249 res <- step (adaptState gst) st
2254 done acc' >>= \r -> return $ Yield r (GO_EMPTY_PAT s)
2255 Skip s -> return $ Skip (GO_EMPTY_PAT s)
2258 stepOuter _ GO_DONE = return Stop
2260data SplitState s arr
2262 | SplitBuffering s arr
2263 | SplitSplitting s arr
2264 | SplitYielding arr (SplitState s arr)
2267-- XXX An alternative approach would be to use a partial fold (Fold m a b) to
2268-- split using a splitBy like combinator. The Fold would consume upto the
2269-- separator and return any leftover which can then be fed to the next fold.
2271-- We can revisit this once we have partial folds/parsers.
2273-- | Performs infix separator style splitting.
2274{-# INLINE_NORMAL splitInnerBy #-}
2277 => (f a -> m (f a, Maybe (f a))) -- splitter
2278 -> (f a -> f a -> m (f a)) -- joiner
2281splitInnerBy splitter joiner (Stream step1 state1) =
2282 (Stream step (SplitInitial state1))
2286 {-# INLINE_LATE step #-}
2287 step gst (SplitInitial st) = do
2291 (x1, mx2) <- splitter x
2292 return $ case mx2 of
2293 Nothing -> Skip (SplitBuffering s x1)
2294 Just x2 -> Skip (SplitYielding x1 (SplitSplitting s x2))
2295 Skip s -> return $ Skip (SplitInitial s)
2296 Stop -> return $ Stop
2298 step gst (SplitBuffering st buf) = do
2302 (x1, mx2) <- splitter x
2303 buf' <- joiner buf x1
2304 return $ case mx2 of
2305 Nothing -> Skip (SplitBuffering s buf')
2306 Just x2 -> Skip (SplitYielding buf' (SplitSplitting s x2))
2307 Skip s -> return $ Skip (SplitBuffering s buf)
2308 Stop -> return $ Skip (SplitYielding buf SplitFinishing)
2310 step _ (SplitSplitting st buf) = do
2311 (x1, mx2) <- splitter buf
2312 return $ case mx2 of
2313 Nothing -> Skip $ SplitBuffering st x1
2314 Just x2 -> Skip $ SplitYielding x1 (SplitSplitting st x2)
2316 step _ (SplitYielding x next) = return $ Yield x next
2317 step _ SplitFinishing = return $ Stop
2319-- | Performs infix separator style splitting.
2320{-# INLINE_NORMAL splitInnerBySuffix #-}
2322 :: (Monad m, Eq (f a), Monoid (f a))
2323 => (f a -> m (f a, Maybe (f a))) -- splitter
2324 -> (f a -> f a -> m (f a)) -- joiner
2327splitInnerBySuffix splitter joiner (Stream step1 state1) =
2328 (Stream step (SplitInitial state1))
2332 {-# INLINE_LATE step #-}
2333 step gst (SplitInitial st) = do
2337 (x1, mx2) <- splitter x
2338 return $ case mx2 of
2339 Nothing -> Skip (SplitBuffering s x1)
2340 Just x2 -> Skip (SplitYielding x1 (SplitSplitting s x2))
2341 Skip s -> return $ Skip (SplitInitial s)
2342 Stop -> return $ Stop
2344 step gst (SplitBuffering st buf) = do
2348 (x1, mx2) <- splitter x
2349 buf' <- joiner buf x1
2350 return $ case mx2 of
2351 Nothing -> Skip (SplitBuffering s buf')
2352 Just x2 -> Skip (SplitYielding buf' (SplitSplitting s x2))
2353 Skip s -> return $ Skip (SplitBuffering s buf)
2357 else Skip (SplitYielding buf SplitFinishing)
2359 step _ (SplitSplitting st buf) = do
2360 (x1, mx2) <- splitter buf
2361 return $ case mx2 of
2362 Nothing -> Skip $ SplitBuffering st x1
2363 Just x2 -> Skip $ SplitYielding x1 (SplitSplitting st x2)
2365 step _ (SplitYielding x next) = return $ Yield x next
2366 step _ SplitFinishing = return $ Stop
2368------------------------------------------------------------------------------
2370------------------------------------------------------------------------------
2372{-# INLINE_NORMAL isPrefixOf #-}
2373isPrefixOf :: (Eq a, Monad m) => Stream m a -> Stream m a -> m Bool
2374isPrefixOf (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
2376 go (sa, sb, Nothing) = do
2377 r <- stepa defState sa
2379 Yield x sa' -> go (sa', sb, Just x)
2380 Skip sa' -> go (sa', sb, Nothing)
2383 go (sa, sb, Just x) = do
2384 r <- stepb defState sb
2388 then go (sa, sb', Nothing)
2390 Skip sb' -> go (sa, sb', Just x)
2391 Stop -> return False
2393{-# INLINE_NORMAL isSubsequenceOf #-}
2394isSubsequenceOf :: (Eq a, Monad m) => Stream m a -> Stream m a -> m Bool
2395isSubsequenceOf (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
2397 go (sa, sb, Nothing) = do
2398 r <- stepa defState sa
2400 Yield x sa' -> go (sa', sb, Just x)
2401 Skip sa' -> go (sa', sb, Nothing)
2404 go (sa, sb, Just x) = do
2405 r <- stepb defState sb
2409 then go (sa, sb', Nothing)
2410 else go (sa, sb', Just x)
2411 Skip sb' -> go (sa, sb', Just x)
2412 Stop -> return False
2414{-# INLINE_NORMAL stripPrefix #-}
2417 => Stream m a -> Stream m a -> m (Maybe (Stream m a))
2418stripPrefix (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
2420 go (sa, sb, Nothing) = do
2421 r <- stepa defState sa
2423 Yield x sa' -> go (sa', sb, Just x)
2424 Skip sa' -> go (sa', sb, Nothing)
2425 Stop -> return $ Just (Stream stepb sb)
2427 go (sa, sb, Just x) = do
2428 r <- stepb defState sb
2432 then go (sa, sb', Nothing)
2434 Skip sb' -> go (sa, sb', Just x)
2435 Stop -> return Nothing
2437------------------------------------------------------------------------------
2439------------------------------------------------------------------------------
2441-- | Execute a monadic action for each element of the 'Stream'
2442{-# INLINE_NORMAL mapM_ #-}
2443mapM_ :: Monad m => (a -> m b) -> Stream m a -> m ()
2444mapM_ m = drain . mapM m
2446-------------------------------------------------------------------------------
2447-- Stream transformations using Unfolds
2448-------------------------------------------------------------------------------
2450-- Define a unique structure to use in inspection testing
2451data ConcatMapUState o i =
2453 | ConcatMapUInner o i
2455-- | @concatMapU unfold stream@ uses @unfold@ to map the input stream elements
2456-- to streams and then flattens the generated streams into a single output
2459-- This is like 'concatMap' but uses an unfold with an explicit state to
2460-- generate the stream instead of a 'Stream' type generator. This allows better
2461-- optimization via fusion. This can be many times more efficient than
2464{-# INLINE_NORMAL concatMapU #-}
2465concatMapU :: Monad m => Unfold m a b -> Stream m a -> Stream m b
2466concatMapU (Unfold istep inject) (Stream ostep ost) =
2467 Stream step (ConcatMapUOuter ost)
2469 {-# INLINE_LATE step #-}
2470 step gst (ConcatMapUOuter o) = do
2471 r <- ostep (adaptState gst) o
2475 i `seq` return (Skip (ConcatMapUInner o' i))
2476 Skip o' -> return $ Skip (ConcatMapUOuter o')
2477 Stop -> return $ Stop
2479 step _ (ConcatMapUInner o i) = do
2482 Yield x i' -> Yield x (ConcatMapUInner o i')
2483 Skip i' -> Skip (ConcatMapUInner o i')
2484 Stop -> Skip (ConcatMapUOuter o)
2486data ConcatUnfoldInterleaveState o i =
2487 ConcatUnfoldInterleaveOuter o [i]
2488 | ConcatUnfoldInterleaveInner o [i]
2489 | ConcatUnfoldInterleaveInnerL [i] [i]
2490 | ConcatUnfoldInterleaveInnerR [i] [i]
2492-- XXX use arrays to store state instead of lists.
2493-- XXX In general we can use different scheduling strategies e.g. how to
2494-- schedule the outer vs inner loop or assigning weights to different streams
2495-- or outer and inner loops.
2497-- After a yield, switch to the next stream. Do not switch streams on Skip.
2498-- Yield from outer stream switches to the inner stream.
2500-- There are two choices here, (1) exhaust the outer stream first and then
2501-- start yielding from the inner streams, this is much simpler to implement,
2502-- (2) yield at least one element from an inner stream before going back to
2503-- outer stream and opening the next stream from it.
2505-- Ideally, we need some scheduling bias to inner streams vs outer stream.
2506-- Maybe we can configure the behavior.
2508{-# INLINE_NORMAL concatUnfoldInterleave #-}
2509concatUnfoldInterleave :: Monad m => Unfold m a b -> Stream m a -> Stream m b
2510concatUnfoldInterleave (Unfold istep inject) (Stream ostep ost) =
2511 Stream step (ConcatUnfoldInterleaveOuter ost [])
2513 {-# INLINE_LATE step #-}
2514 step gst (ConcatUnfoldInterleaveOuter o ls) = do
2515 r <- ostep (adaptState gst) o
2519 i `seq` return (Skip (ConcatUnfoldInterleaveInner o' (i : ls)))
2520 Skip o' -> return $ Skip (ConcatUnfoldInterleaveOuter o' ls)
2521 Stop -> return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2523 step _ (ConcatUnfoldInterleaveInner _ []) = undefined
2524 step _ (ConcatUnfoldInterleaveInner o (st:ls)) = do
2527 Yield x s -> Yield x (ConcatUnfoldInterleaveOuter o (s:ls))
2528 Skip s -> Skip (ConcatUnfoldInterleaveInner o (s:ls))
2529 Stop -> Skip (ConcatUnfoldInterleaveOuter o ls)
2531 step _ (ConcatUnfoldInterleaveInnerL [] []) = return Stop
2532 step _ (ConcatUnfoldInterleaveInnerL [] rs) =
2533 return $ Skip (ConcatUnfoldInterleaveInnerR [] rs)
2535 step _ (ConcatUnfoldInterleaveInnerL (st:ls) rs) = do
2538 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerL ls (s:rs))
2539 Skip s -> Skip (ConcatUnfoldInterleaveInnerL (s:ls) rs)
2540 Stop -> Skip (ConcatUnfoldInterleaveInnerL ls rs)
2542 step _ (ConcatUnfoldInterleaveInnerR [] []) = return Stop
2543 step _ (ConcatUnfoldInterleaveInnerR ls []) =
2544 return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2546 step _ (ConcatUnfoldInterleaveInnerR ls (st:rs)) = do
2549 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerR (s:ls) rs)
2550 Skip s -> Skip (ConcatUnfoldInterleaveInnerR ls (s:rs))
2551 Stop -> Skip (ConcatUnfoldInterleaveInnerR ls rs)
2553-- XXX In general we can use different scheduling strategies e.g. how to
2554-- schedule the outer vs inner loop or assigning weights to different streams
2555-- or outer and inner loops.
2557-- This could be inefficient if the tasks are too small.
2559-- Compared to concatUnfoldInterleave this one switches streams on Skips.
2561{-# INLINE_NORMAL concatUnfoldRoundrobin #-}
2562concatUnfoldRoundrobin :: Monad m => Unfold m a b -> Stream m a -> Stream m b
2563concatUnfoldRoundrobin (Unfold istep inject) (Stream ostep ost) =
2564 Stream step (ConcatUnfoldInterleaveOuter ost [])
2566 {-# INLINE_LATE step #-}
2567 step gst (ConcatUnfoldInterleaveOuter o ls) = do
2568 r <- ostep (adaptState gst) o
2572 i `seq` return (Skip (ConcatUnfoldInterleaveInner o' (i : ls)))
2573 Skip o' -> return $ Skip (ConcatUnfoldInterleaveInner o' ls)
2574 Stop -> return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2576 step _ (ConcatUnfoldInterleaveInner o []) =
2577 return $ Skip (ConcatUnfoldInterleaveOuter o [])
2579 step _ (ConcatUnfoldInterleaveInner o (st:ls)) = do
2582 Yield x s -> Yield x (ConcatUnfoldInterleaveOuter o (s:ls))
2583 Skip s -> Skip (ConcatUnfoldInterleaveOuter o (s:ls))
2584 Stop -> Skip (ConcatUnfoldInterleaveOuter o ls)
2586 step _ (ConcatUnfoldInterleaveInnerL [] []) = return Stop
2587 step _ (ConcatUnfoldInterleaveInnerL [] rs) =
2588 return $ Skip (ConcatUnfoldInterleaveInnerR [] rs)
2590 step _ (ConcatUnfoldInterleaveInnerL (st:ls) rs) = do
2593 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerL ls (s:rs))
2594 Skip s -> Skip (ConcatUnfoldInterleaveInnerL ls (s:rs))
2595 Stop -> Skip (ConcatUnfoldInterleaveInnerL ls rs)
2597 step _ (ConcatUnfoldInterleaveInnerR [] []) = return Stop
2598 step _ (ConcatUnfoldInterleaveInnerR ls []) =
2599 return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2601 step _ (ConcatUnfoldInterleaveInnerR ls (st:rs)) = do
2604 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerR (s:ls) rs)
2605 Skip s -> Skip (ConcatUnfoldInterleaveInnerR (s:ls) rs)
2606 Stop -> Skip (ConcatUnfoldInterleaveInnerR ls rs)
2608data AppendState s1 s2 = AppendFirst s1 | AppendSecond s2
2610-- Note that this could be much faster compared to the CPS stream. However, as
2611-- the number of streams being composed increases this may become expensive.
2612-- Need to see where the breaking point is between the two.
2614{-# INLINE_NORMAL append #-}
2615append :: Monad m => Stream m a -> Stream m a -> Stream m a
2616append (Stream step1 state1) (Stream step2 state2) =
2617 Stream step (AppendFirst state1)
2621 {-# INLINE_LATE step #-}
2622 step gst (AppendFirst st) = do
2625 Yield a s -> Yield a (AppendFirst s)
2626 Skip s -> Skip (AppendFirst s)
2627 Stop -> Skip (AppendSecond state2)
2629 step gst (AppendSecond st) = do
2632 Yield a s -> Yield a (AppendSecond s)
2633 Skip s -> Skip (AppendSecond s)
2636data InterleaveState s1 s2 = InterleaveFirst s1 s2 | InterleaveSecond s1 s2
2637 | InterleaveSecondOnly s2 | InterleaveFirstOnly s1
2639{-# INLINE_NORMAL interleave #-}
2640interleave :: Monad m => Stream m a -> Stream m a -> Stream m a
2641interleave (Stream step1 state1) (Stream step2 state2) =
2642 Stream step (InterleaveFirst state1 state2)
2646 {-# INLINE_LATE step #-}
2647 step gst (InterleaveFirst st1 st2) = do
2650 Yield a s -> Yield a (InterleaveSecond s st2)
2651 Skip s -> Skip (InterleaveFirst s st2)
2652 Stop -> Skip (InterleaveSecondOnly st2)
2654 step gst (InterleaveSecond st1 st2) = do
2657 Yield a s -> Yield a (InterleaveFirst st1 s)
2658 Skip s -> Skip (InterleaveSecond st1 s)
2659 Stop -> Skip (InterleaveFirstOnly st1)
2661 step gst (InterleaveFirstOnly st1) = do
2664 Yield a s -> Yield a (InterleaveFirstOnly s)
2665 Skip s -> Skip (InterleaveFirstOnly s)
2668 step gst (InterleaveSecondOnly st2) = do
2671 Yield a s -> Yield a (InterleaveSecondOnly s)
2672 Skip s -> Skip (InterleaveSecondOnly s)
2675{-# INLINE_NORMAL interleaveMin #-}
2676interleaveMin :: Monad m => Stream m a -> Stream m a -> Stream m a
2677interleaveMin (Stream step1 state1) (Stream step2 state2) =
2678 Stream step (InterleaveFirst state1 state2)
2682 {-# INLINE_LATE step #-}
2683 step gst (InterleaveFirst st1 st2) = do
2686 Yield a s -> Yield a (InterleaveSecond s st2)
2687 Skip s -> Skip (InterleaveFirst s st2)
2690 step gst (InterleaveSecond st1 st2) = do
2693 Yield a s -> Yield a (InterleaveFirst st1 s)
2694 Skip s -> Skip (InterleaveSecond st1 s)
2697 step _ (InterleaveFirstOnly _) = undefined
2698 step _ (InterleaveSecondOnly _) = undefined
2700{-# INLINE_NORMAL interleaveSuffix #-}
2701interleaveSuffix :: Monad m => Stream m a -> Stream m a -> Stream m a
2702interleaveSuffix (Stream step1 state1) (Stream step2 state2) =
2703 Stream step (InterleaveFirst state1 state2)
2707 {-# INLINE_LATE step #-}
2708 step gst (InterleaveFirst st1 st2) = do
2711 Yield a s -> Yield a (InterleaveSecond s st2)
2712 Skip s -> Skip (InterleaveFirst s st2)
2715 step gst (InterleaveSecond st1 st2) = do
2718 Yield a s -> Yield a (InterleaveFirst st1 s)
2719 Skip s -> Skip (InterleaveSecond st1 s)
2720 Stop -> Skip (InterleaveFirstOnly st1)
2722 step gst (InterleaveFirstOnly st1) = do
2725 Yield a s -> Yield a (InterleaveFirstOnly s)
2726 Skip s -> Skip (InterleaveFirstOnly s)
2729 step _ (InterleaveSecondOnly _) = undefined
2731data InterleaveInfixState s1 s2 a
2732 = InterleaveInfixFirst s1 s2
2733 | InterleaveInfixSecondBuf s1 s2
2734 | InterleaveInfixSecondYield s1 s2 a
2735 | InterleaveInfixFirstYield s1 s2 a
2736 | InterleaveInfixFirstOnly s1
2738{-# INLINE_NORMAL interleaveInfix #-}
2739interleaveInfix :: Monad m => Stream m a -> Stream m a -> Stream m a
2740interleaveInfix (Stream step1 state1) (Stream step2 state2) =
2741 Stream step (InterleaveInfixFirst state1 state2)
2745 {-# INLINE_LATE step #-}
2746 step gst (InterleaveInfixFirst st1 st2) = do
2749 Yield a s -> Yield a (InterleaveInfixSecondBuf s st2)
2750 Skip s -> Skip (InterleaveInfixFirst s st2)
2753 step gst (InterleaveInfixSecondBuf st1 st2) = do
2756 Yield a s -> Skip (InterleaveInfixSecondYield st1 s a)
2757 Skip s -> Skip (InterleaveInfixSecondBuf st1 s)
2758 Stop -> Skip (InterleaveInfixFirstOnly st1)
2760 step gst (InterleaveInfixSecondYield st1 st2 x) = do
2763 Yield a s -> Yield x (InterleaveInfixFirstYield s st2 a)
2764 Skip s -> Skip (InterleaveInfixSecondYield s st2 x)
2767 step _ (InterleaveInfixFirstYield st1 st2 x) = do
2768 return $ Yield x (InterleaveInfixSecondBuf st1 st2)
2770 step gst (InterleaveInfixFirstOnly st1) = do
2773 Yield a s -> Yield a (InterleaveInfixFirstOnly s)
2774 Skip s -> Skip (InterleaveInfixFirstOnly s)
2777{-# INLINE_NORMAL roundRobin #-}
2778roundRobin :: Monad m => Stream m a -> Stream m a -> Stream m a
2779roundRobin (Stream step1 state1) (Stream step2 state2) =
2780 Stream step (InterleaveFirst state1 state2)
2784 {-# INLINE_LATE step #-}
2785 step gst (InterleaveFirst st1 st2) = do
2788 Yield a s -> Yield a (InterleaveSecond s st2)
2789 Skip s -> Skip (InterleaveSecond s st2)
2790 Stop -> Skip (InterleaveSecondOnly st2)
2792 step gst (InterleaveSecond st1 st2) = do
2795 Yield a s -> Yield a (InterleaveFirst st1 s)
2796 Skip s -> Skip (InterleaveFirst st1 s)
2797 Stop -> Skip (InterleaveFirstOnly st1)
2799 step gst (InterleaveSecondOnly st2) = do
2802 Yield a s -> Yield a (InterleaveSecondOnly s)
2803 Skip s -> Skip (InterleaveSecondOnly s)
2806 step gst (InterleaveFirstOnly st1) = do
2809 Yield a s -> Yield a (InterleaveFirstOnly s)
2810 Skip s -> Skip (InterleaveFirstOnly s)
2813data ICUState s1 s2 i1 i2 =
2818 | ICUFirstInner s1 s2 i1
2819 | ICUSecondInner s1 s2 i2
2820 | ICUFirstOnlyInner s1 i1
2821 | ICUSecondOnlyInner s2 i2
2823-- | Interleave streams (full streams, not the elements) unfolded from two
2824-- input streams and concat. Stop when the first stream stops. If the second
2825-- stream ends before the first one then first stream still keeps running alone
2826-- without any interleaving with the second stream.
2828-- [a1, a2, ... an] [b1, b2 ...]
2829-- => [streamA1, streamA2, ... streamAn] [streamB1, streamB2, ...]
2830-- => [streamA1, streamB1, streamA2...StreamAn, streamBn]
2831-- => [a11, a12, ...a1j, b11, b12, ...b1k, a21, a22, ...]
2833{-# INLINE_NORMAL gintercalateSuffix #-}
2836 => Unfold m a c -> Stream m a -> Unfold m b c -> Stream m b -> Stream m c
2838 (Unfold istep1 inject1) (Stream step1 state1)
2839 (Unfold istep2 inject2) (Stream step2 state2) =
2840 Stream step (ICUFirst state1 state2)
2844 {-# INLINE_LATE step #-}
2845 step gst (ICUFirst s1 s2) = do
2846 r <- step1 (adaptState gst) s1
2850 i `seq` return (Skip (ICUFirstInner s s2 i))
2851 Skip s -> return $ Skip (ICUFirst s s2)
2854 step gst (ICUFirstOnly s1) = do
2855 r <- step1 (adaptState gst) s1
2859 i `seq` return (Skip (ICUFirstOnlyInner s i))
2860 Skip s -> return $ Skip (ICUFirstOnly s)
2863 step _ (ICUFirstInner s1 s2 i1) = do
2866 Yield x i' -> Yield x (ICUFirstInner s1 s2 i')
2867 Skip i' -> Skip (ICUFirstInner s1 s2 i')
2868 Stop -> Skip (ICUSecond s1 s2)
2870 step _ (ICUFirstOnlyInner s1 i1) = do
2873 Yield x i' -> Yield x (ICUFirstOnlyInner s1 i')
2874 Skip i' -> Skip (ICUFirstOnlyInner s1 i')
2875 Stop -> Skip (ICUFirstOnly s1)
2877 step gst (ICUSecond s1 s2) = do
2878 r <- step2 (adaptState gst) s2
2882 i `seq` return (Skip (ICUSecondInner s1 s i))
2883 Skip s -> return $ Skip (ICUSecond s1 s)
2884 Stop -> return $ Skip (ICUFirstOnly s1)
2886 step _ (ICUSecondInner s1 s2 i2) = do
2889 Yield x i' -> Yield x (ICUSecondInner s1 s2 i')
2890 Skip i' -> Skip (ICUSecondInner s1 s2 i')
2891 Stop -> Skip (ICUFirst s1 s2)
2893 step _ (ICUSecondOnly _s2) = undefined
2894 step _ (ICUSecondOnlyInner _s2 _i2) = undefined
2896data InterposeSuffixState s1 i1 =
2897 InterposeSuffixFirst s1
2898 -- | InterposeSuffixFirstYield s1 i1
2899 | InterposeSuffixFirstInner s1 i1
2900 | InterposeSuffixSecond s1
2902-- Note that if an unfolded layer turns out to be nil we still emit the
2903-- separator effect. An alternate behavior could be to emit the separator
2904-- effect only if at least one element has been yielded by the unfolding.
2905-- However, that becomes a bit complicated, so we have chosen the former
2907{-# INLINE_NORMAL interposeSuffix #-}
2910 => m c -> Unfold m b c -> Stream m b -> Stream m c
2913 (Unfold istep1 inject1) (Stream step1 state1) =
2914 Stream step (InterposeSuffixFirst state1)
2918 {-# INLINE_LATE step #-}
2919 step gst (InterposeSuffixFirst s1) = do
2920 r <- step1 (adaptState gst) s1
2924 i `seq` return (Skip (InterposeSuffixFirstInner s i))
2925 -- i `seq` return (Skip (InterposeSuffixFirstYield s i))
2926 Skip s -> return $ Skip (InterposeSuffixFirst s)
2930 step _ (InterposeSuffixFirstYield s1 i1) = do
2933 Yield x i' -> Yield x (InterposeSuffixFirstInner s1 i')
2934 Skip i' -> Skip (InterposeSuffixFirstYield s1 i')
2935 Stop -> Skip (InterposeSuffixFirst s1)
2938 step _ (InterposeSuffixFirstInner s1 i1) = do
2941 Yield x i' -> Yield x (InterposeSuffixFirstInner s1 i')
2942 Skip i' -> Skip (InterposeSuffixFirstInner s1 i')
2943 Stop -> Skip (InterposeSuffixSecond s1)
2945 step _ (InterposeSuffixSecond s1) = do
2947 return $ Yield r (InterposeSuffixFirst s1)
2949data ICALState s1 s2 i1 i2 a =
2951 -- | ICALFirstYield s1 s2 i1
2952 | ICALFirstInner s1 s2 i1
2954 | ICALFirstOnlyInner s1 i1
2955 | ICALSecondInject s1 s2
2956 | ICALFirstInject s1 s2 i2
2957 -- | ICALFirstBuf s1 s2 i1 i2
2958 | ICALSecondInner s1 s2 i1 i2
2959 -- -- | ICALSecondInner s1 s2 i1 i2 a
2960 -- -- | ICALFirstResume s1 s2 i1 i2 a
2962-- | Interleave streams (full streams, not the elements) unfolded from two
2963-- input streams and concat. Stop when the first stream stops. If the second
2964-- stream ends before the first one then first stream still keeps running alone
2965-- without any interleaving with the second stream.
2967-- [a1, a2, ... an] [b1, b2 ...]
2968-- => [streamA1, streamA2, ... streamAn] [streamB1, streamB2, ...]
2969-- => [streamA1, streamB1, streamA2...StreamAn, streamBn]
2970-- => [a11, a12, ...a1j, b11, b12, ...b1k, a21, a22, ...]
2972{-# INLINE_NORMAL gintercalate #-}
2975 => Unfold m a c -> Stream m a -> Unfold m b c -> Stream m b -> Stream m c
2977 (Unfold istep1 inject1) (Stream step1 state1)
2978 (Unfold istep2 inject2) (Stream step2 state2) =
2979 Stream step (ICALFirst state1 state2)
2983 {-# INLINE_LATE step #-}
2984 step gst (ICALFirst s1 s2) = do
2985 r <- step1 (adaptState gst) s1
2989 i `seq` return (Skip (ICALFirstInner s s2 i))
2990 -- i `seq` return (Skip (ICALFirstYield s s2 i))
2991 Skip s -> return $ Skip (ICALFirst s s2)
2995 step _ (ICALFirstYield s1 s2 i1) = do
2998 Yield x i' -> Yield x (ICALFirstInner s1 s2 i')
2999 Skip i' -> Skip (ICALFirstYield s1 s2 i')
3000 Stop -> Skip (ICALFirst s1 s2)
3003 step _ (ICALFirstInner s1 s2 i1) = do
3006 Yield x i' -> Yield x (ICALFirstInner s1 s2 i')
3007 Skip i' -> Skip (ICALFirstInner s1 s2 i')
3008 Stop -> Skip (ICALSecondInject s1 s2)
3010 step gst (ICALFirstOnly s1) = do
3011 r <- step1 (adaptState gst) s1
3015 i `seq` return (Skip (ICALFirstOnlyInner s i))
3016 Skip s -> return $ Skip (ICALFirstOnly s)
3019 step _ (ICALFirstOnlyInner s1 i1) = do
3022 Yield x i' -> Yield x (ICALFirstOnlyInner s1 i')
3023 Skip i' -> Skip (ICALFirstOnlyInner s1 i')
3024 Stop -> Skip (ICALFirstOnly s1)
3026 -- We inject the second stream even before checking if the first stream
3027 -- would yield any more elements. There is no clear choice whether we
3028 -- should do this before or after that. Doing it after may make the state
3029 -- machine a bit simpler though.
3030 step gst (ICALSecondInject s1 s2) = do
3031 r <- step2 (adaptState gst) s2
3035 i `seq` return (Skip (ICALFirstInject s1 s i))
3036 Skip s -> return $ Skip (ICALSecondInject s1 s)
3037 Stop -> return $ Skip (ICALFirstOnly s1)
3039 step gst (ICALFirstInject s1 s2 i2) = do
3040 r <- step1 (adaptState gst) s1
3044 i `seq` return (Skip (ICALSecondInner s s2 i i2))
3045 -- i `seq` return (Skip (ICALFirstBuf s s2 i i2))
3046 Skip s -> return $ Skip (ICALFirstInject s s2 i2)
3050 step _ (ICALFirstBuf s1 s2 i1 i2) = do
3053 Yield x i' -> Skip (ICALSecondInner s1 s2 i' i2 x)
3054 Skip i' -> Skip (ICALFirstBuf s1 s2 i' i2)
3057 step _ (ICALSecondInner s1 s2 i1 i2 v) = do
3060 Yield x i' -> Yield x (ICALSecondInner s1 s2 i1 i' v)
3061 Skip i' -> Skip (ICALSecondInner s1 s2 i1 i' v)
3062 Stop -> Skip (ICALFirstResume s1 s2 i1 i2 v)
3065 step _ (ICALSecondInner s1 s2 i1 i2) = do
3068 Yield x i' -> Yield x (ICALSecondInner s1 s2 i1 i')
3069 Skip i' -> Skip (ICALSecondInner s1 s2 i1 i')
3070 Stop -> Skip (ICALFirstInner s1 s2 i1)
3071 -- Stop -> Skip (ICALFirstResume s1 s2 i1 i2)
3074 step _ (ICALFirstResume s1 s2 i1 i2 x) = do
3075 return $ Yield x (ICALFirstInner s1 s2 i1 i2)
3078data InterposeState s1 i1 a =
3080 -- | InterposeFirstYield s1 i1
3081 | InterposeFirstInner s1 i1
3082 | InterposeFirstInject s1
3083 -- | InterposeFirstBuf s1 i1
3084 | InterposeSecondYield s1 i1
3085 -- -- | InterposeSecondYield s1 i1 a
3086 -- -- | InterposeFirstResume s1 i1 a
3088-- Note that this only interposes the pure values, we may run many effects to
3089-- generate those values as some effects may not generate anything (Skip).
3090{-# INLINE_NORMAL interpose #-}
3091interpose :: Monad m => m c -> Unfold m b c -> Stream m b -> Stream m c
3094 (Unfold istep1 inject1) (Stream step1 state1) =
3095 Stream step (InterposeFirst state1)
3099 {-# INLINE_LATE step #-}
3100 step gst (InterposeFirst s1) = do
3101 r <- step1 (adaptState gst) s1
3105 i `seq` return (Skip (InterposeFirstInner s i))
3106 -- i `seq` return (Skip (InterposeFirstYield s i))
3107 Skip s -> return $ Skip (InterposeFirst s)
3111 step _ (InterposeFirstYield s1 i1) = do
3114 Yield x i' -> Yield x (InterposeFirstInner s1 i')
3115 Skip i' -> Skip (InterposeFirstYield s1 i')
3116 Stop -> Skip (InterposeFirst s1)
3119 step _ (InterposeFirstInner s1 i1) = do
3122 Yield x i' -> Yield x (InterposeFirstInner s1 i')
3123 Skip i' -> Skip (InterposeFirstInner s1 i')
3124 Stop -> Skip (InterposeFirstInject s1)
3126 step gst (InterposeFirstInject s1) = do
3127 r <- step1 (adaptState gst) s1
3131 -- i `seq` return (Skip (InterposeFirstBuf s i))
3132 i `seq` return (Skip (InterposeSecondYield s i))
3133 Skip s -> return $ Skip (InterposeFirstInject s)
3137 step _ (InterposeFirstBuf s1 i1) = do
3140 Yield x i' -> Skip (InterposeSecondYield s1 i' x)
3141 Skip i' -> Skip (InterposeFirstBuf s1 i')
3146 step _ (InterposeSecondYield s1 i1 v) = do
3148 return $ Yield r (InterposeFirstResume s1 i1 v)
3150 step _ (InterposeSecondYield s1 i1) = do
3152 return $ Yield r (InterposeFirstInner s1 i1)
3155 step _ (InterposeFirstResume s1 i1 v) = do
3156 return $ Yield v (InterposeFirstInner s1 i1)
3159------------------------------------------------------------------------------
3161------------------------------------------------------------------------------
3163data GbracketState s1 s2 v
3165 | GBracketNormal s1 v
3166 | GBracketException s2
3168-- | The most general bracketing and exception combinator. All other
3169-- combinators can be expressed in terms of this combinator. This can also be
3170-- used for cases which are not covered by the standard combinators.
3174{-# INLINE_NORMAL gbracket #-}
3178 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
3179 -> (c -> m d) -- ^ after, on normal stop
3180 -> (c -> e -> Stream m b) -- ^ on exception
3181 -> (c -> Stream m b) -- ^ stream generator
3183gbracket bef exc aft fexc fnormal =
3184 Stream step GBracketInit
3188 {-# INLINE_LATE step #-}
3189 step _ GBracketInit = do
3191 return $ Skip $ GBracketNormal (fnormal r) r
3193 step gst (GBracketNormal (UnStream step1 st) v) = do
3194 res <- exc $ step1 gst st
3196 Right r -> case r of
3198 return $ Yield x (GBracketNormal (Stream step1 s) v)
3199 Skip s -> return $ Skip (GBracketNormal (Stream step1 s) v)
3200 Stop -> aft v >> return Stop
3201 Left e -> return $ Skip (GBracketException (fexc v e))
3202 step gst (GBracketException (UnStream step1 st)) = do
3205 Yield x s -> return $ Yield x (GBracketException (Stream step1 s))
3206 Skip s -> return $ Skip (GBracketException (Stream step1 s))
3209-- | Create an IORef holding a finalizer that is called automatically when the
3210-- IORef is garbage collected. The IORef can be written to with a 'Nothing'
3211-- value to deactivate the finalizer.
3212newFinalizedIORef :: (MonadIO m, MonadBaseControl IO m)
3213 => m a -> m (IORef (Maybe (IO ())))
3214newFinalizedIORef finalizer = do
3215 mrun <- captureMonadState
3216 ref <- liftIO $ newIORef $ Just $ liftIO $ void $ do
3217 _ <- runInIO mrun finalizer
3220 res <- readIORef ref
3222 Nothing -> return ()
3224 _ <- liftIO $ mkWeakIORef ref finalizer1
3227-- | Run the finalizer stored in an IORef and deactivate it so that it is run
3230runIORefFinalizer :: MonadIO m => IORef (Maybe (IO ())) -> m ()
3231runIORefFinalizer ref = liftIO $ do
3232 res <- readIORef ref
3234 Nothing -> return ()
3235 Just f -> writeIORef ref Nothing >> f
3237-- | Deactivate the finalizer stored in an IORef without running it.
3239clearIORefFinalizer :: MonadIO m => IORef (Maybe (IO ())) -> m ()
3240clearIORefFinalizer ref = liftIO $ writeIORef ref Nothing
3242data GbracketIOState s1 s2 v wref
3244 | GBracketIONormal s1 v wref
3245 | GBracketIOException s2
3247-- | Like gbracket but also uses a finalizer to make sure when the stream is
3248-- garbage collected we run the finalizing action. This requires a MonadIO and
3249-- MonadBaseControl IO constraint.
3251-- | The most general bracketing and exception combinator. All other
3252-- combinators can be expressed in terms of this combinator. This can also be
3253-- used for cases which are not covered by the standard combinators.
3257{-# INLINE_NORMAL gbracketIO #-}
3259 :: (MonadIO m, MonadBaseControl IO m)
3261 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
3262 -> (c -> m d) -- ^ after, on normal stop or GC
3263 -> (c -> e -> Stream m b) -- ^ on exception
3264 -> (c -> Stream m b) -- ^ stream generator
3266gbracketIO bef exc aft fexc fnormal =
3267 Stream step GBracketIOInit
3271 -- If the stream is never evaluated the "aft" action will never be
3272 -- called. For that to occur we will need the user of this API to pass a
3273 -- weak pointer to us.
3274 {-# INLINE_LATE step #-}
3275 step _ GBracketIOInit = do
3276 -- We mask asynchronous exceptions to make the execution
3277 -- of 'bef' and the registration of 'aft' atomic.
3278 -- A similar thing is done in the resourcet package: https://git.io/JvKV3
3279 -- Tutorial: https://markkarpov.com/tutorial/exceptions.html
3280 (r, ref) <- liftBaseOp_ mask_ $ do
3282 ref <- newFinalizedIORef (aft r)
3284 return $ Skip $ GBracketIONormal (fnormal r) r ref
3286 step gst (GBracketIONormal (UnStream step1 st) v ref) = do
3287 res <- exc $ step1 gst st
3289 Right r -> case r of
3291 return $ Yield x (GBracketIONormal (Stream step1 s) v ref)
3293 return $ Skip (GBracketIONormal (Stream step1 s) v ref)
3295 runIORefFinalizer ref
3298 clearIORefFinalizer ref
3299 return $ Skip (GBracketIOException (fexc v e))
3300 step gst (GBracketIOException (UnStream step1 st)) = do
3304 return $ Yield x (GBracketIOException (Stream step1 s))
3305 Skip s -> return $ Skip (GBracketIOException (Stream step1 s))
3308-- | Run a side effect before the stream yields its first element.
3309{-# INLINE_NORMAL before #-}
3310before :: Monad m => m b -> Stream m a -> Stream m a
3311before action (Stream step state) = Stream step' Nothing
3315 {-# INLINE_LATE step' #-}
3316 step' _ Nothing = action >> return (Skip (Just state))
3318 step' gst (Just st) = do
3321 Yield x s -> return $ Yield x (Just s)
3322 Skip s -> return $ Skip (Just s)
3325-- | Run a side effect whenever the stream stops normally.
3326{-# INLINE_NORMAL after #-}
3327after :: Monad m => m b -> Stream m a -> Stream m a
3328after action (Stream step state) = Stream step' state
3332 {-# INLINE_LATE step' #-}
3336 Yield x s -> return $ Yield x s
3337 Skip s -> return $ Skip s
3338 Stop -> action >> return Stop
3340{-# INLINE_NORMAL afterIO #-}
3341afterIO :: (MonadIO m, MonadBaseControl IO m)
3342 => m b -> Stream m a -> Stream m a
3343afterIO action (Stream step state) = Stream step' Nothing
3347 {-# INLINE_LATE step' #-}
3348 step' _ Nothing = do
3349 ref <- newFinalizedIORef action
3350 return $ Skip $ Just (state, ref)
3351 step' gst (Just (st, ref)) = do
3354 Yield x s -> return $ Yield x (Just (s, ref))
3355 Skip s -> return $ Skip (Just (s, ref))
3357 runIORefFinalizer ref
3360-- XXX These combinators are expensive due to the call to
3361-- onException/handle/try on each step. Therefore, when possible, they should
3362-- be called in an outer loop where we perform less iterations. For example, we
3363-- cannot call them on each iteration in a char stream, instead we can call
3364-- them when doing an IO on an array.
3366-- XXX For high performance error checks in busy streams we may need another
3367-- Error constructor in step.
3369-- | Run a side effect whenever the stream aborts due to an exception. The
3370-- exception is not caught, simply rethrown.
3371{-# INLINE_NORMAL onException #-}
3372onException :: MonadCatch m => m b -> Stream m a -> Stream m a
3373onException action str =
3374 gbracket (return ()) MC.try return
3375 (\_ (e :: MC.SomeException) -> nilM (action >> MC.throwM e))
3378{-# INLINE_NORMAL _onException #-}
3379_onException :: MonadCatch m => m b -> Stream m a -> Stream m a
3380_onException action (Stream step state) = Stream step' state
3384 {-# INLINE_LATE step' #-}
3386 res <- step gst st `MC.onException` action
3388 Yield x s -> return $ Yield x s
3389 Skip s -> return $ Skip s
3392-- XXX bracket is like concatMap, it generates a stream and then flattens it.
3393-- Like concatMap it has 10x worse performance compared to linear fused
3396-- | Run the first action before the stream starts and remember its output,
3397-- generate a stream using the output, run the second action providing the
3398-- remembered value as an argument whenever the stream ends normally or due to
3400{-# INLINE_NORMAL bracket #-}
3401bracket :: MonadCatch m => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
3402bracket bef aft bet =
3403 gbracket bef MC.try aft
3404 (\a (e :: SomeException) -> nilM (aft a >> MC.throwM e)) bet
3406{-# INLINE_NORMAL bracketIO #-}
3407bracketIO :: (MonadAsync m, MonadCatch m)
3408 => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
3409bracketIO bef aft bet =
3410 gbracketIO bef MC.try aft
3411 (\a (e :: SomeException) -> nilM (aft a >> MC.throwM e)) bet
3413data BracketState s v = BracketInit | BracketRun s v
3415{-# INLINE_NORMAL _bracket #-}
3416_bracket :: MonadCatch m => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
3417_bracket bef aft bet = Stream step' BracketInit
3421 {-# INLINE_LATE step' #-}
3422 step' _ BracketInit = bef >>= \x -> return (Skip (BracketRun (bet x) x))
3424 -- NOTE: It is important to use UnStream instead of the Stream pattern
3425 -- here, otherwise we get huge perf degradation, see note in concatMap.
3426 step' gst (BracketRun (UnStream step state) v) = do
3427 -- res <- step gst state `MC.onException` aft v
3428 res <- MC.try $ step gst state
3430 Left (e :: SomeException) -> aft v >> MC.throwM e >> return Stop
3431 Right r -> case r of
3432 Yield x s -> return $ Yield x (BracketRun (Stream step s) v)
3433 Skip s -> return $ Skip (BracketRun (Stream step s) v)
3434 Stop -> aft v >> return Stop
3436-- | Run a side effect whenever the stream stops normally or aborts due to an
3438{-# INLINE finally #-}
3439finally :: MonadCatch m => m b -> Stream m a -> Stream m a
3440-- finally action xs = after action $ onException action xs
3441finally action xs = bracket (return ()) (\_ -> action) (const xs)
3443{-# INLINE finallyIO #-}
3444finallyIO :: (MonadAsync m, MonadCatch m) => m b -> Stream m a -> Stream m a
3445finallyIO action xs = bracketIO (return ()) (\_ -> action) (const xs)
3447-- | When evaluating a stream if an exception occurs, stream evaluation aborts
3448-- and the specified exception handler is run with the exception as argument.
3449{-# INLINE_NORMAL handle #-}
3450handle :: (MonadCatch m, Exception e)
3451 => (e -> Stream m a) -> Stream m a -> Stream m a
3453 gbracket (return ()) MC.try return (\_ e -> f e) (\_ -> str)
3455{-# INLINE_NORMAL _handle #-}
3456_handle :: (MonadCatch m, Exception e)
3457 => (e -> Stream m a) -> Stream m a -> Stream m a
3458_handle f (Stream step state) = Stream step' (Left state)
3462 {-# INLINE_LATE step' #-}
3463 step' gst (Left st) = do
3464 res <- MC.try $ step gst st
3466 Left e -> return $ Skip $ Right (f e)
3467 Right r -> case r of
3468 Yield x s -> return $ Yield x (Left s)
3469 Skip s -> return $ Skip (Left s)
3472 step' gst (Right (UnStream step1 st)) = do
3475 Yield x s -> return $ Yield x (Right (Stream step1 s))
3476 Skip s -> return $ Skip (Right (Stream step1 s))
3479-------------------------------------------------------------------------------
3480-- General transformation
3481-------------------------------------------------------------------------------
3483{-# INLINE_NORMAL transform #-}
3484transform :: Monad m => Pipe m a b -> Stream m a -> Stream m b
3485transform (Pipe pstep1 pstep2 pstate) (Stream step state) =
3486 Stream step' (Consume pstate, state)
3490 {-# INLINE_LATE step' #-}
3492 step' gst (Consume pst, st) = pst `seq` do
3493 r <- step (adaptState gst) st
3498 Pipe.Yield b pst' -> return $ Yield b (pst', s)
3499 Pipe.Continue pst' -> return $ Skip (pst', s)
3500 Skip s -> return $ Skip (Consume pst, s)
3503 step' _ (Produce pst, st) = pst `seq` do
3506 Pipe.Yield b pst' -> return $ Yield b (pst', st)
3507 Pipe.Continue pst' -> return $ Skip (pst', st)
3509------------------------------------------------------------------------------
3510-- Transformation by Folding (Scans)
3511------------------------------------------------------------------------------
3513------------------------------------------------------------------------------
3515------------------------------------------------------------------------------
3517-- XXX Is a prescan useful, discarding the last step does not sound useful? I
3518-- am not sure about the utility of this function, so this is implemented but
3519-- not exposed. We can expose it if someone provides good reasons why this is
3522-- XXX We have to execute the stream one step ahead to know that we are at the
3523-- last step. The vector implementation of prescan executes the last fold step
3524-- but does not yield the result. This means we have executed the effect but
3525-- discarded value. This does not sound right. In this implementation we are
3526-- not executing the last fold step.
3527{-# INLINE_NORMAL prescanlM' #-}
3528prescanlM' :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
3529prescanlM' f mz (Stream step state) = Stream step' (state, mz)
3531 {-# INLINE_LATE step' #-}
3532 step' gst (st, prev) = do
3533 r <- step (adaptState gst) st
3537 return $ Yield acc (s, f acc x)
3538 Skip s -> return $ Skip (s, prev)
3541{-# INLINE prescanl' #-}
3542prescanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
3543prescanl' f z = prescanlM' (\a b -> return (f a b)) (return z)
3545------------------------------------------------------------------------------
3546-- Monolithic postscans (postscan followed by a map)
3547------------------------------------------------------------------------------
3549-- The performance of a modular postscan followed by a map seems to be
3550-- equivalent to this monolithic scan followed by map therefore we may not need
3551-- this implementation. We just have it for performance comparison and in case
3552-- modular version does not perform well in some situation.
3554{-# INLINE_NORMAL postscanlMx' #-}
3555postscanlMx' :: Monad m
3556 => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> Stream m b
3557postscanlMx' fstep begin done (Stream step state) = do
3558 Stream step' (state, begin)
3560 {-# INLINE_LATE step' #-}
3561 step' gst (st, acc) = do
3562 r <- step (adaptState gst) st
3568 v `seq` y `seq` return (Yield v (s, return y))
3569 Skip s -> return $ Skip (s, acc)
3572{-# INLINE_NORMAL postscanlx' #-}
3573postscanlx' :: Monad m
3574 => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> Stream m b
3575postscanlx' fstep begin done s =
3576 postscanlMx' (\b a -> return (fstep b a)) (return begin) (return . done) s
3578-- XXX do we need consM strict to evaluate the begin value?
3579{-# INLINE scanlMx' #-}
3581 => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> Stream m b
3582scanlMx' fstep begin done s =
3583 (begin >>= \x -> x `seq` done x) `consM` postscanlMx' fstep begin done s
3585{-# INLINE scanlx' #-}
3587 => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> Stream m b
3588scanlx' fstep begin done s =
3589 scanlMx' (\b a -> return (fstep b a)) (return begin) (return . done) s
3591------------------------------------------------------------------------------
3593------------------------------------------------------------------------------
3595{-# INLINE_NORMAL postscanlM' #-}
3596postscanlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
3597postscanlM' fstep begin (Stream step state) =
3598 begin `seq` Stream step' (state, begin)
3600 {-# INLINE_LATE step' #-}
3601 step' gst (st, acc) = acc `seq` do
3602 r <- step (adaptState gst) st
3606 y `seq` return (Yield y (s, y))
3607 Skip s -> return $ Skip (s, acc)
3610{-# INLINE_NORMAL postscanl' #-}
3611postscanl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
3612postscanl' f = postscanlM' (\a b -> return (f a b))
3614-- We can possibly have the "done" function as a Maybe to provide an option to
3615-- emit or not emit the accumulator when the stream stops.
3617-- TBD: use a single Yield point
3619{-# INLINE_NORMAL postscanlMAfter' #-}
3620postscanlMAfter' :: Monad m
3621 => (b -> a -> m b) -> m b -> (b -> m b) -> Stream m a -> Stream m b
3622postscanlMAfter' fstep initial done (Stream step1 state1) = do
3623 Stream step (Just (state1, initial))
3627 {-# INLINE_LATE step #-}
3628 step gst (Just (st, acc)) = do
3629 r <- step1 (adaptState gst) st
3634 y `seq` return (Yield y (Just (s, return y)))
3635 Skip s -> return $ Skip $ Just (s, acc)
3636 Stop -> acc >>= done >>= \res -> return (Yield res Nothing)
3637 step _ Nothing = return Stop
3639{-# INLINE_NORMAL postscanlM #-}
3640postscanlM :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
3641postscanlM fstep begin (Stream step state) = Stream step' (state, begin)
3643 {-# INLINE_LATE step' #-}
3644 step' gst (st, acc) = do
3645 r <- step (adaptState gst) st
3649 return (Yield y (s, y))
3650 Skip s -> return $ Skip (s, acc)
3653{-# INLINE_NORMAL postscanl #-}
3654postscanl :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
3655postscanl f = postscanlM (\a b -> return (f a b))
3657{-# INLINE_NORMAL scanlM' #-}
3658scanlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
3659scanlM' fstep begin s = begin `seq` (begin `cons` postscanlM' fstep begin s)
3661{-# INLINE scanlMAfter' #-}
3662scanlMAfter' :: Monad m
3663 => (b -> a -> m b) -> m b -> (b -> m b) -> Stream m a -> Stream m b
3664scanlMAfter' fstep initial done s =
3665 (initial >>= \x -> x `seq` return x) `consM`
3666 postscanlMAfter' fstep initial done s
3668{-# INLINE scanl' #-}
3669scanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
3670scanl' f = scanlM' (\a b -> return (f a b))
3672{-# INLINE_NORMAL scanlM #-}
3673scanlM :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
3674scanlM fstep begin s = begin `cons` postscanlM fstep begin s
3677scanl :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
3678scanl f = scanlM (\a b -> return (f a b))
3680{-# INLINE_NORMAL scanl1M #-}
3681scanl1M :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
3682scanl1M fstep (Stream step state) = Stream step' (state, Nothing)
3684 {-# INLINE_LATE step' #-}
3685 step' gst (st, Nothing) = do
3688 Yield x s -> return $ Yield x (s, Just x)
3689 Skip s -> return $ Skip (s, Nothing)
3692 step' gst (st, Just acc) = do
3697 return $ Yield z (s, Just z)
3698 Skip s -> return $ Skip (s, Just acc)
3701{-# INLINE scanl1 #-}
3702scanl1 :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
3703scanl1 f = scanl1M (\x y -> return (f x y))
3705{-# INLINE_NORMAL scanl1M' #-}
3706scanl1M' :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
3707scanl1M' fstep (Stream step state) = Stream step' (state, Nothing)
3709 {-# INLINE_LATE step' #-}
3710 step' gst (st, Nothing) = do
3713 Yield x s -> x `seq` return $ Yield x (s, Just x)
3714 Skip s -> return $ Skip (s, Nothing)
3717 step' gst (st, Just acc) = acc `seq` do
3722 z `seq` return $ Yield z (s, Just z)
3723 Skip s -> return $ Skip (s, Just acc)
3726{-# INLINE scanl1' #-}
3727scanl1' :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
3728scanl1' f = scanl1M' (\x y -> return (f x y))
3730------------------------------------------------------------------------------
3732------------------------------------------------------------------------------
3734data RollingMapState s a = RollingMapInit s | RollingMapGo s a
3736{-# INLINE rollingMapM #-}
3737rollingMapM :: Monad m => (a -> a -> m b) -> Stream m a -> Stream m b
3738rollingMapM f (Stream step1 state1) = Stream step (RollingMapInit state1)
3740 step gst (RollingMapInit st) = do
3741 r <- step1 (adaptState gst) st
3743 Yield x s -> Skip $ RollingMapGo s x
3744 Skip s -> Skip $ RollingMapInit s
3747 step gst (RollingMapGo s1 x1) = do
3748 r <- step1 (adaptState gst) s1
3752 return $ Yield res $ RollingMapGo s x
3753 Skip s -> return $ Skip $ RollingMapGo s x1
3754 Stop -> return $ Stop
3756{-# INLINE rollingMap #-}
3757rollingMap :: Monad m => (a -> a -> b) -> Stream m a -> Stream m b
3758rollingMap f = rollingMapM (\x y -> return $ f x y)
3760------------------------------------------------------------------------------
3761-- Tapping/Distributing
3762------------------------------------------------------------------------------
3765tap :: Monad m => Fold m a b -> Stream m a -> Stream m a
3766tap (Fold fstep initial extract) (Stream step state) = Stream step' Nothing
3770 step' _ Nothing = do
3772 return $ Skip (Just (r, state))
3774 step' gst (Just (acc, st)) = acc `seq` do
3779 return $ Yield x (Just (acc', s))
3780 Skip s -> return $ Skip (Just (acc, s))
3785{-# INLINE_NORMAL tapOffsetEvery #-}
3786tapOffsetEvery :: Monad m
3787 => Int -> Int -> Fold m a b -> Stream m a -> Stream m a
3788tapOffsetEvery offset n (Fold fstep initial extract) (Stream step state) =
3789 Stream step' Nothing
3793 {-# INLINE_LATE step' #-}
3794 step' _ Nothing = do
3796 return $ Skip (Just (r, state, offset `mod` n))
3798 step' gst (Just (acc, st, count)) | count <= 0 = do
3802 !acc' <- fstep acc x
3803 return $ Yield x (Just (acc', s, n - 1))
3804 Skip s -> return $ Skip (Just (acc, s, count))
3809 step' gst (Just (acc, st, count)) = do
3812 Yield x s -> return $ Yield x (Just (acc, s, count - 1))
3813 Skip s -> return $ Skip (Just (acc, s, count))
3818{-# INLINE_NORMAL pollCounts #-}
3822 -> (Stream m Int -> Stream m Int)
3826pollCounts predicate transf fld (Stream step state) = Stream step' Nothing
3829 {-# INLINE_LATE step' #-}
3830 step' _ Nothing = do
3831 -- As long as we are using an "Int" for counts lockfree reads from
3832 -- Var should work correctly on both 32-bit and 64-bit machines.
3833 -- However, an Int on a 32-bit machine may overflow quickly.
3834 countVar <- liftIO $ newVar (0 :: Int)
3836 $ void $ runFold fld
3837 $ transf $ fromPrimVar countVar
3838 return $ Skip (Just (countVar, tid, state))
3840 step' gst (Just (countVar, tid, st)) = do
3844 when (predicate x) $ liftIO $ modifyVar' countVar (+ 1)
3845 return $ Yield x (Just (countVar, tid, s))
3846 Skip s -> return $ Skip (Just (countVar, tid, s))
3848 liftIO $ killThread tid
3851{-# INLINE_NORMAL tapRate #-}
3853 (MonadAsync m, MonadCatch m)
3858tapRate samplingRate action (Stream step state) = Stream step' Nothing
3860 {-# NOINLINE loop #-}
3861 loop countVar prev = do
3864 (do liftIO $ threadDelay (round $ samplingRate * 1000000)
3865 i <- liftIO $ readVar countVar
3866 let !diff = i - prev
3869 (\(e :: AsyncException) -> do
3870 i <- liftIO $ readVar countVar
3871 let !diff = i - prev
3873 throwM (MC.toException e))
3876 {-# INLINE_LATE step' #-}
3877 step' _ Nothing = do
3878 countVar <- liftIO $ newVar 0
3879 tid <- fork $ loop countVar 0
3880 ref <- liftIO $ newIORef ()
3881 _ <- liftIO $ mkWeakIORef ref (killThread tid)
3882 return $ Skip (Just (countVar, tid, state, ref))
3884 step' gst (Just (countVar, tid, st, ref)) = do
3888 liftIO $ modifyVar' countVar (+ 1)
3889 return $ Yield x (Just (countVar, tid, s, ref))
3890 Skip s -> return $ Skip (Just (countVar, tid, s, ref))
3892 liftIO $ killThread tid
3896-------------------------------------------------------------------------------
3898-------------------------------------------------------------------------------
3900{-# INLINE_NORMAL takeWhileM #-}
3901takeWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
3902takeWhileM f (Stream step state) = Stream step' state
3904 {-# INLINE_LATE step' #-}
3910 return $ if b then Yield x s else Stop
3911 Skip s -> return $ Skip s
3914{-# INLINE takeWhile #-}
3915takeWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
3916takeWhile f = takeWhileM (return . f)
3918{-# INLINE_NORMAL drop #-}
3919drop :: Monad m => Int -> Stream m a -> Stream m a
3920drop n (Stream step state) = Stream step' (state, Just n)
3922 {-# INLINE_LATE step' #-}
3923 step' gst (st, Just i)
3928 Yield _ s -> Skip (s, Just (i - 1))
3929 Skip s -> Skip (s, Just i)
3931 | otherwise = return $ Skip (st, Nothing)
3933 step' gst (st, Nothing) = do
3937 Yield x s -> Yield x (s, Nothing)
3938 Skip s -> Skip (s, Nothing)
3941data DropWhileState s a
3943 | DropWhileYield a s
3946{-# INLINE_NORMAL dropWhileM #-}
3947dropWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
3948dropWhileM f (Stream step state) = Stream step' (DropWhileDrop state)
3950 {-# INLINE_LATE step' #-}
3951 step' gst (DropWhileDrop st) = do
3957 then return $ Skip (DropWhileDrop s)
3958 else return $ Skip (DropWhileYield x s)
3959 Skip s -> return $ Skip (DropWhileDrop s)
3962 step' gst (DropWhileNext st) = do
3965 Yield x s -> return $ Skip (DropWhileYield x s)
3966 Skip s -> return $ Skip (DropWhileNext s)
3969 step' _ (DropWhileYield x st) = return $ Yield x (DropWhileNext st)
3971{-# INLINE dropWhile #-}
3972dropWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
3973dropWhile f = dropWhileM (return . f)
3975{-# INLINE_NORMAL filterM #-}
3976filterM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
3977filterM f (Stream step state) = Stream step' state
3979 {-# INLINE_LATE step' #-}
3988 Skip s -> return $ Skip s
3991{-# INLINE filter #-}
3992filter :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
3993filter f = filterM (return . f)
3995{-# INLINE_NORMAL uniq #-}
3996uniq :: (Eq a, Monad m) => Stream m a -> Stream m a
3997uniq (Stream step state) = Stream step' (Nothing, state)
3999 {-# INLINE_LATE step' #-}
4000 step' gst (Nothing, st) = do
4003 Yield x s -> return $ Yield x (Just x, s)
4004 Skip s -> return $ Skip (Nothing, s)
4006 step' gst (Just x, st) = do
4009 Yield y s | x == y -> return $ Skip (Just x, s)
4010 | otherwise -> return $ Yield y (Just y, s)
4011 Skip s -> return $ Skip (Just x, s)
4014------------------------------------------------------------------------------
4015-- Transformation by Mapping
4016------------------------------------------------------------------------------
4018{-# INLINE_NORMAL sequence #-}
4019sequence :: Monad m => Stream m (m a) -> Stream m a
4020sequence (Stream step state) = Stream step' state
4022 {-# INLINE_LATE step' #-}
4024 r <- step (adaptState gst) st
4026 Yield x s -> x >>= \a -> return (Yield a s)
4027 Skip s -> return $ Skip s
4030------------------------------------------------------------------------------
4032------------------------------------------------------------------------------
4034data LoopState x s = FirstYield s
4035 | InterspersingYield s
4038{-# INLINE_NORMAL intersperseM #-}
4039intersperseM :: Monad m => m a -> Stream m a -> Stream m a
4040intersperseM m (Stream step state) = Stream step' (FirstYield state)
4042 {-# INLINE_LATE step' #-}
4043 step' gst (FirstYield st) = do
4047 Yield x s -> Skip (YieldAndCarry x s)
4048 Skip s -> Skip (FirstYield s)
4051 step' gst (InterspersingYield st) = do
4056 return $ Yield a (YieldAndCarry x s)
4057 Skip s -> return $ Skip $ InterspersingYield s
4060 step' _ (YieldAndCarry x st) = return $ Yield x (InterspersingYield st)
4062{-# INLINE_NORMAL intersperseM_ #-}
4063intersperseM_ :: Monad m => m b -> Stream m a -> Stream m a
4064intersperseM_ m (Stream step1 state1) = Stream step (Left (pure (), state1))
4066 {-# INLINE_LATE step #-}
4067 step gst (Left (eff, st)) = do
4070 Yield x s -> eff >> return (Yield x (Right s))
4071 Skip s -> return $ Skip (Left (eff, s))
4074 step _ (Right st) = return $ Skip $ Left (void m, st)
4079 | SuffixYield a (SuffixState s a)
4081{-# INLINE_NORMAL intersperseSuffix #-}
4082intersperseSuffix :: forall m a. Monad m => m a -> Stream m a -> Stream m a
4083intersperseSuffix action (Stream step state) = Stream step' (SuffixElem state)
4085 {-# INLINE_LATE step' #-}
4086 step' gst (SuffixElem st) = do
4089 Yield x s -> Skip (SuffixYield x (SuffixSuffix s))
4090 Skip s -> Skip (SuffixElem s)
4093 step' _ (SuffixSuffix st) = do
4094 action >>= \r -> return $ Skip (SuffixYield r (SuffixElem st))
4096 step' _ (SuffixYield x next) = return $ Yield x next
4098{-# INLINE_NORMAL intersperseSuffix_ #-}
4099intersperseSuffix_ :: Monad m => m b -> Stream m a -> Stream m a
4100intersperseSuffix_ m (Stream step1 state1) = Stream step (Left state1)
4102 {-# INLINE_LATE step #-}
4103 step gst (Left st) = do
4106 Yield x s -> return $ Yield x (Right s)
4107 Skip s -> return $ Skip $ Left s
4110 step _ (Right st) = m >> return (Skip (Left st))
4112data SuffixSpanState s a
4113 = SuffixSpanElem s Int
4114 | SuffixSpanSuffix s
4115 | SuffixSpanYield a (SuffixSpanState s a)
4119-- | intersperse after every n items
4120{-# INLINE_NORMAL intersperseSuffixBySpan #-}
4121intersperseSuffixBySpan :: forall m a. Monad m
4122 => Int -> m a -> Stream m a -> Stream m a
4123intersperseSuffixBySpan n action (Stream step state) =
4124 Stream step' (SuffixSpanElem state n)
4126 {-# INLINE_LATE step' #-}
4127 step' gst (SuffixSpanElem st i) | i > 0 = do
4130 Yield x s -> Skip (SuffixSpanYield x (SuffixSpanElem s (i - 1)))
4131 Skip s -> Skip (SuffixSpanElem s i)
4132 Stop -> if i == n then Stop else Skip SuffixSpanLast
4133 step' _ (SuffixSpanElem st _) = return $ Skip (SuffixSpanSuffix st)
4135 step' _ (SuffixSpanSuffix st) = do
4136 action >>= \r -> return $ Skip (SuffixSpanYield r (SuffixSpanElem st n))
4138 step' _ (SuffixSpanLast) = do
4139 action >>= \r -> return $ Skip (SuffixSpanYield r SuffixSpanStop)
4141 step' _ (SuffixSpanYield x next) = return $ Yield x next
4143 step' _ (SuffixSpanStop) = return Stop
4145{-# INLINE intersperse #-}
4146intersperse :: Monad m => a -> Stream m a -> Stream m a
4147intersperse a = intersperseM (return a)
4149{-# INLINE_NORMAL insertBy #-}
4150insertBy :: Monad m => (a -> a -> Ordering) -> a -> Stream m a -> Stream m a
4151insertBy cmp a (Stream step state) = Stream step' (state, False, Nothing)
4153 {-# INLINE_LATE step' #-}
4154 step' gst (st, False, _) = do
4157 Yield x s -> case cmp a x of
4158 GT -> return $ Yield x (s, False, Nothing)
4159 _ -> return $ Yield a (s, True, Just x)
4160 Skip s -> return $ Skip (s, False, Nothing)
4161 Stop -> return $ Yield a (st, True, Nothing)
4163 step' _ (_, True, Nothing) = return Stop
4165 step' gst (st, True, Just prev) = do
4168 Yield x s -> return $ Yield prev (s, True, Just x)
4169 Skip s -> return $ Skip (s, True, Just prev)
4170 Stop -> return $ Yield prev (st, True, Nothing)
4172------------------------------------------------------------------------------
4174------------------------------------------------------------------------------
4176{-# INLINE_NORMAL deleteBy #-}
4177deleteBy :: Monad m => (a -> a -> Bool) -> a -> Stream m a -> Stream m a
4178deleteBy eq x (Stream step state) = Stream step' (state, False)
4180 {-# INLINE_LATE step' #-}
4181 step' gst (st, False) = do
4184 Yield y s -> return $
4185 if eq x y then Skip (s, True) else Yield y (s, False)
4186 Skip s -> return $ Skip (s, False)
4189 step' gst (st, True) = do
4192 Yield y s -> return $ Yield y (s, True)
4193 Skip s -> return $ Skip (s, True)
4196------------------------------------------------------------------------------
4197-- Transformation by Map and Filter
4198------------------------------------------------------------------------------
4200-- XXX Will this always fuse properly?
4201{-# INLINE_NORMAL mapMaybe #-}
4202mapMaybe :: Monad m => (a -> Maybe b) -> Stream m a -> Stream m b
4203mapMaybe f = fmap fromJust . filter isJust . map f
4205{-# INLINE_NORMAL mapMaybeM #-}
4206mapMaybeM :: Monad m => (a -> m (Maybe b)) -> Stream m a -> Stream m b
4207mapMaybeM f = fmap fromJust . filter isJust . mapM f
4209------------------------------------------------------------------------------
4211------------------------------------------------------------------------------
4213{-# INLINE_NORMAL indexed #-}
4214indexed :: Monad m => Stream m a -> Stream m (Int, a)
4215indexed (Stream step state) = Stream step' (state, 0)
4217 {-# INLINE_LATE step' #-}
4218 step' gst (st, i) = i `seq` do
4219 r <- step (adaptState gst) st
4221 Yield x s -> return $ Yield (i, x) (s, i+1)
4222 Skip s -> return $ Skip (s, i)
4225{-# INLINE_NORMAL indexedR #-}
4226indexedR :: Monad m => Int -> Stream m a -> Stream m (Int, a)
4227indexedR m (Stream step state) = Stream step' (state, m)
4229 {-# INLINE_LATE step' #-}
4230 step' gst (st, i) = i `seq` do
4231 r <- step (adaptState gst) st
4233 Yield x s -> let i' = i - 1
4234 in return $ Yield (i, x) (s, i')
4235 Skip s -> return $ Skip (s, i)
4238{-# INLINE_NORMAL zipWithM #-}
4240 => (a -> b -> m c) -> Stream m a -> Stream m b -> Stream m c
4241zipWithM f (Stream stepa ta) (Stream stepb tb) = Stream step (ta, tb, Nothing)
4243 {-# INLINE_LATE step #-}
4244 step gst (sa, sb, Nothing) = do
4245 r <- stepa (adaptState gst) sa
4248 Yield x sa' -> Skip (sa', sb, Just x)
4249 Skip sa' -> Skip (sa', sb, Nothing)
4252 step gst (sa, sb, Just x) = do
4253 r <- stepb (adaptState gst) sb
4257 return $ Yield z (sa, sb', Nothing)
4258 Skip sb' -> return $ Skip (sa, sb', Just x)
4261#if __GLASGOW_HASKELL__ >= 801
4262{-# RULES "zipWithM xs xs"
4263 forall f xs. zipWithM @Identity f xs xs = mapM (\x -> f x x) xs #-}
4266{-# INLINE zipWith #-}
4267zipWith :: Monad m => (a -> b -> c) -> Stream m a -> Stream m b -> Stream m c
4268zipWith f = zipWithM (\a b -> return (f a b))
4270------------------------------------------------------------------------------
4272------------------------------------------------------------------------------
4274{-# INLINE_NORMAL mergeByM #-}
4277 => (a -> a -> m Ordering) -> Stream m a -> Stream m a -> Stream m a
4278mergeByM cmp (Stream stepa ta) (Stream stepb tb) =
4279 Stream step (Just ta, Just tb, Nothing, Nothing)
4281 {-# INLINE_LATE step #-}
4283 -- one of the values is missing, and the corresponding stream is running
4284 step gst (Just sa, sb, Nothing, b) = do
4287 Yield a sa' -> Skip (Just sa', sb, Just a, b)
4288 Skip sa' -> Skip (Just sa', sb, Nothing, b)
4289 Stop -> Skip (Nothing, sb, Nothing, b)
4291 step gst (sa, Just sb, a, Nothing) = do
4294 Yield b sb' -> Skip (sa, Just sb', a, Just b)
4295 Skip sb' -> Skip (sa, Just sb', a, Nothing)
4296 Stop -> Skip (sa, Nothing, a, Nothing)
4298 -- both the values are available
4299 step _ (sa, sb, Just a, Just b) = do
4301 return $ case res of
4302 GT -> Yield b (sa, sb, Just a, Nothing)
4303 _ -> Yield a (sa, sb, Nothing, Just b)
4305 -- one of the values is missing, corresponding stream is done
4306 step _ (Nothing, sb, Nothing, Just b) =
4307 return $ Yield b (Nothing, sb, Nothing, Nothing)
4309 step _ (sa, Nothing, Just a, Nothing) =
4310 return $ Yield a (sa, Nothing, Nothing, Nothing)
4312 step _ (Nothing, Nothing, Nothing, Nothing) = return Stop
4314{-# INLINE mergeBy #-}
4317 => (a -> a -> Ordering) -> Stream m a -> Stream m a -> Stream m a
4318mergeBy cmp = mergeByM (\a b -> return $ cmp a b)
4320------------------------------------------------------------------------------
4321-- Transformation comprehensions
4322------------------------------------------------------------------------------
4324{-# INLINE_NORMAL the #-}
4325the :: (Eq a, Monad m) => Stream m a -> m (Maybe a)
4326the (Stream step state) = go state
4329 r <- step defState st
4331 Yield x s -> go' x s
4333 Stop -> return Nothing
4335 r <- step defState st
4337 Yield x s | x == n -> go' n s
4338 | otherwise -> return Nothing
4340 Stop -> return (Just n)
4342{-# INLINE runFold #-}
4343runFold :: (Monad m) => Fold m a b -> Stream m a -> m b
4344runFold (Fold step begin done) = foldlMx' step begin done
4346-------------------------------------------------------------------------------
4347-- Concurrent application and fold
4348-------------------------------------------------------------------------------
4350-- XXX These functions should be moved to Stream/Parallel.hs
4352-- Using StreamD the worker stream producing code can fuse with the code to
4353-- queue output to the SVar giving some perf boost.
4355-- Note that StreamD can only be used in limited situations, specifically, we
4356-- cannot implement joinStreamVarPar using this.
4358-- XXX make sure that the SVar passed is a Parallel style SVar.
4360-- | Fold the supplied stream to the SVar asynchronously using Parallel
4361-- concurrency style.
4362-- {-# INLINE_NORMAL toSVarParallel #-}
4363{-# INLINE toSVarParallel #-}
4364toSVarParallel :: MonadAsync m
4365 => State t m a -> SVar t m a -> Stream m a -> m ()
4366toSVarParallel st sv xs =
4367 if svarInspectMode sv
4371 case getYieldLimit st of
4372 Nothing -> doFork (work Nothing)
4374 (handleChildException sv)
4375 Just _ -> doFork (workLim Nothing)
4377 (handleChildException sv)
4382 {-# NOINLINE work #-}
4383 work info = (runFold (FL.toParallelSVar sv info) xs)
4385 {-# NOINLINE workLim #-}
4386 workLim info = runFold (FL.toParallelSVarLimited sv info) xs
4388 {-# NOINLINE forkWithDiag #-}
4390 -- We do not use workerCount in case of ParallelVar but still there is
4391 -- no harm in maintaining it correctly.
4392 liftIO $ atomicModifyIORefCAS_ (workerCount sv) $ \n -> n + 1
4394 -- This allocation matters when significant number of workers are being
4395 -- sent. We allocate it only when needed. The overhead increases by 4x.
4397 case yieldRateInfo sv of
4398 Nothing -> return Nothing
4399 Just _ -> liftIO $ do
4400 cntRef <- newIORef 0
4401 t <- getTime Monotonic
4402 lat <- newIORef (0, t)
4403 return $ Just WorkerInfo
4404 { workerYieldMax = 0
4405 , workerYieldCount = cntRef
4406 , workerLatencyStart = lat
4409 case getYieldLimit st of
4410 Nothing -> doFork (work winfo)
4412 (handleChildException sv)
4413 Just _ -> doFork (workLim winfo)
4415 (handleChildException sv)
4418{-# INLINE_NORMAL mkParallelD #-}
4419mkParallelD :: MonadAsync m => Stream m a -> Stream m a
4420mkParallelD m = Stream step Nothing
4423 step gst Nothing = do
4424 sv <- newParallelVar StopNone gst
4425 toSVarParallel gst sv m
4426 -- XXX use unfold instead?
4427 return $ Skip $ Just $ fromSVar sv
4429 step gst (Just (UnStream step1 st)) = do
4432 Yield a s -> Yield a (Just $ Stream step1 s)
4433 Skip s -> Skip (Just $ Stream step1 s)
4436-- Compare with mkAsync. mkAsync uses an Async style SVar whereas this uses a
4437-- parallel style SVar for evaluation. Currently, parallel style cannot use
4438-- rate control whereas Async style can use rate control. In async style SVar
4439-- the worker thread terminates when the buffer is full whereas in Parallel
4442-- | Make the stream producer and consumer run concurrently by introducing a
4443-- buffer between them. The producer thread evaluates the input stream until
4444-- the buffer fills, it blocks if the buffer is full until there is space in
4445-- the buffer. The consumer consumes the stream lazily from the buffer.
4449{-# INLINE_NORMAL mkParallel #-}
4450mkParallel :: (K.IsStream t, MonadAsync m) => t m a -> t m a
4451mkParallel = fromStreamD . mkParallelD . toStreamD
4453-- Note: we can use another API with two callbacks stop and yield if we want
4454-- the callback to be able to indicate end of stream.
4456-- | Generates a callback and a stream pair. The callback returned is used to
4457-- queue values to the stream. The stream is infinite, there is no way for the
4458-- callback to indicate that it is done now.
4462{-# INLINE_NORMAL newCallbackStream #-}
4463newCallbackStream :: (K.IsStream t, MonadAsync m) => m ((a -> m ()), t m a)
4464newCallbackStream = do
4465 sv <- newParallelVar StopNone defState
4467 -- XXX Add our own thread-id to the SVar as we can not know the callback's
4468 -- thread-id and the callback is not run in a managed worker. We need to
4469 -- handle this better.
4470 liftIO myThreadId >>= modifyThread sv
4472 let callback a = liftIO $ void $ send sv (ChildYield a)
4473 -- XXX we can return an SVar and then the consumer can unfold from the
4475 return (callback, fromStreamD (fromSVar sv))
4477-------------------------------------------------------------------------------
4479-------------------------------------------------------------------------------
4481-- | Create an SVar with a fold consumer that will fold any elements sent to it
4482-- using the supplied fold function.
4483{-# INLINE newFoldSVar #-}
4484newFoldSVar :: MonadAsync m => State t m a -> Fold m a b -> m (SVar t m a)
4485newFoldSVar stt f = do
4486 -- Buffer size for the SVar is derived from the current state
4487 sv <- newParallelVar StopAny (adaptState stt)
4488 -- Add the producer thread-id to the SVar.
4489 liftIO myThreadId >>= modifyThread sv
4490 void $ doFork (work sv) (svarMrun sv) (handleFoldException sv)
4495 {-# NOINLINE work #-}
4496 work sv = void $ runFold f $ fromProducer sv
4498data TapState sv st = TapInit | Tapping sv st | TapDone st
4500{-# INLINE_NORMAL tapAsync #-}
4501tapAsync :: MonadAsync m => Fold m a b -> Stream m a -> Stream m a
4502tapAsync f (Stream step1 state1) = Stream step TapInit
4506 -- In general, a Stop event would come equipped with the result
4507 -- of the fold. It is not used here but it would be useful in
4508 -- applicative and distribute.
4509 done <- fromConsumer svr
4510 when (not done) $ do
4511 liftIO $ withDiagMVar svr "teeToSVar: waiting to drain"
4512 $ takeMVar (outputDoorBellFromConsumer svr)
4516 liftIO $ sendStop svr Nothing
4517 -- drain/wait until a stop event arrives from the fold.
4520 {-# INLINE_LATE step #-}
4521 step gst TapInit = do
4522 sv <- newFoldSVar gst f
4523 return $ Skip (Tapping sv state1)
4525 step gst (Tapping sv st) = do
4529 done <- pushToFold sv a
4532 -- XXX we do not need to wait synchronously here
4534 return $ Yield a (TapDone s)
4535 else return $ Yield a (Tapping sv s)
4536 Skip s -> return $ Skip (Tapping sv s)
4541 step gst (TapDone st) = do
4544 Yield a s -> Yield a (TapDone s)
4545 Skip s -> Skip (TapDone s)
4548-- XXX Exported from Array again as this fold is specific to Array
4549-- | Take last 'n' elements from the stream and discard the rest.
4551lastN :: (Storable a, MonadIO m) => Int -> Fold m a (Array a)
4553 | n <= 0 = fmap (const mempty) FL.drain
4554 | otherwise = Fold step initial done
4556 step (Tuple3' rb rh i) a = do
4557 rh1 <- liftIO $ RB.unsafeInsert rb rh a
4558 return $ Tuple3' rb rh1 (i + 1)
4559 initial = fmap (\(a, b) -> Tuple3' a b (0 :: Int)) $ liftIO $ RB.new n
4560 done (Tuple3' rb rh i) = do
4561 arr <- liftIO $ A.newArray n
4562 foldFunc i rh snoc' arr rb
4563 snoc' b a = liftIO $ A.unsafeSnoc b a
4565 | i < n = RB.unsafeFoldRingM
4566 | otherwise = RB.unsafeFoldRingFullM
4568------------------------------------------------------------------------------
4570------------------------------------------------------------------------------
4572-- XXX using getTime in the loop can be pretty expensive especially for
4573-- computations where iterations are lightweight. We have the following
4576-- 1) Run a timeout thread updating a flag asynchronously and check that
4577-- flag here, that way we can have a cheap termination check.
4579-- 2) Use COARSE clock to get time with lower resolution but more efficiently.
4581-- 3) Use rdtscp/rdtsc to get time directly from the processor, compute the
4582-- termination value of rdtsc in the beginning and then in each iteration just
4583-- get rdtsc and check if we should terminate.
4587 | TakeByTimeCheck st s
4588 | TakeByTimeYield st s
4590{-# INLINE_NORMAL takeByTime #-}
4591takeByTime :: (MonadIO m, TimeUnit64 t) => t -> Stream m a -> Stream m a
4592takeByTime duration (Stream step1 state1) = Stream step (TakeByTimeInit state1)
4595 lim = toRelTime64 duration
4597 {-# INLINE_LATE step #-}
4598 step _ (TakeByTimeInit _) | lim == 0 = return Stop
4599 step _ (TakeByTimeInit st) = do
4600 t0 <- liftIO $ getTime Monotonic
4601 return $ Skip (TakeByTimeYield st t0)
4602 step _ (TakeByTimeCheck st t0) = do
4603 t <- liftIO $ getTime Monotonic
4605 if diffAbsTime64 t t0 > lim
4607 else Skip (TakeByTimeYield st t0)
4608 step gst (TakeByTimeYield st t0) = do
4611 Yield x s -> Yield x (TakeByTimeCheck s t0)
4612 Skip s -> Skip (TakeByTimeCheck s t0)
4615data DropByTime st s x
4617 | DropByTimeGen st s
4618 | DropByTimeCheck st s x
4619 | DropByTimeYield st
4621{-# INLINE_NORMAL dropByTime #-}
4622dropByTime :: (MonadIO m, TimeUnit64 t) => t -> Stream m a -> Stream m a
4623dropByTime duration (Stream step1 state1) = Stream step (DropByTimeInit state1)
4626 lim = toRelTime64 duration
4628 {-# INLINE_LATE step #-}
4629 step _ (DropByTimeInit st) = do
4630 t0 <- liftIO $ getTime Monotonic
4631 return $ Skip (DropByTimeGen st t0)
4632 step gst (DropByTimeGen st t0) = do
4635 Yield x s -> Skip (DropByTimeCheck s t0 x)
4636 Skip s -> Skip (DropByTimeGen s t0)
4638 step _ (DropByTimeCheck st t0 x) = do
4639 t <- liftIO $ getTime Monotonic
4640 if diffAbsTime64 t t0 <= lim
4641 then return $ Skip $ DropByTimeGen st t0
4642 else return $ Yield x $ DropByTimeYield st
4643 step gst (DropByTimeYield st) = do
4646 Yield x s -> Yield x (DropByTimeYield s)
4647 Skip s -> Skip (DropByTimeYield s)
4650{-# INLINE updateTimeVar #-}
4651updateTimeVar :: Var IO Int64 -> IO ()
4652updateTimeVar timeVar = do
4653 MicroSecond64 t <- fromAbsTime <$> getTime Monotonic
4654 modifyVar' timeVar (const t)
4656{-# INLINE updateWithDelay #-}
4657updateWithDelay :: RealFrac a => a -> Var IO Int64 -> IO ()
4658updateWithDelay precision timeVar = do
4659 threadDelay (delayTime precision)
4660 updateTimeVar timeVar
4664 -- Keep the minimum at least a millisecond to avoid high CPU usage
4665 {-# INLINE delayTime #-}
4667 | g' >= fromIntegral (maxBound :: Int) = maxBound
4669 | otherwise = round g'
4673 g' = g * 10 ^ (6 :: Int)
4675-- XXX we should move this to stream generation section of this file. Also, the
4676-- take/drop combinators above should be moved to filtering section.
4678{-# INLINE_NORMAL times #-}
4679times :: MonadAsync m => Double -> Stream m (AbsTime, RelTime64)
4680times g = Stream step Nothing
4684 {-# INLINE_LATE step #-}
4686 -- XXX note that this is safe only on a 64-bit machine. On a 32-bit
4687 -- machine a 64-bit 'Var' cannot be read consistently without a lock
4688 -- while another thread is writing to it.
4689 timeVar <- liftIO $ newVar (0 :: Int64)
4690 liftIO $ updateTimeVar timeVar
4691 tid <- forkManaged $ liftIO $ forever (updateWithDelay g timeVar)
4692 a <- liftIO $ readVar timeVar
4693 return $ Skip $ Just (timeVar, tid, a)
4695 step _ s@(Just (timeVar, _, t0)) = do
4696 a <- liftIO $ readVar timeVar
4697 -- XXX we can perhaps use an AbsTime64 using a 64 bit Int for
4698 -- efficiency. or maybe we can use a representation using Double for
4699 -- floating precision time
4700 return $ Yield (toAbsTime (MicroSecond64 t0),
4701 (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