4-- Module : Streamly.Internal.Data.Stream.StreamD
5-- Copyright : (c) 2018 Harendra Kumar
6-- (c) Roman Leshchinskiy 2008-2010
7-- (c) The University of Glasgow, 2009
10-- Maintainer : streamly@composewell.com
11-- Stability : experimental
14-- Direct style re-implementation of CPS style stream in StreamK module. The
15-- symbol or suffix 'D' in this module denotes the "Direct" style. GHC is able
16-- to INLINE and fuse direct style better, providing better performance than
20-- import qualified Streamly.Internal.Data.Stream.StreamD as D
23-- Some of the functions in this file have been adapted from the vector
24-- library, https://hackage.haskell.org/package/vector.
26module Streamly.Internal.Data.Stream.StreamD
30 , Stream (Stream, UnStream)
46 -- ** Specialized Generation
47 -- | Generate a monadic stream from a seed.
60 , enumerateFromStepIntegral
61 , enumerateFromIntegral
62 , enumerateFromThenIntegral
63 , enumerateFromToIntegral
64 , enumerateFromThenToIntegral
66 , enumerateFromStepNum
69 , enumerateFromToFractional
70 , enumerateFromThenToFractional
76 -- | Transform an input structure into a stream.
77 -- | Direct style stream does not support @fromFoldable@.
111 -- ** Specialized Folds
138 -- ** Flattening nested streams
141 , ConcatMapUState (..)
143 , ConcatUnfoldInterleaveState (..)
144 , concatUnfoldInterleave
145 , concatUnfoldRoundrobin
148 , InterleaveState(..)
153 , roundRobin -- interleaveFair?/ParallelFair
186 -- | Transform a stream into another type.
203 -- ** By folding (scans)
254 , intersperseSuffixBySpan
278 -- * Transformation comprehensions
298 -- * Concurrent Application
307import Control.Concurrent (killThread, myThreadId, takeMVar, threadDelay)
308import Control.Exception
309 (assert, Exception, SomeException, AsyncException, fromException, mask_)
310import Control.Monad (void, when, forever)
311import Control.Monad.Catch (MonadCatch, MonadThrow, throwM)
312import Control.Monad.IO.Class (MonadIO(..))
313import Control.Monad.Trans.Reader (ReaderT)
314import Control.Monad.Trans.State.Strict (StateT)
315import Control.Monad.Trans.Class (MonadTrans(lift))
316import Control.Monad.Trans.Control (MonadBaseControl, liftBaseOp_, control)
317import Data.Bits (shiftR, shiftL, (.|.), (.&.))
318import Data.Functor.Identity (Identity(..))
319import Data.Int (Int64)
320import Data.IORef (newIORef, readIORef, mkWeakIORef, writeIORef, IORef)
321import Data.Maybe (fromJust, isJust, isNothing)
322import Data.Word (Word32)
323import Foreign.Ptr (Ptr)
324import Foreign.Storable (Storable(..))
325import GHC.Types (SPEC(..))
326import System.Mem (performMajorGC)
327import Fusion.Plugin.Types (Fuse(..))
328import Streamly.Internal.Data.IORef.Prim (Prim)
329import Streamly.Internal.Data.Time.Units
330 (TimeUnit64, toRelTime64, diffAbsTime64, RelTime64)
331import Streamly.Internal.Data.Atomics (atomicModifyIORefCAS_)
332import Streamly.Internal.Memory.Array.Types (Array(..))
333import Streamly.Internal.Data.Fold.Types (Fold(..))
334import Streamly.Internal.Data.Parser (ParseError(..))
335import Streamly.Internal.Data.Pipe.Types (Pipe(..), PipeState(..))
336import Streamly.Internal.Data.Time.Clock (Clock(Monotonic), getTime)
337import Streamly.Internal.Data.Time.Units
338 (MicroSecond64(..), fromAbsTime, toAbsTime, AbsTime)
339import Streamly.Internal.Data.Unfold.Types (Unfold(..))
340import Streamly.Internal.Data.Tuple.Strict (Tuple3'(..))
341import Streamly.Internal.Data.Stream.SVar (fromConsumer, pushToFold)
343import qualified Streamly.Internal.Data.IORef.Prim as Prim
344import qualified Streamly.Internal.Data.Pipe.Types as Pipe
345import qualified Streamly.Internal.Memory.Array.Types as A
346import qualified Streamly.Internal.Memory.Mutable.Array.Types as MA
347import qualified Streamly.Internal.Data.Fold as FL
348import qualified Streamly.Memory.Ring as RB
349import qualified Streamly.Internal.Data.Stream.StreamK as K
350import qualified Streamly.Internal.Data.Parser as PR
351import qualified Streamly.Internal.Data.Parser.ParserD as PRD
352import qualified Control.Monad.Catch as MC
353import qualified Control.Monad.Trans.Reader as Reader
354import qualified Control.Monad.Trans.State.Strict as State
355import qualified Prelude
358 ( map, mapM, mapM_, repeat, foldr, last, take, filter
359 , takeWhile, drop, dropWhile, all, any, maximum, minimum, elem
360 , notElem, null, head, tail, zipWith, lookup, foldr1, sequence
361 , (!!), scanl, scanl1, concatMap, replicate, enumFromTo, concat
362 , reverse, iterate, splitAt)
363import Streamly.Internal.Data.Stream.StreamD.Type
364import Streamly.Internal.Data.SVar
366------------------------------------------------------------------------------
368------------------------------------------------------------------------------
370-- | An empty 'Stream'.
371{-# INLINE_NORMAL nil #-}
372nil :: Monad m => Stream m a
373nil = Stream (\_ _ -> return Stop) ()
375-- | An empty 'Stream' with a side effect.
376{-# INLINE_NORMAL nilM #-}
377nilM :: Monad m => m b -> Stream m a
378nilM m = Stream (\_ _ -> m >> return Stop) ()
380{-# INLINE_NORMAL consM #-}
381consM :: Monad m => m a -> Stream m a -> Stream m a
382consM m (Stream step state) = Stream step1 Nothing
384 {-# INLINE_LATE step1 #-}
385 step1 _ Nothing = m >>= \x -> return $ Yield x (Just state)
386 step1 gst (Just st) = do
390 Yield a s -> Yield a (Just s)
391 Skip s -> Skip (Just s)
394-- XXX implement in terms of consM?
395-- cons x = consM (return x)
397-- | Can fuse but has O(n^2) complexity.
398{-# INLINE_NORMAL cons #-}
399cons :: Monad m => a -> Stream m a -> Stream m a
400cons x (Stream step state) = Stream step1 Nothing
402 {-# INLINE_LATE step1 #-}
403 step1 _ Nothing = return $ Yield x (Just state)
404 step1 gst (Just st) = do
408 Yield a s -> Yield a (Just s)
409 Skip s -> Skip (Just s)
412-------------------------------------------------------------------------------
414-------------------------------------------------------------------------------
416-- Does not fuse, has the same performance as the StreamK version.
417{-# INLINE_NORMAL uncons #-}
418uncons :: Monad m => Stream m a -> m (Maybe (a, Stream m a))
419uncons (UnStream step state) = go state
422 r <- step defState st
424 Yield x s -> return $ Just (x, Stream step s)
426 Stop -> return Nothing
428------------------------------------------------------------------------------
429-- Generation by unfold
430------------------------------------------------------------------------------
432{-# INLINE_NORMAL unfoldrM #-}
433unfoldrM :: Monad m => (s -> m (Maybe (a, s))) -> s -> Stream m a
434unfoldrM next state = Stream step state
436 {-# INLINE_LATE step #-}
440 Just (x, s) -> Yield x s
443{-# INLINE_LATE unfoldr #-}
444unfoldr :: Monad m => (s -> Maybe (a, s)) -> s -> Stream m a
445unfoldr f = unfoldrM (return . f)
447-- | Convert an 'Unfold' into a 'Stream' by supplying it a seed.
449{-# INLINE_NORMAL unfold #-}
450unfold :: Monad m => Unfold m a b -> a -> Stream m b
451unfold (Unfold ustep inject) seed = Stream step Nothing
453 {-# INLINE_LATE step #-}
454 step _ Nothing = inject seed >>= return . Skip . Just
455 step _ (Just st) = do
458 Yield x s -> Yield x (Just s)
459 Skip s -> Skip (Just s)
462------------------------------------------------------------------------------
463-- Specialized Generation
464------------------------------------------------------------------------------
466{-# INLINE_NORMAL repeatM #-}
467repeatM :: Monad m => m a -> Stream m a
468repeatM x = Stream (\_ _ -> x >>= \r -> return $ Yield r ()) ()
470{-# INLINE_NORMAL repeat #-}
471repeat :: Monad m => a -> Stream m a
472repeat x = Stream (\_ _ -> return $ Yield x ()) ()
474{-# INLINE_NORMAL iterateM #-}
475iterateM :: Monad m => (a -> m a) -> m a -> Stream m a
476iterateM step = Stream (\_ st -> st >>= \x -> return $ Yield x (step x))
478{-# INLINE_NORMAL iterate #-}
479iterate :: Monad m => (a -> a) -> a -> Stream m a
480iterate step st = iterateM (return . step) (return st)
482{-# INLINE_NORMAL replicateM #-}
483replicateM :: forall m a. Monad m => Int -> m a -> Stream m a
484replicateM n p = Stream step n
486 {-# INLINE_LATE step #-}
488 | i <= 0 = return Stop
491 return $ Yield x (i - 1)
493{-# INLINE_NORMAL replicate #-}
494replicate :: Monad m => Int -> a -> Stream m a
495replicate n x = replicateM n (return x)
497-- This would not work properly for floats, therefore we put an Integral
499-- | Can be used to enumerate unbounded integrals. This does not check for
500-- overflow or underflow for bounded integrals.
501{-# INLINE_NORMAL enumerateFromStepIntegral #-}
502enumerateFromStepIntegral :: (Integral a, Monad m) => a -> a -> Stream m a
503enumerateFromStepIntegral from stride =
504 from `seq` stride `seq` Stream step from
506 {-# INLINE_LATE step #-}
507 step _ !x = return $ Yield x $! (x + stride)
509-- We are assuming that "to" is constrained by the type to be within
511{-# INLINE enumerateFromToIntegral #-}
512enumerateFromToIntegral :: (Monad m, Integral a) => a -> a -> Stream m a
513enumerateFromToIntegral from to =
514 takeWhile (<= to) $ enumerateFromStepIntegral from 1
516{-# INLINE enumerateFromIntegral #-}
517enumerateFromIntegral :: (Monad m, Integral a, Bounded a) => a -> Stream m a
518enumerateFromIntegral from = enumerateFromToIntegral from maxBound
520data EnumState a = EnumInit | EnumYield a a a | EnumStop
522{-# INLINE_NORMAL enumerateFromThenToIntegralUp #-}
523enumerateFromThenToIntegralUp
524 :: (Monad m, Integral a)
525 => a -> a -> a -> Stream m a
526enumerateFromThenToIntegralUp from next to = Stream step EnumInit
528 {-# INLINE_LATE step #-}
534 else Yield from EnumStop
535 else -- from <= next <= to
536 let stride = next - from
537 in Skip $ EnumYield from stride (to - stride)
539 step _ (EnumYield x stride toMinus) =
542 then Yield x EnumStop
543 else Yield x $ EnumYield (x + stride) stride toMinus
545 step _ EnumStop = return Stop
547{-# INLINE_NORMAL enumerateFromThenToIntegralDn #-}
548enumerateFromThenToIntegralDn
549 :: (Monad m, Integral a)
550 => a -> a -> a -> Stream m a
551enumerateFromThenToIntegralDn from next to = Stream step EnumInit
553 {-# INLINE_LATE step #-}
555 return $ if to > next
558 else Yield from EnumStop
559 else -- from >= next >= to
560 let stride = next - from
561 in Skip $ EnumYield from stride (to - stride)
563 step _ (EnumYield x stride toMinus) =
566 then Yield x EnumStop
567 else Yield x $ EnumYield (x + stride) stride toMinus
569 step _ EnumStop = return Stop
571{-# INLINE_NORMAL enumerateFromThenToIntegral #-}
572enumerateFromThenToIntegral
573 :: (Monad m, Integral a)
574 => a -> a -> a -> Stream m a
575enumerateFromThenToIntegral from next to
576 | next >= from = enumerateFromThenToIntegralUp from next to
577 | otherwise = enumerateFromThenToIntegralDn from next to
579{-# INLINE_NORMAL enumerateFromThenIntegral #-}
580enumerateFromThenIntegral
581 :: (Monad m, Integral a, Bounded a)
582 => a -> a -> Stream m a
583enumerateFromThenIntegral from next =
585 then enumerateFromThenToIntegralUp from next maxBound
586 else enumerateFromThenToIntegralDn from next minBound
588-- For floating point numbers if the increment is less than the precision then
589-- it just gets lost. Therefore we cannot always increment it correctly by just
591-- 9007199254740992 + 1 + 1 :: Double => 9.007199254740992e15
592-- 9007199254740992 + 2 :: Double => 9.007199254740994e15
594-- Instead we accumulate the increment counter and compute the increment
595-- every time before adding it to the starting number.
597-- This works for Integrals as well as floating point numbers, but
598-- enumerateFromStepIntegral is faster for integrals.
599{-# INLINE_NORMAL enumerateFromStepNum #-}
600enumerateFromStepNum :: (Monad m, Num a) => a -> a -> Stream m a
601enumerateFromStepNum from stride = Stream step 0
603 {-# INLINE_LATE step #-}
604 step _ !i = return $ (Yield $! (from + i * stride)) $! (i + 1)
606{-# INLINE_NORMAL numFrom #-}
607numFrom :: (Monad m, Num a) => a -> Stream m a
608numFrom from = enumerateFromStepNum from 1
610{-# INLINE_NORMAL numFromThen #-}
611numFromThen :: (Monad m, Num a) => a -> a -> Stream m a
612numFromThen from next = enumerateFromStepNum from (next - from)
614-- We cannot write a general function for Num. The only way to write code
615-- portable between the two is to use a 'Real' constraint and convert between
616-- Fractional and Integral using fromRational which is horribly slow.
617{-# INLINE_NORMAL enumerateFromToFractional #-}
618enumerateFromToFractional
619 :: (Monad m, Fractional a, Ord a)
620 => a -> a -> Stream m a
621enumerateFromToFractional from to =
622 takeWhile (<= to + 1 / 2) $ enumerateFromStepNum from 1
624{-# INLINE_NORMAL enumerateFromThenToFractional #-}
625enumerateFromThenToFractional
626 :: (Monad m, Fractional a, Ord a)
627 => a -> a -> a -> Stream m a
628enumerateFromThenToFractional from next to =
629 takeWhile predicate $ numFromThen from next
631 mid = (next - from) / 2
632 predicate | next >= from = (<= to + mid)
633 | otherwise = (>= to + mid)
635-------------------------------------------------------------------------------
636-- Generation by Conversion
637-------------------------------------------------------------------------------
639{-# INLINE_NORMAL fromIndicesM #-}
640fromIndicesM :: Monad m => (Int -> m a) -> Stream m a
641fromIndicesM gen = Stream step 0
643 {-# INLINE_LATE step #-}
646 return $ Yield x (i + 1)
648{-# INLINE fromIndices #-}
649fromIndices :: Monad m => (Int -> a) -> Stream m a
650fromIndices gen = fromIndicesM (return . gen)
652{-# INLINE_NORMAL generateM #-}
653generateM :: Monad m => Int -> (Int -> m a) -> Stream m a
654generateM n gen = n `seq` Stream step 0
656 {-# INLINE_LATE step #-}
657 step _ i | i < n = do
659 return $ Yield x (i + 1)
660 | otherwise = return Stop
662{-# INLINE generate #-}
663generate :: Monad m => Int -> (Int -> a) -> Stream m a
664generate n gen = generateM n (return . gen)
666-- XXX we need the MonadAsync constraint because of a rewrite rule.
667-- | Convert a list of monadic actions to a 'Stream'
668{-# INLINE_LATE fromListM #-}
669fromListM :: MonadAsync m => [m a] -> Stream m a
670fromListM = Stream step
672 {-# INLINE_LATE step #-}
673 step _ (m:ms) = m >>= \x -> return $ Yield x ms
674 step _ [] = return Stop
676{-# INLINE toStreamD #-}
677toStreamD :: (K.IsStream t, Monad m) => t m a -> Stream m a
678toStreamD = fromStreamK . K.toStream
680{-# INLINE_NORMAL fromPrimIORef #-}
681fromPrimIORef :: (MonadIO m, Prim a) => Prim.IORef a -> Stream m a
682fromPrimIORef var = Stream step ()
684 {-# INLINE_LATE step #-}
685 step _ () = liftIO (Prim.readIORef var) >>= \x -> return $ Yield x ()
687-------------------------------------------------------------------------------
688-- Generation from SVar
689-------------------------------------------------------------------------------
691data FromSVarState t m a =
693 | FromSVarRead (SVar t m a)
694 | FromSVarLoop (SVar t m a) [ChildEvent a]
695 | FromSVarDone (SVar t m a)
697{-# INLINE_NORMAL fromSVar #-}
698fromSVar :: (MonadAsync m) => SVar t m a -> Stream m a
699fromSVar svar = Stream step FromSVarInit
702 {-# INLINE_LATE step #-}
703 step _ FromSVarInit = do
704 ref <- liftIO $ newIORef ()
705 _ <- liftIO $ mkWeakIORef ref hook
706 -- when this copy of svar gets garbage collected "ref" will get
707 -- garbage collected and our GC hook will be called.
708 let sv = svar{svarRef = Just ref}
709 return $ Skip (FromSVarRead sv)
713 {-# NOINLINE hook #-}
715 when (svarInspectMode svar) $ do
716 r <- liftIO $ readIORef (svarStopTime (svarStats svar))
718 printSVar svar "SVar Garbage Collected"
720 -- If there are any SVars referenced by this SVar a GC will prompt
721 -- them to be cleaned up quickly.
722 when (svarInspectMode svar) performMajorGC
724 step _ (FromSVarRead sv) = do
725 list <- readOutputQ sv
726 -- Reversing the output is important to guarantee that we process the
727 -- outputs in the same order as they were generated by the constituent
729 return $ Skip $ FromSVarLoop sv (Prelude.reverse list)
731 step _ (FromSVarLoop sv []) = do
732 done <- postProcess sv
733 return $ Skip $ if done
734 then (FromSVarDone sv)
735 else (FromSVarRead sv)
737 step _ (FromSVarLoop sv (ev : es)) = do
739 ChildYield a -> return $ Yield a (FromSVarLoop sv es)
740 ChildStop tid e -> do
744 stop <- shouldStop tid
747 liftIO (cleanupSVar sv)
748 return $ Skip (FromSVarDone sv)
749 else return $ Skip (FromSVarLoop sv es)
751 case fromException ex of
753 return $ Skip (FromSVarLoop sv es)
754 Nothing -> liftIO (cleanupSVar sv) >> throwM ex
758 case svarStopStyle sv of
759 StopNone -> return False
760 StopAny -> return True
762 sid <- liftIO $ readIORef (svarStopBy sv)
763 return $ if tid == sid then True else False
765 step _ (FromSVarDone sv) = do
766 when (svarInspectMode sv) $ do
767 t <- liftIO $ getTime Monotonic
768 liftIO $ writeIORef (svarStopTime (svarStats sv)) (Just t)
769 liftIO $ printSVar sv "SVar Done"
772-------------------------------------------------------------------------------
773-- Process events received by a fold consumer from a stream producer
774-------------------------------------------------------------------------------
776{-# INLINE_NORMAL fromProducer #-}
777fromProducer :: (MonadAsync m) => SVar t m a -> Stream m a
778fromProducer svar = Stream step (FromSVarRead svar)
781 {-# INLINE_LATE step #-}
782 step _ (FromSVarRead sv) = do
783 list <- readOutputQ sv
784 -- Reversing the output is important to guarantee that we process the
785 -- outputs in the same order as they were generated by the constituent
787 return $ Skip $ FromSVarLoop sv (Prelude.reverse list)
789 step _ (FromSVarLoop sv []) = return $ Skip $ FromSVarRead sv
790 step _ (FromSVarLoop sv (ev : es)) = do
792 ChildYield a -> return $ Yield a (FromSVarLoop sv es)
793 ChildStop tid e -> do
797 sendStopToProducer sv
798 return $ Skip (FromSVarDone sv)
799 Just _ -> error "Bug: fromProducer: received exception"
801 step _ (FromSVarDone sv) = do
802 when (svarInspectMode sv) $ do
803 t <- liftIO $ getTime Monotonic
804 liftIO $ writeIORef (svarStopTime (svarStats sv)) (Just t)
805 liftIO $ printSVar sv "SVar Done"
808 step _ FromSVarInit = undefined
810-------------------------------------------------------------------------------
811-- Hoisting the inner monad
812-------------------------------------------------------------------------------
814{-# INLINE_NORMAL hoist #-}
815hoist :: Monad n => (forall x. m x -> n x) -> Stream m a -> Stream n a
816hoist f (Stream step state) = (Stream step' state)
818 {-# INLINE_LATE step' #-}
820 r <- f $ step (adaptState gst) st
822 Yield x s -> Yield x s
826{-# INLINE generally #-}
827generally :: Monad m => Stream Identity a -> Stream m a
828generally = hoist (return . runIdentity)
830{-# INLINE_NORMAL liftInner #-}
831liftInner :: (Monad m, MonadTrans t, Monad (t m))
832 => Stream m a -> Stream (t m) a
833liftInner (Stream step state) = Stream step' state
835 {-# INLINE_LATE step' #-}
837 r <- lift $ step (adaptState gst) st
839 Yield x s -> Yield x s
843{-# INLINE_NORMAL runReaderT #-}
844runReaderT :: Monad m => m s -> Stream (ReaderT s m) a -> Stream m a
845runReaderT env (Stream step state) = Stream step' (state, env)
847 {-# INLINE_LATE step' #-}
848 step' gst (st, action) = do
850 r <- Reader.runReaderT (step (adaptState gst) st) sv
852 Yield x s -> Yield x (s, return sv)
853 Skip s -> Skip (s, return sv)
856{-# INLINE_NORMAL evalStateT #-}
857evalStateT :: Monad m => m s -> Stream (StateT s m) a -> Stream m a
858evalStateT initial (Stream step state) = Stream step' (state, initial)
860 {-# INLINE_LATE step' #-}
861 step' gst (st, action) = do
863 (r, sv') <- State.runStateT (step (adaptState gst) st) sv
865 Yield x s -> Yield x (s, return sv')
866 Skip s -> Skip (s, return sv')
869{-# INLINE_NORMAL runStateT #-}
870runStateT :: Monad m => m s -> Stream (StateT s m) a -> Stream m (s, a)
871runStateT initial (Stream step state) = Stream step' (state, initial)
873 {-# INLINE_LATE step' #-}
874 step' gst (st, action) = do
876 (r, sv') <- State.runStateT (step (adaptState gst) st) sv
878 Yield x s -> Yield (sv', x) (s, return sv')
879 Skip s -> Skip (s, return sv')
882------------------------------------------------------------------------------
883-- Elimination by Folds
884------------------------------------------------------------------------------
886------------------------------------------------------------------------------
888------------------------------------------------------------------------------
890{-# INLINE_NORMAL foldr1 #-}
891foldr1 :: Monad m => (a -> a -> a) -> Stream m a -> m (Maybe a)
895 Nothing -> return Nothing
896 Just (h, t) -> fmap Just (foldr f h t)
898------------------------------------------------------------------------------
900------------------------------------------------------------------------------
902{-# INLINE_NORMAL foldlT #-}
903foldlT :: (Monad m, Monad (s m), MonadTrans s)
904 => (s m b -> a -> s m b) -> s m b -> Stream m a -> s m b
905foldlT fstep begin (Stream step state) = go SPEC begin state
908 r <- lift $ step defState st
910 Yield x s -> go SPEC (fstep acc x) s
911 Skip s -> go SPEC acc s
914-- Note, this is going to have horrible performance, because of the nature of
915-- the stream type (i.e. direct stream vs CPS). Its only for reference, it is
916-- likely be practically unusable.
917{-# INLINE_NORMAL foldlS #-}
919 => (Stream m b -> a -> Stream m b) -> Stream m b -> Stream m a -> Stream m b
920foldlS fstep begin (Stream step state) = Stream step' (Left (state, begin))
922 step' gst (Left (st, acc)) = do
923 r <- step (adaptState gst) st
925 Yield x s -> Skip (Left (s, fstep acc x))
926 Skip s -> Skip (Left (s, acc))
927 Stop -> Skip (Right acc)
929 step' gst (Right (Stream stp stt)) = do
930 r <- stp (adaptState gst) stt
932 Yield x s -> Yield x (Right (Stream stp s))
933 Skip s -> Skip (Right (Stream stp s))
936------------------------------------------------------------------------------
938------------------------------------------------------------------------------
940-- Inlined definition. Without the inline "serially/parser/take" benchmark
941-- degrades and parseMany does not fuse. Even using "inline" at the callsite
943{-# INLINE splitAt #-}
944splitAt :: Int -> [a] -> ([a],[a])
947 | otherwise = splitAt' n ls
949 splitAt' :: Int -> [a] -> ([a], [a])
950 splitAt' _ [] = ([], [])
951 splitAt' 1 (x:xs) = ([x], xs)
952 splitAt' m (x:xs) = (x:xs', xs'')
954 (xs', xs'') = splitAt' (m - 1) xs
956-- | Run a 'Parse' over a stream.
957{-# INLINE_NORMAL parselMx' #-}
960 => (s -> a -> m (PR.Step s b))
965parselMx' pstep initial extract (Stream step state) = do
966 initial >>= go SPEC state []
970 -- XXX currently we are using a dumb list based approach for backtracking
971 -- buffer. This can be replaced by a sliding/ring buffer using Data.Array.
972 -- That will allow us more efficient random back and forth movement.
974 go !_ st buf !pst = do
975 r <- step defState st
980 PR.Partial 0 pst1 -> go SPEC s [] pst1
981 PR.Partial n pst1 -> do
982 assert (n <= length (x:buf)) (return ())
983 let src0 = Prelude.take n (x:buf)
984 src = Prelude.reverse src0
985 gobuf SPEC s [] src pst1
986 PR.Continue 0 pst1 -> go SPEC s (x:buf) pst1
987 PR.Continue n pst1 -> do
988 assert (n <= length (x:buf)) (return ())
989 let (src0, buf1) = splitAt n (x:buf)
990 src = Prelude.reverse src0
991 gobuf SPEC s buf1 src pst1
992 PR.Done _ b -> return b
993 PR.Error err -> throwM $ ParseError err
994 Skip s -> go SPEC s buf pst
997 gobuf !_ s buf [] !pst = go SPEC s buf pst
998 gobuf !_ s buf (x:xs) !pst = do
1001 PR.Partial 0 pst1 ->
1002 gobuf SPEC s [] xs pst1
1003 PR.Partial n pst1 -> do
1004 assert (n <= length (x:buf)) (return ())
1005 let src0 = Prelude.take n (x:buf)
1006 src = Prelude.reverse src0 ++ xs
1007 gobuf SPEC s [] src pst1
1008 PR.Continue 0 pst1 -> gobuf SPEC s (x:buf) xs pst1
1009 PR.Continue n pst1 -> do
1010 assert (n <= length (x:buf)) (return ())
1011 let (src0, buf1) = splitAt n (x:buf)
1012 src = Prelude.reverse src0 ++ xs
1013 gobuf SPEC s buf1 src pst1
1014 PR.Done _ b -> return b
1015 PR.Error err -> throwM $ ParseError err
1017------------------------------------------------------------------------------
1019------------------------------------------------------------------------------
1021{-# ANN type ParseChunksState Fuse #-}
1022data ParseChunksState x inpBuf st pst =
1023 ParseChunksInit inpBuf st
1024 | ParseChunksInitLeftOver inpBuf
1025 | ParseChunksStream st inpBuf pst
1026 | ParseChunksBuf inpBuf st inpBuf pst
1027 | ParseChunksYield x (ParseChunksState x inpBuf st pst)
1029{-# INLINE_NORMAL parseMany #-}
1035parseMany (PRD.Parser pstep initial extract) (Stream step state) =
1036 Stream stepOuter (ParseChunksInit [] state)
1040 {-# INLINE_LATE stepOuter #-}
1041 -- Buffer is empty, go to stream processing loop
1042 stepOuter _ (ParseChunksInit [] st) = do
1043 initial >>= return . Skip . ParseChunksStream st []
1045 -- Buffer is not empty, go to buffered processing loop
1046 stepOuter _ (ParseChunksInit src st) = do
1047 initial >>= return . Skip . ParseChunksBuf src st []
1049 -- XXX we just discard any leftover input at the end
1050 stepOuter _ (ParseChunksInitLeftOver _) = return Stop
1052 -- Buffer is empty process elements from the stream
1053 stepOuter gst (ParseChunksStream st buf pst) = do
1054 r <- step (adaptState gst) st
1059 PR.Partial 0 pst1 ->
1060 return $ Skip $ ParseChunksStream s [] pst1
1061 PR.Partial n pst1 -> do
1062 assert (n <= length (x:buf)) (return ())
1063 let src0 = Prelude.take n (x:buf)
1064 src = Prelude.reverse src0
1065 return $ Skip $ ParseChunksBuf src s [] pst1
1066 -- PR.Continue 0 pst1 ->
1067 -- return $ Skip $ ParseChunksStream s (x:buf) pst1
1068 PR.Continue n pst1 -> do
1069 assert (n <= length (x:buf)) (return ())
1070 let (src0, buf1) = splitAt n (x:buf)
1071 src = Prelude.reverse src0
1072 return $ Skip $ ParseChunksBuf src s buf1 pst1
1073 -- XXX Specialize for Stop 0 common case?
1075 assert (n <= length (x:buf)) (return ())
1076 let src = Prelude.reverse (Prelude.take n (x:buf))
1078 ParseChunksYield b (ParseChunksInit src s)
1079 PR.Error err -> throwM $ ParseError err
1080 Skip s -> return $ Skip $ ParseChunksStream s buf pst
1083 let src = Prelude.reverse buf
1085 ParseChunksYield b (ParseChunksInitLeftOver src)
1087 -- go back to stream processing mode
1088 stepOuter _ (ParseChunksBuf [] s buf pst) =
1089 return $ Skip $ ParseChunksStream s buf pst
1091 -- buffered processing loop
1092 stepOuter _ (ParseChunksBuf (x:xs) s buf pst) = do
1095 PR.Partial 0 pst1 ->
1096 return $ Skip $ ParseChunksBuf xs s [] pst1
1097 PR.Partial n pst1 -> do
1098 assert (n <= length (x:buf)) (return ())
1099 let src0 = Prelude.take n (x:buf)
1100 src = Prelude.reverse src0 ++ xs
1101 return $ Skip $ ParseChunksBuf src s [] pst1
1102 -- PR.Continue 0 pst1 -> return $ Skip $ ParseChunksBuf xs s (x:buf) pst1
1103 PR.Continue n pst1 -> do
1104 assert (n <= length (x:buf)) (return ())
1105 let (src0, buf1) = splitAt n (x:buf)
1106 src = Prelude.reverse src0 ++ xs
1107 return $ Skip $ ParseChunksBuf src s buf1 pst1
1108 -- XXX Specialize for Stop 0 common case?
1110 assert (n <= length (x:buf)) (return ())
1111 let src = Prelude.reverse (Prelude.take n (x:buf)) ++ xs
1112 return $ Skip $ ParseChunksYield b (ParseChunksInit src s)
1113 PR.Error err -> throwM $ ParseError err
1115 stepOuter _ (ParseChunksYield a next) = return $ Yield a next
1117{-# ANN type ConcatParseState Fuse #-}
1118data ConcatParseState x inpBuf st p =
1119 ConcatParseInit inpBuf st p
1120 | ConcatParseInitLeftOver inpBuf
1121 | ConcatParseStream st inpBuf p
1122 | ConcatParseBuf inpBuf st inpBuf p
1123 | ConcatParseYield x (ConcatParseState x inpBuf st p)
1125{-# INLINE_NORMAL parseIterate #-}
1128 => (b -> PRD.Parser m a b)
1132parseIterate func seed (Stream step state) =
1133 Stream stepOuter (ConcatParseInit [] state (func seed))
1137 {-# INLINE_LATE stepOuter #-}
1138 -- Buffer is empty, go to stream processing loop
1139 stepOuter _ (ConcatParseInit [] st (PRD.Parser pstep initial extract)) = do
1140 initial >>= \r -> return $ Skip $ ConcatParseStream st []
1141 (PRD.Parser pstep (return r) extract)
1143 -- Buffer is not empty, go to buffered processing loop
1144 stepOuter _ (ConcatParseInit src st
1145 (PRD.Parser pstep initial extract)) = do
1146 initial >>= \r -> return $ Skip $ ConcatParseBuf src st []
1147 (PRD.Parser pstep (return r) extract)
1149 -- XXX we just discard any leftover input at the end
1150 stepOuter _ (ConcatParseInitLeftOver _) = return Stop
1152 -- Buffer is empty process elements from the stream
1153 stepOuter gst (ConcatParseStream st buf
1154 p@(PRD.Parser pstep initial extract)) = do
1155 r <- step (adaptState gst) st
1161 PR.Partial 0 pst1 ->
1162 return $ Skip $ ConcatParseStream s []
1163 (PRD.Parser pstep (return pst1) extract)
1164 PR.Partial n pst1 -> do
1165 assert (n <= length (x:buf)) (return ())
1166 let src0 = Prelude.take n (x:buf)
1167 src = Prelude.reverse src0
1168 return $ Skip $ ConcatParseBuf src s []
1169 (PRD.Parser pstep (return pst1) extract)
1170 -- PR.Continue 0 pst1 ->
1171 -- return $ Skip $ ConcatParseStream s (x:buf) pst1
1172 PR.Continue n pst1 -> do
1173 assert (n <= length (x:buf)) (return ())
1174 let (src0, buf1) = splitAt n (x:buf)
1175 src = Prelude.reverse src0
1176 return $ Skip $ ConcatParseBuf src s buf1
1177 (PRD.Parser pstep (return pst1) extract)
1178 -- XXX Specialize for Stop 0 common case?
1180 assert (n <= length (x:buf)) (return ())
1181 let src = Prelude.reverse (Prelude.take n (x:buf))
1183 ConcatParseYield b (ConcatParseInit src s (func b))
1184 PR.Error err -> throwM $ ParseError err
1185 Skip s -> return $ Skip $ ConcatParseStream s buf p
1189 let src = Prelude.reverse buf
1191 ConcatParseYield b (ConcatParseInitLeftOver src)
1193 -- go back to stream processing mode
1194 stepOuter _ (ConcatParseBuf [] s buf p) =
1195 return $ Skip $ ConcatParseStream s buf p
1197 -- buffered processing loop
1198 stepOuter _ (ConcatParseBuf (x:xs) s buf
1199 (PRD.Parser pstep initial extract)) = do
1203 PR.Partial 0 pst1 ->
1204 return $ Skip $ ConcatParseBuf xs s []
1205 (PRD.Parser pstep (return pst1) extract)
1206 PR.Partial n pst1 -> do
1207 assert (n <= length (x:buf)) (return ())
1208 let src0 = Prelude.take n (x:buf)
1209 src = Prelude.reverse src0 ++ xs
1210 return $ Skip $ ConcatParseBuf src s []
1211 (PRD.Parser pstep (return pst1) extract)
1212 -- PR.Continue 0 pst1 -> return $ Skip $ ConcatParseBuf xs s (x:buf) pst1
1213 PR.Continue n pst1 -> do
1214 assert (n <= length (x:buf)) (return ())
1215 let (src0, buf1) = splitAt n (x:buf)
1216 src = Prelude.reverse src0 ++ xs
1217 return $ Skip $ ConcatParseBuf src s buf1
1218 (PRD.Parser pstep (return pst1) extract)
1219 -- XXX Specialize for Stop 0 common case?
1221 assert (n <= length (x:buf)) (return ())
1222 let src = Prelude.reverse (Prelude.take n (x:buf)) ++ xs
1223 return $ Skip $ ConcatParseYield b
1224 (ConcatParseInit src s (func b))
1225 PR.Error err -> throwM $ ParseError err
1227 stepOuter _ (ConcatParseYield a next) = return $ Yield a next
1229------------------------------------------------------------------------------
1231------------------------------------------------------------------------------
1233-- | Run a streaming composition, discard the results.
1234{-# INLINE_LATE drain #-}
1235drain :: Monad m => Stream m a -> m ()
1236-- drain = foldrM (\_ xs -> xs) (return ())
1237drain (Stream step state) = go SPEC state
1240 r <- step defState st
1242 Yield _ s -> go SPEC s
1246{-# INLINE_NORMAL null #-}
1247null :: Monad m => Stream m a -> m Bool
1248null m = foldrM (\_ _ -> return False) (return True) m
1250{-# INLINE_NORMAL head #-}
1251head :: Monad m => Stream m a -> m (Maybe a)
1252head m = foldrM (\x _ -> return (Just x)) (return Nothing) m
1254{-# INLINE_NORMAL headElse #-}
1255headElse :: Monad m => a -> Stream m a -> m a
1256headElse a m = foldrM (\x _ -> return x) (return a) m
1258-- Does not fuse, has the same performance as the StreamK version.
1259{-# INLINE_NORMAL tail #-}
1260tail :: Monad m => Stream m a -> m (Maybe (Stream m a))
1261tail (UnStream step state) = go state
1264 r <- step defState st
1266 Yield _ s -> return (Just $ Stream step s)
1268 Stop -> return Nothing
1270-- XXX will it fuse? need custom impl?
1271{-# INLINE_NORMAL last #-}
1272last :: Monad m => Stream m a -> m (Maybe a)
1273last = foldl' (\_ y -> Just y) Nothing
1275{-# INLINE_NORMAL elem #-}
1276elem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
1277-- elem e m = foldrM (\x xs -> if x == e then return True else xs) (return False) m
1278elem e (Stream step state) = go state
1281 r <- step defState st
1284 | x == e -> return True
1287 Stop -> return False
1289{-# INLINE_NORMAL notElem #-}
1290notElem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
1291notElem e s = fmap not (elem e s)
1293{-# INLINE_NORMAL all #-}
1294all :: Monad m => (a -> Bool) -> Stream m a -> m Bool
1295-- all p m = foldrM (\x xs -> if p x then xs else return False) (return True) m
1296all p (Stream step state) = go state
1299 r <- step defState st
1303 | otherwise -> return False
1307{-# INLINE_NORMAL any #-}
1308any :: Monad m => (a -> Bool) -> Stream m a -> m Bool
1309-- any p m = foldrM (\x xs -> if p x then return True else xs) (return False) m
1310any p (Stream step state) = go state
1313 r <- step defState st
1316 | p x -> return True
1319 Stop -> return False
1321{-# INLINE_NORMAL maximum #-}
1322maximum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
1323maximum (Stream step state) = go Nothing state
1326 r <- step defState st
1328 Yield x s -> go (Just x) s
1329 Skip s -> go Nothing s
1330 Stop -> return Nothing
1331 go (Just acc) st = do
1332 r <- step defState st
1335 | acc <= x -> go (Just x) s
1336 | otherwise -> go (Just acc) s
1337 Skip s -> go (Just acc) s
1338 Stop -> return (Just acc)
1340{-# INLINE_NORMAL maximumBy #-}
1341maximumBy :: Monad m => (a -> a -> Ordering) -> Stream m a -> m (Maybe a)
1342maximumBy cmp (Stream step state) = go Nothing state
1345 r <- step defState st
1347 Yield x s -> go (Just x) s
1348 Skip s -> go Nothing s
1349 Stop -> return Nothing
1350 go (Just acc) st = do
1351 r <- step defState st
1353 Yield x s -> case cmp acc x of
1354 GT -> go (Just acc) s
1356 Skip s -> go (Just acc) s
1357 Stop -> return (Just acc)
1359{-# INLINE_NORMAL minimum #-}
1360minimum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
1361minimum (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 acc) s
1374 | otherwise -> go (Just x) s
1375 Skip s -> go (Just acc) s
1376 Stop -> return (Just acc)
1378{-# INLINE_NORMAL minimumBy #-}
1379minimumBy :: Monad m => (a -> a -> Ordering) -> Stream m a -> m (Maybe a)
1380minimumBy 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
1393 _ -> go (Just acc) s
1394 Skip s -> go (Just acc) s
1395 Stop -> return (Just acc)
1397{-# INLINE_NORMAL (!!) #-}
1398(!!) :: (Monad m) => Stream m a -> Int -> m (Maybe a)
1399(Stream step state) !! i = go i state
1402 r <- step defState st
1404 Yield x s | n < 0 -> return Nothing
1405 | n == 0 -> return $ Just x
1406 | otherwise -> go (n - 1) s
1408 Stop -> return Nothing
1410{-# INLINE_NORMAL lookup #-}
1411lookup :: (Monad m, Eq a) => a -> Stream m (a, b) -> m (Maybe b)
1412lookup e m = foldrM (\(a, b) xs -> if e == a then return (Just b) else xs)
1415{-# INLINE_NORMAL findM #-}
1416findM :: Monad m => (a -> m Bool) -> Stream m a -> m (Maybe a)
1417findM p m = foldrM (\x xs -> p x >>= \r -> if r then return (Just x) else xs)
1421find :: Monad m => (a -> Bool) -> Stream m a -> m (Maybe a)
1422find p = findM (return . p)
1424{-# INLINE_NORMAL findIndices #-}
1425findIndices :: Monad m => (a -> Bool) -> Stream m a -> Stream m Int
1426findIndices p (Stream step state) = Stream step' (state, 0)
1428 {-# INLINE_LATE step' #-}
1429 step' gst (st, i) = i `seq` do
1430 r <- step (adaptState gst) st
1432 Yield x s -> if p x then Yield i (s, i+1) else Skip (s, i+1)
1433 Skip s -> Skip (s, i)
1436{-# INLINE toListRev #-}
1437toListRev :: Monad m => Stream m a -> m [a]
1438toListRev = foldl' (flip (:)) []
1440-- We can implement reverse as:
1442-- > reverse = foldlS (flip cons) nil
1444-- However, this implementation is unusable because of the horrible performance
1445-- of cons. So we just convert it to a list first and then stream from the
1448-- XXX Maybe we can use an Array instead of a list here?
1449{-# INLINE_NORMAL reverse #-}
1450reverse :: Monad m => Stream m a -> Stream m a
1451reverse m = Stream step Nothing
1453 {-# INLINE_LATE step #-}
1456 return $ Skip (Just xs)
1457 step _ (Just (x:xs)) = return $ Yield x (Just xs)
1458 step _ (Just []) = return Stop
1460-- Much faster reverse for Storables
1461{-# INLINE_NORMAL reverse' #-}
1462reverse' :: forall m a. (MonadIO m, Storable a) => Stream m a -> Stream m a
1464-- This commented implementation copies the whole stream into one single array
1465-- and then streams from that array, this is 3-4x faster than the chunked code
1466-- that follows. Though this could be problematic due to unbounded large
1467-- allocations. We need to figure out why the chunked code is slower and if we
1468-- can optimize the chunked code to work as fast as this one. It may be a
1470import Foreign.ForeignPtr (touchForeignPtr)
1471import Foreign.ForeignPtr.Unsafe (unsafeForeignPtrToPtr)
1472import Foreign.Ptr (Ptr, plusPtr)
1473reverse' m = Stream step Nothing
1475 {-# INLINE_LATE step #-}
1477 arr <- A.fromStreamD m
1478 let p = aEnd arr `plusPtr` negate (sizeOf (undefined :: a))
1479 return $ Skip $ Just (aStart arr, p)
1481 step _ (Just (start, p)) | p < unsafeForeignPtrToPtr start = return Stop
1483 step _ (Just (start, p)) = do
1484 let !x = A.unsafeInlineIO $ do
1486 touchForeignPtr start
1488 next = p `plusPtr` negate (sizeOf (undefined :: a))
1489 return $ Yield x (Just (start, next))
1496 $ A.fromStreamDArraysOf A.defaultChunkSize m
1499------------------------------------------------------------------------------
1500-- Grouping/Splitting
1501------------------------------------------------------------------------------
1503{-# INLINE_NORMAL splitSuffixBy' #-}
1504splitSuffixBy' :: Monad m
1505 => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1506splitSuffixBy' predicate f (Stream step state) =
1507 Stream (stepOuter f) (Just state)
1511 {-# INLINE_LATE stepOuter #-}
1512 stepOuter (Fold fstep initial done) gst (Just st) = do
1513 res <- step (adaptState gst) st
1519 then done acc' >>= \val -> return $ Yield val (Just s)
1522 Skip s -> return $ Skip $ Just s
1528 res <- step (adaptState gst) stt
1533 then done acc' >>= \val -> return $ Yield val (Just s)
1535 Skip s -> go SPEC s acc
1536 Stop -> done acc >>= \val -> return $ Yield val Nothing
1538 stepOuter _ _ Nothing = return Stop
1540{-# INLINE_NORMAL groupsBy #-}
1546groupsBy cmp f (Stream step state) = Stream (stepOuter f) (Just state, Nothing)
1550 {-# INLINE_LATE stepOuter #-}
1551 stepOuter (Fold fstep initial done) gst (Just st, Nothing) = do
1552 res <- step (adaptState gst) st
1559 Skip s -> return $ Skip $ (Just s, Nothing)
1564 go !_ prev stt !acc = do
1565 res <- step (adaptState gst) stt
1572 else done acc >>= \r -> return $ Yield r (Just s, Just x)
1573 Skip s -> go SPEC prev s acc
1574 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1576 stepOuter (Fold fstep initial done) gst (Just st, Just prev) = do
1578 acc' <- fstep acc prev
1583 -- XXX code duplicated from the previous equation
1585 res <- step (adaptState gst) stt
1592 else done acc >>= \r -> return $ Yield r (Just s, Just x)
1593 Skip s -> go SPEC s acc
1594 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1596 stepOuter _ _ (Nothing,_) = return Stop
1598{-# INLINE_NORMAL groupsRollingBy #-}
1599groupsRollingBy :: Monad m
1604groupsRollingBy cmp f (Stream step state) =
1605 Stream (stepOuter f) (Just state, Nothing)
1608 {-# INLINE_LATE stepOuter #-}
1609 stepOuter (Fold fstep initial done) gst (Just st, Nothing) = do
1610 res <- step (adaptState gst) st
1617 Skip s -> return $ Skip $ (Just s, Nothing)
1621 go !_ prev stt !acc = do
1622 res <- step (adaptState gst) stt
1630 done acc >>= \r -> return $ Yield r (Just s, Just x)
1631 Skip s -> go SPEC prev s acc
1632 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1634 stepOuter (Fold fstep initial done) gst (Just st, Just prev') = do
1636 acc' <- fstep acc prev'
1637 go SPEC prev' st acc'
1640 go !_ prevv stt !acc = do
1641 res <- step (adaptState gst) stt
1648 else done acc >>= \r -> return $ Yield r (Just s, Just x)
1649 Skip s -> go SPEC prevv s acc
1650 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1652 stepOuter _ _ (Nothing, _) = return Stop
1654{-# INLINE_NORMAL splitBy #-}
1655splitBy :: Monad m => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1656splitBy predicate f (Stream step state) = Stream (step' f) (Just state)
1660 {-# INLINE_LATE step' #-}
1661 step' (Fold fstep initial done) gst (Just st) = initial >>= go SPEC st
1666 res <- step (adaptState gst) stt
1670 then done acc >>= \r -> return $ Yield r (Just s)
1674 Skip s -> go SPEC s acc
1675 Stop -> done acc >>= \r -> return $ Yield r Nothing
1677 step' _ _ Nothing = return Stop
1679-- XXX requires -funfolding-use-threshold=150 in lines-unlines benchmark
1680{-# INLINE_NORMAL splitSuffixBy #-}
1681splitSuffixBy :: Monad m
1682 => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1683splitSuffixBy predicate f (Stream step state) = Stream (step' f) (Just state)
1687 {-# INLINE_LATE step' #-}
1688 step' (Fold fstep initial done) gst (Just st) = do
1689 res <- step (adaptState gst) st
1694 then done acc >>= \val -> return $ Yield val (Just s)
1699 Skip s -> return $ Skip $ Just s
1705 res <- step (adaptState gst) stt
1709 then done acc >>= \r -> return $ Yield r (Just s)
1713 Skip s -> go SPEC s acc
1714 Stop -> done acc >>= \r -> return $ Yield r Nothing
1716 step' _ _ Nothing = return Stop
1718{-# INLINE_NORMAL wordsBy #-}
1719wordsBy :: Monad m => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1720wordsBy predicate f (Stream step state) = Stream (stepOuter f) (Just state)
1724 {-# INLINE_LATE stepOuter #-}
1725 stepOuter (Fold fstep initial done) gst (Just st) = do
1726 res <- step (adaptState gst) st
1730 then return $ Skip (Just s)
1736 Skip s -> return $ Skip $ Just s
1742 res <- step (adaptState gst) stt
1746 then done acc >>= \r -> return $ Yield r (Just s)
1750 Skip s -> go SPEC s acc
1751 Stop -> done acc >>= \r -> return $ Yield r Nothing
1753 stepOuter _ _ Nothing = return Stop
1755-- String search algorithms:
1756-- http://www-igm.univ-mlv.fr/~lecroq/string/index.html
1759-- TODO can we unify the splitting operations using a splitting configuration
1760-- like in the split package.
1762data SplitStyle = Infix | Suffix | Prefix deriving (Eq, Show)
1764data SplitOptions = SplitOptions
1765 { style :: SplitStyle
1766 , withSep :: Bool -- ^ keep the separators in output
1767 -- , compact :: Bool -- ^ treat multiple consecutive separators as one
1768 -- , trimHead :: Bool -- ^ drop blank at head
1769 -- , trimTail :: Bool -- ^ drop blank at tail
1773data SplitOnState s a =
1778 | GO_KARP_RABIN s !(RB.Ring a) !(Ptr a)
1781{-# INLINE_NORMAL splitOn #-}
1783 :: forall m a b. (MonadIO m, Storable a, Enum a, Eq a)
1788splitOn patArr (Fold fstep initial done) (Stream step state) =
1789 Stream stepOuter GO_START
1793 patLen = A.length patArr
1794 maxIndex = patLen - 1
1795 elemBits = sizeOf (undefined :: a) * 8
1797 {-# INLINE_LATE stepOuter #-}
1798 stepOuter _ GO_START =
1800 then return $ Skip $ GO_EMPTY_PAT state
1803 r <- liftIO $ (A.unsafeIndexIO patArr 0)
1804 return $ Skip $ GO_SINGLE_PAT state r
1805 else if sizeOf (undefined :: a) * patLen
1806 <= sizeOf (undefined :: Word)
1807 then return $ Skip $ GO_SHORT_PAT state
1809 (rb, rhead) <- liftIO $ RB.new patLen
1810 return $ Skip $ GO_KARP_RABIN state rb rhead
1812 stepOuter gst (GO_SINGLE_PAT stt pat) = initial >>= go SPEC stt
1817 res <- step (adaptState gst) st
1823 return $ Yield r (GO_SINGLE_PAT s pat)
1824 else fstep acc x >>= go SPEC s
1825 Skip s -> go SPEC s acc
1826 Stop -> done acc >>= \r -> return $ Yield r GO_DONE
1828 stepOuter gst (GO_SHORT_PAT stt) = initial >>= go0 SPEC 0 (0 :: Word) stt
1833 mask = (1 `shiftL` (elemBits * patLen)) - 1
1835 addToWord wrd a = (wrd `shiftL` elemBits) .|. fromIntegral (fromEnum a)
1838 patWord = mask .&. A.foldl' addToWord 0 patArr
1840 go0 !_ !idx wrd st !acc = do
1841 res <- step (adaptState gst) st
1844 let wrd' = addToWord wrd x
1847 if wrd' .&. mask == patWord
1850 return $ Yield r (GO_SHORT_PAT s)
1851 else go1 SPEC wrd' s acc
1852 else go0 SPEC (idx + 1) wrd' s acc
1853 Skip s -> go0 SPEC idx wrd s acc
1856 then go2 wrd idx acc
1858 done acc' >>= \r -> return $ Yield r GO_DONE
1861 go1 !_ wrd st !acc = do
1862 res <- step (adaptState gst) st
1865 let wrd' = addToWord wrd x
1866 old = (mask .&. wrd) `shiftR` (elemBits * (patLen - 1))
1867 acc' <- fstep acc (toEnum $ fromIntegral old)
1868 if wrd' .&. mask == patWord
1869 then done acc' >>= \r -> return $ Yield r (GO_SHORT_PAT s)
1870 else go1 SPEC wrd' s acc'
1871 Skip s -> go1 SPEC wrd s acc
1873 acc' <- go2 wrd patLen acc
1874 done acc' >>= \r -> return $ Yield r GO_DONE
1876 go2 !wrd !n !acc | n > 0 = do
1877 let old = (mask .&. wrd) `shiftR` (elemBits * (n - 1))
1878 fstep acc (toEnum $ fromIntegral old) >>= go2 wrd (n - 1)
1879 go2 _ _ acc = return acc
1881 stepOuter gst (GO_KARP_RABIN stt rb rhead) = do
1882 initial >>= go0 SPEC 0 rhead stt
1886 k = 2891336453 :: Word32
1888 addCksum cksum a = cksum * k + fromIntegral (fromEnum a)
1889 deltaCksum cksum old new =
1890 addCksum cksum new - coeff * fromIntegral (fromEnum old)
1892 -- XXX shall we use a random starting hash or 1 instead of 0?
1893 patHash = A.foldl' addCksum 0 patArr
1896 go0 !_ !idx !rh st !acc = do
1897 res <- step (adaptState gst) st
1900 rh' <- liftIO $ RB.unsafeInsert rb rh x
1903 let fold = RB.unsafeFoldRing (RB.ringBound rb)
1904 let !ringHash = fold addCksum 0 rb
1905 if ringHash == patHash
1906 then go2 SPEC ringHash rh' s acc
1907 else go1 SPEC ringHash rh' s acc
1908 else go0 SPEC (idx + 1) rh' s acc
1909 Skip s -> go0 SPEC idx rh s acc
1911 !acc' <- if idx /= 0
1912 then RB.unsafeFoldRingM rh fstep acc rb
1914 done acc' >>= \r -> return $ Yield r GO_DONE
1916 -- XXX Theoretically this code can do 4 times faster if GHC generates
1917 -- optimal code. If we use just "(cksum' == patHash)" condition it goes
1918 -- 4x faster, as soon as we add the "RB.unsafeEqArray rb v" condition
1919 -- the generated code changes drastically and becomes 4x slower. Need
1920 -- to investigate what is going on with GHC.
1922 go1 !_ !cksum !rh st !acc = do
1923 res <- step (adaptState gst) st
1926 old <- liftIO $ peek rh
1927 let cksum' = deltaCksum cksum old x
1928 acc' <- fstep acc old
1930 if (cksum' == patHash)
1932 rh' <- liftIO (RB.unsafeInsert rb rh x)
1933 go2 SPEC cksum' rh' s acc'
1935 rh' <- liftIO (RB.unsafeInsert rb rh x)
1936 go1 SPEC cksum' rh' s acc'
1937 Skip s -> go1 SPEC cksum rh s acc
1939 acc' <- RB.unsafeFoldRingFullM rh fstep acc rb
1940 done acc' >>= \r -> return $ Yield r GO_DONE
1942 go2 !_ !cksum' !rh' s !acc' = do
1943 if RB.unsafeEqArray rb rh' patArr
1946 return $ Yield r (GO_KARP_RABIN s rb rhead)
1947 else go1 SPEC cksum' rh' s acc'
1949 stepOuter gst (GO_EMPTY_PAT st) = do
1950 res <- step (adaptState gst) st
1955 done acc' >>= \r -> return $ Yield r (GO_EMPTY_PAT s)
1956 Skip s -> return $ Skip (GO_EMPTY_PAT s)
1959 stepOuter _ GO_DONE = return Stop
1961{-# INLINE_NORMAL splitSuffixOn #-}
1963 :: forall m a b. (MonadIO m, Storable a, Enum a, Eq a)
1969splitSuffixOn withSep patArr (Fold fstep initial done)
1970 (Stream step state) =
1971 Stream stepOuter GO_START
1975 patLen = A.length patArr
1976 maxIndex = patLen - 1
1977 elemBits = sizeOf (undefined :: a) * 8
1979 {-# INLINE_LATE stepOuter #-}
1980 stepOuter _ GO_START =
1982 then return $ Skip $ GO_EMPTY_PAT state
1985 r <- liftIO $ (A.unsafeIndexIO patArr 0)
1986 return $ Skip $ GO_SINGLE_PAT state r
1987 else if sizeOf (undefined :: a) * patLen
1988 <= sizeOf (undefined :: Word)
1989 then return $ Skip $ GO_SHORT_PAT state
1991 (rb, rhead) <- liftIO $ RB.new patLen
1992 return $ Skip $ GO_KARP_RABIN state rb rhead
1994 stepOuter gst (GO_SINGLE_PAT stt pat) = do
1995 -- This first part is the only difference between splitOn and
1997 -- If the last element is a separator do not issue a blank segment.
1998 res <- step (adaptState gst) stt
2004 acc' <- if withSep then fstep acc x else return acc
2005 done acc' >>= \r -> return $ Yield r (GO_SINGLE_PAT s pat)
2006 else fstep acc x >>= go SPEC s
2007 Skip s -> return $ Skip $ (GO_SINGLE_PAT s pat)
2012 -- This is identical for splitOn and splitSuffixOn
2014 res <- step (adaptState gst) st
2019 acc' <- if withSep then fstep acc x else return acc
2021 return $ Yield r (GO_SINGLE_PAT s pat)
2022 else fstep acc x >>= go SPEC s
2023 Skip s -> go SPEC s acc
2024 Stop -> done acc >>= \r -> return $ Yield r GO_DONE
2026 stepOuter gst (GO_SHORT_PAT stt) = do
2028 -- Call "initial" only if the stream yields an element, otherwise we
2029 -- may call "initial" but never yield anything. initial may produce a
2030 -- side effect, therefore we will end up doing and discard a side
2035 res <- step (adaptState gst) stt
2039 let wrd' = addToWord wrd x
2040 acc' <- if withSep then fstep acc x else return acc
2043 if wrd' .&. mask == patWord
2044 then done acc' >>= \r -> return $ Yield r (GO_SHORT_PAT s)
2045 else go0 SPEC (idx + 1) wrd' s acc'
2046 else go0 SPEC (idx + 1) wrd' s acc'
2047 Skip s -> return $ Skip (GO_SHORT_PAT s)
2053 mask = (1 `shiftL` (elemBits * patLen)) - 1
2055 addToWord wrd a = (wrd `shiftL` elemBits) .|. fromIntegral (fromEnum a)
2058 patWord = mask .&. A.foldl' addToWord 0 patArr
2060 go0 !_ !idx wrd st !acc = do
2061 res <- step (adaptState gst) st
2064 let wrd' = addToWord wrd x
2065 acc' <- if withSep then fstep acc x else return acc
2068 if wrd' .&. mask == patWord
2071 return $ Yield r (GO_SHORT_PAT s)
2072 else go1 SPEC wrd' s acc'
2073 else go0 SPEC (idx + 1) wrd' s acc'
2074 Skip s -> go0 SPEC idx wrd s acc
2076 if (idx == maxIndex) && (wrd .&. mask == patWord)
2079 acc' <- if idx /= 0 && not withSep
2080 then go2 wrd idx acc
2082 done acc' >>= \r -> return $ Yield r GO_DONE
2085 go1 !_ wrd st !acc = do
2086 res <- step (adaptState gst) st
2089 let wrd' = addToWord wrd x
2090 old = (mask .&. wrd) `shiftR` (elemBits * (patLen - 1))
2093 else fstep acc (toEnum $ fromIntegral old)
2094 if wrd' .&. mask == patWord
2095 then done acc' >>= \r -> return $ Yield r (GO_SHORT_PAT s)
2096 else go1 SPEC wrd' s acc'
2097 Skip s -> go1 SPEC wrd s acc
2099 -- If the last sequence is a separator do not issue a blank
2101 if wrd .&. mask == patWord
2106 else go2 wrd patLen acc
2107 done acc' >>= \r -> return $ Yield r GO_DONE
2109 go2 !wrd !n !acc | n > 0 = do
2110 let old = (mask .&. wrd) `shiftR` (elemBits * (n - 1))
2111 fstep acc (toEnum $ fromIntegral old) >>= go2 wrd (n - 1)
2112 go2 _ _ acc = return acc
2114 stepOuter gst (GO_KARP_RABIN stt rb rhead) = do
2116 res <- step (adaptState gst) stt
2120 acc' <- if withSep then fstep acc x else return acc
2121 rh' <- liftIO (RB.unsafeInsert rb rhead x)
2124 let fold = RB.unsafeFoldRing (RB.ringBound rb)
2125 let !ringHash = fold addCksum 0 rb
2126 if ringHash == patHash
2127 then go2 SPEC ringHash rh' s acc'
2128 else go0 SPEC (idx + 1) rh' s acc'
2129 else go0 SPEC (idx + 1) rh' s acc'
2130 Skip s -> return $ Skip (GO_KARP_RABIN s rb rhead)
2135 k = 2891336453 :: Word32
2137 addCksum cksum a = cksum * k + fromIntegral (fromEnum a)
2138 deltaCksum cksum old new =
2139 addCksum cksum new - coeff * fromIntegral (fromEnum old)
2141 -- XXX shall we use a random starting hash or 1 instead of 0?
2142 patHash = A.foldl' addCksum 0 patArr
2145 go0 !_ !idx !rh st !acc = do
2146 res <- step (adaptState gst) st
2149 acc' <- if withSep then fstep acc x else return acc
2150 rh' <- liftIO (RB.unsafeInsert rb rh x)
2153 let fold = RB.unsafeFoldRing (RB.ringBound rb)
2154 let !ringHash = fold addCksum 0 rb
2155 if ringHash == patHash
2156 then go2 SPEC ringHash rh' s acc'
2157 else go1 SPEC ringHash rh' s acc'
2158 else go0 SPEC (idx + 1) rh' s acc'
2159 Skip s -> go0 SPEC idx rh s acc
2161 -- do not issue a blank segment when we end at pattern
2162 if (idx == maxIndex) && RB.unsafeEqArray rb rh patArr
2165 !acc' <- if idx /= 0 && not withSep
2166 then RB.unsafeFoldRingM rh fstep acc rb
2168 done acc' >>= \r -> return $ Yield r GO_DONE
2170 -- XXX Theoretically this code can do 4 times faster if GHC generates
2171 -- optimal code. If we use just "(cksum' == patHash)" condition it goes
2172 -- 4x faster, as soon as we add the "RB.unsafeEqArray rb v" condition
2173 -- the generated code changes drastically and becomes 4x slower. Need
2174 -- to investigate what is going on with GHC.
2176 go1 !_ !cksum !rh st !acc = do
2177 res <- step (adaptState gst) st
2180 old <- liftIO $ peek rh
2181 let cksum' = deltaCksum cksum old x
2186 if (cksum' == patHash)
2188 rh' <- liftIO (RB.unsafeInsert rb rh x)
2189 go2 SPEC cksum' rh' s acc'
2191 rh' <- liftIO (RB.unsafeInsert rb rh x)
2192 go1 SPEC cksum' rh' s acc'
2193 Skip s -> go1 SPEC cksum rh s acc
2195 if RB.unsafeEqArray rb rh patArr
2200 else RB.unsafeFoldRingFullM rh fstep acc rb
2201 done acc' >>= \r -> return $ Yield r GO_DONE
2203 go2 !_ !cksum' !rh' s !acc' = do
2204 if RB.unsafeEqArray rb rh' patArr
2207 return $ Yield r (GO_KARP_RABIN s rb rhead)
2208 else go1 SPEC cksum' rh' s acc'
2210 stepOuter gst (GO_EMPTY_PAT st) = do
2211 res <- step (adaptState gst) st
2216 done acc' >>= \r -> return $ Yield r (GO_EMPTY_PAT s)
2217 Skip s -> return $ Skip (GO_EMPTY_PAT s)
2220 stepOuter _ GO_DONE = return Stop
2222data SplitState s arr
2224 | SplitBuffering s arr
2225 | SplitSplitting s arr
2226 | SplitYielding arr (SplitState s arr)
2229-- XXX An alternative approach would be to use a partial fold (Fold m a b) to
2230-- split using a splitBy like combinator. The Fold would consume upto the
2231-- separator and return any leftover which can then be fed to the next fold.
2233-- We can revisit this once we have partial folds/parsers.
2235-- | Performs infix separator style splitting.
2236{-# INLINE_NORMAL splitInnerBy #-}
2239 => (f a -> m (f a, Maybe (f a))) -- splitter
2240 -> (f a -> f a -> m (f a)) -- joiner
2243splitInnerBy splitter joiner (Stream step1 state1) =
2244 (Stream step (SplitInitial state1))
2248 {-# INLINE_LATE step #-}
2249 step gst (SplitInitial st) = do
2253 (x1, mx2) <- splitter x
2254 return $ case mx2 of
2255 Nothing -> Skip (SplitBuffering s x1)
2256 Just x2 -> Skip (SplitYielding x1 (SplitSplitting s x2))
2257 Skip s -> return $ Skip (SplitInitial s)
2258 Stop -> return $ Stop
2260 step gst (SplitBuffering st buf) = do
2264 (x1, mx2) <- splitter x
2265 buf' <- joiner buf x1
2266 return $ case mx2 of
2267 Nothing -> Skip (SplitBuffering s buf')
2268 Just x2 -> Skip (SplitYielding buf' (SplitSplitting s x2))
2269 Skip s -> return $ Skip (SplitBuffering s buf)
2270 Stop -> return $ Skip (SplitYielding buf SplitFinishing)
2272 step _ (SplitSplitting st buf) = do
2273 (x1, mx2) <- splitter buf
2274 return $ case mx2 of
2275 Nothing -> Skip $ SplitBuffering st x1
2276 Just x2 -> Skip $ SplitYielding x1 (SplitSplitting st x2)
2278 step _ (SplitYielding x next) = return $ Yield x next
2279 step _ SplitFinishing = return $ Stop
2281-- | Performs infix separator style splitting.
2282{-# INLINE_NORMAL splitInnerBySuffix #-}
2284 :: (Monad m, Eq (f a), Monoid (f a))
2285 => (f a -> m (f a, Maybe (f a))) -- splitter
2286 -> (f a -> f a -> m (f a)) -- joiner
2289splitInnerBySuffix splitter joiner (Stream step1 state1) =
2290 (Stream step (SplitInitial state1))
2294 {-# INLINE_LATE step #-}
2295 step gst (SplitInitial st) = do
2299 (x1, mx2) <- splitter x
2300 return $ case mx2 of
2301 Nothing -> Skip (SplitBuffering s x1)
2302 Just x2 -> Skip (SplitYielding x1 (SplitSplitting s x2))
2303 Skip s -> return $ Skip (SplitInitial s)
2304 Stop -> return $ Stop
2306 step gst (SplitBuffering st buf) = do
2310 (x1, mx2) <- splitter x
2311 buf' <- joiner buf x1
2312 return $ case mx2 of
2313 Nothing -> Skip (SplitBuffering s buf')
2314 Just x2 -> Skip (SplitYielding buf' (SplitSplitting s x2))
2315 Skip s -> return $ Skip (SplitBuffering s buf)
2319 else Skip (SplitYielding buf SplitFinishing)
2321 step _ (SplitSplitting st buf) = do
2322 (x1, mx2) <- splitter buf
2323 return $ case mx2 of
2324 Nothing -> Skip $ SplitBuffering st x1
2325 Just x2 -> Skip $ SplitYielding x1 (SplitSplitting st x2)
2327 step _ (SplitYielding x next) = return $ Yield x next
2328 step _ SplitFinishing = return $ Stop
2330------------------------------------------------------------------------------
2332------------------------------------------------------------------------------
2334{-# INLINE_NORMAL isPrefixOf #-}
2335isPrefixOf :: (Eq a, Monad m) => Stream m a -> Stream m a -> m Bool
2336isPrefixOf (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
2338 go (sa, sb, Nothing) = do
2339 r <- stepa defState sa
2341 Yield x sa' -> go (sa', sb, Just x)
2342 Skip sa' -> go (sa', sb, Nothing)
2345 go (sa, sb, Just x) = do
2346 r <- stepb defState sb
2350 then go (sa, sb', Nothing)
2352 Skip sb' -> go (sa, sb', Just x)
2353 Stop -> return False
2355{-# INLINE_NORMAL isSubsequenceOf #-}
2356isSubsequenceOf :: (Eq a, Monad m) => Stream m a -> Stream m a -> m Bool
2357isSubsequenceOf (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
2359 go (sa, sb, Nothing) = do
2360 r <- stepa defState sa
2362 Yield x sa' -> go (sa', sb, Just x)
2363 Skip sa' -> go (sa', sb, Nothing)
2366 go (sa, sb, Just x) = do
2367 r <- stepb defState sb
2371 then go (sa, sb', Nothing)
2372 else go (sa, sb', Just x)
2373 Skip sb' -> go (sa, sb', Just x)
2374 Stop -> return False
2376{-# INLINE_NORMAL stripPrefix #-}
2379 => Stream m a -> Stream m a -> m (Maybe (Stream m a))
2380stripPrefix (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
2382 go (sa, sb, Nothing) = do
2383 r <- stepa defState sa
2385 Yield x sa' -> go (sa', sb, Just x)
2386 Skip sa' -> go (sa', sb, Nothing)
2387 Stop -> return $ Just (Stream stepb sb)
2389 go (sa, sb, Just x) = do
2390 r <- stepb defState sb
2394 then go (sa, sb', Nothing)
2396 Skip sb' -> go (sa, sb', Just x)
2397 Stop -> return Nothing
2399------------------------------------------------------------------------------
2401------------------------------------------------------------------------------
2403-- | Execute a monadic action for each element of the 'Stream'
2404{-# INLINE_NORMAL mapM_ #-}
2405mapM_ :: Monad m => (a -> m b) -> Stream m a -> m ()
2406mapM_ m = drain . mapM m
2408-------------------------------------------------------------------------------
2409-- Stream transformations using Unfolds
2410-------------------------------------------------------------------------------
2412-- Define a unique structure to use in inspection testing
2413data ConcatMapUState o i =
2415 | ConcatMapUInner o i
2417-- | @concatMapU unfold stream@ uses @unfold@ to map the input stream elements
2418-- to streams and then flattens the generated streams into a single output
2421-- This is like 'concatMap' but uses an unfold with an explicit state to
2422-- generate the stream instead of a 'Stream' type generator. This allows better
2423-- optimization via fusion. This can be many times more efficient than
2426{-# INLINE_NORMAL concatMapU #-}
2427concatMapU :: Monad m => Unfold m a b -> Stream m a -> Stream m b
2428concatMapU (Unfold istep inject) (Stream ostep ost) =
2429 Stream step (ConcatMapUOuter ost)
2431 {-# INLINE_LATE step #-}
2432 step gst (ConcatMapUOuter o) = do
2433 r <- ostep (adaptState gst) o
2437 i `seq` return (Skip (ConcatMapUInner o' i))
2438 Skip o' -> return $ Skip (ConcatMapUOuter o')
2439 Stop -> return $ Stop
2441 step _ (ConcatMapUInner o i) = do
2444 Yield x i' -> Yield x (ConcatMapUInner o i')
2445 Skip i' -> Skip (ConcatMapUInner o i')
2446 Stop -> Skip (ConcatMapUOuter o)
2448data ConcatUnfoldInterleaveState o i =
2449 ConcatUnfoldInterleaveOuter o [i]
2450 | ConcatUnfoldInterleaveInner o [i]
2451 | ConcatUnfoldInterleaveInnerL [i] [i]
2452 | ConcatUnfoldInterleaveInnerR [i] [i]
2454-- XXX use arrays to store state instead of lists.
2455-- XXX In general we can use different scheduling strategies e.g. how to
2456-- schedule the outer vs inner loop or assigning weights to different streams
2457-- or outer and inner loops.
2459-- After a yield, switch to the next stream. Do not switch streams on Skip.
2460-- Yield from outer stream switches to the inner stream.
2462-- There are two choices here, (1) exhaust the outer stream first and then
2463-- start yielding from the inner streams, this is much simpler to implement,
2464-- (2) yield at least one element from an inner stream before going back to
2465-- outer stream and opening the next stream from it.
2467-- Ideally, we need some scheduling bias to inner streams vs outer stream.
2468-- Maybe we can configure the behavior.
2470{-# INLINE_NORMAL concatUnfoldInterleave #-}
2471concatUnfoldInterleave :: Monad m => Unfold m a b -> Stream m a -> Stream m b
2472concatUnfoldInterleave (Unfold istep inject) (Stream ostep ost) =
2473 Stream step (ConcatUnfoldInterleaveOuter ost [])
2475 {-# INLINE_LATE step #-}
2476 step gst (ConcatUnfoldInterleaveOuter o ls) = do
2477 r <- ostep (adaptState gst) o
2481 i `seq` return (Skip (ConcatUnfoldInterleaveInner o' (i : ls)))
2482 Skip o' -> return $ Skip (ConcatUnfoldInterleaveOuter o' ls)
2483 Stop -> return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2485 step _ (ConcatUnfoldInterleaveInner _ []) = undefined
2486 step _ (ConcatUnfoldInterleaveInner o (st:ls)) = do
2489 Yield x s -> Yield x (ConcatUnfoldInterleaveOuter o (s:ls))
2490 Skip s -> Skip (ConcatUnfoldInterleaveInner o (s:ls))
2491 Stop -> Skip (ConcatUnfoldInterleaveOuter o ls)
2493 step _ (ConcatUnfoldInterleaveInnerL [] []) = return Stop
2494 step _ (ConcatUnfoldInterleaveInnerL [] rs) =
2495 return $ Skip (ConcatUnfoldInterleaveInnerR [] rs)
2497 step _ (ConcatUnfoldInterleaveInnerL (st:ls) rs) = do
2500 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerL ls (s:rs))
2501 Skip s -> Skip (ConcatUnfoldInterleaveInnerL (s:ls) rs)
2502 Stop -> Skip (ConcatUnfoldInterleaveInnerL ls rs)
2504 step _ (ConcatUnfoldInterleaveInnerR [] []) = return Stop
2505 step _ (ConcatUnfoldInterleaveInnerR ls []) =
2506 return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2508 step _ (ConcatUnfoldInterleaveInnerR ls (st:rs)) = do
2511 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerR (s:ls) rs)
2512 Skip s -> Skip (ConcatUnfoldInterleaveInnerR ls (s:rs))
2513 Stop -> Skip (ConcatUnfoldInterleaveInnerR ls rs)
2515-- XXX In general we can use different scheduling strategies e.g. how to
2516-- schedule the outer vs inner loop or assigning weights to different streams
2517-- or outer and inner loops.
2519-- This could be inefficient if the tasks are too small.
2521-- Compared to concatUnfoldInterleave this one switches streams on Skips.
2523{-# INLINE_NORMAL concatUnfoldRoundrobin #-}
2524concatUnfoldRoundrobin :: Monad m => Unfold m a b -> Stream m a -> Stream m b
2525concatUnfoldRoundrobin (Unfold istep inject) (Stream ostep ost) =
2526 Stream step (ConcatUnfoldInterleaveOuter ost [])
2528 {-# INLINE_LATE step #-}
2529 step gst (ConcatUnfoldInterleaveOuter o ls) = do
2530 r <- ostep (adaptState gst) o
2534 i `seq` return (Skip (ConcatUnfoldInterleaveInner o' (i : ls)))
2535 Skip o' -> return $ Skip (ConcatUnfoldInterleaveInner o' ls)
2536 Stop -> return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2538 step _ (ConcatUnfoldInterleaveInner o []) =
2539 return $ Skip (ConcatUnfoldInterleaveOuter o [])
2541 step _ (ConcatUnfoldInterleaveInner o (st:ls)) = do
2544 Yield x s -> Yield x (ConcatUnfoldInterleaveOuter o (s:ls))
2545 Skip s -> Skip (ConcatUnfoldInterleaveOuter o (s:ls))
2546 Stop -> Skip (ConcatUnfoldInterleaveOuter o ls)
2548 step _ (ConcatUnfoldInterleaveInnerL [] []) = return Stop
2549 step _ (ConcatUnfoldInterleaveInnerL [] rs) =
2550 return $ Skip (ConcatUnfoldInterleaveInnerR [] rs)
2552 step _ (ConcatUnfoldInterleaveInnerL (st:ls) rs) = do
2555 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerL ls (s:rs))
2556 Skip s -> Skip (ConcatUnfoldInterleaveInnerL ls (s:rs))
2557 Stop -> Skip (ConcatUnfoldInterleaveInnerL ls rs)
2559 step _ (ConcatUnfoldInterleaveInnerR [] []) = return Stop
2560 step _ (ConcatUnfoldInterleaveInnerR ls []) =
2561 return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2563 step _ (ConcatUnfoldInterleaveInnerR ls (st:rs)) = do
2566 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerR (s:ls) rs)
2567 Skip s -> Skip (ConcatUnfoldInterleaveInnerR (s:ls) rs)
2568 Stop -> Skip (ConcatUnfoldInterleaveInnerR ls rs)
2570data AppendState s1 s2 = AppendFirst s1 | AppendSecond s2
2572-- Note that this could be much faster compared to the CPS stream. However, as
2573-- the number of streams being composed increases this may become expensive.
2574-- Need to see where the breaking point is between the two.
2576{-# INLINE_NORMAL append #-}
2577append :: Monad m => Stream m a -> Stream m a -> Stream m a
2578append (Stream step1 state1) (Stream step2 state2) =
2579 Stream step (AppendFirst state1)
2583 {-# INLINE_LATE step #-}
2584 step gst (AppendFirst st) = do
2587 Yield a s -> Yield a (AppendFirst s)
2588 Skip s -> Skip (AppendFirst s)
2589 Stop -> Skip (AppendSecond state2)
2591 step gst (AppendSecond st) = do
2594 Yield a s -> Yield a (AppendSecond s)
2595 Skip s -> Skip (AppendSecond s)
2598data InterleaveState s1 s2 = InterleaveFirst s1 s2 | InterleaveSecond s1 s2
2599 | InterleaveSecondOnly s2 | InterleaveFirstOnly s1
2601{-# INLINE_NORMAL interleave #-}
2602interleave :: Monad m => Stream m a -> Stream m a -> Stream m a
2603interleave (Stream step1 state1) (Stream step2 state2) =
2604 Stream step (InterleaveFirst state1 state2)
2608 {-# INLINE_LATE step #-}
2609 step gst (InterleaveFirst st1 st2) = do
2612 Yield a s -> Yield a (InterleaveSecond s st2)
2613 Skip s -> Skip (InterleaveFirst s st2)
2614 Stop -> Skip (InterleaveSecondOnly st2)
2616 step gst (InterleaveSecond st1 st2) = do
2619 Yield a s -> Yield a (InterleaveFirst st1 s)
2620 Skip s -> Skip (InterleaveSecond st1 s)
2621 Stop -> Skip (InterleaveFirstOnly st1)
2623 step gst (InterleaveFirstOnly st1) = do
2626 Yield a s -> Yield a (InterleaveFirstOnly s)
2627 Skip s -> Skip (InterleaveFirstOnly s)
2630 step gst (InterleaveSecondOnly st2) = do
2633 Yield a s -> Yield a (InterleaveSecondOnly s)
2634 Skip s -> Skip (InterleaveSecondOnly s)
2637{-# INLINE_NORMAL interleaveMin #-}
2638interleaveMin :: Monad m => Stream m a -> Stream m a -> Stream m a
2639interleaveMin (Stream step1 state1) (Stream step2 state2) =
2640 Stream step (InterleaveFirst state1 state2)
2644 {-# INLINE_LATE step #-}
2645 step gst (InterleaveFirst st1 st2) = do
2648 Yield a s -> Yield a (InterleaveSecond s st2)
2649 Skip s -> Skip (InterleaveFirst s st2)
2652 step gst (InterleaveSecond st1 st2) = do
2655 Yield a s -> Yield a (InterleaveFirst st1 s)
2656 Skip s -> Skip (InterleaveSecond st1 s)
2659 step _ (InterleaveFirstOnly _) = undefined
2660 step _ (InterleaveSecondOnly _) = undefined
2662{-# INLINE_NORMAL interleaveSuffix #-}
2663interleaveSuffix :: Monad m => Stream m a -> Stream m a -> Stream m a
2664interleaveSuffix (Stream step1 state1) (Stream step2 state2) =
2665 Stream step (InterleaveFirst state1 state2)
2669 {-# INLINE_LATE step #-}
2670 step gst (InterleaveFirst st1 st2) = do
2673 Yield a s -> Yield a (InterleaveSecond s st2)
2674 Skip s -> Skip (InterleaveFirst s st2)
2677 step gst (InterleaveSecond st1 st2) = do
2680 Yield a s -> Yield a (InterleaveFirst st1 s)
2681 Skip s -> Skip (InterleaveSecond st1 s)
2682 Stop -> Skip (InterleaveFirstOnly st1)
2684 step gst (InterleaveFirstOnly st1) = do
2687 Yield a s -> Yield a (InterleaveFirstOnly s)
2688 Skip s -> Skip (InterleaveFirstOnly s)
2691 step _ (InterleaveSecondOnly _) = undefined
2693data InterleaveInfixState s1 s2 a
2694 = InterleaveInfixFirst s1 s2
2695 | InterleaveInfixSecondBuf s1 s2
2696 | InterleaveInfixSecondYield s1 s2 a
2697 | InterleaveInfixFirstYield s1 s2 a
2698 | InterleaveInfixFirstOnly s1
2700{-# INLINE_NORMAL interleaveInfix #-}
2701interleaveInfix :: Monad m => Stream m a -> Stream m a -> Stream m a
2702interleaveInfix (Stream step1 state1) (Stream step2 state2) =
2703 Stream step (InterleaveInfixFirst state1 state2)
2707 {-# INLINE_LATE step #-}
2708 step gst (InterleaveInfixFirst st1 st2) = do
2711 Yield a s -> Yield a (InterleaveInfixSecondBuf s st2)
2712 Skip s -> Skip (InterleaveInfixFirst s st2)
2715 step gst (InterleaveInfixSecondBuf st1 st2) = do
2718 Yield a s -> Skip (InterleaveInfixSecondYield st1 s a)
2719 Skip s -> Skip (InterleaveInfixSecondBuf st1 s)
2720 Stop -> Skip (InterleaveInfixFirstOnly st1)
2722 step gst (InterleaveInfixSecondYield st1 st2 x) = do
2725 Yield a s -> Yield x (InterleaveInfixFirstYield s st2 a)
2726 Skip s -> Skip (InterleaveInfixSecondYield s st2 x)
2729 step _ (InterleaveInfixFirstYield st1 st2 x) = do
2730 return $ Yield x (InterleaveInfixSecondBuf st1 st2)
2732 step gst (InterleaveInfixFirstOnly st1) = do
2735 Yield a s -> Yield a (InterleaveInfixFirstOnly s)
2736 Skip s -> Skip (InterleaveInfixFirstOnly s)
2739{-# INLINE_NORMAL roundRobin #-}
2740roundRobin :: Monad m => Stream m a -> Stream m a -> Stream m a
2741roundRobin (Stream step1 state1) (Stream step2 state2) =
2742 Stream step (InterleaveFirst state1 state2)
2746 {-# INLINE_LATE step #-}
2747 step gst (InterleaveFirst st1 st2) = do
2750 Yield a s -> Yield a (InterleaveSecond s st2)
2751 Skip s -> Skip (InterleaveSecond s st2)
2752 Stop -> Skip (InterleaveSecondOnly st2)
2754 step gst (InterleaveSecond st1 st2) = do
2757 Yield a s -> Yield a (InterleaveFirst st1 s)
2758 Skip s -> Skip (InterleaveFirst st1 s)
2759 Stop -> Skip (InterleaveFirstOnly st1)
2761 step gst (InterleaveSecondOnly st2) = do
2764 Yield a s -> Yield a (InterleaveSecondOnly s)
2765 Skip s -> Skip (InterleaveSecondOnly s)
2768 step gst (InterleaveFirstOnly st1) = do
2771 Yield a s -> Yield a (InterleaveFirstOnly s)
2772 Skip s -> Skip (InterleaveFirstOnly s)
2775data ICUState s1 s2 i1 i2 =
2780 | ICUFirstInner s1 s2 i1
2781 | ICUSecondInner s1 s2 i2
2782 | ICUFirstOnlyInner s1 i1
2783 | ICUSecondOnlyInner s2 i2
2785-- | Interleave streams (full streams, not the elements) unfolded from two
2786-- input streams and concat. Stop when the first stream stops. If the second
2787-- stream ends before the first one then first stream still keeps running alone
2788-- without any interleaving with the second stream.
2790-- [a1, a2, ... an] [b1, b2 ...]
2791-- => [streamA1, streamA2, ... streamAn] [streamB1, streamB2, ...]
2792-- => [streamA1, streamB1, streamA2...StreamAn, streamBn]
2793-- => [a11, a12, ...a1j, b11, b12, ...b1k, a21, a22, ...]
2795{-# INLINE_NORMAL gintercalateSuffix #-}
2798 => Unfold m a c -> Stream m a -> Unfold m b c -> Stream m b -> Stream m c
2800 (Unfold istep1 inject1) (Stream step1 state1)
2801 (Unfold istep2 inject2) (Stream step2 state2) =
2802 Stream step (ICUFirst state1 state2)
2806 {-# INLINE_LATE step #-}
2807 step gst (ICUFirst s1 s2) = do
2808 r <- step1 (adaptState gst) s1
2812 i `seq` return (Skip (ICUFirstInner s s2 i))
2813 Skip s -> return $ Skip (ICUFirst s s2)
2816 step gst (ICUFirstOnly s1) = do
2817 r <- step1 (adaptState gst) s1
2821 i `seq` return (Skip (ICUFirstOnlyInner s i))
2822 Skip s -> return $ Skip (ICUFirstOnly s)
2825 step _ (ICUFirstInner s1 s2 i1) = do
2828 Yield x i' -> Yield x (ICUFirstInner s1 s2 i')
2829 Skip i' -> Skip (ICUFirstInner s1 s2 i')
2830 Stop -> Skip (ICUSecond s1 s2)
2832 step _ (ICUFirstOnlyInner s1 i1) = do
2835 Yield x i' -> Yield x (ICUFirstOnlyInner s1 i')
2836 Skip i' -> Skip (ICUFirstOnlyInner s1 i')
2837 Stop -> Skip (ICUFirstOnly s1)
2839 step gst (ICUSecond s1 s2) = do
2840 r <- step2 (adaptState gst) s2
2844 i `seq` return (Skip (ICUSecondInner s1 s i))
2845 Skip s -> return $ Skip (ICUSecond s1 s)
2846 Stop -> return $ Skip (ICUFirstOnly s1)
2848 step _ (ICUSecondInner s1 s2 i2) = do
2851 Yield x i' -> Yield x (ICUSecondInner s1 s2 i')
2852 Skip i' -> Skip (ICUSecondInner s1 s2 i')
2853 Stop -> Skip (ICUFirst s1 s2)
2855 step _ (ICUSecondOnly _s2) = undefined
2856 step _ (ICUSecondOnlyInner _s2 _i2) = undefined
2858data InterposeSuffixState s1 i1 =
2859 InterposeSuffixFirst s1
2860 -- | InterposeSuffixFirstYield s1 i1
2861 | InterposeSuffixFirstInner s1 i1
2862 | InterposeSuffixSecond s1
2864-- Note that if an unfolded layer turns out to be nil we still emit the
2865-- separator effect. An alternate behavior could be to emit the separator
2866-- effect only if at least one element has been yielded by the unfolding.
2867-- However, that becomes a bit complicated, so we have chosen the former
2869{-# INLINE_NORMAL interposeSuffix #-}
2872 => m c -> Unfold m b c -> Stream m b -> Stream m c
2875 (Unfold istep1 inject1) (Stream step1 state1) =
2876 Stream step (InterposeSuffixFirst state1)
2880 {-# INLINE_LATE step #-}
2881 step gst (InterposeSuffixFirst s1) = do
2882 r <- step1 (adaptState gst) s1
2886 i `seq` return (Skip (InterposeSuffixFirstInner s i))
2887 -- i `seq` return (Skip (InterposeSuffixFirstYield s i))
2888 Skip s -> return $ Skip (InterposeSuffixFirst s)
2892 step _ (InterposeSuffixFirstYield s1 i1) = do
2895 Yield x i' -> Yield x (InterposeSuffixFirstInner s1 i')
2896 Skip i' -> Skip (InterposeSuffixFirstYield s1 i')
2897 Stop -> Skip (InterposeSuffixFirst s1)
2900 step _ (InterposeSuffixFirstInner s1 i1) = do
2903 Yield x i' -> Yield x (InterposeSuffixFirstInner s1 i')
2904 Skip i' -> Skip (InterposeSuffixFirstInner s1 i')
2905 Stop -> Skip (InterposeSuffixSecond s1)
2907 step _ (InterposeSuffixSecond s1) = do
2909 return $ Yield r (InterposeSuffixFirst s1)
2911data ICALState s1 s2 i1 i2 a =
2913 -- | ICALFirstYield s1 s2 i1
2914 | ICALFirstInner s1 s2 i1
2916 | ICALFirstOnlyInner s1 i1
2917 | ICALSecondInject s1 s2
2918 | ICALFirstInject s1 s2 i2
2919 -- | ICALFirstBuf s1 s2 i1 i2
2920 | ICALSecondInner s1 s2 i1 i2
2921 -- -- | ICALSecondInner s1 s2 i1 i2 a
2922 -- -- | ICALFirstResume s1 s2 i1 i2 a
2924-- | Interleave streams (full streams, not the elements) unfolded from two
2925-- input streams and concat. Stop when the first stream stops. If the second
2926-- stream ends before the first one then first stream still keeps running alone
2927-- without any interleaving with the second stream.
2929-- [a1, a2, ... an] [b1, b2 ...]
2930-- => [streamA1, streamA2, ... streamAn] [streamB1, streamB2, ...]
2931-- => [streamA1, streamB1, streamA2...StreamAn, streamBn]
2932-- => [a11, a12, ...a1j, b11, b12, ...b1k, a21, a22, ...]
2934{-# INLINE_NORMAL gintercalate #-}
2937 => Unfold m a c -> Stream m a -> Unfold m b c -> Stream m b -> Stream m c
2939 (Unfold istep1 inject1) (Stream step1 state1)
2940 (Unfold istep2 inject2) (Stream step2 state2) =
2941 Stream step (ICALFirst state1 state2)
2945 {-# INLINE_LATE step #-}
2946 step gst (ICALFirst s1 s2) = do
2947 r <- step1 (adaptState gst) s1
2951 i `seq` return (Skip (ICALFirstInner s s2 i))
2952 -- i `seq` return (Skip (ICALFirstYield s s2 i))
2953 Skip s -> return $ Skip (ICALFirst s s2)
2957 step _ (ICALFirstYield s1 s2 i1) = do
2960 Yield x i' -> Yield x (ICALFirstInner s1 s2 i')
2961 Skip i' -> Skip (ICALFirstYield s1 s2 i')
2962 Stop -> Skip (ICALFirst s1 s2)
2965 step _ (ICALFirstInner s1 s2 i1) = do
2968 Yield x i' -> Yield x (ICALFirstInner s1 s2 i')
2969 Skip i' -> Skip (ICALFirstInner s1 s2 i')
2970 Stop -> Skip (ICALSecondInject s1 s2)
2972 step gst (ICALFirstOnly s1) = do
2973 r <- step1 (adaptState gst) s1
2977 i `seq` return (Skip (ICALFirstOnlyInner s i))
2978 Skip s -> return $ Skip (ICALFirstOnly s)
2981 step _ (ICALFirstOnlyInner s1 i1) = do
2984 Yield x i' -> Yield x (ICALFirstOnlyInner s1 i')
2985 Skip i' -> Skip (ICALFirstOnlyInner s1 i')
2986 Stop -> Skip (ICALFirstOnly s1)
2988 -- We inject the second stream even before checking if the first stream
2989 -- would yield any more elements. There is no clear choice whether we
2990 -- should do this before or after that. Doing it after may make the state
2991 -- machine a bit simpler though.
2992 step gst (ICALSecondInject s1 s2) = do
2993 r <- step2 (adaptState gst) s2
2997 i `seq` return (Skip (ICALFirstInject s1 s i))
2998 Skip s -> return $ Skip (ICALSecondInject s1 s)
2999 Stop -> return $ Skip (ICALFirstOnly s1)
3001 step gst (ICALFirstInject s1 s2 i2) = do
3002 r <- step1 (adaptState gst) s1
3006 i `seq` return (Skip (ICALSecondInner s s2 i i2))
3007 -- i `seq` return (Skip (ICALFirstBuf s s2 i i2))
3008 Skip s -> return $ Skip (ICALFirstInject s s2 i2)
3012 step _ (ICALFirstBuf s1 s2 i1 i2) = do
3015 Yield x i' -> Skip (ICALSecondInner s1 s2 i' i2 x)
3016 Skip i' -> Skip (ICALFirstBuf s1 s2 i' i2)
3019 step _ (ICALSecondInner s1 s2 i1 i2 v) = do
3022 Yield x i' -> Yield x (ICALSecondInner s1 s2 i1 i' v)
3023 Skip i' -> Skip (ICALSecondInner s1 s2 i1 i' v)
3024 Stop -> Skip (ICALFirstResume s1 s2 i1 i2 v)
3027 step _ (ICALSecondInner s1 s2 i1 i2) = do
3030 Yield x i' -> Yield x (ICALSecondInner s1 s2 i1 i')
3031 Skip i' -> Skip (ICALSecondInner s1 s2 i1 i')
3032 Stop -> Skip (ICALFirstInner s1 s2 i1)
3033 -- Stop -> Skip (ICALFirstResume s1 s2 i1 i2)
3036 step _ (ICALFirstResume s1 s2 i1 i2 x) = do
3037 return $ Yield x (ICALFirstInner s1 s2 i1 i2)
3040data InterposeState s1 i1 a =
3042 -- | InterposeFirstYield s1 i1
3043 | InterposeFirstInner s1 i1
3044 | InterposeFirstInject s1
3045 -- | InterposeFirstBuf s1 i1
3046 | InterposeSecondYield s1 i1
3047 -- -- | InterposeSecondYield s1 i1 a
3048 -- -- | InterposeFirstResume s1 i1 a
3050-- Note that this only interposes the pure values, we may run many effects to
3051-- generate those values as some effects may not generate anything (Skip).
3052{-# INLINE_NORMAL interpose #-}
3053interpose :: Monad m => m c -> Unfold m b c -> Stream m b -> Stream m c
3056 (Unfold istep1 inject1) (Stream step1 state1) =
3057 Stream step (InterposeFirst state1)
3061 {-# INLINE_LATE step #-}
3062 step gst (InterposeFirst s1) = do
3063 r <- step1 (adaptState gst) s1
3067 i `seq` return (Skip (InterposeFirstInner s i))
3068 -- i `seq` return (Skip (InterposeFirstYield s i))
3069 Skip s -> return $ Skip (InterposeFirst s)
3073 step _ (InterposeFirstYield s1 i1) = do
3076 Yield x i' -> Yield x (InterposeFirstInner s1 i')
3077 Skip i' -> Skip (InterposeFirstYield s1 i')
3078 Stop -> Skip (InterposeFirst s1)
3081 step _ (InterposeFirstInner s1 i1) = do
3084 Yield x i' -> Yield x (InterposeFirstInner s1 i')
3085 Skip i' -> Skip (InterposeFirstInner s1 i')
3086 Stop -> Skip (InterposeFirstInject s1)
3088 step gst (InterposeFirstInject s1) = do
3089 r <- step1 (adaptState gst) s1
3093 -- i `seq` return (Skip (InterposeFirstBuf s i))
3094 i `seq` return (Skip (InterposeSecondYield s i))
3095 Skip s -> return $ Skip (InterposeFirstInject s)
3099 step _ (InterposeFirstBuf s1 i1) = do
3102 Yield x i' -> Skip (InterposeSecondYield s1 i' x)
3103 Skip i' -> Skip (InterposeFirstBuf s1 i')
3108 step _ (InterposeSecondYield s1 i1 v) = do
3110 return $ Yield r (InterposeFirstResume s1 i1 v)
3112 step _ (InterposeSecondYield s1 i1) = do
3114 return $ Yield r (InterposeFirstInner s1 i1)
3117 step _ (InterposeFirstResume s1 i1 v) = do
3118 return $ Yield v (InterposeFirstInner s1 i1)
3121------------------------------------------------------------------------------
3123------------------------------------------------------------------------------
3125data GbracketState s1 s2 v
3127 | GBracketNormal s1 v
3128 | GBracketException s2
3130-- | The most general bracketing and exception combinator. All other
3131-- combinators can be expressed in terms of this combinator. This can also be
3132-- used for cases which are not covered by the standard combinators.
3136{-# INLINE_NORMAL gbracket_ #-}
3140 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
3141 -> (c -> m d) -- ^ after, on normal stop
3142 -> (c -> e -> Stream m b -> Stream m b) -- ^ on exception
3143 -> (c -> Stream m b) -- ^ stream generator
3145gbracket_ bef exc aft fexc fnormal =
3146 Stream step GBracketInit
3150 {-# INLINE_LATE step #-}
3151 step _ GBracketInit = do
3153 return $ Skip $ GBracketNormal (fnormal r) r
3155 step gst (GBracketNormal (UnStream step1 st) v) = do
3156 res <- exc $ step1 gst st
3158 Right r -> case r of
3160 return $ Yield x (GBracketNormal (Stream step1 s) v)
3161 Skip s -> return $ Skip (GBracketNormal (Stream step1 s) v)
3162 Stop -> aft v >> return Stop
3164 return $ Skip (GBracketException (fexc v e (UnStream step1 st)))
3165 step gst (GBracketException (UnStream step1 st)) = do
3168 Yield x s -> return $ Yield x (GBracketException (Stream step1 s))
3169 Skip s -> return $ Skip (GBracketException (Stream step1 s))
3172------------------------------------------------------------------------------
3174------------------------------------------------------------------------------
3176-- | Make a finalizer from a monadic action @m a@ that can run in IO monad.
3177mkIOFinalizer :: MonadBaseControl IO m => m b -> m (IO ())
3179 mrun <- captureMonadState
3185-- | Run an IO action stored in a finalized IORef.
3186runIORefFinalizerGC :: IORef (Maybe (IO ())) -> IO ()
3187runIORefFinalizerGC ref = do
3188 res <- readIORef ref
3190 Nothing -> return ()
3193-- | Create an IORef holding a finalizer that is called automatically when
3194-- the IORef is garbage collected. The IORef can be written to with a 'Nothing'
3195-- value to deactivate the finalizer.
3197-- Note: The finalizer is always run with the state of the monad captured at
3198-- the time of calling newFinalizedIORef. To run it on garbage collection we
3199-- have no option but to take a snapshot of the monadic state at some point of
3200-- time. For normal case we could run it with the current state of the monad
3201-- but we want to keep both the cases consistent.
3203newFinalizedIORef :: (MonadIO m, MonadBaseControl IO m)
3204 => m a -> m (IORef (Maybe (IO ())))
3205newFinalizedIORef finalizer = do
3206 f <- mkIOFinalizer finalizer
3207 ref <- liftIO $ newIORef $ Just f
3208 _ <- liftIO $ mkWeakIORef ref (runIORefFinalizerGC ref)
3211-- | Run the finalizer stored in an IORef and deactivate it so that it is run
3212-- only once. Note, the action runs with async exceptions masked.
3214runIORefFinalizer :: MonadIO m => IORef (Maybe (IO ())) -> m ()
3215runIORefFinalizer ref = liftIO $ do
3216 res <- readIORef ref
3218 Nothing -> return ()
3220 -- if an async exception comes after writing 'Nothing' then the
3221 -- finalizing action will never be run. We need to do this
3222 -- atomically wrt async exceptions.
3224 writeIORef ref Nothing
3227-- | Run an action clearing the finalizer IORef atomically wrt async
3228-- exceptions. The action is run with async exceptions masked.
3229withIORefFinalizer :: MonadBaseControl IO m
3230 => IORef (Maybe (IO ())) -> m a -> m a
3231withIORefFinalizer ref action = do
3232 control $ \runinio ->
3234 writeIORef ref Nothing
3237------------------------------------------------------------------------------
3239data GbracketIOState s1 s2 v wref
3241 | GBracketIONormal s1 v wref
3242 | GBracketIOException s2
3244-- | The most general bracketing and exception combinator. All other
3245-- combinators can be expressed in terms of this combinator. This can also be
3246-- used for cases which are not covered by the standard combinators.
3248-- | It uses a finalizer to make sure we run the finalizing action
3249-- when the stream is garbage collected.
3253{-# INLINE_NORMAL gbracket #-}
3255 :: (MonadIO m, MonadBaseControl IO m)
3257 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
3258 -> (c -> m d) -- ^ after, on normal stop or GC
3259 -> (c -> e -> Stream m b -> m (Stream m b)) -- ^ on exception
3260 -> (c -> Stream m b) -- ^ stream generator
3262gbracket bef exc aft fexc fnormal =
3263 Stream step GBracketIOInit
3267 -- If the stream is never evaluated the "aft" action will never be
3268 -- called. For that to occur we will need the user of this API to pass a
3269 -- weak pointer to us.
3270 {-# INLINE_LATE step #-}
3271 step _ GBracketIOInit = do
3272 -- We mask asynchronous exceptions to make the execution
3273 -- of 'bef' and the registration of 'aft' atomic.
3274 -- A similar thing is done in the resourcet package: https://git.io/JvKV3
3275 -- Tutorial: https://markkarpov.com/tutorial/exceptions.html
3276 (r, ref) <- liftBaseOp_ mask_ $ do
3278 ref <- newFinalizedIORef (aft r)
3280 return $ Skip $ GBracketIONormal (fnormal r) r ref
3282 step gst (GBracketIONormal (UnStream step1 st) v ref) = do
3283 res <- exc $ step1 gst st
3285 Right r -> case r of
3287 return $ Yield x (GBracketIONormal (Stream step1 s) v ref)
3289 return $ Skip (GBracketIONormal (Stream step1 s) v ref)
3291 runIORefFinalizer ref
3294 -- Clearing of finalizer and running of exception handler must
3295 -- be atomic wrt async exceptions. Otherwise if we have cleared
3296 -- the finalizer and have not run the exception handler then we
3297 -- may leak the resource.
3298 stream <- withIORefFinalizer ref (fexc v e (UnStream step1 st))
3299 return $ Skip (GBracketIOException stream)
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 after #-}
3341after :: (MonadIO m, MonadBaseControl IO m)
3342 => m b -> Stream m a -> Stream m a
3343after 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-- | Run the first action before the stream starts and remember its output,
3407-- generate a stream using the output, run the second action providing the
3408-- remembered value as an argument whenever the stream ends normally or due to
3409-- an exception or gets garbage collected.
3411{-# INLINE_NORMAL bracket #-}
3412bracket :: (MonadAsync m, MonadCatch m)
3413 => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
3414bracket bef aft bet =
3415 gbracket bef MC.try aft
3416 (\a (e :: SomeException) _ -> aft a >> return (nilM (MC.throwM e))) bet
3418data BracketState s v = BracketInit | BracketRun s v
3420{-# INLINE_NORMAL _bracket #-}
3421_bracket :: MonadCatch m => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
3422_bracket bef aft bet = Stream step' BracketInit
3426 {-# INLINE_LATE step' #-}
3427 step' _ BracketInit = bef >>= \x -> return (Skip (BracketRun (bet x) x))
3429 -- NOTE: It is important to use UnStream instead of the Stream pattern
3430 -- here, otherwise we get huge perf degradation, see note in concatMap.
3431 step' gst (BracketRun (UnStream step state) v) = do
3432 -- res <- step gst state `MC.onException` aft v
3433 res <- MC.try $ step gst state
3435 Left (e :: SomeException) -> aft v >> MC.throwM e >> return Stop
3436 Right r -> case r of
3437 Yield x s -> return $ Yield x (BracketRun (Stream step s) v)
3438 Skip s -> return $ Skip (BracketRun (Stream step s) v)
3439 Stop -> aft v >> return Stop
3441-- | Run a side effect whenever the stream stops normally or aborts due to an
3443{-# INLINE finally_ #-}
3444finally_ :: MonadCatch m => m b -> Stream m a -> Stream m a
3445-- finally action xs = after action $ onException action xs
3446finally_ action xs = bracket_ (return ()) (\_ -> action) (const xs)
3448-- | Run a side effect whenever the stream stops normally or aborts due to an
3449-- exception or gets garbage collected.
3451{-# INLINE finally #-}
3452finally :: (MonadAsync m, MonadCatch m) => m b -> Stream m a -> Stream m a
3453finally action xs = bracket (return ()) (\_ -> action) (const xs)
3455-- | When evaluating a stream if an exception occurs, stream
3456-- evaluation aborts and the specified exception handler is run with
3457-- the exception and the Stream which threw the exception as argument.
3459{-# INLINE_NORMAL ghandle #-}
3460ghandle :: (MonadCatch m, Exception e)
3461 => (e -> Stream m a -> Stream m a) -> Stream m a -> Stream m a
3463 gbracket_ (return ()) MC.try return (\_ -> f) (\_ -> str)
3465-- | When evaluating a stream if an exception occurs, stream evaluation aborts
3466-- and the specified exception handler is run with the exception as argument.
3467{-# INLINE_NORMAL handle #-}
3468handle :: (MonadCatch m, Exception e)
3469 => (e -> Stream m a) -> Stream m a -> Stream m a
3471 gbracket_ (return ()) MC.try return (\_ e _ -> f e) (\_ -> str)
3473{-# INLINE_NORMAL _handle #-}
3474_handle :: (MonadCatch m, Exception e)
3475 => (e -> Stream m a) -> Stream m a -> Stream m a
3476_handle f (Stream step state) = Stream step' (Left state)
3480 {-# INLINE_LATE step' #-}
3481 step' gst (Left st) = do
3482 res <- MC.try $ step gst st
3484 Left e -> return $ Skip $ Right (f e)
3485 Right r -> case r of
3486 Yield x s -> return $ Yield x (Left s)
3487 Skip s -> return $ Skip (Left s)
3490 step' gst (Right (UnStream step1 st)) = do
3493 Yield x s -> return $ Yield x (Right (Stream step1 s))
3494 Skip s -> return $ Skip (Right (Stream step1 s))
3497-------------------------------------------------------------------------------
3498-- General transformation
3499-------------------------------------------------------------------------------
3501{-# INLINE_NORMAL transform #-}
3502transform :: Monad m => Pipe m a b -> Stream m a -> Stream m b
3503transform (Pipe pstep1 pstep2 pstate) (Stream step state) =
3504 Stream step' (Consume pstate, state)
3508 {-# INLINE_LATE step' #-}
3510 step' gst (Consume pst, st) = pst `seq` do
3511 r <- step (adaptState gst) st
3516 Pipe.Yield b pst' -> return $ Yield b (pst', s)
3517 Pipe.Continue pst' -> return $ Skip (pst', s)
3518 Skip s -> return $ Skip (Consume pst, s)
3521 step' _ (Produce pst, st) = pst `seq` do
3524 Pipe.Yield b pst' -> return $ Yield b (pst', st)
3525 Pipe.Continue pst' -> return $ Skip (pst', st)
3527------------------------------------------------------------------------------
3528-- Transformation by Folding (Scans)
3529------------------------------------------------------------------------------
3531------------------------------------------------------------------------------
3533------------------------------------------------------------------------------
3535-- XXX Is a prescan useful, discarding the last step does not sound useful? I
3536-- am not sure about the utility of this function, so this is implemented but
3537-- not exposed. We can expose it if someone provides good reasons why this is
3540-- XXX We have to execute the stream one step ahead to know that we are at the
3541-- last step. The vector implementation of prescan executes the last fold step
3542-- but does not yield the result. This means we have executed the effect but
3543-- discarded value. This does not sound right. In this implementation we are
3544-- not executing the last fold step.
3545{-# INLINE_NORMAL prescanlM' #-}
3546prescanlM' :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
3547prescanlM' f mz (Stream step state) = Stream step' (state, mz)
3549 {-# INLINE_LATE step' #-}
3550 step' gst (st, prev) = do
3551 r <- step (adaptState gst) st
3555 return $ Yield acc (s, f acc x)
3556 Skip s -> return $ Skip (s, prev)
3559{-# INLINE prescanl' #-}
3560prescanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
3561prescanl' f z = prescanlM' (\a b -> return (f a b)) (return z)
3563------------------------------------------------------------------------------
3564-- Monolithic postscans (postscan followed by a map)
3565------------------------------------------------------------------------------
3567-- The performance of a modular postscan followed by a map seems to be
3568-- equivalent to this monolithic scan followed by map therefore we may not need
3569-- this implementation. We just have it for performance comparison and in case
3570-- modular version does not perform well in some situation.
3572{-# INLINE_NORMAL postscanlMx' #-}
3573postscanlMx' :: Monad m
3574 => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> Stream m b
3575postscanlMx' fstep begin done (Stream step state) = do
3576 Stream step' (state, begin)
3578 {-# INLINE_LATE step' #-}
3579 step' gst (st, acc) = do
3580 r <- step (adaptState gst) st
3586 v `seq` y `seq` return (Yield v (s, return y))
3587 Skip s -> return $ Skip (s, acc)
3590{-# INLINE_NORMAL postscanlx' #-}
3591postscanlx' :: Monad m
3592 => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> Stream m b
3593postscanlx' fstep begin done s =
3594 postscanlMx' (\b a -> return (fstep b a)) (return begin) (return . done) s
3596-- XXX do we need consM strict to evaluate the begin value?
3597{-# INLINE scanlMx' #-}
3599 => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> Stream m b
3600scanlMx' fstep begin done s =
3601 (begin >>= \x -> x `seq` done x) `consM` postscanlMx' fstep begin done s
3603{-# INLINE scanlx' #-}
3605 => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> Stream m b
3606scanlx' fstep begin done s =
3607 scanlMx' (\b a -> return (fstep b a)) (return begin) (return . done) s
3609------------------------------------------------------------------------------
3611------------------------------------------------------------------------------
3613{-# INLINE_NORMAL postscanlM' #-}
3614postscanlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
3615postscanlM' fstep begin (Stream step state) =
3616 begin `seq` Stream step' (state, begin)
3618 {-# INLINE_LATE step' #-}
3619 step' gst (st, acc) = acc `seq` do
3620 r <- step (adaptState gst) st
3624 y `seq` return (Yield y (s, y))
3625 Skip s -> return $ Skip (s, acc)
3628{-# INLINE_NORMAL postscanl' #-}
3629postscanl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
3630postscanl' f = postscanlM' (\a b -> return (f a b))
3632-- We can possibly have the "done" function as a Maybe to provide an option to
3633-- emit or not emit the accumulator when the stream stops.
3635-- TBD: use a single Yield point
3637{-# INLINE_NORMAL postscanlMAfter' #-}
3638postscanlMAfter' :: Monad m
3639 => (b -> a -> m b) -> m b -> (b -> m b) -> Stream m a -> Stream m b
3640postscanlMAfter' fstep initial done (Stream step1 state1) = do
3641 Stream step (Just (state1, initial))
3645 {-# INLINE_LATE step #-}
3646 step gst (Just (st, acc)) = do
3647 r <- step1 (adaptState gst) st
3652 y `seq` return (Yield y (Just (s, return y)))
3653 Skip s -> return $ Skip $ Just (s, acc)
3654 Stop -> acc >>= done >>= \res -> return (Yield res Nothing)
3655 step _ Nothing = return Stop
3657{-# INLINE_NORMAL postscanlM #-}
3658postscanlM :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
3659postscanlM fstep begin (Stream step state) = Stream step' (state, begin)
3661 {-# INLINE_LATE step' #-}
3662 step' gst (st, acc) = do
3663 r <- step (adaptState gst) st
3667 return (Yield y (s, y))
3668 Skip s -> return $ Skip (s, acc)
3671{-# INLINE_NORMAL postscanl #-}
3672postscanl :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
3673postscanl f = postscanlM (\a b -> return (f a b))
3675{-# INLINE_NORMAL scanlM' #-}
3676scanlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
3677scanlM' fstep begin s = begin `seq` (begin `cons` postscanlM' fstep begin s)
3679{-# INLINE scanlMAfter' #-}
3680scanlMAfter' :: Monad m
3681 => (b -> a -> m b) -> m b -> (b -> m b) -> Stream m a -> Stream m b
3682scanlMAfter' fstep initial done s =
3683 (initial >>= \x -> x `seq` return x) `consM`
3684 postscanlMAfter' fstep initial done s
3686{-# INLINE scanl' #-}
3687scanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
3688scanl' f = scanlM' (\a b -> return (f a b))
3690{-# INLINE_NORMAL scanlM #-}
3691scanlM :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
3692scanlM fstep begin s = begin `cons` postscanlM fstep begin s
3695scanl :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
3696scanl f = scanlM (\a b -> return (f a b))
3698{-# INLINE_NORMAL scanl1M #-}
3699scanl1M :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
3700scanl1M fstep (Stream step state) = Stream step' (state, Nothing)
3702 {-# INLINE_LATE step' #-}
3703 step' gst (st, Nothing) = do
3706 Yield x s -> return $ Yield x (s, Just x)
3707 Skip s -> return $ Skip (s, Nothing)
3710 step' gst (st, Just acc) = do
3715 return $ Yield z (s, Just z)
3716 Skip s -> return $ Skip (s, Just acc)
3719{-# INLINE scanl1 #-}
3720scanl1 :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
3721scanl1 f = scanl1M (\x y -> return (f x y))
3723{-# INLINE_NORMAL scanl1M' #-}
3724scanl1M' :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
3725scanl1M' fstep (Stream step state) = Stream step' (state, Nothing)
3727 {-# INLINE_LATE step' #-}
3728 step' gst (st, Nothing) = do
3731 Yield x s -> x `seq` return $ Yield x (s, Just x)
3732 Skip s -> return $ Skip (s, Nothing)
3735 step' gst (st, Just acc) = acc `seq` do
3740 z `seq` return $ Yield z (s, Just z)
3741 Skip s -> return $ Skip (s, Just acc)
3744{-# INLINE scanl1' #-}
3745scanl1' :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
3746scanl1' f = scanl1M' (\x y -> return (f x y))
3748------------------------------------------------------------------------------
3750------------------------------------------------------------------------------
3752data RollingMapState s a = RollingMapInit s | RollingMapGo s a
3754{-# INLINE rollingMapM #-}
3755rollingMapM :: Monad m => (a -> a -> m b) -> Stream m a -> Stream m b
3756rollingMapM f (Stream step1 state1) = Stream step (RollingMapInit state1)
3758 step gst (RollingMapInit st) = do
3759 r <- step1 (adaptState gst) st
3761 Yield x s -> Skip $ RollingMapGo s x
3762 Skip s -> Skip $ RollingMapInit s
3765 step gst (RollingMapGo s1 x1) = do
3766 r <- step1 (adaptState gst) s1
3770 return $ Yield res $ RollingMapGo s x
3771 Skip s -> return $ Skip $ RollingMapGo s x1
3772 Stop -> return $ Stop
3774{-# INLINE rollingMap #-}
3775rollingMap :: Monad m => (a -> a -> b) -> Stream m a -> Stream m b
3776rollingMap f = rollingMapM (\x y -> return $ f x y)
3778------------------------------------------------------------------------------
3779-- Tapping/Distributing
3780------------------------------------------------------------------------------
3783tap :: Monad m => Fold m a b -> Stream m a -> Stream m a
3784tap (Fold fstep initial extract) (Stream step state) = Stream step' Nothing
3788 step' _ Nothing = do
3790 return $ Skip (Just (r, state))
3792 step' gst (Just (acc, st)) = acc `seq` do
3797 return $ Yield x (Just (acc', s))
3798 Skip s -> return $ Skip (Just (acc, s))
3803{-# INLINE_NORMAL tapOffsetEvery #-}
3804tapOffsetEvery :: Monad m
3805 => Int -> Int -> Fold m a b -> Stream m a -> Stream m a
3806tapOffsetEvery offset n (Fold fstep initial extract) (Stream step state) =
3807 Stream step' Nothing
3811 {-# INLINE_LATE step' #-}
3812 step' _ Nothing = do
3814 return $ Skip (Just (r, state, offset `mod` n))
3816 step' gst (Just (acc, st, count)) | count <= 0 = do
3820 !acc' <- fstep acc x
3821 return $ Yield x (Just (acc', s, n - 1))
3822 Skip s -> return $ Skip (Just (acc, s, count))
3827 step' gst (Just (acc, st, count)) = do
3830 Yield x s -> return $ Yield x (Just (acc, s, count - 1))
3831 Skip s -> return $ Skip (Just (acc, s, count))
3836{-# INLINE_NORMAL pollCounts #-}
3840 -> (Stream m Int -> Stream m Int)
3844pollCounts predicate transf fld (Stream step state) = Stream step' Nothing
3847 {-# INLINE_LATE step' #-}
3848 step' _ Nothing = do
3849 -- As long as we are using an "Int" for counts lockfree reads from
3850 -- Var should work correctly on both 32-bit and 64-bit machines.
3851 -- However, an Int on a 32-bit machine may overflow quickly.
3852 countVar <- liftIO $ Prim.newIORef (0 :: Int)
3854 $ void $ runFold fld
3855 $ transf $ fromPrimIORef countVar
3856 return $ Skip (Just (countVar, tid, state))
3858 step' gst (Just (countVar, tid, st)) = do
3862 when (predicate x) $ liftIO $ Prim.modifyIORef' countVar (+ 1)
3863 return $ Yield x (Just (countVar, tid, s))
3864 Skip s -> return $ Skip (Just (countVar, tid, s))
3866 liftIO $ killThread tid
3869{-# INLINE_NORMAL tapRate #-}
3871 (MonadAsync m, MonadCatch m)
3876tapRate samplingRate action (Stream step state) = Stream step' Nothing
3878 {-# NOINLINE loop #-}
3879 loop countVar prev = do
3882 (do liftIO $ threadDelay (round $ samplingRate * 1000000)
3883 i <- liftIO $ Prim.readIORef countVar
3884 let !diff = i - prev
3887 (\(e :: AsyncException) -> do
3888 i <- liftIO $ Prim.readIORef countVar
3889 let !diff = i - prev
3891 throwM (MC.toException e))
3894 {-# INLINE_LATE step' #-}
3895 step' _ Nothing = do
3896 countVar <- liftIO $ Prim.newIORef 0
3897 tid <- fork $ loop countVar 0
3898 ref <- liftIO $ newIORef ()
3899 _ <- liftIO $ mkWeakIORef ref (killThread tid)
3900 return $ Skip (Just (countVar, tid, state, ref))
3902 step' gst (Just (countVar, tid, st, ref)) = do
3906 liftIO $ Prim.modifyIORef' countVar (+ 1)
3907 return $ Yield x (Just (countVar, tid, s, ref))
3908 Skip s -> return $ Skip (Just (countVar, tid, s, ref))
3910 liftIO $ killThread tid
3914-------------------------------------------------------------------------------
3916-------------------------------------------------------------------------------
3918{-# INLINE_NORMAL takeWhileM #-}
3919takeWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
3920takeWhileM f (Stream step state) = Stream step' state
3922 {-# INLINE_LATE step' #-}
3928 return $ if b then Yield x s else Stop
3929 Skip s -> return $ Skip s
3932{-# INLINE takeWhile #-}
3933takeWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
3934takeWhile f = takeWhileM (return . f)
3936{-# INLINE_NORMAL drop #-}
3937drop :: Monad m => Int -> Stream m a -> Stream m a
3938drop n (Stream step state) = Stream step' (state, Just n)
3940 {-# INLINE_LATE step' #-}
3941 step' gst (st, Just i)
3946 Yield _ s -> Skip (s, Just (i - 1))
3947 Skip s -> Skip (s, Just i)
3949 | otherwise = return $ Skip (st, Nothing)
3951 step' gst (st, Nothing) = do
3955 Yield x s -> Yield x (s, Nothing)
3956 Skip s -> Skip (s, Nothing)
3959data DropWhileState s a
3961 | DropWhileYield a s
3964{-# INLINE_NORMAL dropWhileM #-}
3965dropWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
3966dropWhileM f (Stream step state) = Stream step' (DropWhileDrop state)
3968 {-# INLINE_LATE step' #-}
3969 step' gst (DropWhileDrop st) = do
3975 then return $ Skip (DropWhileDrop s)
3976 else return $ Skip (DropWhileYield x s)
3977 Skip s -> return $ Skip (DropWhileDrop s)
3980 step' gst (DropWhileNext st) = do
3983 Yield x s -> return $ Skip (DropWhileYield x s)
3984 Skip s -> return $ Skip (DropWhileNext s)
3987 step' _ (DropWhileYield x st) = return $ Yield x (DropWhileNext st)
3989{-# INLINE dropWhile #-}
3990dropWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
3991dropWhile f = dropWhileM (return . f)
3993{-# INLINE_NORMAL filterM #-}
3994filterM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
3995filterM f (Stream step state) = Stream step' state
3997 {-# INLINE_LATE step' #-}
4006 Skip s -> return $ Skip s
4009{-# INLINE filter #-}
4010filter :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
4011filter f = filterM (return . f)
4013{-# INLINE_NORMAL uniq #-}
4014uniq :: (Eq a, Monad m) => Stream m a -> Stream m a
4015uniq (Stream step state) = Stream step' (Nothing, state)
4017 {-# INLINE_LATE step' #-}
4018 step' gst (Nothing, st) = do
4021 Yield x s -> return $ Yield x (Just x, s)
4022 Skip s -> return $ Skip (Nothing, s)
4024 step' gst (Just x, st) = do
4027 Yield y s | x == y -> return $ Skip (Just x, s)
4028 | otherwise -> return $ Yield y (Just y, s)
4029 Skip s -> return $ Skip (Just x, s)
4032------------------------------------------------------------------------------
4033-- Transformation by Mapping
4034------------------------------------------------------------------------------
4036{-# INLINE_NORMAL sequence #-}
4037sequence :: Monad m => Stream m (m a) -> Stream m a
4038sequence (Stream step state) = Stream step' state
4040 {-# INLINE_LATE step' #-}
4042 r <- step (adaptState gst) st
4044 Yield x s -> x >>= \a -> return (Yield a s)
4045 Skip s -> return $ Skip s
4048------------------------------------------------------------------------------
4050------------------------------------------------------------------------------
4052data LoopState x s = FirstYield s
4053 | InterspersingYield s
4056{-# INLINE_NORMAL intersperseM #-}
4057intersperseM :: Monad m => m a -> Stream m a -> Stream m a
4058intersperseM m (Stream step state) = Stream step' (FirstYield state)
4060 {-# INLINE_LATE step' #-}
4061 step' gst (FirstYield st) = do
4065 Yield x s -> Skip (YieldAndCarry x s)
4066 Skip s -> Skip (FirstYield s)
4069 step' gst (InterspersingYield st) = do
4074 return $ Yield a (YieldAndCarry x s)
4075 Skip s -> return $ Skip $ InterspersingYield s
4078 step' _ (YieldAndCarry x st) = return $ Yield x (InterspersingYield st)
4080{-# INLINE_NORMAL intersperseM_ #-}
4081intersperseM_ :: Monad m => m b -> Stream m a -> Stream m a
4082intersperseM_ m (Stream step1 state1) = Stream step (Left (pure (), state1))
4084 {-# INLINE_LATE step #-}
4085 step gst (Left (eff, st)) = do
4088 Yield x s -> eff >> return (Yield x (Right s))
4089 Skip s -> return $ Skip (Left (eff, s))
4092 step _ (Right st) = return $ Skip $ Left (void m, st)
4097 | SuffixYield a (SuffixState s a)
4099{-# INLINE_NORMAL intersperseSuffix #-}
4100intersperseSuffix :: forall m a. Monad m => m a -> Stream m a -> Stream m a
4101intersperseSuffix action (Stream step state) = Stream step' (SuffixElem state)
4103 {-# INLINE_LATE step' #-}
4104 step' gst (SuffixElem st) = do
4107 Yield x s -> Skip (SuffixYield x (SuffixSuffix s))
4108 Skip s -> Skip (SuffixElem s)
4111 step' _ (SuffixSuffix st) = do
4112 action >>= \r -> return $ Skip (SuffixYield r (SuffixElem st))
4114 step' _ (SuffixYield x next) = return $ Yield x next
4116{-# INLINE_NORMAL intersperseSuffix_ #-}
4117intersperseSuffix_ :: Monad m => m b -> Stream m a -> Stream m a
4118intersperseSuffix_ m (Stream step1 state1) = Stream step (Left state1)
4120 {-# INLINE_LATE step #-}
4121 step gst (Left st) = do
4124 Yield x s -> return $ Yield x (Right s)
4125 Skip s -> return $ Skip $ Left s
4128 step _ (Right st) = m >> return (Skip (Left st))
4130data SuffixSpanState s a
4131 = SuffixSpanElem s Int
4132 | SuffixSpanSuffix s
4133 | SuffixSpanYield a (SuffixSpanState s a)
4137-- | intersperse after every n items
4138{-# INLINE_NORMAL intersperseSuffixBySpan #-}
4139intersperseSuffixBySpan :: forall m a. Monad m
4140 => Int -> m a -> Stream m a -> Stream m a
4141intersperseSuffixBySpan n action (Stream step state) =
4142 Stream step' (SuffixSpanElem state n)
4144 {-# INLINE_LATE step' #-}
4145 step' gst (SuffixSpanElem st i) | i > 0 = do
4148 Yield x s -> Skip (SuffixSpanYield x (SuffixSpanElem s (i - 1)))
4149 Skip s -> Skip (SuffixSpanElem s i)
4150 Stop -> if i == n then Stop else Skip SuffixSpanLast
4151 step' _ (SuffixSpanElem st _) = return $ Skip (SuffixSpanSuffix st)
4153 step' _ (SuffixSpanSuffix st) = do
4154 action >>= \r -> return $ Skip (SuffixSpanYield r (SuffixSpanElem st n))
4156 step' _ (SuffixSpanLast) = do
4157 action >>= \r -> return $ Skip (SuffixSpanYield r SuffixSpanStop)
4159 step' _ (SuffixSpanYield x next) = return $ Yield x next
4161 step' _ (SuffixSpanStop) = return Stop
4163{-# INLINE intersperse #-}
4164intersperse :: Monad m => a -> Stream m a -> Stream m a
4165intersperse a = intersperseM (return a)
4167{-# INLINE_NORMAL insertBy #-}
4168insertBy :: Monad m => (a -> a -> Ordering) -> a -> Stream m a -> Stream m a
4169insertBy cmp a (Stream step state) = Stream step' (state, False, Nothing)
4171 {-# INLINE_LATE step' #-}
4172 step' gst (st, False, _) = do
4175 Yield x s -> case cmp a x of
4176 GT -> return $ Yield x (s, False, Nothing)
4177 _ -> return $ Yield a (s, True, Just x)
4178 Skip s -> return $ Skip (s, False, Nothing)
4179 Stop -> return $ Yield a (st, True, Nothing)
4181 step' _ (_, True, Nothing) = return Stop
4183 step' gst (st, True, Just prev) = do
4186 Yield x s -> return $ Yield prev (s, True, Just x)
4187 Skip s -> return $ Skip (s, True, Just prev)
4188 Stop -> return $ Yield prev (st, True, Nothing)
4190------------------------------------------------------------------------------
4192------------------------------------------------------------------------------
4194{-# INLINE_NORMAL deleteBy #-}
4195deleteBy :: Monad m => (a -> a -> Bool) -> a -> Stream m a -> Stream m a
4196deleteBy eq x (Stream step state) = Stream step' (state, False)
4198 {-# INLINE_LATE step' #-}
4199 step' gst (st, False) = do
4202 Yield y s -> return $
4203 if eq x y then Skip (s, True) else Yield y (s, False)
4204 Skip s -> return $ Skip (s, False)
4207 step' gst (st, True) = do
4210 Yield y s -> return $ Yield y (s, True)
4211 Skip s -> return $ Skip (s, True)
4214------------------------------------------------------------------------------
4215-- Transformation by Map and Filter
4216------------------------------------------------------------------------------
4218-- XXX Will this always fuse properly?
4219{-# INLINE_NORMAL mapMaybe #-}
4220mapMaybe :: Monad m => (a -> Maybe b) -> Stream m a -> Stream m b
4221mapMaybe f = fmap fromJust . filter isJust . map f
4223{-# INLINE_NORMAL mapMaybeM #-}
4224mapMaybeM :: Monad m => (a -> m (Maybe b)) -> Stream m a -> Stream m b
4225mapMaybeM f = fmap fromJust . filter isJust . mapM f
4227------------------------------------------------------------------------------
4229------------------------------------------------------------------------------
4231{-# INLINE_NORMAL indexed #-}
4232indexed :: Monad m => Stream m a -> Stream m (Int, a)
4233indexed (Stream step state) = Stream step' (state, 0)
4235 {-# INLINE_LATE step' #-}
4236 step' gst (st, i) = i `seq` do
4237 r <- step (adaptState gst) st
4239 Yield x s -> return $ Yield (i, x) (s, i+1)
4240 Skip s -> return $ Skip (s, i)
4243{-# INLINE_NORMAL indexedR #-}
4244indexedR :: Monad m => Int -> Stream m a -> Stream m (Int, a)
4245indexedR m (Stream step state) = Stream step' (state, m)
4247 {-# INLINE_LATE step' #-}
4248 step' gst (st, i) = i `seq` do
4249 r <- step (adaptState gst) st
4251 Yield x s -> let i' = i - 1
4252 in return $ Yield (i, x) (s, i')
4253 Skip s -> return $ Skip (s, i)
4256{-# INLINE_NORMAL zipWithM #-}
4258 => (a -> b -> m c) -> Stream m a -> Stream m b -> Stream m c
4259zipWithM f (Stream stepa ta) (Stream stepb tb) = Stream step (ta, tb, Nothing)
4261 {-# INLINE_LATE step #-}
4262 step gst (sa, sb, Nothing) = do
4263 r <- stepa (adaptState gst) sa
4266 Yield x sa' -> Skip (sa', sb, Just x)
4267 Skip sa' -> Skip (sa', sb, Nothing)
4270 step gst (sa, sb, Just x) = do
4271 r <- stepb (adaptState gst) sb
4275 return $ Yield z (sa, sb', Nothing)
4276 Skip sb' -> return $ Skip (sa, sb', Just x)
4279#if __GLASGOW_HASKELL__ >= 801
4280{-# RULES "zipWithM xs xs"
4281 forall f xs. zipWithM @Identity f xs xs = mapM (\x -> f x x) xs #-}
4284{-# INLINE zipWith #-}
4285zipWith :: Monad m => (a -> b -> c) -> Stream m a -> Stream m b -> Stream m c
4286zipWith f = zipWithM (\a b -> return (f a b))
4288------------------------------------------------------------------------------
4290------------------------------------------------------------------------------
4292{-# INLINE_NORMAL mergeByM #-}
4295 => (a -> a -> m Ordering) -> Stream m a -> Stream m a -> Stream m a
4296mergeByM cmp (Stream stepa ta) (Stream stepb tb) =
4297 Stream step (Just ta, Just tb, Nothing, Nothing)
4299 {-# INLINE_LATE step #-}
4301 -- one of the values is missing, and the corresponding stream is running
4302 step gst (Just sa, sb, Nothing, b) = do
4305 Yield a sa' -> Skip (Just sa', sb, Just a, b)
4306 Skip sa' -> Skip (Just sa', sb, Nothing, b)
4307 Stop -> Skip (Nothing, sb, Nothing, b)
4309 step gst (sa, Just sb, a, Nothing) = do
4312 Yield b sb' -> Skip (sa, Just sb', a, Just b)
4313 Skip sb' -> Skip (sa, Just sb', a, Nothing)
4314 Stop -> Skip (sa, Nothing, a, Nothing)
4316 -- both the values are available
4317 step _ (sa, sb, Just a, Just b) = do
4319 return $ case res of
4320 GT -> Yield b (sa, sb, Just a, Nothing)
4321 _ -> Yield a (sa, sb, Nothing, Just b)
4323 -- one of the values is missing, corresponding stream is done
4324 step _ (Nothing, sb, Nothing, Just b) =
4325 return $ Yield b (Nothing, sb, Nothing, Nothing)
4327 step _ (sa, Nothing, Just a, Nothing) =
4328 return $ Yield a (sa, Nothing, Nothing, Nothing)
4330 step _ (Nothing, Nothing, Nothing, Nothing) = return Stop
4332{-# INLINE mergeBy #-}
4335 => (a -> a -> Ordering) -> Stream m a -> Stream m a -> Stream m a
4336mergeBy cmp = mergeByM (\a b -> return $ cmp a b)
4338------------------------------------------------------------------------------
4339-- Transformation comprehensions
4340------------------------------------------------------------------------------
4342{-# INLINE_NORMAL the #-}
4343the :: (Eq a, Monad m) => Stream m a -> m (Maybe a)
4344the (Stream step state) = go state
4347 r <- step defState st
4349 Yield x s -> go' x s
4351 Stop -> return Nothing
4353 r <- step defState st
4355 Yield x s | x == n -> go' n s
4356 | otherwise -> return Nothing
4358 Stop -> return (Just n)
4360{-# INLINE runFold #-}
4361runFold :: (Monad m) => Fold m a b -> Stream m a -> m b
4362runFold (Fold step begin done) = foldlMx' step begin done
4364-------------------------------------------------------------------------------
4365-- Concurrent application and fold
4366-------------------------------------------------------------------------------
4368-- XXX These functions should be moved to Stream/Parallel.hs
4370-- Using StreamD the worker stream producing code can fuse with the code to
4371-- queue output to the SVar giving some perf boost.
4373-- Note that StreamD can only be used in limited situations, specifically, we
4374-- cannot implement joinStreamVarPar using this.
4376-- XXX make sure that the SVar passed is a Parallel style SVar.
4378-- | Fold the supplied stream to the SVar asynchronously using Parallel
4379-- concurrency style.
4380-- {-# INLINE_NORMAL toSVarParallel #-}
4381{-# INLINE toSVarParallel #-}
4382toSVarParallel :: MonadAsync m
4383 => State t m a -> SVar t m a -> Stream m a -> m ()
4384toSVarParallel st sv xs =
4385 if svarInspectMode sv
4389 case getYieldLimit st of
4390 Nothing -> doFork (work Nothing)
4392 (handleChildException sv)
4393 Just _ -> doFork (workLim Nothing)
4395 (handleChildException sv)
4400 {-# NOINLINE work #-}
4401 work info = (runFold (FL.toParallelSVar sv info) xs)
4403 {-# NOINLINE workLim #-}
4404 workLim info = runFold (FL.toParallelSVarLimited sv info) xs
4406 {-# NOINLINE forkWithDiag #-}
4408 -- We do not use workerCount in case of ParallelVar but still there is
4409 -- no harm in maintaining it correctly.
4410 liftIO $ atomicModifyIORefCAS_ (workerCount sv) $ \n -> n + 1
4412 -- This allocation matters when significant number of workers are being
4413 -- sent. We allocate it only when needed. The overhead increases by 4x.
4415 case yieldRateInfo sv of
4416 Nothing -> return Nothing
4417 Just _ -> liftIO $ do
4418 cntRef <- newIORef 0
4419 t <- getTime Monotonic
4420 lat <- newIORef (0, t)
4421 return $ Just WorkerInfo
4422 { workerYieldMax = 0
4423 , workerYieldCount = cntRef
4424 , workerLatencyStart = lat
4427 case getYieldLimit st of
4428 Nothing -> doFork (work winfo)
4430 (handleChildException sv)
4431 Just _ -> doFork (workLim winfo)
4433 (handleChildException sv)
4436{-# INLINE_NORMAL mkParallelD #-}
4437mkParallelD :: MonadAsync m => Stream m a -> Stream m a
4438mkParallelD m = Stream step Nothing
4441 step gst Nothing = do
4442 sv <- newParallelVar StopNone gst
4443 toSVarParallel gst sv m
4444 -- XXX use unfold instead?
4445 return $ Skip $ Just $ fromSVar sv
4447 step gst (Just (UnStream step1 st)) = do
4450 Yield a s -> Yield a (Just $ Stream step1 s)
4451 Skip s -> Skip (Just $ Stream step1 s)
4454-- Compare with mkAsync. mkAsync uses an Async style SVar whereas this uses a
4455-- parallel style SVar for evaluation. Currently, parallel style cannot use
4456-- rate control whereas Async style can use rate control. In async style SVar
4457-- the worker thread terminates when the buffer is full whereas in Parallel
4460-- | Make the stream producer and consumer run concurrently by introducing a
4461-- buffer between them. The producer thread evaluates the input stream until
4462-- the buffer fills, it blocks if the buffer is full until there is space in
4463-- the buffer. The consumer consumes the stream lazily from the buffer.
4467{-# INLINE_NORMAL mkParallel #-}
4468mkParallel :: (K.IsStream t, MonadAsync m) => t m a -> t m a
4469mkParallel = fromStreamD . mkParallelD . toStreamD
4471-- Note: we can use another API with two callbacks stop and yield if we want
4472-- the callback to be able to indicate end of stream.
4474-- | Generates a callback and a stream pair. The callback returned is used to
4475-- queue values to the stream. The stream is infinite, there is no way for the
4476-- callback to indicate that it is done now.
4480{-# INLINE_NORMAL newCallbackStream #-}
4481newCallbackStream :: (K.IsStream t, MonadAsync m) => m ((a -> m ()), t m a)
4482newCallbackStream = do
4483 sv <- newParallelVar StopNone defState
4485 -- XXX Add our own thread-id to the SVar as we can not know the callback's
4486 -- thread-id and the callback is not run in a managed worker. We need to
4487 -- handle this better.
4488 liftIO myThreadId >>= modifyThread sv
4490 let callback a = liftIO $ void $ send sv (ChildYield a)
4491 -- XXX we can return an SVar and then the consumer can unfold from the
4493 return (callback, fromStreamD (fromSVar sv))
4495-------------------------------------------------------------------------------
4497-------------------------------------------------------------------------------
4499-- | Create an SVar with a fold consumer that will fold any elements sent to it
4500-- using the supplied fold function.
4501{-# INLINE newFoldSVar #-}
4502newFoldSVar :: MonadAsync m => State t m a -> Fold m a b -> m (SVar t m a)
4503newFoldSVar stt f = do
4504 -- Buffer size for the SVar is derived from the current state
4505 sv <- newParallelVar StopAny (adaptState stt)
4506 -- Add the producer thread-id to the SVar.
4507 liftIO myThreadId >>= modifyThread sv
4508 void $ doFork (work sv) (svarMrun sv) (handleFoldException sv)
4513 {-# NOINLINE work #-}
4514 work sv = void $ runFold f $ fromProducer sv
4516data TapState sv st = TapInit | Tapping sv st | TapDone st
4518{-# INLINE_NORMAL tapAsync #-}
4519tapAsync :: MonadAsync m => Fold m a b -> Stream m a -> Stream m a
4520tapAsync f (Stream step1 state1) = Stream step TapInit
4524 -- In general, a Stop event would come equipped with the result
4525 -- of the fold. It is not used here but it would be useful in
4526 -- applicative and distribute.
4527 done <- fromConsumer svr
4528 when (not done) $ do
4529 liftIO $ withDiagMVar svr "teeToSVar: waiting to drain"
4530 $ takeMVar (outputDoorBellFromConsumer svr)
4534 liftIO $ sendStop svr Nothing
4535 -- drain/wait until a stop event arrives from the fold.
4538 {-# INLINE_LATE step #-}
4539 step gst TapInit = do
4540 sv <- newFoldSVar gst f
4541 return $ Skip (Tapping sv state1)
4543 step gst (Tapping sv st) = do
4547 done <- pushToFold sv a
4550 -- XXX we do not need to wait synchronously here
4552 return $ Yield a (TapDone s)
4553 else return $ Yield a (Tapping sv s)
4554 Skip s -> return $ Skip (Tapping sv s)
4559 step gst (TapDone st) = do
4562 Yield a s -> Yield a (TapDone s)
4563 Skip s -> Skip (TapDone s)
4566-- XXX Exported from Array again as this fold is specific to Array
4567-- | Take last 'n' elements from the stream and discard the rest.
4569lastN :: (Storable a, MonadIO m) => Int -> Fold m a (Array a)
4571 | n <= 0 = fmap (const mempty) FL.drain
4572 | otherwise = A.unsafeFreeze <$> Fold step initial done
4574 step (Tuple3' rb rh i) a = do
4575 rh1 <- liftIO $ RB.unsafeInsert rb rh a
4576 return $ Tuple3' rb rh1 (i + 1)
4577 initial = fmap (\(a, b) -> Tuple3' a b (0 :: Int)) $ liftIO $ RB.new n
4578 done (Tuple3' rb rh i) = do
4579 arr <- liftIO $ MA.newArray n
4580 foldFunc i rh snoc' arr rb
4581 snoc' b a = liftIO $ MA.unsafeSnoc b a
4583 | i < n = RB.unsafeFoldRingM
4584 | otherwise = RB.unsafeFoldRingFullM
4586------------------------------------------------------------------------------
4588------------------------------------------------------------------------------
4590-- XXX using getTime in the loop can be pretty expensive especially for
4591-- computations where iterations are lightweight. We have the following
4594-- 1) Run a timeout thread updating a flag asynchronously and check that
4595-- flag here, that way we can have a cheap termination check.
4597-- 2) Use COARSE clock to get time with lower resolution but more efficiently.
4599-- 3) Use rdtscp/rdtsc to get time directly from the processor, compute the
4600-- termination value of rdtsc in the beginning and then in each iteration just
4601-- get rdtsc and check if we should terminate.
4605 | TakeByTimeCheck st s
4606 | TakeByTimeYield st s
4608{-# INLINE_NORMAL takeByTime #-}
4609takeByTime :: (MonadIO m, TimeUnit64 t) => t -> Stream m a -> Stream m a
4610takeByTime duration (Stream step1 state1) = Stream step (TakeByTimeInit state1)
4613 lim = toRelTime64 duration
4615 {-# INLINE_LATE step #-}
4616 step _ (TakeByTimeInit _) | lim == 0 = return Stop
4617 step _ (TakeByTimeInit st) = do
4618 t0 <- liftIO $ getTime Monotonic
4619 return $ Skip (TakeByTimeYield st t0)
4620 step _ (TakeByTimeCheck st t0) = do
4621 t <- liftIO $ getTime Monotonic
4623 if diffAbsTime64 t t0 > lim
4625 else Skip (TakeByTimeYield st t0)
4626 step gst (TakeByTimeYield st t0) = do
4629 Yield x s -> Yield x (TakeByTimeCheck s t0)
4630 Skip s -> Skip (TakeByTimeCheck s t0)
4633data DropByTime st s x
4635 | DropByTimeGen st s
4636 | DropByTimeCheck st s x
4637 | DropByTimeYield st
4639{-# INLINE_NORMAL dropByTime #-}
4640dropByTime :: (MonadIO m, TimeUnit64 t) => t -> Stream m a -> Stream m a
4641dropByTime duration (Stream step1 state1) = Stream step (DropByTimeInit state1)
4644 lim = toRelTime64 duration
4646 {-# INLINE_LATE step #-}
4647 step _ (DropByTimeInit st) = do
4648 t0 <- liftIO $ getTime Monotonic
4649 return $ Skip (DropByTimeGen st t0)
4650 step gst (DropByTimeGen st t0) = do
4653 Yield x s -> Skip (DropByTimeCheck s t0 x)
4654 Skip s -> Skip (DropByTimeGen s t0)
4656 step _ (DropByTimeCheck st t0 x) = do
4657 t <- liftIO $ getTime Monotonic
4658 if diffAbsTime64 t t0 <= lim
4659 then return $ Skip $ DropByTimeGen st t0
4660 else return $ Yield x $ DropByTimeYield st
4661 step gst (DropByTimeYield st) = do
4664 Yield x s -> Yield x (DropByTimeYield s)
4665 Skip s -> Skip (DropByTimeYield s)
4668{-# INLINE updateTimeVar #-}
4669updateTimeVar :: Prim.IORef Int64 -> IO ()
4670updateTimeVar timeVar = do
4671 MicroSecond64 t <- fromAbsTime <$> getTime Monotonic
4672 Prim.modifyIORef' timeVar (const t)
4674{-# INLINE updateWithDelay #-}
4675updateWithDelay :: RealFrac a => a -> Prim.IORef Int64 -> IO ()
4676updateWithDelay precision timeVar = do
4677 threadDelay (delayTime precision)
4678 updateTimeVar timeVar
4682 -- Keep the minimum at least a millisecond to avoid high CPU usage
4683 {-# INLINE delayTime #-}
4685 | g' >= fromIntegral (maxBound :: Int) = maxBound
4687 | otherwise = round g'
4691 g' = g * 10 ^ (6 :: Int)
4693-- XXX we should move this to stream generation section of this file. Also, the
4694-- take/drop combinators above should be moved to filtering section.
4696{-# INLINE_NORMAL times #-}
4697times :: MonadAsync m => Double -> Stream m (AbsTime, RelTime64)
4698times g = Stream step Nothing
4702 {-# INLINE_LATE step #-}
4704 -- XXX note that this is safe only on a 64-bit machine. On a 32-bit
4705 -- machine a 64-bit 'Var' cannot be read consistently without a lock
4706 -- while another thread is writing to it.
4707 timeVar <- liftIO $ Prim.newIORef (0 :: Int64)
4708 liftIO $ updateTimeVar timeVar
4709 tid <- forkManaged $ liftIO $ forever (updateWithDelay g timeVar)
4710 a <- liftIO $ Prim.readIORef timeVar
4711 return $ Skip $ Just (timeVar, tid, a)
4713 step _ s@(Just (timeVar, _, t0)) = do
4714 a <- liftIO $ Prim.readIORef timeVar
4715 -- XXX we can perhaps use an AbsTime64 using a 64 bit Int for
4716 -- efficiency. or maybe we can use a representation using Double for
4717 -- floating precision time
4718 return $ Yield (toAbsTime (MicroSecond64 t0),
4719 (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