4-- Module : Streamly.Internal.Data.Stream.StreamD
5-- Copyright : (c) 2018 Composewell Technologies
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.Storable (Storable(..))
324import GHC.Exts (SpecConstrAnnotation(..))
325import GHC.Types (SPEC(..))
326import System.Mem (performMajorGC)
327import Fusion.Plugin.Types (Fuse(..))
328import Streamly.Internal.Data.IORef.Prim (Prim)
329import Streamly.Internal.Data.Time.Units
330 (TimeUnit64, toRelTime64, diffAbsTime64, RelTime64)
331import Streamly.Internal.Data.Atomics (atomicModifyIORefCAS_)
332import Streamly.Internal.Data.Array.Storable.Foreign.Types (Array(..))
333import Streamly.Internal.Data.Fold.Types (Fold(..))
334import Streamly.Internal.Data.Parser (ParseError(..))
335import Streamly.Internal.Data.Pipe.Types (Pipe(..), PipeState(..))
336import Streamly.Internal.Data.Time.Clock (Clock(Monotonic), getTime)
337import Streamly.Internal.Data.Time.Units
338 (MicroSecond64(..), fromAbsTime, toAbsTime, AbsTime)
339import Streamly.Internal.Data.Unfold.Types (Unfold(..))
340import Streamly.Internal.Data.Tuple.Strict (Tuple3'(..))
341import Streamly.Internal.Data.Stream.SVar (fromConsumer, pushToFold)
343import qualified Streamly.Internal.Data.IORef.Prim as Prim
344import qualified Streamly.Internal.Data.Pipe.Types as Pipe
345import qualified Streamly.Internal.Data.Array.Storable.Foreign.Types as A
346import qualified Streamly.Internal.Data.Array.Storable.Foreign.Mut.Types as MA
347import qualified Streamly.Internal.Data.Fold as FL
348import qualified Streamly.Memory.Ring as RB
349import qualified Streamly.Internal.Data.Stream.StreamK as K
350import qualified Streamly.Internal.Data.Parser as PR
351import qualified Streamly.Internal.Data.Parser.ParserD as PRD
352import qualified Control.Monad.Catch as MC
353import qualified Control.Monad.Trans.Reader as Reader
354import qualified Control.Monad.Trans.State.Strict as State
355import qualified Prelude
358 ( map, mapM, mapM_, repeat, foldr, last, take, filter
359 , takeWhile, drop, dropWhile, all, any, maximum, minimum, elem
360 , notElem, null, head, tail, zipWith, lookup, foldr1, sequence
361 , (!!), scanl, scanl1, concatMap, replicate, enumFromTo, concat
362 , reverse, iterate, splitAt)
363import Streamly.Internal.Data.Stream.StreamD.Type
364import Streamly.Internal.Data.SVar
366------------------------------------------------------------------------------
368------------------------------------------------------------------------------
370-- | An empty 'Stream'.
371{-# INLINE_NORMAL nil #-}
372nil :: Monad m => Stream m a
373nil = Stream (\_ _ -> return Stop) ()
375-- | An empty 'Stream' with a side effect.
376{-# INLINE_NORMAL nilM #-}
377nilM :: Monad m => m b -> Stream m a
378nilM m = Stream (\_ _ -> m >> return Stop) ()
380{-# INLINE_NORMAL consM #-}
381consM :: Monad m => m a -> Stream m a -> Stream m a
382consM m (Stream step state) = Stream step1 Nothing
384 {-# INLINE_LATE step1 #-}
385 step1 _ Nothing = m >>= \x -> return $ Yield x (Just state)
386 step1 gst (Just st) = do
390 Yield a s -> Yield a (Just s)
391 Skip s -> Skip (Just s)
394-- XXX implement in terms of consM?
395-- cons x = consM (return x)
397-- | Can fuse but has O(n^2) complexity.
398{-# INLINE_NORMAL cons #-}
399cons :: Monad m => a -> Stream m a -> Stream m a
400cons x (Stream step state) = Stream step1 Nothing
402 {-# INLINE_LATE step1 #-}
403 step1 _ Nothing = return $ Yield x (Just state)
404 step1 gst (Just st) = do
408 Yield a s -> Yield a (Just s)
409 Skip s -> Skip (Just s)
412-------------------------------------------------------------------------------
414-------------------------------------------------------------------------------
416-- Does not fuse, has the same performance as the StreamK version.
417{-# INLINE_NORMAL uncons #-}
418uncons :: Monad m => Stream m a -> m (Maybe (a, Stream m a))
419uncons (UnStream step state) = go state
422 r <- step defState st
424 Yield x s -> return $ Just (x, Stream step s)
426 Stop -> return Nothing
428------------------------------------------------------------------------------
429-- Generation by unfold
430------------------------------------------------------------------------------
432{-# INLINE_NORMAL unfoldrM #-}
433unfoldrM :: Monad m => (s -> m (Maybe (a, s))) -> s -> Stream m a
434unfoldrM next state = Stream step state
436 {-# INLINE_LATE step #-}
440 Just (x, s) -> Yield x s
443{-# INLINE_LATE unfoldr #-}
444unfoldr :: Monad m => (s -> Maybe (a, s)) -> s -> Stream m a
445unfoldr f = unfoldrM (return . f)
447-- | Convert an 'Unfold' into a 'Stream' by supplying it a seed.
449{-# INLINE_NORMAL unfold #-}
450unfold :: Monad m => Unfold m a b -> a -> Stream m b
451unfold (Unfold ustep inject) seed = Stream step Nothing
453 {-# INLINE_LATE step #-}
454 step _ Nothing = inject seed >>= return . Skip . Just
455 step _ (Just st) = do
458 Yield x s -> Yield x (Just s)
459 Skip s -> Skip (Just s)
462------------------------------------------------------------------------------
463-- Specialized Generation
464------------------------------------------------------------------------------
466{-# INLINE_NORMAL repeatM #-}
467repeatM :: Monad m => m a -> Stream m a
468repeatM x = Stream (\_ _ -> x >>= \r -> return $ Yield r ()) ()
470{-# INLINE_NORMAL repeat #-}
471repeat :: Monad m => a -> Stream m a
472repeat x = Stream (\_ _ -> return $ Yield x ()) ()
474{-# INLINE_NORMAL iterateM #-}
475iterateM :: Monad m => (a -> m a) -> m a -> Stream m a
476iterateM step = Stream (\_ st -> st >>= \x -> return $ Yield x (step x))
478{-# INLINE_NORMAL iterate #-}
479iterate :: Monad m => (a -> a) -> a -> Stream m a
480iterate step st = iterateM (return . step) (return st)
482{-# INLINE_NORMAL replicateM #-}
483replicateM :: forall m a. Monad m => Int -> m a -> Stream m a
484replicateM n p = Stream step n
486 {-# INLINE_LATE step #-}
488 | i <= 0 = return Stop
491 return $ Yield x (i - 1)
493{-# INLINE_NORMAL replicate #-}
494replicate :: Monad m => Int -> a -> Stream m a
495replicate n x = replicateM n (return x)
497-- This would not work properly for floats, therefore we put an Integral
499-- | Can be used to enumerate unbounded integrals. This does not check for
500-- overflow or underflow for bounded integrals.
501{-# INLINE_NORMAL enumerateFromStepIntegral #-}
502enumerateFromStepIntegral :: (Integral a, Monad m) => a -> a -> Stream m a
503enumerateFromStepIntegral from stride =
504 from `seq` stride `seq` Stream step from
506 {-# INLINE_LATE step #-}
507 step _ !x = return $ Yield x $! (x + stride)
509-- We are assuming that "to" is constrained by the type to be within
511{-# INLINE enumerateFromToIntegral #-}
512enumerateFromToIntegral :: (Monad m, Integral a) => a -> a -> Stream m a
513enumerateFromToIntegral from to =
514 takeWhile (<= to) $ enumerateFromStepIntegral from 1
516{-# INLINE enumerateFromIntegral #-}
517enumerateFromIntegral :: (Monad m, Integral a, Bounded a) => a -> Stream m a
518enumerateFromIntegral from = enumerateFromToIntegral from maxBound
520data EnumState a = EnumInit | EnumYield a a a | EnumStop
522{-# INLINE_NORMAL enumerateFromThenToIntegralUp #-}
523enumerateFromThenToIntegralUp
524 :: (Monad m, Integral a)
525 => a -> a -> a -> Stream m a
526enumerateFromThenToIntegralUp from next to = Stream step EnumInit
528 {-# INLINE_LATE step #-}
534 else Yield from EnumStop
535 else -- from <= next <= to
536 let stride = next - from
537 in Skip $ EnumYield from stride (to - stride)
539 step _ (EnumYield x stride toMinus) =
542 then Yield x EnumStop
543 else Yield x $ EnumYield (x + stride) stride toMinus
545 step _ EnumStop = return Stop
547{-# INLINE_NORMAL enumerateFromThenToIntegralDn #-}
548enumerateFromThenToIntegralDn
549 :: (Monad m, Integral a)
550 => a -> a -> a -> Stream m a
551enumerateFromThenToIntegralDn from next to = Stream step EnumInit
553 {-# INLINE_LATE step #-}
555 return $ if to > next
558 else Yield from EnumStop
559 else -- from >= next >= to
560 let stride = next - from
561 in Skip $ EnumYield from stride (to - stride)
563 step _ (EnumYield x stride toMinus) =
566 then Yield x EnumStop
567 else Yield x $ EnumYield (x + stride) stride toMinus
569 step _ EnumStop = return Stop
571{-# INLINE_NORMAL enumerateFromThenToIntegral #-}
572enumerateFromThenToIntegral
573 :: (Monad m, Integral a)
574 => a -> a -> a -> Stream m a
575enumerateFromThenToIntegral from next to
576 | next >= from = enumerateFromThenToIntegralUp from next to
577 | otherwise = enumerateFromThenToIntegralDn from next to
579{-# INLINE_NORMAL enumerateFromThenIntegral #-}
580enumerateFromThenIntegral
581 :: (Monad m, Integral a, Bounded a)
582 => a -> a -> Stream m a
583enumerateFromThenIntegral from next =
585 then enumerateFromThenToIntegralUp from next maxBound
586 else enumerateFromThenToIntegralDn from next minBound
588-- For floating point numbers if the increment is less than the precision then
589-- it just gets lost. Therefore we cannot always increment it correctly by just
591-- 9007199254740992 + 1 + 1 :: Double => 9.007199254740992e15
592-- 9007199254740992 + 2 :: Double => 9.007199254740994e15
594-- Instead we accumulate the increment counter and compute the increment
595-- every time before adding it to the starting number.
597-- This works for Integrals as well as floating point numbers, but
598-- enumerateFromStepIntegral is faster for integrals.
599{-# INLINE_NORMAL enumerateFromStepNum #-}
600enumerateFromStepNum :: (Monad m, Num a) => a -> a -> Stream m a
601enumerateFromStepNum from stride = Stream step 0
603 {-# INLINE_LATE step #-}
604 step _ !i = return $ (Yield $! (from + i * stride)) $! (i + 1)
606{-# INLINE_NORMAL numFrom #-}
607numFrom :: (Monad m, Num a) => a -> Stream m a
608numFrom from = enumerateFromStepNum from 1
610{-# INLINE_NORMAL numFromThen #-}
611numFromThen :: (Monad m, Num a) => a -> a -> Stream m a
612numFromThen from next = enumerateFromStepNum from (next - from)
614-- We cannot write a general function for Num. The only way to write code
615-- portable between the two is to use a 'Real' constraint and convert between
616-- Fractional and Integral using fromRational which is horribly slow.
617{-# INLINE_NORMAL enumerateFromToFractional #-}
618enumerateFromToFractional
619 :: (Monad m, Fractional a, Ord a)
620 => a -> a -> Stream m a
621enumerateFromToFractional from to =
622 takeWhile (<= to + 1 / 2) $ enumerateFromStepNum from 1
624{-# INLINE_NORMAL enumerateFromThenToFractional #-}
625enumerateFromThenToFractional
626 :: (Monad m, Fractional a, Ord a)
627 => a -> a -> a -> Stream m a
628enumerateFromThenToFractional from next to =
629 takeWhile predicate $ numFromThen from next
631 mid = (next - from) / 2
632 predicate | next >= from = (<= to + mid)
633 | otherwise = (>= to + mid)
635-------------------------------------------------------------------------------
636-- Generation by Conversion
637-------------------------------------------------------------------------------
639{-# INLINE_NORMAL fromIndicesM #-}
640fromIndicesM :: Monad m => (Int -> m a) -> Stream m a
641fromIndicesM gen = Stream step 0
643 {-# INLINE_LATE step #-}
646 return $ Yield x (i + 1)
648{-# INLINE fromIndices #-}
649fromIndices :: Monad m => (Int -> a) -> Stream m a
650fromIndices gen = fromIndicesM (return . gen)
652{-# INLINE_NORMAL generateM #-}
653generateM :: Monad m => Int -> (Int -> m a) -> Stream m a
654generateM n gen = n `seq` Stream step 0
656 {-# INLINE_LATE step #-}
657 step _ i | i < n = do
659 return $ Yield x (i + 1)
660 | otherwise = return Stop
662{-# INLINE generate #-}
663generate :: Monad m => Int -> (Int -> a) -> Stream m a
664generate n gen = generateM n (return . gen)
666-- XXX we need the MonadAsync constraint because of a rewrite rule.
667-- | Convert a list of monadic actions to a 'Stream'
668{-# INLINE_LATE fromListM #-}
669fromListM :: MonadAsync m => [m a] -> Stream m a
670fromListM = Stream step
672 {-# INLINE_LATE step #-}
673 step _ (m:ms) = m >>= \x -> return $ Yield x ms
674 step _ [] = return Stop
676{-# INLINE toStreamD #-}
677toStreamD :: (K.IsStream t, Monad m) => t m a -> Stream m a
678toStreamD = fromStreamK . K.toStream
680{-# INLINE_NORMAL fromPrimIORef #-}
681fromPrimIORef :: (MonadIO m, Prim a) => Prim.IORef a -> Stream m a
682fromPrimIORef var = Stream step ()
684 {-# INLINE_LATE step #-}
685 step _ () = liftIO (Prim.readIORef var) >>= \x -> return $ Yield x ()
687-------------------------------------------------------------------------------
688-- Generation from SVar
689-------------------------------------------------------------------------------
691data FromSVarState t m a =
693 | FromSVarRead (SVar t m a)
694 | FromSVarLoop (SVar t m a) [ChildEvent a]
695 | FromSVarDone (SVar t m a)
697{-# INLINE_NORMAL fromSVar #-}
698fromSVar :: (MonadAsync m) => SVar t m a -> Stream m a
699fromSVar svar = Stream step FromSVarInit
702 {-# INLINE_LATE step #-}
703 step _ FromSVarInit = do
704 ref <- liftIO $ newIORef ()
705 _ <- liftIO $ mkWeakIORef ref hook
706 -- when this copy of svar gets garbage collected "ref" will get
707 -- garbage collected and our GC hook will be called.
708 let sv = svar{svarRef = Just ref}
709 return $ Skip (FromSVarRead sv)
713 {-# NOINLINE hook #-}
715 when (svarInspectMode svar) $ do
716 r <- liftIO $ readIORef (svarStopTime (svarStats svar))
718 printSVar svar "SVar Garbage Collected"
720 -- If there are any SVars referenced by this SVar a GC will prompt
721 -- them to be cleaned up quickly.
722 when (svarInspectMode svar) performMajorGC
724 step _ (FromSVarRead sv) = do
725 list <- readOutputQ sv
726 -- Reversing the output is important to guarantee that we process the
727 -- outputs in the same order as they were generated by the constituent
729 return $ Skip $ FromSVarLoop sv (Prelude.reverse list)
731 step _ (FromSVarLoop sv []) = do
732 done <- postProcess sv
733 return $ Skip $ if done
734 then (FromSVarDone sv)
735 else (FromSVarRead sv)
737 step _ (FromSVarLoop sv (ev : es)) = do
739 ChildYield a -> return $ Yield a (FromSVarLoop sv es)
740 ChildStop tid e -> do
744 stop <- shouldStop tid
747 liftIO (cleanupSVar sv)
748 return $ Skip (FromSVarDone sv)
749 else return $ Skip (FromSVarLoop sv es)
751 case fromException ex of
753 return $ Skip (FromSVarLoop sv es)
754 Nothing -> liftIO (cleanupSVar sv) >> throwM ex
758 case svarStopStyle sv of
759 StopNone -> return False
760 StopAny -> return True
762 sid <- liftIO $ readIORef (svarStopBy sv)
763 return $ if tid == sid then True else False
765 step _ (FromSVarDone sv) = do
766 when (svarInspectMode sv) $ do
767 t <- liftIO $ getTime Monotonic
768 liftIO $ writeIORef (svarStopTime (svarStats sv)) (Just t)
769 liftIO $ printSVar sv "SVar Done"
772-------------------------------------------------------------------------------
773-- Process events received by a fold consumer from a stream producer
774-------------------------------------------------------------------------------
776{-# INLINE_NORMAL fromProducer #-}
777fromProducer :: (MonadAsync m) => SVar t m a -> Stream m a
778fromProducer svar = Stream step (FromSVarRead svar)
781 {-# INLINE_LATE step #-}
782 step _ (FromSVarRead sv) = do
783 list <- readOutputQ sv
784 -- Reversing the output is important to guarantee that we process the
785 -- outputs in the same order as they were generated by the constituent
787 return $ Skip $ FromSVarLoop sv (Prelude.reverse list)
789 step _ (FromSVarLoop sv []) = return $ Skip $ FromSVarRead sv
790 step _ (FromSVarLoop sv (ev : es)) = do
792 ChildYield a -> return $ Yield a (FromSVarLoop sv es)
793 ChildStop tid e -> do
797 sendStopToProducer sv
798 return $ Skip (FromSVarDone sv)
799 Just _ -> error "Bug: fromProducer: received exception"
801 step _ (FromSVarDone sv) = do
802 when (svarInspectMode sv) $ do
803 t <- liftIO $ getTime Monotonic
804 liftIO $ writeIORef (svarStopTime (svarStats sv)) (Just t)
805 liftIO $ printSVar sv "SVar Done"
808 step _ FromSVarInit = undefined
810-------------------------------------------------------------------------------
811-- Hoisting the inner monad
812-------------------------------------------------------------------------------
814{-# INLINE_NORMAL hoist #-}
815hoist :: Monad n => (forall x. m x -> n x) -> Stream m a -> Stream n a
816hoist f (Stream step state) = (Stream step' state)
818 {-# INLINE_LATE step' #-}
820 r <- f $ step (adaptState gst) st
822 Yield x s -> Yield x s
826{-# INLINE generally #-}
827generally :: Monad m => Stream Identity a -> Stream m a
828generally = hoist (return . runIdentity)
830{-# INLINE_NORMAL liftInner #-}
831liftInner :: (Monad m, MonadTrans t, Monad (t m))
832 => Stream m a -> Stream (t m) a
833liftInner (Stream step state) = Stream step' state
835 {-# INLINE_LATE step' #-}
837 r <- lift $ step (adaptState gst) st
839 Yield x s -> Yield x s
843{-# INLINE_NORMAL runReaderT #-}
844runReaderT :: Monad m => m s -> Stream (ReaderT s m) a -> Stream m a
845runReaderT env (Stream step state) = Stream step' (state, env)
847 {-# INLINE_LATE step' #-}
848 step' gst (st, action) = do
850 r <- Reader.runReaderT (step (adaptState gst) st) sv
852 Yield x s -> Yield x (s, return sv)
853 Skip s -> Skip (s, return sv)
856{-# INLINE_NORMAL evalStateT #-}
857evalStateT :: Monad m => m s -> Stream (StateT s m) a -> Stream m a
858evalStateT initial (Stream step state) = Stream step' (state, initial)
860 {-# INLINE_LATE step' #-}
861 step' gst (st, action) = do
863 (r, sv') <- State.runStateT (step (adaptState gst) st) sv
865 Yield x s -> Yield x (s, return sv')
866 Skip s -> Skip (s, return sv')
869{-# INLINE_NORMAL runStateT #-}
870runStateT :: Monad m => m s -> Stream (StateT s m) a -> Stream m (s, a)
871runStateT initial (Stream step state) = Stream step' (state, initial)
873 {-# INLINE_LATE step' #-}
874 step' gst (st, action) = do
876 (r, sv') <- State.runStateT (step (adaptState gst) st) sv
878 Yield x s -> Yield (sv', x) (s, return sv')
879 Skip s -> Skip (s, return sv')
882------------------------------------------------------------------------------
883-- Elimination by Folds
884------------------------------------------------------------------------------
886------------------------------------------------------------------------------
888------------------------------------------------------------------------------
890{-# INLINE_NORMAL foldr1 #-}
891foldr1 :: Monad m => (a -> a -> a) -> Stream m a -> m (Maybe a)
895 Nothing -> return Nothing
896 Just (h, t) -> fmap Just (foldr f h t)
898------------------------------------------------------------------------------
900------------------------------------------------------------------------------
902{-# INLINE_NORMAL foldlT #-}
903foldlT :: (Monad m, Monad (s m), MonadTrans s)
904 => (s m b -> a -> s m b) -> s m b -> Stream m a -> s m b
905foldlT fstep begin (Stream step state) = go SPEC begin state
908 r <- lift $ step defState st
910 Yield x s -> go SPEC (fstep acc x) s
911 Skip s -> go SPEC acc s
914-- Note, this is going to have horrible performance, because of the nature of
915-- the stream type (i.e. direct stream vs CPS). Its only for reference, it is
916-- likely be practically unusable.
917{-# INLINE_NORMAL foldlS #-}
919 => (Stream m b -> a -> Stream m b) -> Stream m b -> Stream m a -> Stream m b
920foldlS fstep begin (Stream step state) = Stream step' (Left (state, begin))
922 step' gst (Left (st, acc)) = do
923 r <- step (adaptState gst) st
925 Yield x s -> Skip (Left (s, fstep acc x))
926 Skip s -> Skip (Left (s, acc))
927 Stop -> Skip (Right acc)
929 step' gst (Right (Stream stp stt)) = do
930 r <- stp (adaptState gst) stt
932 Yield x s -> Yield x (Right (Stream stp s))
933 Skip s -> Skip (Right (Stream stp s))
936------------------------------------------------------------------------------
938------------------------------------------------------------------------------
940-- Inlined definition. Without the inline "serially/parser/take" benchmark
941-- degrades and parseMany does not fuse. Even using "inline" at the callsite
943{-# INLINE splitAt #-}
944splitAt :: Int -> [a] -> ([a],[a])
947 | otherwise = splitAt' n ls
949 splitAt' :: Int -> [a] -> ([a], [a])
950 splitAt' _ [] = ([], [])
951 splitAt' 1 (x:xs) = ([x], xs)
952 splitAt' m (x:xs) = (x:xs', xs'')
954 (xs', xs'') = splitAt' (m - 1) xs
956-- GHC parser does not accept {-# ANN type [] NoSpecConstr #-}, so we need
958{-# ANN type List NoSpecConstr #-}
959newtype List a = List {getList :: [a]}
961-- | Run a 'Parse' over a stream.
962{-# INLINE_NORMAL parselMx' #-}
965 => (s -> a -> m (PR.Step s b))
970parselMx' pstep initial extract (Stream step state) = do
971 initial >>= go SPEC state (List [])
975 -- XXX currently we are using a dumb list based approach for backtracking
976 -- buffer. This can be replaced by a sliding/ring buffer using Data.Array.
977 -- That will allow us more efficient random back and forth movement.
979 go !_ st buf !pst = do
980 r <- step defState st
985 PR.Partial 0 pst1 -> go SPEC s (List []) pst1
986 PR.Partial n pst1 -> do
987 assert (n <= length (x:getList buf)) (return ())
988 let src0 = Prelude.take n (x:getList buf)
989 src = Prelude.reverse src0
990 gobuf SPEC s (List []) (List src) pst1
991 PR.Continue 0 pst1 -> go SPEC s (List (x:getList buf)) pst1
992 PR.Continue n pst1 -> do
993 assert (n <= length (x:getList buf)) (return ())
994 let (src0, buf1) = splitAt n (x:getList buf)
995 src = Prelude.reverse src0
996 gobuf SPEC s (List buf1) (List src) pst1
997 PR.Done _ b -> return b
998 PR.Error err -> throwM $ ParseError err
999 Skip s -> go SPEC s buf pst
1002 gobuf !_ s buf (List []) !pst = go SPEC s buf pst
1003 gobuf !_ s buf (List (x:xs)) !pst = do
1006 PR.Partial 0 pst1 ->
1007 gobuf SPEC s (List []) (List xs) pst1
1008 PR.Partial n pst1 -> do
1009 assert (n <= length (x:getList buf)) (return ())
1010 let src0 = Prelude.take n (x:getList buf)
1011 src = Prelude.reverse src0 ++ xs
1012 gobuf SPEC s (List []) (List src) pst1
1013 PR.Continue 0 pst1 ->
1014 gobuf SPEC s (List (x:getList buf)) (List xs) pst1
1015 PR.Continue n pst1 -> do
1016 assert (n <= length (x:getList buf)) (return ())
1017 let (src0, buf1) = splitAt n (x:getList buf)
1018 src = Prelude.reverse src0 ++ xs
1019 gobuf SPEC s (List buf1) (List src) pst1
1020 PR.Done _ b -> return b
1021 PR.Error err -> throwM $ ParseError err
1023------------------------------------------------------------------------------
1025------------------------------------------------------------------------------
1027{-# ANN type ParseChunksState Fuse #-}
1028data ParseChunksState x inpBuf st pst =
1029 ParseChunksInit inpBuf st
1030 | ParseChunksInitLeftOver inpBuf
1031 | ParseChunksStream st inpBuf !pst
1032 | ParseChunksBuf inpBuf st inpBuf !pst
1033 | ParseChunksYield x (ParseChunksState x inpBuf st pst)
1035{-# INLINE_NORMAL parseMany #-}
1041parseMany (PRD.Parser pstep initial extract) (Stream step state) =
1042 Stream stepOuter (ParseChunksInit [] state)
1046 {-# INLINE_LATE stepOuter #-}
1047 -- Buffer is empty, get the first element from the stream, initialize the
1048 -- fold and then go to stream processing loop.
1049 stepOuter gst (ParseChunksInit [] st) = do
1050 r <- step (adaptState gst) st
1052 Yield x s -> initial >>= return . Skip . ParseChunksBuf [x] s []
1053 Skip s -> return $ Skip $ ParseChunksInit [] s
1056 -- Buffer is not empty, go to buffered processing loop
1057 stepOuter _ (ParseChunksInit src st) = do
1058 initial >>= return . Skip . ParseChunksBuf src st []
1060 -- XXX we just discard any leftover input at the end
1061 stepOuter _ (ParseChunksInitLeftOver _) = return Stop
1063 -- Buffer is empty, process elements from the stream
1064 stepOuter gst (ParseChunksStream st buf pst) = do
1065 r <- step (adaptState gst) st
1070 PR.Partial 0 pst1 ->
1071 return $ Skip $ ParseChunksStream s [] pst1
1072 PR.Partial n pst1 -> do
1073 assert (n <= length (x:buf)) (return ())
1074 let src0 = Prelude.take n (x:buf)
1075 src = Prelude.reverse src0
1076 return $ Skip $ ParseChunksBuf src s [] pst1
1077 PR.Continue 0 pst1 ->
1078 return $ Skip $ ParseChunksStream s (x:buf) pst1
1079 PR.Continue n pst1 -> do
1080 assert (n <= length (x:buf)) (return ())
1081 let (src0, buf1) = splitAt n (x:buf)
1082 src = Prelude.reverse src0
1083 return $ Skip $ ParseChunksBuf src s buf1 pst1
1086 ParseChunksYield b (ParseChunksInit [] s)
1088 assert (n <= length (x:buf)) (return ())
1089 let src = Prelude.reverse (Prelude.take n (x:buf))
1091 ParseChunksYield b (ParseChunksInit src s)
1092 PR.Error err -> throwM $ ParseError err
1093 Skip s -> return $ Skip $ ParseChunksStream s buf pst
1096 let src = Prelude.reverse buf
1098 ParseChunksYield b (ParseChunksInitLeftOver src)
1100 -- go back to stream processing mode
1101 stepOuter _ (ParseChunksBuf [] s buf pst) =
1102 return $ Skip $ ParseChunksStream s buf pst
1104 -- buffered processing loop
1105 stepOuter _ (ParseChunksBuf (x:xs) s buf pst) = do
1108 PR.Partial 0 pst1 ->
1109 return $ Skip $ ParseChunksBuf xs s [] pst1
1110 PR.Partial n pst1 -> do
1111 assert (n <= length (x:buf)) (return ())
1112 let src0 = Prelude.take n (x:buf)
1113 src = Prelude.reverse src0 ++ xs
1114 return $ Skip $ ParseChunksBuf src s [] pst1
1115 PR.Continue 0 pst1 ->
1116 return $ Skip $ ParseChunksBuf xs s (x:buf) pst1
1117 PR.Continue n pst1 -> do
1118 assert (n <= length (x:buf)) (return ())
1119 let (src0, buf1) = splitAt n (x:buf)
1120 src = Prelude.reverse src0 ++ xs
1121 return $ Skip $ ParseChunksBuf src s buf1 pst1
1123 return $ Skip $ ParseChunksYield b (ParseChunksInit xs s)
1125 assert (n <= length (x:buf)) (return ())
1126 let src = Prelude.reverse (Prelude.take n (x:buf)) ++ xs
1127 return $ Skip $ ParseChunksYield b (ParseChunksInit src s)
1128 PR.Error err -> throwM $ ParseError err
1130 stepOuter _ (ParseChunksYield a next) = return $ Yield a next
1132{-# ANN type ConcatParseState Fuse #-}
1133data ConcatParseState x inpBuf st p =
1134 ConcatParseInit inpBuf st p
1135 | ConcatParseInitLeftOver inpBuf
1136 | ConcatParseStream st inpBuf p
1137 | ConcatParseBuf inpBuf st inpBuf p
1138 | ConcatParseYield x (ConcatParseState x inpBuf st p)
1140{-# INLINE_NORMAL parseIterate #-}
1143 => (b -> PRD.Parser m a b)
1147parseIterate func seed (Stream step state) =
1148 Stream stepOuter (ConcatParseInit [] state (func seed))
1152 {-# INLINE_LATE stepOuter #-}
1153 -- Buffer is empty, go to stream processing loop
1154 stepOuter _ (ConcatParseInit [] st (PRD.Parser pstep initial extract)) = do
1155 initial >>= \r -> return $ Skip $ ConcatParseStream st []
1156 (PRD.Parser pstep (return r) extract)
1158 -- Buffer is not empty, go to buffered processing loop
1159 stepOuter _ (ConcatParseInit src st
1160 (PRD.Parser pstep initial extract)) = do
1161 initial >>= \r -> return $ Skip $ ConcatParseBuf src st []
1162 (PRD.Parser pstep (return r) extract)
1164 -- XXX we just discard any leftover input at the end
1165 stepOuter _ (ConcatParseInitLeftOver _) = return Stop
1167 -- Buffer is empty process elements from the stream
1168 stepOuter gst (ConcatParseStream st buf
1169 p@(PRD.Parser pstep initial extract)) = do
1170 r <- step (adaptState gst) st
1176 PR.Partial 0 pst1 ->
1177 return $ Skip $ ConcatParseStream s []
1178 (PRD.Parser pstep (return pst1) extract)
1179 PR.Partial n pst1 -> do
1180 assert (n <= length (x:buf)) (return ())
1181 let src0 = Prelude.take n (x:buf)
1182 src = Prelude.reverse src0
1183 return $ Skip $ ConcatParseBuf src s []
1184 (PRD.Parser pstep (return pst1) extract)
1185 -- PR.Continue 0 pst1 ->
1186 -- return $ Skip $ ConcatParseStream s (x:buf) pst1
1187 PR.Continue n pst1 -> do
1188 assert (n <= length (x:buf)) (return ())
1189 let (src0, buf1) = splitAt n (x:buf)
1190 src = Prelude.reverse src0
1191 return $ Skip $ ConcatParseBuf src s buf1
1192 (PRD.Parser pstep (return pst1) extract)
1193 -- XXX Specialize for Stop 0 common case?
1195 assert (n <= length (x:buf)) (return ())
1196 let src = Prelude.reverse (Prelude.take n (x:buf))
1198 ConcatParseYield b (ConcatParseInit src s (func b))
1199 PR.Error err -> throwM $ ParseError err
1200 Skip s -> return $ Skip $ ConcatParseStream s buf p
1204 let src = Prelude.reverse buf
1206 ConcatParseYield b (ConcatParseInitLeftOver src)
1208 -- go back to stream processing mode
1209 stepOuter _ (ConcatParseBuf [] s buf p) =
1210 return $ Skip $ ConcatParseStream s buf p
1212 -- buffered processing loop
1213 stepOuter _ (ConcatParseBuf (x:xs) s buf
1214 (PRD.Parser pstep initial extract)) = do
1218 PR.Partial 0 pst1 ->
1219 return $ Skip $ ConcatParseBuf xs s []
1220 (PRD.Parser pstep (return pst1) extract)
1221 PR.Partial n pst1 -> do
1222 assert (n <= length (x:buf)) (return ())
1223 let src0 = Prelude.take n (x:buf)
1224 src = Prelude.reverse src0 ++ xs
1225 return $ Skip $ ConcatParseBuf src s []
1226 (PRD.Parser pstep (return pst1) extract)
1227 -- PR.Continue 0 pst1 -> return $ Skip $ ConcatParseBuf xs s (x:buf) pst1
1228 PR.Continue n pst1 -> do
1229 assert (n <= length (x:buf)) (return ())
1230 let (src0, buf1) = splitAt n (x:buf)
1231 src = Prelude.reverse src0 ++ xs
1232 return $ Skip $ ConcatParseBuf src s buf1
1233 (PRD.Parser pstep (return pst1) extract)
1234 -- XXX Specialize for Stop 0 common case?
1236 assert (n <= length (x:buf)) (return ())
1237 let src = Prelude.reverse (Prelude.take n (x:buf)) ++ xs
1238 return $ Skip $ ConcatParseYield b
1239 (ConcatParseInit src s (func b))
1240 PR.Error err -> throwM $ ParseError err
1242 stepOuter _ (ConcatParseYield a next) = return $ Yield a next
1244------------------------------------------------------------------------------
1246------------------------------------------------------------------------------
1248-- | Run a streaming composition, discard the results.
1249{-# INLINE_LATE drain #-}
1250drain :: Monad m => Stream m a -> m ()
1251-- drain = foldrM (\_ xs -> xs) (return ())
1252drain (Stream step state) = go SPEC state
1255 r <- step defState st
1257 Yield _ s -> go SPEC s
1261{-# INLINE_NORMAL null #-}
1262null :: Monad m => Stream m a -> m Bool
1263null m = foldrM (\_ _ -> return False) (return True) m
1265{-# INLINE_NORMAL head #-}
1266head :: Monad m => Stream m a -> m (Maybe a)
1267head m = foldrM (\x _ -> return (Just x)) (return Nothing) m
1269{-# INLINE_NORMAL headElse #-}
1270headElse :: Monad m => a -> Stream m a -> m a
1271headElse a m = foldrM (\x _ -> return x) (return a) m
1273-- Does not fuse, has the same performance as the StreamK version.
1274{-# INLINE_NORMAL tail #-}
1275tail :: Monad m => Stream m a -> m (Maybe (Stream m a))
1276tail (UnStream step state) = go state
1279 r <- step defState st
1281 Yield _ s -> return (Just $ Stream step s)
1283 Stop -> return Nothing
1285-- XXX will it fuse? need custom impl?
1286{-# INLINE_NORMAL last #-}
1287last :: Monad m => Stream m a -> m (Maybe a)
1288last = foldl' (\_ y -> Just y) Nothing
1290{-# INLINE_NORMAL elem #-}
1291elem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
1292-- elem e m = foldrM (\x xs -> if x == e then return True else xs) (return False) m
1293elem e (Stream step state) = go state
1296 r <- step defState st
1299 | x == e -> return True
1302 Stop -> return False
1304{-# INLINE_NORMAL notElem #-}
1305notElem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
1306notElem e s = fmap not (elem e s)
1308{-# INLINE_NORMAL all #-}
1309all :: Monad m => (a -> Bool) -> Stream m a -> m Bool
1310-- all p m = foldrM (\x xs -> if p x then xs else return False) (return True) m
1311all p (Stream step state) = go state
1314 r <- step defState st
1318 | otherwise -> return False
1322{-# INLINE_NORMAL any #-}
1323any :: Monad m => (a -> Bool) -> Stream m a -> m Bool
1324-- any p m = foldrM (\x xs -> if p x then return True else xs) (return False) m
1325any p (Stream step state) = go state
1328 r <- step defState st
1331 | p x -> return True
1334 Stop -> return False
1336{-# INLINE_NORMAL maximum #-}
1337maximum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
1338maximum (Stream step state) = go Nothing state
1341 r <- step defState st
1343 Yield x s -> go (Just x) s
1344 Skip s -> go Nothing s
1345 Stop -> return Nothing
1346 go (Just acc) st = do
1347 r <- step defState st
1350 | acc <= x -> go (Just x) s
1351 | otherwise -> go (Just acc) s
1352 Skip s -> go (Just acc) s
1353 Stop -> return (Just acc)
1355{-# INLINE_NORMAL maximumBy #-}
1356maximumBy :: Monad m => (a -> a -> Ordering) -> Stream m a -> m (Maybe a)
1357maximumBy cmp (Stream step state) = go Nothing state
1360 r <- step defState st
1362 Yield x s -> go (Just x) s
1363 Skip s -> go Nothing s
1364 Stop -> return Nothing
1365 go (Just acc) st = do
1366 r <- step defState st
1368 Yield x s -> case cmp acc x of
1369 GT -> go (Just acc) s
1371 Skip s -> go (Just acc) s
1372 Stop -> return (Just acc)
1374{-# INLINE_NORMAL minimum #-}
1375minimum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
1376minimum (Stream step state) = go Nothing state
1379 r <- step defState st
1381 Yield x s -> go (Just x) s
1382 Skip s -> go Nothing s
1383 Stop -> return Nothing
1384 go (Just acc) st = do
1385 r <- step defState st
1388 | acc <= x -> go (Just acc) s
1389 | otherwise -> go (Just x) s
1390 Skip s -> go (Just acc) s
1391 Stop -> return (Just acc)
1393{-# INLINE_NORMAL minimumBy #-}
1394minimumBy :: Monad m => (a -> a -> Ordering) -> Stream m a -> m (Maybe a)
1395minimumBy cmp (Stream step state) = go Nothing state
1398 r <- step defState st
1400 Yield x s -> go (Just x) s
1401 Skip s -> go Nothing s
1402 Stop -> return Nothing
1403 go (Just acc) st = do
1404 r <- step defState st
1406 Yield x s -> case cmp acc x of
1408 _ -> go (Just acc) s
1409 Skip s -> go (Just acc) s
1410 Stop -> return (Just acc)
1412{-# INLINE_NORMAL (!!) #-}
1413(!!) :: (Monad m) => Stream m a -> Int -> m (Maybe a)
1414(Stream step state) !! i = go i state
1417 r <- step defState st
1419 Yield x s | n < 0 -> return Nothing
1420 | n == 0 -> return $ Just x
1421 | otherwise -> go (n - 1) s
1423 Stop -> return Nothing
1425{-# INLINE_NORMAL lookup #-}
1426lookup :: (Monad m, Eq a) => a -> Stream m (a, b) -> m (Maybe b)
1427lookup e m = foldrM (\(a, b) xs -> if e == a then return (Just b) else xs)
1430{-# INLINE_NORMAL findM #-}
1431findM :: Monad m => (a -> m Bool) -> Stream m a -> m (Maybe a)
1432findM p m = foldrM (\x xs -> p x >>= \r -> if r then return (Just x) else xs)
1436find :: Monad m => (a -> Bool) -> Stream m a -> m (Maybe a)
1437find p = findM (return . p)
1439{-# INLINE_NORMAL findIndices #-}
1440findIndices :: Monad m => (a -> Bool) -> Stream m a -> Stream m Int
1441findIndices p (Stream step state) = Stream step' (state, 0)
1443 {-# INLINE_LATE step' #-}
1444 step' gst (st, i) = i `seq` do
1445 r <- step (adaptState gst) st
1447 Yield x s -> if p x then Yield i (s, i+1) else Skip (s, i+1)
1448 Skip s -> Skip (s, i)
1451{-# INLINE toListRev #-}
1452toListRev :: Monad m => Stream m a -> m [a]
1453toListRev = foldl' (flip (:)) []
1455-- We can implement reverse as:
1457-- > reverse = foldlS (flip cons) nil
1459-- However, this implementation is unusable because of the horrible performance
1460-- of cons. So we just convert it to a list first and then stream from the
1463-- XXX Maybe we can use an Array instead of a list here?
1464{-# INLINE_NORMAL reverse #-}
1465reverse :: Monad m => Stream m a -> Stream m a
1466reverse m = Stream step Nothing
1468 {-# INLINE_LATE step #-}
1471 return $ Skip (Just xs)
1472 step _ (Just (x:xs)) = return $ Yield x (Just xs)
1473 step _ (Just []) = return Stop
1475-- Much faster reverse for Storables
1476{-# INLINE_NORMAL reverse' #-}
1477reverse' :: forall m a. (MonadIO m, Storable a) => Stream m a -> Stream m a
1479-- This commented implementation copies the whole stream into one single array
1480-- and then streams from that array, this is 3-4x faster than the chunked code
1481-- that follows. Though this could be problematic due to unbounded large
1482-- allocations. We need to figure out why the chunked code is slower and if we
1483-- can optimize the chunked code to work as fast as this one. It may be a
1485import Foreign.ForeignPtr (touchForeignPtr)
1486import Foreign.ForeignPtr.Unsafe (unsafeForeignPtrToPtr)
1487import Foreign.Ptr (Ptr, plusPtr)
1488reverse' m = Stream step Nothing
1490 {-# INLINE_LATE step #-}
1492 arr <- A.fromStreamD m
1493 let p = aEnd arr `plusPtr` negate (sizeOf (undefined :: a))
1494 return $ Skip $ Just (aStart arr, p)
1496 step _ (Just (start, p)) | p < unsafeForeignPtrToPtr start = return Stop
1498 step _ (Just (start, p)) = do
1499 let !x = A.unsafeInlineIO $ do
1501 touchForeignPtr start
1503 next = p `plusPtr` negate (sizeOf (undefined :: a))
1504 return $ Yield x (Just (start, next))
1511 $ A.fromStreamDArraysOf A.defaultChunkSize m
1514------------------------------------------------------------------------------
1515-- Grouping/Splitting
1516------------------------------------------------------------------------------
1518{-# INLINE_NORMAL splitSuffixBy' #-}
1519splitSuffixBy' :: Monad m
1520 => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1521splitSuffixBy' predicate f (Stream step state) =
1522 Stream (stepOuter f) (Just state)
1526 {-# INLINE_LATE stepOuter #-}
1527 stepOuter (Fold fstep initial done) gst (Just st) = do
1528 res <- step (adaptState gst) st
1534 then done acc' >>= \val -> return $ Yield val (Just s)
1537 Skip s -> return $ Skip $ Just s
1543 res <- step (adaptState gst) stt
1548 then done acc' >>= \val -> return $ Yield val (Just s)
1550 Skip s -> go SPEC s acc
1551 Stop -> done acc >>= \val -> return $ Yield val Nothing
1553 stepOuter _ _ Nothing = return Stop
1555{-# INLINE_NORMAL groupsBy #-}
1561groupsBy cmp f (Stream step state) = Stream (stepOuter f) (Just state, Nothing)
1565 {-# INLINE_LATE stepOuter #-}
1566 stepOuter (Fold fstep initial done) gst (Just st, Nothing) = do
1567 res <- step (adaptState gst) st
1574 Skip s -> return $ Skip $ (Just s, Nothing)
1579 go !_ prev stt !acc = do
1580 res <- step (adaptState gst) stt
1587 else done acc >>= \r -> return $ Yield r (Just s, Just x)
1588 Skip s -> go SPEC prev s acc
1589 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1591 stepOuter (Fold fstep initial done) gst (Just st, Just prev) = do
1593 acc' <- fstep acc prev
1598 -- XXX code duplicated from the previous equation
1600 res <- step (adaptState gst) stt
1607 else done acc >>= \r -> return $ Yield r (Just s, Just x)
1608 Skip s -> go SPEC s acc
1609 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1611 stepOuter _ _ (Nothing,_) = return Stop
1613{-# INLINE_NORMAL groupsRollingBy #-}
1614groupsRollingBy :: Monad m
1619groupsRollingBy cmp f (Stream step state) =
1620 Stream (stepOuter f) (Just state, Nothing)
1623 {-# INLINE_LATE stepOuter #-}
1624 stepOuter (Fold fstep initial done) gst (Just st, Nothing) = do
1625 res <- step (adaptState gst) st
1632 Skip s -> return $ Skip $ (Just s, Nothing)
1636 go !_ prev stt !acc = do
1637 res <- step (adaptState gst) stt
1645 done acc >>= \r -> return $ Yield r (Just s, Just x)
1646 Skip s -> go SPEC prev s acc
1647 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1649 stepOuter (Fold fstep initial done) gst (Just st, Just prev') = do
1651 acc' <- fstep acc prev'
1652 go SPEC prev' st acc'
1655 go !_ prevv stt !acc = do
1656 res <- step (adaptState gst) stt
1663 else done acc >>= \r -> return $ Yield r (Just s, Just x)
1664 Skip s -> go SPEC prevv s acc
1665 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1667 stepOuter _ _ (Nothing, _) = return Stop
1669{-# INLINE_NORMAL splitBy #-}
1670splitBy :: Monad m => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1671splitBy predicate f (Stream step state) = Stream (step' f) (Just state)
1675 {-# INLINE_LATE step' #-}
1676 step' (Fold fstep initial done) gst (Just st) = initial >>= go SPEC st
1681 res <- step (adaptState gst) stt
1685 then done acc >>= \r -> return $ Yield r (Just s)
1689 Skip s -> go SPEC s acc
1690 Stop -> done acc >>= \r -> return $ Yield r Nothing
1692 step' _ _ Nothing = return Stop
1694-- XXX requires -funfolding-use-threshold=150 in lines-unlines benchmark
1695{-# INLINE_NORMAL splitSuffixBy #-}
1696splitSuffixBy :: Monad m
1697 => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1698splitSuffixBy predicate f (Stream step state) = Stream (step' f) (Just state)
1702 {-# INLINE_LATE step' #-}
1703 step' (Fold fstep initial done) gst (Just st) = do
1704 res <- step (adaptState gst) st
1709 then done acc >>= \val -> return $ Yield val (Just s)
1714 Skip s -> return $ Skip $ Just s
1720 res <- step (adaptState gst) stt
1724 then done acc >>= \r -> return $ Yield r (Just s)
1728 Skip s -> go SPEC s acc
1729 Stop -> done acc >>= \r -> return $ Yield r Nothing
1731 step' _ _ Nothing = return Stop
1733{-# INLINE_NORMAL wordsBy #-}
1734wordsBy :: Monad m => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1735wordsBy predicate f (Stream step state) = Stream (stepOuter f) (Just state)
1739 {-# INLINE_LATE stepOuter #-}
1740 stepOuter (Fold fstep initial done) gst (Just st) = do
1741 res <- step (adaptState gst) st
1745 then return $ Skip (Just s)
1751 Skip s -> return $ Skip $ Just s
1757 res <- step (adaptState gst) stt
1761 then done acc >>= \r -> return $ Yield r (Just s)
1765 Skip s -> go SPEC s acc
1766 Stop -> done acc >>= \r -> return $ Yield r Nothing
1768 stepOuter _ _ Nothing = return Stop
1770-- String search algorithms:
1771-- http://www-igm.univ-mlv.fr/~lecroq/string/index.html
1774-- TODO can we unify the splitting operations using a splitting configuration
1775-- like in the split package.
1777data SplitStyle = Infix | Suffix | Prefix deriving (Eq, Show)
1779data SplitOptions = SplitOptions
1780 { style :: SplitStyle
1781 , withSep :: Bool -- ^ keep the separators in output
1782 -- , compact :: Bool -- ^ treat multiple consecutive separators as one
1783 -- , trimHead :: Bool -- ^ drop blank at head
1784 -- , trimTail :: Bool -- ^ drop blank at tail
1788{-# ANN type SplitOnSeqState Fuse #-}
1789data SplitOnSeqState rb rh ck w fs s b x =
1791 | SplitOnSeqYield b (SplitOnSeqState rb rh ck w fs s b x)
1796 | SplitOnSeqSingle !fs s x
1798 | SplitOnSeqWordInit s
1799 | SplitOnSeqWordLoop !w s !fs
1800 | SplitOnSeqWordDone Int !fs !w
1802 | SplitOnSeqKRInit Int s rb !rh
1803 | SplitOnSeqKRLoop fs s rb !rh !ck
1804 | SplitOnSeqKRCheck fs s rb !rh
1805 | SplitOnSeqKRDone Int !fs rb !rh
1807{-# INLINE_NORMAL splitOnSeq #-}
1809 :: forall m a b. (MonadIO m, Storable a, Enum a, Eq a)
1814splitOnSeq patArr (Fold fstep initial done) (Stream step state) =
1815 Stream stepOuter SplitOnSeqInit
1819 patLen = A.length patArr
1820 maxIndex = patLen - 1
1821 elemBits = sizeOf (undefined :: a) * 8
1823 -- For word pattern case
1825 wordMask = (1 `shiftL` (elemBits * patLen)) - 1
1828 elemMask = (1 `shiftL` elemBits) - 1
1831 wordPat = wordMask .&. A.foldl' addToWord 0 patArr
1833 addToWord wd a = (wd `shiftL` elemBits) .|. fromIntegral (fromEnum a)
1835 -- For Rabin-Karp search
1836 k = 2891336453 :: Word32
1839 addCksum cksum a = cksum * k + fromIntegral (fromEnum a)
1841 deltaCksum cksum old new =
1842 addCksum cksum new - coeff * fromIntegral (fromEnum old)
1844 -- XXX shall we use a random starting hash or 1 instead of 0?
1845 patHash = A.foldl' addCksum 0 patArr
1847 skip = return . Skip
1849 {-# INLINE_LATE stepOuter #-}
1850 stepOuter _ SplitOnSeqInit =
1852 then return $ Skip $ SplitOnSeqEmpty state
1856 pat <- liftIO $ A.unsafeIndexIO patArr 0
1857 return $ Skip $ SplitOnSeqSingle acc state pat
1858 else if sizeOf (undefined :: a) * patLen
1859 <= sizeOf (undefined :: Word)
1860 then return $ Skip $ SplitOnSeqWordInit state
1862 (rb, rhead) <- liftIO $ RB.new patLen
1863 skip $ SplitOnSeqKRInit 0 state rb rhead
1865 stepOuter _ (SplitOnSeqYield x next) = return $ Yield x next
1867 ---------------------------
1869 ---------------------------
1871 stepOuter gst (SplitOnSeqEmpty st) = do
1872 res <- step (adaptState gst) st
1878 skip $ SplitOnSeqYield r (SplitOnSeqEmpty s)
1879 Skip s -> return $ Skip (SplitOnSeqEmpty s)
1886 stepOuter _ SplitOnSeqDone = return Stop
1892 stepOuter gst (SplitOnSeqSingle fs st pat) = do
1893 res <- step (adaptState gst) st
1900 return $ Skip $ SplitOnSeqYield r (SplitOnSeqSingle fs1 s pat)
1903 return $ Skip $ (SplitOnSeqSingle fs1 s pat)
1904 Skip s -> return $ Skip $ SplitOnSeqSingle fs s pat
1907 return $ Skip $ SplitOnSeqYield r SplitOnSeqDone
1909 ---------------------------
1910 -- Short Pattern - Shift Or
1911 ---------------------------
1913 stepOuter _ (SplitOnSeqWordDone 0 fs _) = do
1915 skip $ SplitOnSeqYield r SplitOnSeqDone
1916 stepOuter _ (SplitOnSeqWordDone n fs wrd) = do
1917 let old = elemMask .&. (wrd `shiftR` (elemBits * (n - 1)))
1918 fs1 <- fstep fs (toEnum $ fromIntegral old)
1919 skip $ SplitOnSeqWordDone (n - 1) fs1 wrd
1921 stepOuter gst (SplitOnSeqWordInit st0) =
1927 go !_ !idx !wrd !st = do
1928 res <- step (adaptState gst) st
1931 let wrd1 = addToWord wrd x
1935 if wrd1 .&. wordMask == wordPat
1938 let next = SplitOnSeqWordInit s
1939 skip $ SplitOnSeqYield r next
1940 else skip $ SplitOnSeqWordLoop wrd1 s fs
1941 else go SPEC (idx + 1) wrd1 s
1942 Skip s -> go SPEC idx wrd s
1946 then skip $ SplitOnSeqWordDone idx fs wrd
1949 skip $ SplitOnSeqYield r SplitOnSeqDone
1951 stepOuter gst (SplitOnSeqWordLoop wrd0 st0 fs0) =
1952 go SPEC wrd0 st0 fs0
1957 go !_ !wrd !st !fs = do
1958 res <- step (adaptState gst) st
1961 let wrd1 = addToWord wrd x
1962 old = (wordMask .&. wrd)
1963 `shiftR` (elemBits * (patLen - 1))
1964 fs1 <- fstep fs (toEnum $ fromIntegral old)
1965 if wrd1 .&. wordMask == wordPat
1968 let next = SplitOnSeqWordInit s
1969 skip $ SplitOnSeqYield r next
1970 else go SPEC wrd1 s fs1
1971 Skip s -> go SPEC wrd s fs
1972 Stop -> skip $ SplitOnSeqWordDone patLen fs wrd
1974 -------------------------------
1975 -- General Pattern - Karp Rabin
1976 -------------------------------
1978 stepOuter gst (SplitOnSeqKRInit idx st rb rh) = do
1979 res <- step (adaptState gst) st
1982 rh1 <- liftIO $ RB.unsafeInsert rb rh x
1985 let fold = RB.unsafeFoldRing (RB.ringBound rb)
1986 let !ringHash = fold addCksum 0 rb
1988 if ringHash == patHash
1989 then skip $ SplitOnSeqKRCheck fs s rb rh1
1990 else skip $ SplitOnSeqKRLoop fs s rb rh1 ringHash
1991 else skip $ SplitOnSeqKRInit (idx + 1) s rb rh1
1992 Skip s -> skip $ SplitOnSeqKRInit idx s rb rh
1995 skip $ SplitOnSeqKRDone idx fs rb (RB.startOf rb)
1997 -- XXX The recursive "go" is more efficient than the state based recursion
1998 -- code commented out below. Perhaps its more efficient because of
1999 -- factoring out "rb" outside the loop.
2001 stepOuter gst (SplitOnSeqKRLoop fs0 st0 rb rh0 cksum0) =
2002 go SPEC fs0 st0 rh0 cksum0
2006 go !_ !fs !st !rh !cksum = do
2007 res <- step (adaptState gst) st
2010 old <- liftIO $ peek rh
2011 let cksum1 = deltaCksum cksum old x
2013 rh1 <- liftIO (RB.unsafeInsert rb rh x)
2014 if (cksum1 == patHash)
2015 then skip $ SplitOnSeqKRCheck fs1 s rb rh1
2016 else go SPEC fs1 s rh1 cksum1
2017 Skip s -> go SPEC fs s rh cksum
2018 Stop -> skip $ SplitOnSeqKRDone patLen fs rb rh
2020 -- XXX The following code is 5 times slower compared to the recursive loop
2021 -- based code above. Need to investigate why. One possibility is that the
2022 -- go loop above does not thread around the ring buffer (rb). This code may
2023 -- be causing the state to bloat and getting allocated on each iteration.
2024 -- We can check the cmm/asm code to confirm. If so a good GHC solution to
2025 -- such problem is needed. One way to avoid this could be to use unboxed
2028 stepOuter gst (SplitOnSeqKRLoop fs st rb rh cksum) = do
2029 res <- step (adaptState gst) st
2032 old <- liftIO $ peek rh
2033 let cksum1 = deltaCksum cksum old x
2035 if (cksum1 == patHash)
2038 skip $ SplitOnSeqYield r $ SplitOnSeqKRInit 0 s rb rh
2040 rh1 <- liftIO (RB.unsafeInsert rb rh x)
2041 skip $ SplitOnSeqKRLoop fs1 s rb rh1 cksum1
2042 Skip s -> skip $ SplitOnSeqKRLoop fs s rb rh cksum
2043 Stop -> skip $ SplitOnSeqKRDone patLen fs rb rh
2046 stepOuter _ (SplitOnSeqKRCheck fs st rb rh) = do
2047 if RB.unsafeEqArray rb rh patArr
2050 let next = SplitOnSeqKRInit 0 st rb (RB.startOf rb)
2051 skip $ SplitOnSeqYield r next
2052 else skip $ SplitOnSeqKRLoop fs st rb rh patHash
2053 stepOuter _ (SplitOnSeqKRDone 0 fs _ _) = do
2055 skip $ SplitOnSeqYield r SplitOnSeqDone
2056 stepOuter _ (SplitOnSeqKRDone n fs rb rh) = do
2057 old <- liftIO $ peek rh
2058 let rh1 = RB.advance rb rh
2060 skip $ SplitOnSeqKRDone (n - 1) fs1 rb rh1
2062{-# ANN type SplitOnSuffixSeqState Fuse #-}
2063data SplitOnSuffixSeqState rb rh ck w fs s b x =
2064 SplitOnSuffixSeqInit
2065 | SplitOnSuffixSeqYield b (SplitOnSuffixSeqState rb rh ck w fs s b x)
2066 | SplitOnSuffixSeqDone
2068 | SplitOnSuffixSeqEmpty s
2070 | SplitOnSuffixSeqSingleInit s x
2071 | SplitOnSuffixSeqSingle !fs s x
2073 | SplitOnSuffixSeqWordInit s
2074 | SplitOnSuffixSeqWordLoop !w s !fs
2075 | SplitOnSuffixSeqWordDone Int !fs !w
2077 | SplitOnSuffixSeqKRInit Int s rb !rh
2078 | SplitOnSuffixSeqKRInit1 !fs s rb !rh
2079 | SplitOnSuffixSeqKRLoop fs s rb !rh !ck
2080 | SplitOnSuffixSeqKRCheck fs s rb !rh
2081 | SplitOnSuffixSeqKRDone Int !fs rb !rh
2083{-# INLINE_NORMAL splitOnSuffixSeq #-}
2085 :: forall m a b. (MonadIO m, Storable a, Enum a, Eq a)
2091splitOnSuffixSeq withSep patArr (Fold fstep initial done) (Stream step state) =
2092 Stream stepOuter SplitOnSuffixSeqInit
2096 patLen = A.length patArr
2097 maxIndex = patLen - 1
2098 elemBits = sizeOf (undefined :: a) * 8
2100 -- For word pattern case
2102 wordMask = (1 `shiftL` (elemBits * patLen)) - 1
2105 elemMask = (1 `shiftL` elemBits) - 1
2108 wordPat = wordMask .&. A.foldl' addToWord 0 patArr
2110 addToWord wd a = (wd `shiftL` elemBits) .|. fromIntegral (fromEnum a)
2112 -- For single element pattern case
2113 {-# INLINE processYieldSingle #-}
2114 processYieldSingle pat x s fs =
2117 fs1 <- if withSep then fstep fs x else return fs
2119 let next = SplitOnSuffixSeqSingleInit s pat
2120 skip $ SplitOnSuffixSeqYield r next
2123 skip $ SplitOnSuffixSeqSingle fs1 s pat
2125 -- For Rabin-Karp search
2126 k = 2891336453 :: Word32
2129 addCksum cksum a = cksum * k + fromIntegral (fromEnum a)
2131 deltaCksum cksum old new =
2132 addCksum cksum new - coeff * fromIntegral (fromEnum old)
2134 -- XXX shall we use a random starting hash or 1 instead of 0?
2135 patHash = A.foldl' addCksum 0 patArr
2137 skip = return . Skip
2139 {-# INLINE_LATE stepOuter #-}
2140 stepOuter _ SplitOnSuffixSeqInit =
2142 then skip $ SplitOnSuffixSeqEmpty state
2145 pat <- liftIO $ A.unsafeIndexIO patArr 0
2146 skip $ SplitOnSuffixSeqSingleInit state pat
2147 else if sizeOf (undefined :: a) * patLen
2148 <= sizeOf (undefined :: Word)
2149 then skip $ SplitOnSuffixSeqWordInit state
2151 (rb, rhead) <- liftIO $ RB.new patLen
2152 skip $ SplitOnSuffixSeqKRInit 0 state rb rhead
2154 stepOuter _ (SplitOnSuffixSeqYield x next) = return $ Yield x next
2156 ---------------------------
2158 ---------------------------
2160 stepOuter gst (SplitOnSuffixSeqEmpty st) = do
2161 res <- step (adaptState gst) st
2167 skip $ SplitOnSuffixSeqYield r (SplitOnSuffixSeqEmpty s)
2168 Skip s -> skip (SplitOnSuffixSeqEmpty s)
2175 stepOuter _ SplitOnSuffixSeqDone = return Stop
2181 stepOuter gst (SplitOnSuffixSeqSingleInit st pat) = do
2182 res <- step (adaptState gst) st
2184 Yield x s -> initial >>= processYieldSingle pat x s
2185 Skip s -> skip $ SplitOnSuffixSeqSingleInit s pat
2188 stepOuter gst (SplitOnSuffixSeqSingle fs st pat) = do
2189 res <- step (adaptState gst) st
2191 Yield x s -> processYieldSingle pat x s fs
2192 Skip s -> skip $ SplitOnSuffixSeqSingle fs s pat
2195 skip $ SplitOnSuffixSeqYield r SplitOnSuffixSeqDone
2197 ---------------------------
2198 -- Short Pattern - Shift Or
2199 ---------------------------
2201 stepOuter _ (SplitOnSuffixSeqWordDone 0 fs _) = do
2203 skip $ SplitOnSuffixSeqYield r SplitOnSuffixSeqDone
2204 stepOuter _ (SplitOnSuffixSeqWordDone n fs wrd) = do
2205 let old = elemMask .&. (wrd `shiftR` (elemBits * (n - 1)))
2206 fs1 <- fstep fs (toEnum $ fromIntegral old)
2207 skip $ SplitOnSuffixSeqWordDone (n - 1) fs1 wrd
2209 stepOuter gst (SplitOnSuffixSeqWordInit st0) = do
2210 res <- step (adaptState gst) st0
2214 let wrd = addToWord 0 x
2215 fs1 <- if withSep then fstep fs x else return fs
2217 Skip s -> skip (SplitOnSuffixSeqWordInit s)
2223 go !_ !idx !wrd !st !fs = do
2224 res <- step (adaptState gst) st
2227 let wrd1 = addToWord wrd x
2228 fs1 <- if withSep then fstep fs x else return fs
2230 then go SPEC (idx + 1) wrd1 s fs1
2231 else if wrd1 .&. wordMask /= wordPat
2232 then skip $ SplitOnSuffixSeqWordLoop wrd1 s fs1
2235 let next = SplitOnSuffixSeqWordInit s
2236 skip $ SplitOnSuffixSeqYield r next
2237 Skip s -> go SPEC idx wrd s fs
2238 Stop -> skip $ SplitOnSuffixSeqWordDone idx fs wrd
2240 stepOuter gst (SplitOnSuffixSeqWordLoop wrd0 st0 fs0) =
2241 go SPEC wrd0 st0 fs0
2246 go !_ !wrd !st !fs = do
2247 res <- step (adaptState gst) st
2250 let wrd1 = addToWord wrd x
2251 old = (wordMask .&. wrd)
2252 `shiftR` (elemBits * (patLen - 1))
2256 else fstep fs (toEnum $ fromIntegral old)
2257 if wrd1 .&. wordMask == wordPat
2260 let next = SplitOnSuffixSeqWordInit s
2261 skip $ SplitOnSuffixSeqYield r next
2262 else go SPEC wrd1 s fs1
2263 Skip s -> go SPEC wrd s fs
2265 if wrd .&. wordMask == wordPat
2270 skip $ SplitOnSuffixSeqYield r SplitOnSuffixSeqDone
2271 else skip $ SplitOnSuffixSeqWordDone patLen fs wrd
2273 -------------------------------
2274 -- General Pattern - Karp Rabin
2275 -------------------------------
2277 stepOuter gst (SplitOnSuffixSeqKRInit idx0 st0 rb rh0) = do
2278 res <- step (adaptState gst) st0
2281 rh1 <- liftIO $ RB.unsafeInsert rb rh0 x
2283 fs1 <- if withSep then fstep fs x else return fs
2284 skip $ SplitOnSuffixSeqKRInit1 fs1 s rb rh1
2285 Skip s -> skip $ SplitOnSuffixSeqKRInit idx0 s rb rh0
2288 stepOuter gst (SplitOnSuffixSeqKRInit1 fs0 st0 rb rh0) = do
2289 go SPEC 1 rh0 st0 fs0
2293 go !_ !idx !rh st !fs = do
2294 res <- step (adaptState gst) st
2297 rh1 <- liftIO (RB.unsafeInsert rb rh x)
2298 fs1 <- if withSep then fstep fs x else return fs
2300 then go SPEC (idx + 1) rh1 s fs1
2302 let fold = RB.unsafeFoldRing (RB.ringBound rb)
2303 !ringHash = fold addCksum 0 rb
2304 in if ringHash == patHash
2305 then SplitOnSuffixSeqKRCheck fs1 s rb rh1
2306 else SplitOnSuffixSeqKRLoop fs1 s rb rh1 ringHash
2307 Skip s -> go SPEC idx rh s fs
2309 -- do not issue a blank segment when we end at pattern
2310 if (idx == maxIndex) && RB.unsafeEqArray rb rh patArr
2315 skip $ SplitOnSuffixSeqYield r SplitOnSuffixSeqDone
2316 else skip $ SplitOnSuffixSeqKRDone idx fs rb (RB.startOf rb)
2318 stepOuter gst (SplitOnSuffixSeqKRLoop fs0 st0 rb rh0 cksum0) =
2319 go SPEC fs0 st0 rh0 cksum0
2323 go !_ !fs !st !rh !cksum = do
2324 res <- step (adaptState gst) st
2327 old <- liftIO $ peek rh
2328 rh1 <- liftIO (RB.unsafeInsert rb rh x)
2329 let cksum1 = deltaCksum cksum old x
2330 fs1 <- if withSep then fstep fs x else fstep fs old
2331 if (cksum1 /= patHash)
2332 then go SPEC fs1 s rh1 cksum1
2333 else skip $ SplitOnSuffixSeqKRCheck fs1 s rb rh1
2334 Skip s -> go SPEC fs s rh cksum
2336 if RB.unsafeEqArray rb rh patArr
2341 skip $ SplitOnSuffixSeqYield r SplitOnSuffixSeqDone
2342 else skip $ SplitOnSuffixSeqKRDone patLen fs rb rh
2344 stepOuter _ (SplitOnSuffixSeqKRCheck fs st rb rh) = do
2345 if RB.unsafeEqArray rb rh patArr
2348 let next = SplitOnSuffixSeqKRInit 0 st rb (RB.startOf rb)
2349 skip $ SplitOnSuffixSeqYield r next
2350 else skip $ SplitOnSuffixSeqKRLoop fs st rb rh patHash
2352 stepOuter _ (SplitOnSuffixSeqKRDone 0 fs _ _) = do
2354 skip $ SplitOnSuffixSeqYield r SplitOnSuffixSeqDone
2355 stepOuter _ (SplitOnSuffixSeqKRDone n fs rb rh) = do
2356 old <- liftIO $ peek rh
2357 let rh1 = RB.advance rb rh
2359 skip $ SplitOnSuffixSeqKRDone (n - 1) fs1 rb rh1
2361data SplitState s arr
2363 | SplitBuffering s arr
2364 | SplitSplitting s arr
2365 | SplitYielding arr (SplitState s arr)
2368-- XXX An alternative approach would be to use a partial fold (Fold m a b) to
2369-- split using a splitBy like combinator. The Fold would consume upto the
2370-- separator and return any leftover which can then be fed to the next fold.
2372-- We can revisit this once we have partial folds/parsers.
2374-- | Performs infix separator style splitting.
2375{-# INLINE_NORMAL splitInnerBy #-}
2378 => (f a -> m (f a, Maybe (f a))) -- splitter
2379 -> (f a -> f a -> m (f a)) -- joiner
2382splitInnerBy splitter joiner (Stream step1 state1) =
2383 (Stream step (SplitInitial state1))
2387 {-# INLINE_LATE step #-}
2388 step gst (SplitInitial st) = do
2392 (x1, mx2) <- splitter x
2393 return $ case mx2 of
2394 Nothing -> Skip (SplitBuffering s x1)
2395 Just x2 -> Skip (SplitYielding x1 (SplitSplitting s x2))
2396 Skip s -> return $ Skip (SplitInitial s)
2397 Stop -> return $ Stop
2399 step gst (SplitBuffering st buf) = do
2403 (x1, mx2) <- splitter x
2404 buf' <- joiner buf x1
2405 return $ case mx2 of
2406 Nothing -> Skip (SplitBuffering s buf')
2407 Just x2 -> Skip (SplitYielding buf' (SplitSplitting s x2))
2408 Skip s -> return $ Skip (SplitBuffering s buf)
2409 Stop -> return $ Skip (SplitYielding buf SplitFinishing)
2411 step _ (SplitSplitting st buf) = do
2412 (x1, mx2) <- splitter buf
2413 return $ case mx2 of
2414 Nothing -> Skip $ SplitBuffering st x1
2415 Just x2 -> Skip $ SplitYielding x1 (SplitSplitting st x2)
2417 step _ (SplitYielding x next) = return $ Yield x next
2418 step _ SplitFinishing = return $ Stop
2420-- | Performs infix separator style splitting.
2421{-# INLINE_NORMAL splitInnerBySuffix #-}
2423 :: (Monad m, Eq (f a), Monoid (f a))
2424 => (f a -> m (f a, Maybe (f a))) -- splitter
2425 -> (f a -> f a -> m (f a)) -- joiner
2428splitInnerBySuffix splitter joiner (Stream step1 state1) =
2429 (Stream step (SplitInitial state1))
2433 {-# INLINE_LATE step #-}
2434 step gst (SplitInitial st) = do
2438 (x1, mx2) <- splitter x
2439 return $ case mx2 of
2440 Nothing -> Skip (SplitBuffering s x1)
2441 Just x2 -> Skip (SplitYielding x1 (SplitSplitting s x2))
2442 Skip s -> return $ Skip (SplitInitial s)
2443 Stop -> return $ Stop
2445 step gst (SplitBuffering st buf) = do
2449 (x1, mx2) <- splitter x
2450 buf' <- joiner buf x1
2451 return $ case mx2 of
2452 Nothing -> Skip (SplitBuffering s buf')
2453 Just x2 -> Skip (SplitYielding buf' (SplitSplitting s x2))
2454 Skip s -> return $ Skip (SplitBuffering s buf)
2458 else Skip (SplitYielding buf SplitFinishing)
2460 step _ (SplitSplitting st buf) = do
2461 (x1, mx2) <- splitter buf
2462 return $ case mx2 of
2463 Nothing -> Skip $ SplitBuffering st x1
2464 Just x2 -> Skip $ SplitYielding x1 (SplitSplitting st x2)
2466 step _ (SplitYielding x next) = return $ Yield x next
2467 step _ SplitFinishing = return $ Stop
2469------------------------------------------------------------------------------
2471------------------------------------------------------------------------------
2473{-# INLINE_NORMAL isPrefixOf #-}
2474isPrefixOf :: (Eq a, Monad m) => Stream m a -> Stream m a -> m Bool
2475isPrefixOf (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
2477 go (sa, sb, Nothing) = do
2478 r <- stepa defState sa
2480 Yield x sa' -> go (sa', sb, Just x)
2481 Skip sa' -> go (sa', sb, Nothing)
2484 go (sa, sb, Just x) = do
2485 r <- stepb defState sb
2489 then go (sa, sb', Nothing)
2491 Skip sb' -> go (sa, sb', Just x)
2492 Stop -> return False
2494{-# INLINE_NORMAL isSubsequenceOf #-}
2495isSubsequenceOf :: (Eq a, Monad m) => Stream m a -> Stream m a -> m Bool
2496isSubsequenceOf (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
2498 go (sa, sb, Nothing) = do
2499 r <- stepa defState sa
2501 Yield x sa' -> go (sa', sb, Just x)
2502 Skip sa' -> go (sa', sb, Nothing)
2505 go (sa, sb, Just x) = do
2506 r <- stepb defState sb
2510 then go (sa, sb', Nothing)
2511 else go (sa, sb', Just x)
2512 Skip sb' -> go (sa, sb', Just x)
2513 Stop -> return False
2515{-# INLINE_NORMAL stripPrefix #-}
2518 => Stream m a -> Stream m a -> m (Maybe (Stream m a))
2519stripPrefix (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
2521 go (sa, sb, Nothing) = do
2522 r <- stepa defState sa
2524 Yield x sa' -> go (sa', sb, Just x)
2525 Skip sa' -> go (sa', sb, Nothing)
2526 Stop -> return $ Just (Stream stepb sb)
2528 go (sa, sb, Just x) = do
2529 r <- stepb defState sb
2533 then go (sa, sb', Nothing)
2535 Skip sb' -> go (sa, sb', Just x)
2536 Stop -> return Nothing
2538------------------------------------------------------------------------------
2540------------------------------------------------------------------------------
2542-- | Execute a monadic action for each element of the 'Stream'
2543{-# INLINE_NORMAL mapM_ #-}
2544mapM_ :: Monad m => (a -> m b) -> Stream m a -> m ()
2545mapM_ m = drain . mapM m
2547-------------------------------------------------------------------------------
2548-- Stream transformations using Unfolds
2549-------------------------------------------------------------------------------
2551-- Define a unique structure to use in inspection testing
2552data ConcatMapUState o i =
2554 | ConcatMapUInner o i
2556-- | @concatMapU unfold stream@ uses @unfold@ to map the input stream elements
2557-- to streams and then flattens the generated streams into a single output
2560-- This is like 'concatMap' but uses an unfold with an explicit state to
2561-- generate the stream instead of a 'Stream' type generator. This allows better
2562-- optimization via fusion. This can be many times more efficient than
2565{-# INLINE_NORMAL concatMapU #-}
2566concatMapU :: Monad m => Unfold m a b -> Stream m a -> Stream m b
2567concatMapU (Unfold istep inject) (Stream ostep ost) =
2568 Stream step (ConcatMapUOuter ost)
2570 {-# INLINE_LATE step #-}
2571 step gst (ConcatMapUOuter o) = do
2572 r <- ostep (adaptState gst) o
2576 i `seq` return (Skip (ConcatMapUInner o' i))
2577 Skip o' -> return $ Skip (ConcatMapUOuter o')
2578 Stop -> return $ Stop
2580 step _ (ConcatMapUInner o i) = do
2583 Yield x i' -> Yield x (ConcatMapUInner o i')
2584 Skip i' -> Skip (ConcatMapUInner o i')
2585 Stop -> Skip (ConcatMapUOuter o)
2587data ConcatUnfoldInterleaveState o i =
2588 ConcatUnfoldInterleaveOuter o [i]
2589 | ConcatUnfoldInterleaveInner o [i]
2590 | ConcatUnfoldInterleaveInnerL [i] [i]
2591 | ConcatUnfoldInterleaveInnerR [i] [i]
2593-- XXX use arrays to store state instead of lists.
2594-- XXX In general we can use different scheduling strategies e.g. how to
2595-- schedule the outer vs inner loop or assigning weights to different streams
2596-- or outer and inner loops.
2598-- After a yield, switch to the next stream. Do not switch streams on Skip.
2599-- Yield from outer stream switches to the inner stream.
2601-- There are two choices here, (1) exhaust the outer stream first and then
2602-- start yielding from the inner streams, this is much simpler to implement,
2603-- (2) yield at least one element from an inner stream before going back to
2604-- outer stream and opening the next stream from it.
2606-- Ideally, we need some scheduling bias to inner streams vs outer stream.
2607-- Maybe we can configure the behavior.
2609{-# INLINE_NORMAL concatUnfoldInterleave #-}
2610concatUnfoldInterleave :: Monad m => Unfold m a b -> Stream m a -> Stream m b
2611concatUnfoldInterleave (Unfold istep inject) (Stream ostep ost) =
2612 Stream step (ConcatUnfoldInterleaveOuter ost [])
2614 {-# INLINE_LATE step #-}
2615 step gst (ConcatUnfoldInterleaveOuter o ls) = do
2616 r <- ostep (adaptState gst) o
2620 i `seq` return (Skip (ConcatUnfoldInterleaveInner o' (i : ls)))
2621 Skip o' -> return $ Skip (ConcatUnfoldInterleaveOuter o' ls)
2622 Stop -> return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2624 step _ (ConcatUnfoldInterleaveInner _ []) = undefined
2625 step _ (ConcatUnfoldInterleaveInner o (st:ls)) = do
2628 Yield x s -> Yield x (ConcatUnfoldInterleaveOuter o (s:ls))
2629 Skip s -> Skip (ConcatUnfoldInterleaveInner o (s:ls))
2630 Stop -> Skip (ConcatUnfoldInterleaveOuter o ls)
2632 step _ (ConcatUnfoldInterleaveInnerL [] []) = return Stop
2633 step _ (ConcatUnfoldInterleaveInnerL [] rs) =
2634 return $ Skip (ConcatUnfoldInterleaveInnerR [] rs)
2636 step _ (ConcatUnfoldInterleaveInnerL (st:ls) rs) = do
2639 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerL ls (s:rs))
2640 Skip s -> Skip (ConcatUnfoldInterleaveInnerL (s:ls) rs)
2641 Stop -> Skip (ConcatUnfoldInterleaveInnerL ls rs)
2643 step _ (ConcatUnfoldInterleaveInnerR [] []) = return Stop
2644 step _ (ConcatUnfoldInterleaveInnerR ls []) =
2645 return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2647 step _ (ConcatUnfoldInterleaveInnerR ls (st:rs)) = do
2650 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerR (s:ls) rs)
2651 Skip s -> Skip (ConcatUnfoldInterleaveInnerR ls (s:rs))
2652 Stop -> Skip (ConcatUnfoldInterleaveInnerR ls rs)
2654-- XXX In general we can use different scheduling strategies e.g. how to
2655-- schedule the outer vs inner loop or assigning weights to different streams
2656-- or outer and inner loops.
2658-- This could be inefficient if the tasks are too small.
2660-- Compared to concatUnfoldInterleave this one switches streams on Skips.
2662{-# INLINE_NORMAL concatUnfoldRoundrobin #-}
2663concatUnfoldRoundrobin :: Monad m => Unfold m a b -> Stream m a -> Stream m b
2664concatUnfoldRoundrobin (Unfold istep inject) (Stream ostep ost) =
2665 Stream step (ConcatUnfoldInterleaveOuter ost [])
2667 {-# INLINE_LATE step #-}
2668 step gst (ConcatUnfoldInterleaveOuter o ls) = do
2669 r <- ostep (adaptState gst) o
2673 i `seq` return (Skip (ConcatUnfoldInterleaveInner o' (i : ls)))
2674 Skip o' -> return $ Skip (ConcatUnfoldInterleaveInner o' ls)
2675 Stop -> return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2677 step _ (ConcatUnfoldInterleaveInner o []) =
2678 return $ Skip (ConcatUnfoldInterleaveOuter o [])
2680 step _ (ConcatUnfoldInterleaveInner o (st:ls)) = do
2683 Yield x s -> Yield x (ConcatUnfoldInterleaveOuter o (s:ls))
2684 Skip s -> Skip (ConcatUnfoldInterleaveOuter o (s:ls))
2685 Stop -> Skip (ConcatUnfoldInterleaveOuter o ls)
2687 step _ (ConcatUnfoldInterleaveInnerL [] []) = return Stop
2688 step _ (ConcatUnfoldInterleaveInnerL [] rs) =
2689 return $ Skip (ConcatUnfoldInterleaveInnerR [] rs)
2691 step _ (ConcatUnfoldInterleaveInnerL (st:ls) rs) = do
2694 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerL ls (s:rs))
2695 Skip s -> Skip (ConcatUnfoldInterleaveInnerL ls (s:rs))
2696 Stop -> Skip (ConcatUnfoldInterleaveInnerL ls rs)
2698 step _ (ConcatUnfoldInterleaveInnerR [] []) = return Stop
2699 step _ (ConcatUnfoldInterleaveInnerR ls []) =
2700 return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2702 step _ (ConcatUnfoldInterleaveInnerR ls (st:rs)) = do
2705 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerR (s:ls) rs)
2706 Skip s -> Skip (ConcatUnfoldInterleaveInnerR (s:ls) rs)
2707 Stop -> Skip (ConcatUnfoldInterleaveInnerR ls rs)
2709data AppendState s1 s2 = AppendFirst s1 | AppendSecond s2
2711-- Note that this could be much faster compared to the CPS stream. However, as
2712-- the number of streams being composed increases this may become expensive.
2713-- Need to see where the breaking point is between the two.
2715{-# INLINE_NORMAL append #-}
2716append :: Monad m => Stream m a -> Stream m a -> Stream m a
2717append (Stream step1 state1) (Stream step2 state2) =
2718 Stream step (AppendFirst state1)
2722 {-# INLINE_LATE step #-}
2723 step gst (AppendFirst st) = do
2726 Yield a s -> Yield a (AppendFirst s)
2727 Skip s -> Skip (AppendFirst s)
2728 Stop -> Skip (AppendSecond state2)
2730 step gst (AppendSecond st) = do
2733 Yield a s -> Yield a (AppendSecond s)
2734 Skip s -> Skip (AppendSecond s)
2737data InterleaveState s1 s2 = InterleaveFirst s1 s2 | InterleaveSecond s1 s2
2738 | InterleaveSecondOnly s2 | InterleaveFirstOnly s1
2740{-# INLINE_NORMAL interleave #-}
2741interleave :: Monad m => Stream m a -> Stream m a -> Stream m a
2742interleave (Stream step1 state1) (Stream step2 state2) =
2743 Stream step (InterleaveFirst state1 state2)
2747 {-# INLINE_LATE step #-}
2748 step gst (InterleaveFirst st1 st2) = do
2751 Yield a s -> Yield a (InterleaveSecond s st2)
2752 Skip s -> Skip (InterleaveFirst s st2)
2753 Stop -> Skip (InterleaveSecondOnly st2)
2755 step gst (InterleaveSecond st1 st2) = do
2758 Yield a s -> Yield a (InterleaveFirst st1 s)
2759 Skip s -> Skip (InterleaveSecond st1 s)
2760 Stop -> Skip (InterleaveFirstOnly st1)
2762 step gst (InterleaveFirstOnly st1) = do
2765 Yield a s -> Yield a (InterleaveFirstOnly s)
2766 Skip s -> Skip (InterleaveFirstOnly s)
2769 step gst (InterleaveSecondOnly st2) = do
2772 Yield a s -> Yield a (InterleaveSecondOnly s)
2773 Skip s -> Skip (InterleaveSecondOnly s)
2776{-# INLINE_NORMAL interleaveMin #-}
2777interleaveMin :: Monad m => Stream m a -> Stream m a -> Stream m a
2778interleaveMin (Stream step1 state1) (Stream step2 state2) =
2779 Stream step (InterleaveFirst state1 state2)
2783 {-# INLINE_LATE step #-}
2784 step gst (InterleaveFirst st1 st2) = do
2787 Yield a s -> Yield a (InterleaveSecond s st2)
2788 Skip s -> Skip (InterleaveFirst s st2)
2791 step gst (InterleaveSecond st1 st2) = do
2794 Yield a s -> Yield a (InterleaveFirst st1 s)
2795 Skip s -> Skip (InterleaveSecond st1 s)
2798 step _ (InterleaveFirstOnly _) = undefined
2799 step _ (InterleaveSecondOnly _) = undefined
2801{-# INLINE_NORMAL interleaveSuffix #-}
2802interleaveSuffix :: Monad m => Stream m a -> Stream m a -> Stream m a
2803interleaveSuffix (Stream step1 state1) (Stream step2 state2) =
2804 Stream step (InterleaveFirst state1 state2)
2808 {-# INLINE_LATE step #-}
2809 step gst (InterleaveFirst st1 st2) = do
2812 Yield a s -> Yield a (InterleaveSecond s st2)
2813 Skip s -> Skip (InterleaveFirst s st2)
2816 step gst (InterleaveSecond st1 st2) = do
2819 Yield a s -> Yield a (InterleaveFirst st1 s)
2820 Skip s -> Skip (InterleaveSecond st1 s)
2821 Stop -> Skip (InterleaveFirstOnly st1)
2823 step gst (InterleaveFirstOnly st1) = do
2826 Yield a s -> Yield a (InterleaveFirstOnly s)
2827 Skip s -> Skip (InterleaveFirstOnly s)
2830 step _ (InterleaveSecondOnly _) = undefined
2832data InterleaveInfixState s1 s2 a
2833 = InterleaveInfixFirst s1 s2
2834 | InterleaveInfixSecondBuf s1 s2
2835 | InterleaveInfixSecondYield s1 s2 a
2836 | InterleaveInfixFirstYield s1 s2 a
2837 | InterleaveInfixFirstOnly s1
2839{-# INLINE_NORMAL interleaveInfix #-}
2840interleaveInfix :: Monad m => Stream m a -> Stream m a -> Stream m a
2841interleaveInfix (Stream step1 state1) (Stream step2 state2) =
2842 Stream step (InterleaveInfixFirst state1 state2)
2846 {-# INLINE_LATE step #-}
2847 step gst (InterleaveInfixFirst st1 st2) = do
2850 Yield a s -> Yield a (InterleaveInfixSecondBuf s st2)
2851 Skip s -> Skip (InterleaveInfixFirst s st2)
2854 step gst (InterleaveInfixSecondBuf st1 st2) = do
2857 Yield a s -> Skip (InterleaveInfixSecondYield st1 s a)
2858 Skip s -> Skip (InterleaveInfixSecondBuf st1 s)
2859 Stop -> Skip (InterleaveInfixFirstOnly st1)
2861 step gst (InterleaveInfixSecondYield st1 st2 x) = do
2864 Yield a s -> Yield x (InterleaveInfixFirstYield s st2 a)
2865 Skip s -> Skip (InterleaveInfixSecondYield s st2 x)
2868 step _ (InterleaveInfixFirstYield st1 st2 x) = do
2869 return $ Yield x (InterleaveInfixSecondBuf st1 st2)
2871 step gst (InterleaveInfixFirstOnly st1) = do
2874 Yield a s -> Yield a (InterleaveInfixFirstOnly s)
2875 Skip s -> Skip (InterleaveInfixFirstOnly s)
2878{-# INLINE_NORMAL roundRobin #-}
2879roundRobin :: Monad m => Stream m a -> Stream m a -> Stream m a
2880roundRobin (Stream step1 state1) (Stream step2 state2) =
2881 Stream step (InterleaveFirst state1 state2)
2885 {-# INLINE_LATE step #-}
2886 step gst (InterleaveFirst st1 st2) = do
2889 Yield a s -> Yield a (InterleaveSecond s st2)
2890 Skip s -> Skip (InterleaveSecond s st2)
2891 Stop -> Skip (InterleaveSecondOnly st2)
2893 step gst (InterleaveSecond st1 st2) = do
2896 Yield a s -> Yield a (InterleaveFirst st1 s)
2897 Skip s -> Skip (InterleaveFirst st1 s)
2898 Stop -> Skip (InterleaveFirstOnly st1)
2900 step gst (InterleaveSecondOnly st2) = do
2903 Yield a s -> Yield a (InterleaveSecondOnly s)
2904 Skip s -> Skip (InterleaveSecondOnly s)
2907 step gst (InterleaveFirstOnly st1) = do
2910 Yield a s -> Yield a (InterleaveFirstOnly s)
2911 Skip s -> Skip (InterleaveFirstOnly s)
2914data ICUState s1 s2 i1 i2 =
2919 | ICUFirstInner s1 s2 i1
2920 | ICUSecondInner s1 s2 i2
2921 | ICUFirstOnlyInner s1 i1
2922 | ICUSecondOnlyInner s2 i2
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 gintercalateSuffix #-}
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 (ICUFirst state1 state2)
2945 {-# INLINE_LATE step #-}
2946 step gst (ICUFirst s1 s2) = do
2947 r <- step1 (adaptState gst) s1
2951 i `seq` return (Skip (ICUFirstInner s s2 i))
2952 Skip s -> return $ Skip (ICUFirst s s2)
2955 step gst (ICUFirstOnly s1) = do
2956 r <- step1 (adaptState gst) s1
2960 i `seq` return (Skip (ICUFirstOnlyInner s i))
2961 Skip s -> return $ Skip (ICUFirstOnly s)
2964 step _ (ICUFirstInner s1 s2 i1) = do
2967 Yield x i' -> Yield x (ICUFirstInner s1 s2 i')
2968 Skip i' -> Skip (ICUFirstInner s1 s2 i')
2969 Stop -> Skip (ICUSecond s1 s2)
2971 step _ (ICUFirstOnlyInner s1 i1) = do
2974 Yield x i' -> Yield x (ICUFirstOnlyInner s1 i')
2975 Skip i' -> Skip (ICUFirstOnlyInner s1 i')
2976 Stop -> Skip (ICUFirstOnly s1)
2978 step gst (ICUSecond s1 s2) = do
2979 r <- step2 (adaptState gst) s2
2983 i `seq` return (Skip (ICUSecondInner s1 s i))
2984 Skip s -> return $ Skip (ICUSecond s1 s)
2985 Stop -> return $ Skip (ICUFirstOnly s1)
2987 step _ (ICUSecondInner s1 s2 i2) = do
2990 Yield x i' -> Yield x (ICUSecondInner s1 s2 i')
2991 Skip i' -> Skip (ICUSecondInner s1 s2 i')
2992 Stop -> Skip (ICUFirst s1 s2)
2994 step _ (ICUSecondOnly _s2) = undefined
2995 step _ (ICUSecondOnlyInner _s2 _i2) = undefined
2997data InterposeSuffixState s1 i1 =
2998 InterposeSuffixFirst s1
2999 -- | InterposeSuffixFirstYield s1 i1
3000 | InterposeSuffixFirstInner s1 i1
3001 | InterposeSuffixSecond s1
3003-- Note that if an unfolded layer turns out to be nil we still emit the
3004-- separator effect. An alternate behavior could be to emit the separator
3005-- effect only if at least one element has been yielded by the unfolding.
3006-- However, that becomes a bit complicated, so we have chosen the former
3008{-# INLINE_NORMAL interposeSuffix #-}
3011 => m c -> Unfold m b c -> Stream m b -> Stream m c
3014 (Unfold istep1 inject1) (Stream step1 state1) =
3015 Stream step (InterposeSuffixFirst state1)
3019 {-# INLINE_LATE step #-}
3020 step gst (InterposeSuffixFirst s1) = do
3021 r <- step1 (adaptState gst) s1
3025 i `seq` return (Skip (InterposeSuffixFirstInner s i))
3026 -- i `seq` return (Skip (InterposeSuffixFirstYield s i))
3027 Skip s -> return $ Skip (InterposeSuffixFirst s)
3031 step _ (InterposeSuffixFirstYield s1 i1) = do
3034 Yield x i' -> Yield x (InterposeSuffixFirstInner s1 i')
3035 Skip i' -> Skip (InterposeSuffixFirstYield s1 i')
3036 Stop -> Skip (InterposeSuffixFirst s1)
3039 step _ (InterposeSuffixFirstInner s1 i1) = do
3042 Yield x i' -> Yield x (InterposeSuffixFirstInner s1 i')
3043 Skip i' -> Skip (InterposeSuffixFirstInner s1 i')
3044 Stop -> Skip (InterposeSuffixSecond s1)
3046 step _ (InterposeSuffixSecond s1) = do
3048 return $ Yield r (InterposeSuffixFirst s1)
3050data ICALState s1 s2 i1 i2 a =
3052 -- | ICALFirstYield s1 s2 i1
3053 | ICALFirstInner s1 s2 i1
3055 | ICALFirstOnlyInner s1 i1
3056 | ICALSecondInject s1 s2
3057 | ICALFirstInject s1 s2 i2
3058 -- | ICALFirstBuf s1 s2 i1 i2
3059 | ICALSecondInner s1 s2 i1 i2
3060 -- -- | ICALSecondInner s1 s2 i1 i2 a
3061 -- -- | ICALFirstResume s1 s2 i1 i2 a
3063-- | Interleave streams (full streams, not the elements) unfolded from two
3064-- input streams and concat. Stop when the first stream stops. If the second
3065-- stream ends before the first one then first stream still keeps running alone
3066-- without any interleaving with the second stream.
3068-- [a1, a2, ... an] [b1, b2 ...]
3069-- => [streamA1, streamA2, ... streamAn] [streamB1, streamB2, ...]
3070-- => [streamA1, streamB1, streamA2...StreamAn, streamBn]
3071-- => [a11, a12, ...a1j, b11, b12, ...b1k, a21, a22, ...]
3073{-# INLINE_NORMAL gintercalate #-}
3076 => Unfold m a c -> Stream m a -> Unfold m b c -> Stream m b -> Stream m c
3078 (Unfold istep1 inject1) (Stream step1 state1)
3079 (Unfold istep2 inject2) (Stream step2 state2) =
3080 Stream step (ICALFirst state1 state2)
3084 {-# INLINE_LATE step #-}
3085 step gst (ICALFirst s1 s2) = do
3086 r <- step1 (adaptState gst) s1
3090 i `seq` return (Skip (ICALFirstInner s s2 i))
3091 -- i `seq` return (Skip (ICALFirstYield s s2 i))
3092 Skip s -> return $ Skip (ICALFirst s s2)
3096 step _ (ICALFirstYield s1 s2 i1) = do
3099 Yield x i' -> Yield x (ICALFirstInner s1 s2 i')
3100 Skip i' -> Skip (ICALFirstYield s1 s2 i')
3101 Stop -> Skip (ICALFirst s1 s2)
3104 step _ (ICALFirstInner s1 s2 i1) = do
3107 Yield x i' -> Yield x (ICALFirstInner s1 s2 i')
3108 Skip i' -> Skip (ICALFirstInner s1 s2 i')
3109 Stop -> Skip (ICALSecondInject s1 s2)
3111 step gst (ICALFirstOnly s1) = do
3112 r <- step1 (adaptState gst) s1
3116 i `seq` return (Skip (ICALFirstOnlyInner s i))
3117 Skip s -> return $ Skip (ICALFirstOnly s)
3120 step _ (ICALFirstOnlyInner s1 i1) = do
3123 Yield x i' -> Yield x (ICALFirstOnlyInner s1 i')
3124 Skip i' -> Skip (ICALFirstOnlyInner s1 i')
3125 Stop -> Skip (ICALFirstOnly s1)
3127 -- We inject the second stream even before checking if the first stream
3128 -- would yield any more elements. There is no clear choice whether we
3129 -- should do this before or after that. Doing it after may make the state
3130 -- machine a bit simpler though.
3131 step gst (ICALSecondInject s1 s2) = do
3132 r <- step2 (adaptState gst) s2
3136 i `seq` return (Skip (ICALFirstInject s1 s i))
3137 Skip s -> return $ Skip (ICALSecondInject s1 s)
3138 Stop -> return $ Skip (ICALFirstOnly s1)
3140 step gst (ICALFirstInject s1 s2 i2) = do
3141 r <- step1 (adaptState gst) s1
3145 i `seq` return (Skip (ICALSecondInner s s2 i i2))
3146 -- i `seq` return (Skip (ICALFirstBuf s s2 i i2))
3147 Skip s -> return $ Skip (ICALFirstInject s s2 i2)
3151 step _ (ICALFirstBuf s1 s2 i1 i2) = do
3154 Yield x i' -> Skip (ICALSecondInner s1 s2 i' i2 x)
3155 Skip i' -> Skip (ICALFirstBuf s1 s2 i' i2)
3158 step _ (ICALSecondInner s1 s2 i1 i2 v) = do
3161 Yield x i' -> Yield x (ICALSecondInner s1 s2 i1 i' v)
3162 Skip i' -> Skip (ICALSecondInner s1 s2 i1 i' v)
3163 Stop -> Skip (ICALFirstResume s1 s2 i1 i2 v)
3166 step _ (ICALSecondInner s1 s2 i1 i2) = do
3169 Yield x i' -> Yield x (ICALSecondInner s1 s2 i1 i')
3170 Skip i' -> Skip (ICALSecondInner s1 s2 i1 i')
3171 Stop -> Skip (ICALFirstInner s1 s2 i1)
3172 -- Stop -> Skip (ICALFirstResume s1 s2 i1 i2)
3175 step _ (ICALFirstResume s1 s2 i1 i2 x) = do
3176 return $ Yield x (ICALFirstInner s1 s2 i1 i2)
3179data InterposeState s1 i1 a =
3181 -- | InterposeFirstYield s1 i1
3182 | InterposeFirstInner s1 i1
3183 | InterposeFirstInject s1
3184 -- | InterposeFirstBuf s1 i1
3185 | InterposeSecondYield s1 i1
3186 -- -- | InterposeSecondYield s1 i1 a
3187 -- -- | InterposeFirstResume s1 i1 a
3189-- Note that this only interposes the pure values, we may run many effects to
3190-- generate those values as some effects may not generate anything (Skip).
3191{-# INLINE_NORMAL interpose #-}
3192interpose :: Monad m => m c -> Unfold m b c -> Stream m b -> Stream m c
3195 (Unfold istep1 inject1) (Stream step1 state1) =
3196 Stream step (InterposeFirst state1)
3200 {-# INLINE_LATE step #-}
3201 step gst (InterposeFirst s1) = do
3202 r <- step1 (adaptState gst) s1
3206 i `seq` return (Skip (InterposeFirstInner s i))
3207 -- i `seq` return (Skip (InterposeFirstYield s i))
3208 Skip s -> return $ Skip (InterposeFirst s)
3212 step _ (InterposeFirstYield s1 i1) = do
3215 Yield x i' -> Yield x (InterposeFirstInner s1 i')
3216 Skip i' -> Skip (InterposeFirstYield s1 i')
3217 Stop -> Skip (InterposeFirst s1)
3220 step _ (InterposeFirstInner s1 i1) = do
3223 Yield x i' -> Yield x (InterposeFirstInner s1 i')
3224 Skip i' -> Skip (InterposeFirstInner s1 i')
3225 Stop -> Skip (InterposeFirstInject s1)
3227 step gst (InterposeFirstInject s1) = do
3228 r <- step1 (adaptState gst) s1
3232 -- i `seq` return (Skip (InterposeFirstBuf s i))
3233 i `seq` return (Skip (InterposeSecondYield s i))
3234 Skip s -> return $ Skip (InterposeFirstInject s)
3238 step _ (InterposeFirstBuf s1 i1) = do
3241 Yield x i' -> Skip (InterposeSecondYield s1 i' x)
3242 Skip i' -> Skip (InterposeFirstBuf s1 i')
3247 step _ (InterposeSecondYield s1 i1 v) = do
3249 return $ Yield r (InterposeFirstResume s1 i1 v)
3251 step _ (InterposeSecondYield s1 i1) = do
3253 return $ Yield r (InterposeFirstInner s1 i1)
3256 step _ (InterposeFirstResume s1 i1 v) = do
3257 return $ Yield v (InterposeFirstInner s1 i1)
3260------------------------------------------------------------------------------
3262------------------------------------------------------------------------------
3264data GbracketState s1 s2 v
3266 | GBracketNormal s1 v
3267 | GBracketException s2
3269-- | Like 'gbracket' but with following differences:
3271-- * alloc action @m c@ runs with async exceptions enabled
3272-- * cleanup action @c -> m d@ won't run if the stream is garbage collected
3273-- after partial evaluation.
3274-- * does not require a 'MonadAsync' constraint.
3276-- /Inhibits stream fusion/
3280{-# INLINE_NORMAL gbracket_ #-}
3284 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
3285 -> (c -> m d) -- ^ after, on normal stop
3286 -> (c -> e -> Stream m b -> Stream m b) -- ^ on exception
3287 -> (c -> Stream m b) -- ^ stream generator
3289gbracket_ bef exc aft fexc fnormal =
3290 Stream step GBracketInit
3294 {-# INLINE_LATE step #-}
3295 step _ GBracketInit = do
3297 return $ Skip $ GBracketNormal (fnormal r) r
3299 step gst (GBracketNormal (UnStream step1 st) v) = do
3300 res <- exc $ step1 gst st
3302 Right r -> case r of
3304 return $ Yield x (GBracketNormal (Stream step1 s) v)
3305 Skip s -> return $ Skip (GBracketNormal (Stream step1 s) v)
3306 Stop -> aft v >> return Stop
3307 -- XXX Do not handle async exceptions, just rethrow them.
3309 return $ Skip (GBracketException (fexc v e (UnStream step1 st)))
3310 step gst (GBracketException (UnStream step1 st)) = do
3313 Yield x s -> return $ Yield x (GBracketException (Stream step1 s))
3314 Skip s -> return $ Skip (GBracketException (Stream step1 s))
3317------------------------------------------------------------------------------
3319------------------------------------------------------------------------------
3321-- | Make a finalizer from a monadic action @m a@ that can run in IO monad.
3322mkIOFinalizer :: MonadBaseControl IO m => m b -> m (IO ())
3324 mrun <- captureMonadState
3330-- | Run an IO action stored in a finalized IORef.
3331runIORefFinalizerGC :: IORef (Maybe (IO ())) -> IO ()
3332runIORefFinalizerGC ref = do
3333 res <- readIORef ref
3335 Nothing -> return ()
3338-- | Create an IORef holding a finalizer that is called automatically when
3339-- the IORef is garbage collected. The IORef can be written to with a 'Nothing'
3340-- value to deactivate the finalizer.
3342-- Note: The finalizer is always run with the state of the monad captured at
3343-- the time of calling newFinalizedIORef. To run it on garbage collection we
3344-- have no option but to take a snapshot of the monadic state at some point of
3345-- time. For normal case we could run it with the current state of the monad
3346-- but we want to keep both the cases consistent.
3348newFinalizedIORef :: (MonadIO m, MonadBaseControl IO m)
3349 => m a -> m (IORef (Maybe (IO ())))
3350newFinalizedIORef finalizer = do
3351 f <- mkIOFinalizer finalizer
3352 ref <- liftIO $ newIORef $ Just f
3353 _ <- liftIO $ mkWeakIORef ref (runIORefFinalizerGC ref)
3356-- | Run the finalizer stored in an IORef and deactivate it so that it is run
3357-- only once. Note, the action runs with async exceptions masked.
3359runIORefFinalizer :: MonadIO m => IORef (Maybe (IO ())) -> m ()
3360runIORefFinalizer ref = liftIO $ do
3361 res <- readIORef ref
3363 Nothing -> return ()
3365 -- if an async exception comes after writing 'Nothing' then the
3366 -- finalizing action will never be run. We need to do this
3367 -- atomically wrt async exceptions.
3369 writeIORef ref Nothing
3372-- | Run an action clearing the finalizer IORef atomically wrt async
3373-- exceptions. The action is run with async exceptions masked.
3374withIORefFinalizer :: MonadBaseControl IO m
3375 => IORef (Maybe (IO ())) -> m a -> m a
3376withIORefFinalizer ref action = do
3377 control $ \runinio ->
3379 writeIORef ref Nothing
3382------------------------------------------------------------------------------
3384data GbracketIOState s1 s2 v wref
3386 | GBracketIONormal s1 v wref
3387 | GBracketIOException s2
3389-- | Run the alloc action @m c@ with async exceptions disabled but keeping
3390-- blocking operations interruptible (see 'Control.Exception.mask'). Use the
3391-- output @c@ as input to @c -> Stream m b@ to generate an output stream. When
3392-- generating the stream use the supplied @try@ operation @forall s. m s -> m
3393-- (Either e s)@ to catch synchronous exceptions. If an exception occurs run
3394-- the exception handler @c -> e -> Stream m b -> m (Stream m b)@.
3396-- The cleanup action @c -> m d@, runs whenever the stream ends normally, due
3397-- to a sync or async exception or if it gets garbage collected after a partial
3398-- lazy evaluation. See 'bracket' for the semantics of the cleanup action.
3400-- 'gbracket' can express all other exception handling combinators.
3402-- /Inhibits stream fusion/
3405{-# INLINE_NORMAL gbracket #-}
3407 :: (MonadIO m, MonadBaseControl IO m)
3409 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
3410 -> (c -> m d) -- ^ after, on normal stop or GC
3411 -> (c -> e -> Stream m b -> m (Stream m b)) -- ^ on exception
3412 -> (c -> Stream m b) -- ^ stream generator
3414gbracket bef exc aft fexc fnormal =
3415 Stream step GBracketIOInit
3419 -- If the stream is never evaluated the "aft" action will never be
3420 -- called. For that to occur we will need the user of this API to pass a
3421 -- weak pointer to us.
3422 {-# INLINE_LATE step #-}
3423 step _ GBracketIOInit = do
3424 -- We mask asynchronous exceptions to make the execution
3425 -- of 'bef' and the registration of 'aft' atomic.
3426 -- A similar thing is done in the resourcet package: https://git.io/JvKV3
3427 -- Tutorial: https://markkarpov.com/tutorial/exceptions.html
3428 (r, ref) <- liftBaseOp_ mask_ $ do
3430 ref <- newFinalizedIORef (aft r)
3432 return $ Skip $ GBracketIONormal (fnormal r) r ref
3434 step gst (GBracketIONormal (UnStream step1 st) v ref) = do
3435 res <- exc $ step1 gst st
3437 Right r -> case r of
3439 return $ Yield x (GBracketIONormal (Stream step1 s) v ref)
3441 return $ Skip (GBracketIONormal (Stream step1 s) v ref)
3443 runIORefFinalizer ref
3445 -- XXX Do not handle async exceptions, just rethrow them.
3447 -- Clearing of finalizer and running of exception handler must
3448 -- be atomic wrt async exceptions. Otherwise if we have cleared
3449 -- the finalizer and have not run the exception handler then we
3450 -- may leak the resource.
3451 stream <- withIORefFinalizer ref (fexc v e (UnStream step1 st))
3452 return $ Skip (GBracketIOException stream)
3453 step gst (GBracketIOException (UnStream step1 st)) = do
3457 return $ Yield x (GBracketIOException (Stream step1 s))
3458 Skip s -> return $ Skip (GBracketIOException (Stream step1 s))
3461-- | See 'Streamly.Internal.Data.Stream.IsStream.after_'.
3463{-# INLINE_NORMAL before #-}
3464before :: Monad m => m b -> Stream m a -> Stream m a
3465before action (Stream step state) = Stream step' Nothing
3469 {-# INLINE_LATE step' #-}
3470 step' _ Nothing = action >> return (Skip (Just state))
3472 step' gst (Just st) = do
3475 Yield x s -> return $ Yield x (Just s)
3476 Skip s -> return $ Skip (Just s)
3479-- | See 'Streamly.Internal.Data.Stream.IsStream.after_'.
3481{-# INLINE_NORMAL after_ #-}
3482after_ :: Monad m => m b -> Stream m a -> Stream m a
3483after_ action (Stream step state) = Stream step' state
3487 {-# INLINE_LATE step' #-}
3491 Yield x s -> return $ Yield x s
3492 Skip s -> return $ Skip s
3493 Stop -> action >> return Stop
3495-- | See 'Streamly.Internal.Data.Stream.IsStream.after'.
3497{-# INLINE_NORMAL after #-}
3498after :: (MonadIO m, MonadBaseControl IO m)
3499 => m b -> Stream m a -> Stream m a
3500after action (Stream step state) = Stream step' Nothing
3504 {-# INLINE_LATE step' #-}
3505 step' _ Nothing = do
3506 ref <- newFinalizedIORef action
3507 return $ Skip $ Just (state, ref)
3508 step' gst (Just (st, ref)) = do
3511 Yield x s -> return $ Yield x (Just (s, ref))
3512 Skip s -> return $ Skip (Just (s, ref))
3514 runIORefFinalizer ref
3517-- XXX For high performance error checks in busy streams we may need another
3518-- Error constructor in step.
3520-- | See 'Streamly.Internal.Data.Stream.IsStream.onException'.
3522{-# INLINE_NORMAL onException #-}
3523onException :: MonadCatch m => m b -> Stream m a -> Stream m a
3524onException action str =
3525 gbracket_ (return ()) MC.try return
3526 (\_ (e :: MC.SomeException) _ -> nilM (action >> MC.throwM e))
3529{-# INLINE_NORMAL _onException #-}
3530_onException :: MonadCatch m => m b -> Stream m a -> Stream m a
3531_onException action (Stream step state) = Stream step' state
3535 {-# INLINE_LATE step' #-}
3537 res <- step gst st `MC.onException` action
3539 Yield x s -> return $ Yield x s
3540 Skip s -> return $ Skip s
3543-- | See 'Streamly.Internal.Data.Stream.IsStream.bracket_'.
3545{-# INLINE_NORMAL bracket_ #-}
3546bracket_ :: MonadCatch m
3547 => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
3548bracket_ bef aft bet =
3549 gbracket_ bef MC.try aft
3550 (\a (e :: SomeException) _ -> nilM (aft a >> MC.throwM e)) bet
3552-- | See 'Streamly.Internal.Data.Stream.IsStream.bracket'.
3554{-# INLINE_NORMAL bracket #-}
3555bracket :: (MonadAsync m, MonadCatch m)
3556 => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
3557bracket bef aft bet =
3558 gbracket bef MC.try aft
3559 (\a (e :: SomeException) _ -> aft a >> return (nilM (MC.throwM e))) bet
3561data BracketState s v = BracketInit | BracketRun s v
3563-- | Alternate (custom) implementation of 'bracket'.
3565{-# INLINE_NORMAL _bracket #-}
3566_bracket :: MonadCatch m
3567 => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
3568_bracket bef aft bet = Stream step' BracketInit
3572 {-# INLINE_LATE step' #-}
3573 step' _ BracketInit = bef >>= \x -> return (Skip (BracketRun (bet x) x))
3575 -- NOTE: It is important to use UnStream instead of the Stream pattern
3576 -- here, otherwise we get huge perf degradation, see note in concatMap.
3577 step' gst (BracketRun (UnStream step state) v) = do
3578 -- res <- step gst state `MC.onException` aft v
3579 res <- MC.try $ step gst state
3581 Left (e :: SomeException) -> aft v >> MC.throwM e >> return Stop
3582 Right r -> case r of
3583 Yield x s -> return $ Yield x (BracketRun (Stream step s) v)
3584 Skip s -> return $ Skip (BracketRun (Stream step s) v)
3585 Stop -> aft v >> return Stop
3587-- | See 'Streamly.Internal.Data.Stream.IsStream.finally_'.
3589{-# INLINE finally_ #-}
3590finally_ :: MonadCatch m => m b -> Stream m a -> Stream m a
3591finally_ action xs = bracket_ (return ()) (\_ -> action) (const xs)
3593-- | See 'Streamly.Internal.Data.Stream.IsStream.finally'.
3595-- finally action xs = after action $ onException action xs
3597{-# INLINE finally #-}
3598finally :: (MonadAsync m, MonadCatch m) => m b -> Stream m a -> Stream m a
3599finally action xs = bracket (return ()) (\_ -> action) (const xs)
3601-- | See 'Streamly.Internal.Data.Stream.IsStream.ghandle'.
3603{-# INLINE_NORMAL ghandle #-}
3604ghandle :: (MonadCatch m, Exception e)
3605 => (e -> Stream m a -> Stream m a) -> Stream m a -> Stream m a
3607 gbracket_ (return ()) MC.try return (\_ -> f) (\_ -> str)
3609-- | See 'Streamly.Internal.Data.Stream.IsStream.handle'.
3611{-# INLINE_NORMAL handle #-}
3612handle :: (MonadCatch m, Exception e)
3613 => (e -> Stream m a) -> Stream m a -> Stream m a
3615 gbracket_ (return ()) MC.try return (\_ e _ -> f e) (\_ -> str)
3617-- | Alternate (custom) implementation of 'handle'.
3619{-# INLINE_NORMAL _handle #-}
3620_handle :: (MonadCatch m, Exception e)
3621 => (e -> Stream m a) -> Stream m a -> Stream m a
3622_handle f (Stream step state) = Stream step' (Left state)
3626 {-# INLINE_LATE step' #-}
3627 step' gst (Left st) = do
3628 res <- MC.try $ step gst st
3630 Left e -> return $ Skip $ Right (f e)
3631 Right r -> case r of
3632 Yield x s -> return $ Yield x (Left s)
3633 Skip s -> return $ Skip (Left s)
3636 step' gst (Right (UnStream step1 st)) = do
3639 Yield x s -> return $ Yield x (Right (Stream step1 s))
3640 Skip s -> return $ Skip (Right (Stream step1 s))
3643-------------------------------------------------------------------------------
3644-- General transformation
3645-------------------------------------------------------------------------------
3647{-# INLINE_NORMAL transform #-}
3648transform :: Monad m => Pipe m a b -> Stream m a -> Stream m b
3649transform (Pipe pstep1 pstep2 pstate) (Stream step state) =
3650 Stream step' (Consume pstate, state)
3654 {-# INLINE_LATE step' #-}
3656 step' gst (Consume pst, st) = pst `seq` do
3657 r <- step (adaptState gst) st
3662 Pipe.Yield b pst' -> return $ Yield b (pst', s)
3663 Pipe.Continue pst' -> return $ Skip (pst', s)
3664 Skip s -> return $ Skip (Consume pst, s)
3667 step' _ (Produce pst, st) = pst `seq` do
3670 Pipe.Yield b pst' -> return $ Yield b (pst', st)
3671 Pipe.Continue pst' -> return $ Skip (pst', st)
3673------------------------------------------------------------------------------
3674-- Transformation by Folding (Scans)
3675------------------------------------------------------------------------------
3677------------------------------------------------------------------------------
3679------------------------------------------------------------------------------
3681-- XXX Is a prescan useful, discarding the last step does not sound useful? I
3682-- am not sure about the utility of this function, so this is implemented but
3683-- not exposed. We can expose it if someone provides good reasons why this is
3686-- XXX We have to execute the stream one step ahead to know that we are at the
3687-- last step. The vector implementation of prescan executes the last fold step
3688-- but does not yield the result. This means we have executed the effect but
3689-- discarded value. This does not sound right. In this implementation we are
3690-- not executing the last fold step.
3691{-# INLINE_NORMAL prescanlM' #-}
3692prescanlM' :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
3693prescanlM' f mz (Stream step state) = Stream step' (state, mz)
3695 {-# INLINE_LATE step' #-}
3696 step' gst (st, prev) = do
3697 r <- step (adaptState gst) st
3701 return $ Yield acc (s, f acc x)
3702 Skip s -> return $ Skip (s, prev)
3705{-# INLINE prescanl' #-}
3706prescanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
3707prescanl' f z = prescanlM' (\a b -> return (f a b)) (return z)
3709------------------------------------------------------------------------------
3710-- Monolithic postscans (postscan followed by a map)
3711------------------------------------------------------------------------------
3713-- The performance of a modular postscan followed by a map seems to be
3714-- equivalent to this monolithic scan followed by map therefore we may not need
3715-- this implementation. We just have it for performance comparison and in case
3716-- modular version does not perform well in some situation.
3718{-# INLINE_NORMAL postscanlMx' #-}
3719postscanlMx' :: Monad m
3720 => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> Stream m b
3721postscanlMx' fstep begin done (Stream step state) = do
3722 Stream step' (state, begin)
3724 {-# INLINE_LATE step' #-}
3725 step' gst (st, acc) = do
3726 r <- step (adaptState gst) st
3732 v `seq` y `seq` return (Yield v (s, return y))
3733 Skip s -> return $ Skip (s, acc)
3736{-# INLINE_NORMAL postscanlx' #-}
3737postscanlx' :: Monad m
3738 => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> Stream m b
3739postscanlx' fstep begin done s =
3740 postscanlMx' (\b a -> return (fstep b a)) (return begin) (return . done) s
3742-- XXX do we need consM strict to evaluate the begin value?
3743{-# INLINE scanlMx' #-}
3745 => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> Stream m b
3746scanlMx' fstep begin done s =
3747 (begin >>= \x -> x `seq` done x) `consM` postscanlMx' fstep begin done s
3749{-# INLINE scanlx' #-}
3751 => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> Stream m b
3752scanlx' fstep begin done s =
3753 scanlMx' (\b a -> return (fstep b a)) (return begin) (return . done) s
3755------------------------------------------------------------------------------
3757------------------------------------------------------------------------------
3759{-# INLINE_NORMAL postscanlM' #-}
3760postscanlM' :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
3761postscanlM' fstep begin (Stream step state) =
3762 Stream step' Nothing
3764 {-# INLINE_LATE step' #-}
3765 step' _ Nothing = do
3767 return $ Skip (Just (state, x))
3769 step' gst (Just (st, acc)) = do
3770 r <- step (adaptState gst) st
3774 return $ Yield y (Just (s, y))
3775 Skip s -> return $ Skip (Just (s, acc))
3778{-# INLINE_NORMAL postscanl' #-}
3779postscanl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
3780postscanl' f seed = postscanlM' (\a b -> return (f a b)) (return seed)
3782-- We can possibly have the "done" function as a Maybe to provide an option to
3783-- emit or not emit the accumulator when the stream stops.
3785-- TBD: use a single Yield point
3787{-# INLINE_NORMAL postscanlMAfter' #-}
3788postscanlMAfter' :: Monad m
3789 => (b -> a -> m b) -> m b -> (b -> m b) -> Stream m a -> Stream m b
3790postscanlMAfter' fstep initial done (Stream step1 state1) = do
3791 Stream step (Just (state1, initial))
3795 {-# INLINE_LATE step #-}
3796 step gst (Just (st, acc)) = do
3797 r <- step1 (adaptState gst) st
3802 y `seq` return (Yield y (Just (s, return y)))
3803 Skip s -> return $ Skip $ Just (s, acc)
3804 Stop -> acc >>= done >>= \res -> return (Yield res Nothing)
3805 step _ Nothing = return Stop
3807{-# INLINE_NORMAL postscanlM #-}
3808postscanlM :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
3809postscanlM fstep begin (Stream step state) = Stream step' Nothing
3811 {-# INLINE_LATE step' #-}
3812 step' _ Nothing = do
3814 return $ Skip (Just (state, r))
3816 step' gst (Just (st, acc)) = do
3817 r <- step (adaptState gst) st
3821 return (Yield y (Just (s, y)))
3822 Skip s -> return $ Skip (Just (s, acc))
3825{-# INLINE_NORMAL postscanl #-}
3826postscanl :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
3827postscanl f seed = postscanlM (\a b -> return (f a b)) (return seed)
3829{-# INLINE_NORMAL scanlM' #-}
3830scanlM' :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
3831scanlM' fstep begin (Stream step state) = Stream step' Nothing
3833 {-# INLINE_LATE step' #-}
3834 step' _ Nothing = do
3836 return $ Yield x (Just (state, x))
3837 step' gst (Just (st, acc)) = do
3838 r <- step (adaptState gst) st
3842 return $ Yield y (Just (s, y))
3843 Skip s -> return $ Skip (Just (s, acc))
3846{-# INLINE scanlMAfter' #-}
3847scanlMAfter' :: Monad m
3848 => (b -> a -> m b) -> m b -> (b -> m b) -> Stream m a -> Stream m b
3849scanlMAfter' fstep initial done s =
3850 (initial >>= \x -> x `seq` return x) `consM`
3851 postscanlMAfter' fstep initial done s
3853{-# INLINE scanl' #-}
3854scanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
3855scanl' f seed = scanlM' (\a b -> return (f a b)) (return seed)
3857{-# INLINE_NORMAL scanlM #-}
3858scanlM :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
3859scanlM fstep begin (Stream step state) = Stream step' Nothing
3861 {-# INLINE_LATE step' #-}
3862 step' _ Nothing = do
3864 return $ Yield x (Just (state, x))
3865 step' gst (Just (st, acc)) = do
3866 r <- step (adaptState gst) st
3870 return $ Yield y (Just (s, y))
3871 Skip s -> return $ Skip (Just (s, acc))
3872 Stop -> return $ Stop
3875scanl :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
3876scanl f seed = scanlM (\a b -> return (f a b)) (return seed)
3878{-# INLINE_NORMAL scanl1M #-}
3879scanl1M :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
3880scanl1M fstep (Stream step state) = Stream step' (state, Nothing)
3882 {-# INLINE_LATE step' #-}
3883 step' gst (st, Nothing) = do
3886 Yield x s -> return $ Yield x (s, Just x)
3887 Skip s -> return $ Skip (s, Nothing)
3890 step' gst (st, Just acc) = do
3895 return $ Yield z (s, Just z)
3896 Skip s -> return $ Skip (s, Just acc)
3899{-# INLINE scanl1 #-}
3900scanl1 :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
3901scanl1 f = scanl1M (\x y -> return (f x y))
3903{-# INLINE_NORMAL scanl1M' #-}
3904scanl1M' :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
3905scanl1M' fstep (Stream step state) = Stream step' (state, Nothing)
3907 {-# INLINE_LATE step' #-}
3908 step' gst (st, Nothing) = do
3911 Yield x s -> x `seq` return $ Yield x (s, Just x)
3912 Skip s -> return $ Skip (s, Nothing)
3915 step' gst (st, Just acc) = acc `seq` do
3920 z `seq` return $ Yield z (s, Just z)
3921 Skip s -> return $ Skip (s, Just acc)
3924{-# INLINE scanl1' #-}
3925scanl1' :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
3926scanl1' f = scanl1M' (\x y -> return (f x y))
3928------------------------------------------------------------------------------
3930------------------------------------------------------------------------------
3932data RollingMapState s a = RollingMapInit s | RollingMapGo s a
3934{-# INLINE rollingMapM #-}
3935rollingMapM :: Monad m => (a -> a -> m b) -> Stream m a -> Stream m b
3936rollingMapM f (Stream step1 state1) = Stream step (RollingMapInit state1)
3938 step gst (RollingMapInit st) = do
3939 r <- step1 (adaptState gst) st
3941 Yield x s -> Skip $ RollingMapGo s x
3942 Skip s -> Skip $ RollingMapInit s
3945 step gst (RollingMapGo s1 x1) = do
3946 r <- step1 (adaptState gst) s1
3950 return $ Yield res $ RollingMapGo s x
3951 Skip s -> return $ Skip $ RollingMapGo s x1
3952 Stop -> return $ Stop
3954{-# INLINE rollingMap #-}
3955rollingMap :: Monad m => (a -> a -> b) -> Stream m a -> Stream m b
3956rollingMap f = rollingMapM (\x y -> return $ f x y)
3958------------------------------------------------------------------------------
3959-- Tapping/Distributing
3960------------------------------------------------------------------------------
3963tap :: Monad m => Fold m a b -> Stream m a -> Stream m a
3964tap (Fold fstep initial extract) (Stream step state) = Stream step' Nothing
3968 step' _ Nothing = do
3970 return $ Skip (Just (r, state))
3972 step' gst (Just (acc, st)) = acc `seq` do
3977 return $ Yield x (Just (acc', s))
3978 Skip s -> return $ Skip (Just (acc, s))
3983{-# INLINE_NORMAL tapOffsetEvery #-}
3984tapOffsetEvery :: Monad m
3985 => Int -> Int -> Fold m a b -> Stream m a -> Stream m a
3986tapOffsetEvery offset n (Fold fstep initial extract) (Stream step state) =
3987 Stream step' Nothing
3991 {-# INLINE_LATE step' #-}
3992 step' _ Nothing = do
3994 return $ Skip (Just (r, state, offset `mod` n))
3996 step' gst (Just (acc, st, count)) | count <= 0 = do
4000 !acc' <- fstep acc x
4001 return $ Yield x (Just (acc', s, n - 1))
4002 Skip s -> return $ Skip (Just (acc, s, count))
4007 step' gst (Just (acc, st, count)) = do
4010 Yield x s -> return $ Yield x (Just (acc, s, count - 1))
4011 Skip s -> return $ Skip (Just (acc, s, count))
4016{-# INLINE_NORMAL pollCounts #-}
4020 -> (Stream m Int -> Stream m Int)
4024pollCounts predicate transf fld (Stream step state) = Stream step' Nothing
4027 {-# INLINE_LATE step' #-}
4028 step' _ Nothing = do
4029 -- As long as we are using an "Int" for counts lockfree reads from
4030 -- Var should work correctly on both 32-bit and 64-bit machines.
4031 -- However, an Int on a 32-bit machine may overflow quickly.
4032 countVar <- liftIO $ Prim.newIORef (0 :: Int)
4034 $ void $ runFold fld
4035 $ transf $ fromPrimIORef countVar
4036 return $ Skip (Just (countVar, tid, state))
4038 step' gst (Just (countVar, tid, st)) = do
4042 when (predicate x) $ liftIO $ Prim.modifyIORef' countVar (+ 1)
4043 return $ Yield x (Just (countVar, tid, s))
4044 Skip s -> return $ Skip (Just (countVar, tid, s))
4046 liftIO $ killThread tid
4049{-# INLINE_NORMAL tapRate #-}
4051 (MonadAsync m, MonadCatch m)
4056tapRate samplingRate action (Stream step state) = Stream step' Nothing
4058 {-# NOINLINE loop #-}
4059 loop countVar prev = do
4062 (do liftIO $ threadDelay (round $ samplingRate * 1000000)
4063 i <- liftIO $ Prim.readIORef countVar
4064 let !diff = i - prev
4067 (\(e :: AsyncException) -> do
4068 i <- liftIO $ Prim.readIORef countVar
4069 let !diff = i - prev
4071 throwM (MC.toException e))
4074 {-# INLINE_LATE step' #-}
4075 step' _ Nothing = do
4076 countVar <- liftIO $ Prim.newIORef 0
4077 tid <- fork $ loop countVar 0
4078 ref <- liftIO $ newIORef ()
4079 _ <- liftIO $ mkWeakIORef ref (killThread tid)
4080 return $ Skip (Just (countVar, tid, state, ref))
4082 step' gst (Just (countVar, tid, st, ref)) = do
4086 liftIO $ Prim.modifyIORef' countVar (+ 1)
4087 return $ Yield x (Just (countVar, tid, s, ref))
4088 Skip s -> return $ Skip (Just (countVar, tid, s, ref))
4090 liftIO $ killThread tid
4094-------------------------------------------------------------------------------
4096-------------------------------------------------------------------------------
4098{-# INLINE_NORMAL takeWhileM #-}
4099takeWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
4100takeWhileM f (Stream step state) = Stream step' state
4102 {-# INLINE_LATE step' #-}
4108 return $ if b then Yield x s else Stop
4109 Skip s -> return $ Skip s
4112{-# INLINE takeWhile #-}
4113takeWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
4114takeWhile f = takeWhileM (return . f)
4116{-# INLINE_NORMAL drop #-}
4117drop :: Monad m => Int -> Stream m a -> Stream m a
4118drop n (Stream step state) = Stream step' (state, Just n)
4120 {-# INLINE_LATE step' #-}
4121 step' gst (st, Just i)
4126 Yield _ s -> Skip (s, Just (i - 1))
4127 Skip s -> Skip (s, Just i)
4129 | otherwise = return $ Skip (st, Nothing)
4131 step' gst (st, Nothing) = do
4135 Yield x s -> Yield x (s, Nothing)
4136 Skip s -> Skip (s, Nothing)
4139data DropWhileState s a
4141 | DropWhileYield a s
4144{-# INLINE_NORMAL dropWhileM #-}
4145dropWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
4146dropWhileM f (Stream step state) = Stream step' (DropWhileDrop state)
4148 {-# INLINE_LATE step' #-}
4149 step' gst (DropWhileDrop st) = do
4155 then return $ Skip (DropWhileDrop s)
4156 else return $ Skip (DropWhileYield x s)
4157 Skip s -> return $ Skip (DropWhileDrop s)
4160 step' gst (DropWhileNext st) = do
4163 Yield x s -> return $ Skip (DropWhileYield x s)
4164 Skip s -> return $ Skip (DropWhileNext s)
4167 step' _ (DropWhileYield x st) = return $ Yield x (DropWhileNext st)
4169{-# INLINE dropWhile #-}
4170dropWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
4171dropWhile f = dropWhileM (return . f)
4173{-# INLINE_NORMAL filterM #-}
4174filterM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
4175filterM f (Stream step state) = Stream step' state
4177 {-# INLINE_LATE step' #-}
4186 Skip s -> return $ Skip s
4189{-# INLINE filter #-}
4190filter :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
4191filter f = filterM (return . f)
4193{-# INLINE_NORMAL uniq #-}
4194uniq :: (Eq a, Monad m) => Stream m a -> Stream m a
4195uniq (Stream step state) = Stream step' (Nothing, state)
4197 {-# INLINE_LATE step' #-}
4198 step' gst (Nothing, st) = do
4201 Yield x s -> return $ Yield x (Just x, s)
4202 Skip s -> return $ Skip (Nothing, s)
4204 step' gst (Just x, st) = do
4207 Yield y s | x == y -> return $ Skip (Just x, s)
4208 | otherwise -> return $ Yield y (Just y, s)
4209 Skip s -> return $ Skip (Just x, s)
4212------------------------------------------------------------------------------
4213-- Transformation by Mapping
4214------------------------------------------------------------------------------
4216{-# INLINE_NORMAL sequence #-}
4217sequence :: Monad m => Stream m (m a) -> Stream m a
4218sequence (Stream step state) = Stream step' state
4220 {-# INLINE_LATE step' #-}
4222 r <- step (adaptState gst) st
4224 Yield x s -> x >>= \a -> return (Yield a s)
4225 Skip s -> return $ Skip s
4228------------------------------------------------------------------------------
4230------------------------------------------------------------------------------
4232data LoopState x s = FirstYield s
4233 | InterspersingYield s
4236{-# INLINE_NORMAL intersperseM #-}
4237intersperseM :: Monad m => m a -> Stream m a -> Stream m a
4238intersperseM m (Stream step state) = Stream step' (FirstYield state)
4240 {-# INLINE_LATE step' #-}
4241 step' gst (FirstYield st) = do
4245 Yield x s -> Skip (YieldAndCarry x s)
4246 Skip s -> Skip (FirstYield s)
4249 step' gst (InterspersingYield st) = do
4254 return $ Yield a (YieldAndCarry x s)
4255 Skip s -> return $ Skip $ InterspersingYield s
4258 step' _ (YieldAndCarry x st) = return $ Yield x (InterspersingYield st)
4260{-# INLINE_NORMAL intersperseM_ #-}
4261intersperseM_ :: Monad m => m b -> Stream m a -> Stream m a
4262intersperseM_ m (Stream step1 state1) = Stream step (Left (pure (), state1))
4264 {-# INLINE_LATE step #-}
4265 step gst (Left (eff, st)) = do
4268 Yield x s -> eff >> return (Yield x (Right s))
4269 Skip s -> return $ Skip (Left (eff, s))
4272 step _ (Right st) = return $ Skip $ Left (void m, st)
4277 | SuffixYield a (SuffixState s a)
4279{-# INLINE_NORMAL intersperseSuffix #-}
4280intersperseSuffix :: forall m a. Monad m => m a -> Stream m a -> Stream m a
4281intersperseSuffix action (Stream step state) = Stream step' (SuffixElem state)
4283 {-# INLINE_LATE step' #-}
4284 step' gst (SuffixElem st) = do
4287 Yield x s -> Skip (SuffixYield x (SuffixSuffix s))
4288 Skip s -> Skip (SuffixElem s)
4291 step' _ (SuffixSuffix st) = do
4292 action >>= \r -> return $ Skip (SuffixYield r (SuffixElem st))
4294 step' _ (SuffixYield x next) = return $ Yield x next
4296{-# INLINE_NORMAL intersperseSuffix_ #-}
4297intersperseSuffix_ :: Monad m => m b -> Stream m a -> Stream m a
4298intersperseSuffix_ m (Stream step1 state1) = Stream step (Left state1)
4300 {-# INLINE_LATE step #-}
4301 step gst (Left st) = do
4304 Yield x s -> return $ Yield x (Right s)
4305 Skip s -> return $ Skip $ Left s
4308 step _ (Right st) = m >> return (Skip (Left st))
4310data SuffixSpanState s a
4311 = SuffixSpanElem s Int
4312 | SuffixSpanSuffix s
4313 | SuffixSpanYield a (SuffixSpanState s a)
4317-- | intersperse after every n items
4318{-# INLINE_NORMAL intersperseSuffixBySpan #-}
4319intersperseSuffixBySpan :: forall m a. Monad m
4320 => Int -> m a -> Stream m a -> Stream m a
4321intersperseSuffixBySpan n action (Stream step state) =
4322 Stream step' (SuffixSpanElem state n)
4324 {-# INLINE_LATE step' #-}
4325 step' gst (SuffixSpanElem st i) | i > 0 = do
4328 Yield x s -> Skip (SuffixSpanYield x (SuffixSpanElem s (i - 1)))
4329 Skip s -> Skip (SuffixSpanElem s i)
4330 Stop -> if i == n then Stop else Skip SuffixSpanLast
4331 step' _ (SuffixSpanElem st _) = return $ Skip (SuffixSpanSuffix st)
4333 step' _ (SuffixSpanSuffix st) = do
4334 action >>= \r -> return $ Skip (SuffixSpanYield r (SuffixSpanElem st n))
4336 step' _ (SuffixSpanLast) = do
4337 action >>= \r -> return $ Skip (SuffixSpanYield r SuffixSpanStop)
4339 step' _ (SuffixSpanYield x next) = return $ Yield x next
4341 step' _ (SuffixSpanStop) = return Stop
4343{-# INLINE intersperse #-}
4344intersperse :: Monad m => a -> Stream m a -> Stream m a
4345intersperse a = intersperseM (return a)
4347{-# INLINE_NORMAL insertBy #-}
4348insertBy :: Monad m => (a -> a -> Ordering) -> a -> Stream m a -> Stream m a
4349insertBy cmp a (Stream step state) = Stream step' (state, False, Nothing)
4351 {-# INLINE_LATE step' #-}
4352 step' gst (st, False, _) = do
4355 Yield x s -> case cmp a x of
4356 GT -> return $ Yield x (s, False, Nothing)
4357 _ -> return $ Yield a (s, True, Just x)
4358 Skip s -> return $ Skip (s, False, Nothing)
4359 Stop -> return $ Yield a (st, True, Nothing)
4361 step' _ (_, True, Nothing) = return Stop
4363 step' gst (st, True, Just prev) = do
4366 Yield x s -> return $ Yield prev (s, True, Just x)
4367 Skip s -> return $ Skip (s, True, Just prev)
4368 Stop -> return $ Yield prev (st, True, Nothing)
4370------------------------------------------------------------------------------
4372------------------------------------------------------------------------------
4374{-# INLINE_NORMAL deleteBy #-}
4375deleteBy :: Monad m => (a -> a -> Bool) -> a -> Stream m a -> Stream m a
4376deleteBy eq x (Stream step state) = Stream step' (state, False)
4378 {-# INLINE_LATE step' #-}
4379 step' gst (st, False) = do
4382 Yield y s -> return $
4383 if eq x y then Skip (s, True) else Yield y (s, False)
4384 Skip s -> return $ Skip (s, False)
4387 step' gst (st, True) = do
4390 Yield y s -> return $ Yield y (s, True)
4391 Skip s -> return $ Skip (s, True)
4394------------------------------------------------------------------------------
4395-- Transformation by Map and Filter
4396------------------------------------------------------------------------------
4398-- XXX Will this always fuse properly?
4399{-# INLINE_NORMAL mapMaybe #-}
4400mapMaybe :: Monad m => (a -> Maybe b) -> Stream m a -> Stream m b
4401mapMaybe f = fmap fromJust . filter isJust . map f
4403{-# INLINE_NORMAL mapMaybeM #-}
4404mapMaybeM :: Monad m => (a -> m (Maybe b)) -> Stream m a -> Stream m b
4405mapMaybeM f = fmap fromJust . filter isJust . mapM f
4407------------------------------------------------------------------------------
4409------------------------------------------------------------------------------
4411{-# INLINE_NORMAL indexed #-}
4412indexed :: Monad m => Stream m a -> Stream m (Int, a)
4413indexed (Stream step state) = Stream step' (state, 0)
4415 {-# INLINE_LATE step' #-}
4416 step' gst (st, i) = i `seq` do
4417 r <- step (adaptState gst) st
4419 Yield x s -> return $ Yield (i, x) (s, i+1)
4420 Skip s -> return $ Skip (s, i)
4423{-# INLINE_NORMAL indexedR #-}
4424indexedR :: Monad m => Int -> Stream m a -> Stream m (Int, a)
4425indexedR m (Stream step state) = Stream step' (state, m)
4427 {-# INLINE_LATE step' #-}
4428 step' gst (st, i) = i `seq` do
4429 r <- step (adaptState gst) st
4431 Yield x s -> let i' = i - 1
4432 in return $ Yield (i, x) (s, i')
4433 Skip s -> return $ Skip (s, i)
4436{-# INLINE_NORMAL zipWithM #-}
4438 => (a -> b -> m c) -> Stream m a -> Stream m b -> Stream m c
4439zipWithM f (Stream stepa ta) (Stream stepb tb) = Stream step (ta, tb, Nothing)
4441 {-# INLINE_LATE step #-}
4442 step gst (sa, sb, Nothing) = do
4443 r <- stepa (adaptState gst) sa
4446 Yield x sa' -> Skip (sa', sb, Just x)
4447 Skip sa' -> Skip (sa', sb, Nothing)
4450 step gst (sa, sb, Just x) = do
4451 r <- stepb (adaptState gst) sb
4455 return $ Yield z (sa, sb', Nothing)
4456 Skip sb' -> return $ Skip (sa, sb', Just x)
4459#if __GLASGOW_HASKELL__ >= 801
4460{-# RULES "zipWithM xs xs"
4461 forall f xs. zipWithM @Identity f xs xs = mapM (\x -> f x x) xs #-}
4464{-# INLINE zipWith #-}
4465zipWith :: Monad m => (a -> b -> c) -> Stream m a -> Stream m b -> Stream m c
4466zipWith f = zipWithM (\a b -> return (f a b))
4468------------------------------------------------------------------------------
4470------------------------------------------------------------------------------
4472{-# INLINE_NORMAL mergeByM #-}
4475 => (a -> a -> m Ordering) -> Stream m a -> Stream m a -> Stream m a
4476mergeByM cmp (Stream stepa ta) (Stream stepb tb) =
4477 Stream step (Just ta, Just tb, Nothing, Nothing)
4479 {-# INLINE_LATE step #-}
4481 -- one of the values is missing, and the corresponding stream is running
4482 step gst (Just sa, sb, Nothing, b) = do
4485 Yield a sa' -> Skip (Just sa', sb, Just a, b)
4486 Skip sa' -> Skip (Just sa', sb, Nothing, b)
4487 Stop -> Skip (Nothing, sb, Nothing, b)
4489 step gst (sa, Just sb, a, Nothing) = do
4492 Yield b sb' -> Skip (sa, Just sb', a, Just b)
4493 Skip sb' -> Skip (sa, Just sb', a, Nothing)
4494 Stop -> Skip (sa, Nothing, a, Nothing)
4496 -- both the values are available
4497 step _ (sa, sb, Just a, Just b) = do
4499 return $ case res of
4500 GT -> Yield b (sa, sb, Just a, Nothing)
4501 _ -> Yield a (sa, sb, Nothing, Just b)
4503 -- one of the values is missing, corresponding stream is done
4504 step _ (Nothing, sb, Nothing, Just b) =
4505 return $ Yield b (Nothing, sb, Nothing, Nothing)
4507 step _ (sa, Nothing, Just a, Nothing) =
4508 return $ Yield a (sa, Nothing, Nothing, Nothing)
4510 step _ (Nothing, Nothing, Nothing, Nothing) = return Stop
4512{-# INLINE mergeBy #-}
4515 => (a -> a -> Ordering) -> Stream m a -> Stream m a -> Stream m a
4516mergeBy cmp = mergeByM (\a b -> return $ cmp a b)
4518------------------------------------------------------------------------------
4519-- Transformation comprehensions
4520------------------------------------------------------------------------------
4522{-# INLINE_NORMAL the #-}
4523the :: (Eq a, Monad m) => Stream m a -> m (Maybe a)
4524the (Stream step state) = go state
4527 r <- step defState st
4529 Yield x s -> go' x s
4531 Stop -> return Nothing
4533 r <- step defState st
4535 Yield x s | x == n -> go' n s
4536 | otherwise -> return Nothing
4538 Stop -> return (Just n)
4540{-# INLINE runFold #-}
4541runFold :: (Monad m) => Fold m a b -> Stream m a -> m b
4542runFold (Fold step begin done) = foldlMx' step begin done
4544-------------------------------------------------------------------------------
4545-- Concurrent application and fold
4546-------------------------------------------------------------------------------
4548-- XXX These functions should be moved to Stream/Parallel.hs
4550-- Using StreamD the worker stream producing code can fuse with the code to
4551-- queue output to the SVar giving some perf boost.
4553-- Note that StreamD can only be used in limited situations, specifically, we
4554-- cannot implement joinStreamVarPar using this.
4556-- XXX make sure that the SVar passed is a Parallel style SVar.
4558-- | Fold the supplied stream to the SVar asynchronously using Parallel
4559-- concurrency style.
4560-- {-# INLINE_NORMAL toSVarParallel #-}
4561{-# INLINE toSVarParallel #-}
4562toSVarParallel :: MonadAsync m
4563 => State t m a -> SVar t m a -> Stream m a -> m ()
4564toSVarParallel st sv xs =
4565 if svarInspectMode sv
4569 case getYieldLimit st of
4570 Nothing -> doFork (work Nothing)
4572 (handleChildException sv)
4573 Just _ -> doFork (workLim Nothing)
4575 (handleChildException sv)
4580 {-# NOINLINE work #-}
4581 work info = (runFold (FL.toParallelSVar sv info) xs)
4583 {-# NOINLINE workLim #-}
4584 workLim info = runFold (FL.toParallelSVarLimited sv info) xs
4586 {-# NOINLINE forkWithDiag #-}
4588 -- We do not use workerCount in case of ParallelVar but still there is
4589 -- no harm in maintaining it correctly.
4590 liftIO $ atomicModifyIORefCAS_ (workerCount sv) $ \n -> n + 1
4592 -- This allocation matters when significant number of workers are being
4593 -- sent. We allocate it only when needed. The overhead increases by 4x.
4595 case yieldRateInfo sv of
4596 Nothing -> return Nothing
4597 Just _ -> liftIO $ do
4598 cntRef <- newIORef 0
4599 t <- getTime Monotonic
4600 lat <- newIORef (0, t)
4601 return $ Just WorkerInfo
4602 { workerYieldMax = 0
4603 , workerYieldCount = cntRef
4604 , workerLatencyStart = lat
4607 case getYieldLimit st of
4608 Nothing -> doFork (work winfo)
4610 (handleChildException sv)
4611 Just _ -> doFork (workLim winfo)
4613 (handleChildException sv)
4616{-# INLINE_NORMAL mkParallelD #-}
4617mkParallelD :: MonadAsync m => Stream m a -> Stream m a
4618mkParallelD m = Stream step Nothing
4621 step gst Nothing = do
4622 sv <- newParallelVar StopNone gst
4623 toSVarParallel gst sv m
4624 -- XXX use unfold instead?
4625 return $ Skip $ Just $ fromSVar sv
4627 step gst (Just (UnStream step1 st)) = do
4630 Yield a s -> Yield a (Just $ Stream step1 s)
4631 Skip s -> Skip (Just $ Stream step1 s)
4634-- Compare with mkAsync. mkAsync uses an Async style SVar whereas this uses a
4635-- parallel style SVar for evaluation. Currently, parallel style cannot use
4636-- rate control whereas Async style can use rate control. In async style SVar
4637-- the worker thread terminates when the buffer is full whereas in Parallel
4640-- | Make the stream producer and consumer run concurrently by introducing a
4641-- buffer between them. The producer thread evaluates the input stream until
4642-- the buffer fills, it blocks if the buffer is full until there is space in
4643-- the buffer. The consumer consumes the stream lazily from the buffer.
4647{-# INLINE_NORMAL mkParallel #-}
4648mkParallel :: (K.IsStream t, MonadAsync m) => t m a -> t m a
4649mkParallel = fromStreamD . mkParallelD . toStreamD
4651-- Note: we can use another API with two callbacks stop and yield if we want
4652-- the callback to be able to indicate end of stream.
4654-- | Generates a callback and a stream pair. The callback returned is used to
4655-- queue values to the stream. The stream is infinite, there is no way for the
4656-- callback to indicate that it is done now.
4660{-# INLINE_NORMAL newCallbackStream #-}
4661newCallbackStream :: (K.IsStream t, MonadAsync m) => m ((a -> m ()), t m a)
4662newCallbackStream = do
4663 sv <- newParallelVar StopNone defState
4665 -- XXX Add our own thread-id to the SVar as we can not know the callback's
4666 -- thread-id and the callback is not run in a managed worker. We need to
4667 -- handle this better.
4668 liftIO myThreadId >>= modifyThread sv
4670 let callback a = liftIO $ void $ send sv (ChildYield a)
4671 -- XXX we can return an SVar and then the consumer can unfold from the
4673 return (callback, fromStreamD (fromSVar sv))
4675-------------------------------------------------------------------------------
4677-------------------------------------------------------------------------------
4679-- | Create an SVar with a fold consumer that will fold any elements sent to it
4680-- using the supplied fold function.
4681{-# INLINE newFoldSVar #-}
4682newFoldSVar :: MonadAsync m => State t m a -> Fold m a b -> m (SVar t m a)
4683newFoldSVar stt f = do
4684 -- Buffer size for the SVar is derived from the current state
4685 sv <- newParallelVar StopAny (adaptState stt)
4686 -- Add the producer thread-id to the SVar.
4687 liftIO myThreadId >>= modifyThread sv
4688 void $ doFork (work sv) (svarMrun sv) (handleFoldException sv)
4693 {-# NOINLINE work #-}
4694 work sv = void $ runFold f $ fromProducer sv
4696data TapState sv st = TapInit | Tapping sv st | TapDone st
4698{-# INLINE_NORMAL tapAsync #-}
4699tapAsync :: MonadAsync m => Fold m a b -> Stream m a -> Stream m a
4700tapAsync f (Stream step1 state1) = Stream step TapInit
4704 -- In general, a Stop event would come equipped with the result
4705 -- of the fold. It is not used here but it would be useful in
4706 -- applicative and distribute.
4707 done <- fromConsumer svr
4708 when (not done) $ do
4709 liftIO $ withDiagMVar svr "teeToSVar: waiting to drain"
4710 $ takeMVar (outputDoorBellFromConsumer svr)
4714 liftIO $ sendStop svr Nothing
4715 -- drain/wait until a stop event arrives from the fold.
4718 {-# INLINE_LATE step #-}
4719 step gst TapInit = do
4720 sv <- newFoldSVar gst f
4721 return $ Skip (Tapping sv state1)
4723 step gst (Tapping sv st) = do
4727 done <- pushToFold sv a
4730 -- XXX we do not need to wait synchronously here
4732 return $ Yield a (TapDone s)
4733 else return $ Yield a (Tapping sv s)
4734 Skip s -> return $ Skip (Tapping sv s)
4739 step gst (TapDone st) = do
4742 Yield a s -> Yield a (TapDone s)
4743 Skip s -> Skip (TapDone s)
4746-- XXX Exported from Array again as this fold is specific to Array
4747-- | Take last 'n' elements from the stream and discard the rest.
4749lastN :: (Storable a, MonadIO m) => Int -> Fold m a (Array a)
4751 | n <= 0 = fmap (const mempty) FL.drain
4752 | otherwise = A.unsafeFreeze <$> Fold step initial done
4754 step (Tuple3' rb rh i) a = do
4755 rh1 <- liftIO $ RB.unsafeInsert rb rh a
4756 return $ Tuple3' rb rh1 (i + 1)
4757 initial = fmap (\(a, b) -> Tuple3' a b (0 :: Int)) $ liftIO $ RB.new n
4758 done (Tuple3' rb rh i) = do
4759 arr <- liftIO $ MA.newArray n
4760 foldFunc i rh snoc' arr rb
4761 snoc' b a = liftIO $ MA.unsafeSnoc b a
4763 | i < n = RB.unsafeFoldRingM
4764 | otherwise = RB.unsafeFoldRingFullM
4766------------------------------------------------------------------------------
4768------------------------------------------------------------------------------
4770-- XXX using getTime in the loop can be pretty expensive especially for
4771-- computations where iterations are lightweight. We have the following
4774-- 1) Run a timeout thread updating a flag asynchronously and check that
4775-- flag here, that way we can have a cheap termination check.
4777-- 2) Use COARSE clock to get time with lower resolution but more efficiently.
4779-- 3) Use rdtscp/rdtsc to get time directly from the processor, compute the
4780-- termination value of rdtsc in the beginning and then in each iteration just
4781-- get rdtsc and check if we should terminate.
4785 | TakeByTimeCheck st s
4786 | TakeByTimeYield st s
4788{-# INLINE_NORMAL takeByTime #-}
4789takeByTime :: (MonadIO m, TimeUnit64 t) => t -> Stream m a -> Stream m a
4790takeByTime duration (Stream step1 state1) = Stream step (TakeByTimeInit state1)
4793 lim = toRelTime64 duration
4795 {-# INLINE_LATE step #-}
4796 step _ (TakeByTimeInit _) | lim == 0 = return Stop
4797 step _ (TakeByTimeInit st) = do
4798 t0 <- liftIO $ getTime Monotonic
4799 return $ Skip (TakeByTimeYield st t0)
4800 step _ (TakeByTimeCheck st t0) = do
4801 t <- liftIO $ getTime Monotonic
4803 if diffAbsTime64 t t0 > lim
4805 else Skip (TakeByTimeYield st t0)
4806 step gst (TakeByTimeYield st t0) = do
4809 Yield x s -> Yield x (TakeByTimeCheck s t0)
4810 Skip s -> Skip (TakeByTimeCheck s t0)
4813data DropByTime st s x
4815 | DropByTimeGen st s
4816 | DropByTimeCheck st s x
4817 | DropByTimeYield st
4819{-# INLINE_NORMAL dropByTime #-}
4820dropByTime :: (MonadIO m, TimeUnit64 t) => t -> Stream m a -> Stream m a
4821dropByTime duration (Stream step1 state1) = Stream step (DropByTimeInit state1)
4824 lim = toRelTime64 duration
4826 {-# INLINE_LATE step #-}
4827 step _ (DropByTimeInit st) = do
4828 t0 <- liftIO $ getTime Monotonic
4829 return $ Skip (DropByTimeGen st t0)
4830 step gst (DropByTimeGen st t0) = do
4833 Yield x s -> Skip (DropByTimeCheck s t0 x)
4834 Skip s -> Skip (DropByTimeGen s t0)
4836 step _ (DropByTimeCheck st t0 x) = do
4837 t <- liftIO $ getTime Monotonic
4838 if diffAbsTime64 t t0 <= lim
4839 then return $ Skip $ DropByTimeGen st t0
4840 else return $ Yield x $ DropByTimeYield st
4841 step gst (DropByTimeYield st) = do
4844 Yield x s -> Yield x (DropByTimeYield s)
4845 Skip s -> Skip (DropByTimeYield s)
4848{-# INLINE updateTimeVar #-}
4849updateTimeVar :: Prim.IORef Int64 -> IO ()
4850updateTimeVar timeVar = do
4851 MicroSecond64 t <- fromAbsTime <$> getTime Monotonic
4852 Prim.modifyIORef' timeVar (const t)
4854{-# INLINE updateWithDelay #-}
4855updateWithDelay :: RealFrac a => a -> Prim.IORef Int64 -> IO ()
4856updateWithDelay precision timeVar = do
4857 threadDelay (delayTime precision)
4858 updateTimeVar timeVar
4862 -- Keep the minimum at least a millisecond to avoid high CPU usage
4863 {-# INLINE delayTime #-}
4865 | g' >= fromIntegral (maxBound :: Int) = maxBound
4867 | otherwise = round g'
4871 g' = g * 10 ^ (6 :: Int)
4873-- XXX we should move this to stream generation section of this file. Also, the
4874-- take/drop combinators above should be moved to filtering section.
4876{-# INLINE_NORMAL times #-}
4877times :: MonadAsync m => Double -> Stream m (AbsTime, RelTime64)
4878times g = Stream step Nothing
4882 {-# INLINE_LATE step #-}
4884 -- XXX note that this is safe only on a 64-bit machine. On a 32-bit
4885 -- machine a 64-bit 'Var' cannot be read consistently without a lock
4886 -- while another thread is writing to it.
4887 timeVar <- liftIO $ Prim.newIORef (0 :: Int64)
4888 liftIO $ updateTimeVar timeVar
4889 tid <- forkManaged $ liftIO $ forever (updateWithDelay g timeVar)
4890 a <- liftIO $ Prim.readIORef timeVar
4891 return $ Skip $ Just (timeVar, tid, a)
4893 step _ s@(Just (timeVar, _, t0)) = do
4894 a <- liftIO $ Prim.readIORef timeVar
4895 -- XXX we can perhaps use an AbsTime64 using a 64 bit Int for
4896 -- efficiency. or maybe we can use a representation using Double for
4897 -- floating precision time
4898 return $ Yield (toAbsTime (MicroSecond64 t0),
4899 (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