1{-# LANGUAGE BangPatterns #-}
3{-# LANGUAGE ConstraintKinds #-}
4{-# LANGUAGE ExistentialQuantification #-}
5{-# LANGUAGE FlexibleContexts #-}
6{-# LANGUAGE FlexibleInstances #-}
7{-# LANGUAGE MultiParamTypeClasses #-}
8{-# LANGUAGE PatternSynonyms #-}
9{-# LANGUAGE RecordWildCards #-}
10{-# LANGUAGE ScopedTypeVariables #-}
11{-# LANGUAGE ViewPatterns #-}
12{-# LANGUAGE RankNTypes #-}
13{-# LANGUAGE MagicHash #-}
15#if __GLASGOW_HASKELL__ >= 801
16{-# LANGUAGE TypeApplications #-}
22-- Module : Streamly.Internal.Data.Stream.StreamD
23-- Copyright : (c) 2018 Harendra Kumar
24-- (c) Roman Leshchinskiy 2008-2010
25-- (c) The University of Glasgow, 2009
28-- Maintainer : streamly@composewell.com
29-- Stability : experimental
32-- Direct style re-implementation of CPS style stream in StreamK module. The
33-- symbol or suffix 'D' in this module denotes the "Direct" style. GHC is able
34-- to INLINE and fuse direct style better, providing better performance than
38-- import qualified Streamly.Internal.Data.Stream.StreamD as D
41-- Some of the functions in this file have been adapted from the vector
42-- library, https://hackage.haskell.org/package/vector.
44module Streamly.Internal.Data.Stream.StreamD
49#if __GLASGOW_HASKELL__ >= 800
50 , Stream (Stream, UnStream)
70 -- ** Specialized Generation
71 -- | Generate a monadic stream from a seed.
84 , enumerateFromStepIntegral
85 , enumerateFromIntegral
86 , enumerateFromThenIntegral
87 , enumerateFromToIntegral
88 , enumerateFromThenToIntegral
90 , enumerateFromStepNum
93 , enumerateFromToFractional
94 , enumerateFromThenToFractional
100 -- | Transform an input structure into a stream.
101 -- | Direct style stream does not support @fromFoldable@.
131 -- ** Specialized Folds
158 -- ** Flattening nested streams
161 , ConcatMapUState (..)
163 , ConcatUnfoldInterleaveState (..)
164 , concatUnfoldInterleave
165 , concatUnfoldRoundrobin
168 , InterleaveState(..)
173 , roundRobin -- interleaveFair?/ParallelFair
206 -- | Transform a stream into another type.
223 -- ** By folding (scans)
270 , intersperseSuffixBySpan
294 -- * Transformation comprehensions
310 -- * Concurrent Application
318import Control.Concurrent (killThread, myThreadId, takeMVar, threadDelay)
319import Control.Exception
320 (Exception, SomeException, AsyncException, fromException)
321import Control.Monad (void, when, forever)
322import Control.Monad.Catch (MonadCatch, throwM)
323import Control.Monad.IO.Class (MonadIO(..))
324import Control.Monad.Reader (ReaderT)
325import Control.Monad.State.Strict (StateT)
326import Control.Monad.Trans (MonadTrans(lift))
327import Control.Monad.Trans.Control (MonadBaseControl)
328import Data.Bits (shiftR, shiftL, (.|.), (.&.))
329import Data.Functor.Identity (Identity(..))
330import Data.Int (Int64)
331import Data.IORef (newIORef, readIORef, mkWeakIORef, writeIORef, IORef)
332import Data.Maybe (fromJust, isJust, isNothing)
333import Data.Word (Word32)
334import Foreign.Ptr (Ptr)
335import Foreign.Storable (Storable(..))
336import GHC.Types (SPEC(..))
337import System.Mem (performMajorGC)
339 hiding (map, mapM, mapM_, repeat, foldr, last, take, filter,
340 takeWhile, drop, dropWhile, all, any, maximum, minimum, elem,
341 notElem, null, head, tail, zipWith, lookup, foldr1, sequence,
342 (!!), scanl, scanl1, concatMap, replicate, enumFromTo, concat,
345import qualified Control.Monad.Catch as MC
346import qualified Control.Monad.Reader as Reader
347import qualified Control.Monad.State.Strict as State
348import qualified Prelude
350import Streamly.Internal.Mutable.Prim.Var
351 (Prim, Var, readVar, newVar, modifyVar')
352import Streamly.Internal.Data.Time.Units
353 (TimeUnit64, toRelTime64, diffAbsTime64)
355import Streamly.Internal.Data.Atomics (atomicModifyIORefCAS_)
356import Streamly.Internal.Memory.Array.Types (Array(..))
357import Streamly.Internal.Data.Fold.Types (Fold(..))
358import Streamly.Internal.Data.Pipe.Types (Pipe(..), PipeState(..))
359import Streamly.Internal.Data.Time.Clock (Clock(Monotonic), getTime)
360import Streamly.Internal.Data.Time.Units
361 (MicroSecond64(..), fromAbsTime, toAbsTime, AbsTime)
362import Streamly.Internal.Data.Unfold.Types (Unfold(..))
363import Streamly.Internal.Data.Strict (Tuple3'(..))
365import Streamly.Internal.Data.Stream.StreamD.Type
366import Streamly.Internal.Data.SVar
367import Streamly.Internal.Data.Stream.SVar (fromConsumer, pushToFold)
369import qualified Streamly.Internal.Data.Pipe.Types as Pipe
370import qualified Streamly.Internal.Memory.Array.Types as A
371import qualified Streamly.Internal.Data.Fold as FL
372import qualified Streamly.Memory.Ring as RB
373import qualified Streamly.Internal.Data.Stream.StreamK as K
375------------------------------------------------------------------------------
377------------------------------------------------------------------------------
379-- | An empty 'Stream'.
380{-# INLINE_NORMAL nil #-}
381nil :: Monad m => Stream m a
382nil = Stream (\_ _ -> return Stop) ()
384-- | An empty 'Stream' with a side effect.
385{-# INLINE_NORMAL nilM #-}
386nilM :: Monad m => m b -> Stream m a
387nilM m = Stream (\_ _ -> m >> return Stop) ()
389{-# INLINE_NORMAL consM #-}
390consM :: Monad m => m a -> Stream m a -> Stream m a
391consM m (Stream step state) = Stream step1 Nothing
393 {-# INLINE_LATE step1 #-}
394 step1 _ Nothing = m >>= \x -> return $ Yield x (Just state)
395 step1 gst (Just st) = do
399 Yield a s -> Yield a (Just s)
400 Skip s -> Skip (Just s)
403-- XXX implement in terms of consM?
404-- cons x = consM (return x)
406-- | Can fuse but has O(n^2) complexity.
407{-# INLINE_NORMAL cons #-}
408cons :: Monad m => a -> Stream m a -> Stream m a
409cons x (Stream step state) = Stream step1 Nothing
411 {-# INLINE_LATE step1 #-}
412 step1 _ Nothing = return $ Yield x (Just state)
413 step1 gst (Just st) = do
417 Yield a s -> Yield a (Just s)
418 Skip s -> Skip (Just s)
421-------------------------------------------------------------------------------
423-------------------------------------------------------------------------------
425-- Does not fuse, has the same performance as the StreamK version.
426{-# INLINE_NORMAL uncons #-}
427uncons :: Monad m => Stream m a -> m (Maybe (a, Stream m a))
428uncons (UnStream step state) = go state
431 r <- step defState st
433 Yield x s -> return $ Just (x, Stream step s)
435 Stop -> return Nothing
437------------------------------------------------------------------------------
438-- Generation by unfold
439------------------------------------------------------------------------------
441{-# INLINE_NORMAL unfoldrM #-}
442unfoldrM :: Monad m => (s -> m (Maybe (a, s))) -> s -> Stream m a
443unfoldrM next state = Stream step state
445 {-# INLINE_LATE step #-}
449 Just (x, s) -> Yield x s
452{-# INLINE_LATE unfoldr #-}
453unfoldr :: Monad m => (s -> Maybe (a, s)) -> s -> Stream m a
454unfoldr f = unfoldrM (return . f)
456-- | Convert an 'Unfold' into a 'Stream' by supplying it a seed.
458{-# INLINE_NORMAL unfold #-}
459unfold :: Monad m => Unfold m a b -> a -> Stream m b
460unfold (Unfold ustep inject) seed = Stream step Nothing
462 {-# INLINE_LATE step #-}
463 step _ Nothing = inject seed >>= return . Skip . Just
464 step _ (Just st) = do
467 Yield x s -> Yield x (Just s)
468 Skip s -> Skip (Just s)
471------------------------------------------------------------------------------
472-- Specialized Generation
473------------------------------------------------------------------------------
475{-# INLINE_NORMAL repeatM #-}
476repeatM :: Monad m => m a -> Stream m a
477repeatM x = Stream (\_ _ -> x >>= \r -> return $ Yield r ()) ()
479{-# INLINE_NORMAL repeat #-}
480repeat :: Monad m => a -> Stream m a
481repeat x = Stream (\_ _ -> return $ Yield x ()) ()
483{-# INLINE_NORMAL iterateM #-}
484iterateM :: Monad m => (a -> m a) -> m a -> Stream m a
485iterateM step = Stream (\_ st -> st >>= \x -> return $ Yield x (step x))
487{-# INLINE_NORMAL iterate #-}
488iterate :: Monad m => (a -> a) -> a -> Stream m a
489iterate step st = iterateM (return . step) (return st)
491{-# INLINE_NORMAL replicateM #-}
492replicateM :: forall m a. Monad m => Int -> m a -> Stream m a
493replicateM n p = Stream step n
495 {-# INLINE_LATE step #-}
497 | i <= 0 = return Stop
500 return $ Yield x (i - 1)
502{-# INLINE_NORMAL replicate #-}
503replicate :: Monad m => Int -> a -> Stream m a
504replicate n x = replicateM n (return x)
506-- This would not work properly for floats, therefore we put an Integral
508-- | Can be used to enumerate unbounded integrals. This does not check for
509-- overflow or underflow for bounded integrals.
510{-# INLINE_NORMAL enumerateFromStepIntegral #-}
511enumerateFromStepIntegral :: (Integral a, Monad m) => a -> a -> Stream m a
512enumerateFromStepIntegral from stride =
513 from `seq` stride `seq` Stream step from
515 {-# INLINE_LATE step #-}
516 step _ !x = return $ Yield x $! (x + stride)
518-- We are assuming that "to" is constrained by the type to be within
520{-# INLINE enumerateFromToIntegral #-}
521enumerateFromToIntegral :: (Monad m, Integral a) => a -> a -> Stream m a
522enumerateFromToIntegral from to =
523 takeWhile (<= to) $ enumerateFromStepIntegral from 1
525{-# INLINE enumerateFromIntegral #-}
526enumerateFromIntegral :: (Monad m, Integral a, Bounded a) => a -> Stream m a
527enumerateFromIntegral from = enumerateFromToIntegral from maxBound
529data EnumState a = EnumInit | EnumYield a a a | EnumStop
531{-# INLINE_NORMAL enumerateFromThenToIntegralUp #-}
532enumerateFromThenToIntegralUp
533 :: (Monad m, Integral a)
534 => a -> a -> a -> Stream m a
535enumerateFromThenToIntegralUp from next to = Stream step EnumInit
537 {-# INLINE_LATE step #-}
543 else Yield from EnumStop
544 else -- from <= next <= to
545 let stride = next - from
546 in Skip $ EnumYield from stride (to - stride)
548 step _ (EnumYield x stride toMinus) =
551 then Yield x EnumStop
552 else Yield x $ EnumYield (x + stride) stride toMinus
554 step _ EnumStop = return Stop
556{-# INLINE_NORMAL enumerateFromThenToIntegralDn #-}
557enumerateFromThenToIntegralDn
558 :: (Monad m, Integral a)
559 => a -> a -> a -> Stream m a
560enumerateFromThenToIntegralDn from next to = Stream step EnumInit
562 {-# INLINE_LATE step #-}
564 return $ if to > next
567 else Yield from EnumStop
568 else -- from >= next >= to
569 let stride = next - from
570 in Skip $ EnumYield from stride (to - stride)
572 step _ (EnumYield x stride toMinus) =
575 then Yield x EnumStop
576 else Yield x $ EnumYield (x + stride) stride toMinus
578 step _ EnumStop = return Stop
580{-# INLINE_NORMAL enumerateFromThenToIntegral #-}
581enumerateFromThenToIntegral
582 :: (Monad m, Integral a)
583 => a -> a -> a -> Stream m a
584enumerateFromThenToIntegral from next to
585 | next >= from = enumerateFromThenToIntegralUp from next to
586 | otherwise = enumerateFromThenToIntegralDn from next to
588{-# INLINE_NORMAL enumerateFromThenIntegral #-}
589enumerateFromThenIntegral
590 :: (Monad m, Integral a, Bounded a)
591 => a -> a -> Stream m a
592enumerateFromThenIntegral from next =
594 then enumerateFromThenToIntegralUp from next maxBound
595 else enumerateFromThenToIntegralDn from next minBound
597-- For floating point numbers if the increment is less than the precision then
598-- it just gets lost. Therefore we cannot always increment it correctly by just
600-- 9007199254740992 + 1 + 1 :: Double => 9.007199254740992e15
601-- 9007199254740992 + 2 :: Double => 9.007199254740994e15
603-- Instead we accumulate the increment counter and compute the increment
604-- every time before adding it to the starting number.
606-- This works for Integrals as well as floating point numbers, but
607-- enumerateFromStepIntegral is faster for integrals.
608{-# INLINE_NORMAL enumerateFromStepNum #-}
609enumerateFromStepNum :: (Monad m, Num a) => a -> a -> Stream m a
610enumerateFromStepNum from stride = Stream step 0
612 {-# INLINE_LATE step #-}
613 step _ !i = return $ (Yield $! (from + i * stride)) $! (i + 1)
615{-# INLINE_NORMAL numFrom #-}
616numFrom :: (Monad m, Num a) => a -> Stream m a
617numFrom from = enumerateFromStepNum from 1
619{-# INLINE_NORMAL numFromThen #-}
620numFromThen :: (Monad m, Num a) => a -> a -> Stream m a
621numFromThen from next = enumerateFromStepNum from (next - from)
623-- We cannot write a general function for Num. The only way to write code
624-- portable between the two is to use a 'Real' constraint and convert between
625-- Fractional and Integral using fromRational which is horribly slow.
626{-# INLINE_NORMAL enumerateFromToFractional #-}
627enumerateFromToFractional
628 :: (Monad m, Fractional a, Ord a)
629 => a -> a -> Stream m a
630enumerateFromToFractional from to =
631 takeWhile (<= to + 1 / 2) $ enumerateFromStepNum from 1
633{-# INLINE_NORMAL enumerateFromThenToFractional #-}
634enumerateFromThenToFractional
635 :: (Monad m, Fractional a, Ord a)
636 => a -> a -> a -> Stream m a
637enumerateFromThenToFractional from next to =
638 takeWhile predicate $ numFromThen from next
640 mid = (next - from) / 2
641 predicate | next >= from = (<= to + mid)
642 | otherwise = (>= to + mid)
644-------------------------------------------------------------------------------
645-- Generation by Conversion
646-------------------------------------------------------------------------------
648{-# INLINE_NORMAL fromIndicesM #-}
649fromIndicesM :: Monad m => (Int -> m a) -> Stream m a
650fromIndicesM gen = Stream step 0
652 {-# INLINE_LATE step #-}
655 return $ Yield x (i + 1)
657{-# INLINE fromIndices #-}
658fromIndices :: Monad m => (Int -> a) -> Stream m a
659fromIndices gen = fromIndicesM (return . gen)
661{-# INLINE_NORMAL generateM #-}
662generateM :: Monad m => Int -> (Int -> m a) -> Stream m a
663generateM n gen = n `seq` Stream step 0
665 {-# INLINE_LATE step #-}
666 step _ i | i < n = do
668 return $ Yield x (i + 1)
669 | otherwise = return Stop
671{-# INLINE generate #-}
672generate :: Monad m => Int -> (Int -> a) -> Stream m a
673generate n gen = generateM n (return . gen)
675-- XXX we need the MonadAsync constraint because of a rewrite rule.
676-- | Convert a list of monadic actions to a 'Stream'
677{-# INLINE_LATE fromListM #-}
678fromListM :: MonadAsync m => [m a] -> Stream m a
679fromListM = Stream step
681 {-# INLINE_LATE step #-}
682 step _ (m:ms) = m >>= \x -> return $ Yield x ms
683 step _ [] = return Stop
685{-# INLINE toStreamD #-}
686toStreamD :: (K.IsStream t, Monad m) => t m a -> Stream m a
687toStreamD = fromStreamK . K.toStream
689{-# INLINE_NORMAL fromPrimVar #-}
690fromPrimVar :: (MonadIO m, Prim a) => Var IO a -> Stream m a
691fromPrimVar var = Stream step ()
693 {-# INLINE_LATE step #-}
694 step _ () = liftIO (readVar var) >>= \x -> return $ Yield x ()
696-------------------------------------------------------------------------------
697-- Generation from SVar
698-------------------------------------------------------------------------------
700data FromSVarState t m a =
702 | FromSVarRead (SVar t m a)
703 | FromSVarLoop (SVar t m a) [ChildEvent a]
704 | FromSVarDone (SVar t m a)
706{-# INLINE_NORMAL fromSVar #-}
707fromSVar :: (MonadAsync m) => SVar t m a -> Stream m a
708fromSVar svar = Stream step FromSVarInit
711 {-# INLINE_LATE step #-}
712 step _ FromSVarInit = do
713 ref <- liftIO $ newIORef ()
714 _ <- liftIO $ mkWeakIORef ref hook
715 -- when this copy of svar gets garbage collected "ref" will get
716 -- garbage collected and our GC hook will be called.
717 let sv = svar{svarRef = Just ref}
718 return $ Skip (FromSVarRead sv)
722 {-# NOINLINE hook #-}
724 when (svarInspectMode svar) $ do
725 r <- liftIO $ readIORef (svarStopTime (svarStats svar))
727 printSVar svar "SVar Garbage Collected"
729 -- If there are any SVars referenced by this SVar a GC will prompt
730 -- them to be cleaned up quickly.
731 when (svarInspectMode svar) performMajorGC
733 step _ (FromSVarRead sv) = do
734 list <- readOutputQ sv
735 -- Reversing the output is important to guarantee that we process the
736 -- outputs in the same order as they were generated by the constituent
738 return $ Skip $ FromSVarLoop sv (Prelude.reverse list)
740 step _ (FromSVarLoop sv []) = do
741 done <- postProcess sv
742 return $ Skip $ if done
743 then (FromSVarDone sv)
744 else (FromSVarRead sv)
746 step _ (FromSVarLoop sv (ev : es)) = do
748 ChildYield a -> return $ Yield a (FromSVarLoop sv es)
749 ChildStop tid e -> do
753 stop <- shouldStop tid
756 liftIO (cleanupSVar sv)
757 return $ Skip (FromSVarDone sv)
758 else return $ Skip (FromSVarLoop sv es)
760 case fromException ex of
762 return $ Skip (FromSVarLoop sv es)
763 Nothing -> liftIO (cleanupSVar sv) >> throwM ex
767 case svarStopStyle sv of
768 StopNone -> return False
769 StopAny -> return True
771 sid <- liftIO $ readIORef (svarStopBy sv)
772 return $ if tid == sid then True else False
774 step _ (FromSVarDone sv) = do
775 when (svarInspectMode sv) $ do
776 t <- liftIO $ getTime Monotonic
777 liftIO $ writeIORef (svarStopTime (svarStats sv)) (Just t)
778 liftIO $ printSVar sv "SVar Done"
781-------------------------------------------------------------------------------
782-- Process events received by a fold consumer from a stream producer
783-------------------------------------------------------------------------------
785{-# INLINE_NORMAL fromProducer #-}
786fromProducer :: (MonadAsync m) => SVar t m a -> Stream m a
787fromProducer svar = Stream step (FromSVarRead svar)
790 {-# INLINE_LATE step #-}
791 step _ (FromSVarRead sv) = do
792 list <- readOutputQ sv
793 -- Reversing the output is important to guarantee that we process the
794 -- outputs in the same order as they were generated by the constituent
796 return $ Skip $ FromSVarLoop sv (Prelude.reverse list)
798 step _ (FromSVarLoop sv []) = return $ Skip $ FromSVarRead sv
799 step _ (FromSVarLoop sv (ev : es)) = do
801 ChildYield a -> return $ Yield a (FromSVarLoop sv es)
802 ChildStop tid e -> do
806 sendStopToProducer sv
807 return $ Skip (FromSVarDone sv)
808 Just _ -> error "Bug: fromProducer: received exception"
810 step _ (FromSVarDone sv) = do
811 when (svarInspectMode sv) $ do
812 t <- liftIO $ getTime Monotonic
813 liftIO $ writeIORef (svarStopTime (svarStats sv)) (Just t)
814 liftIO $ printSVar sv "SVar Done"
817 step _ FromSVarInit = undefined
819-------------------------------------------------------------------------------
820-- Hoisting the inner monad
821-------------------------------------------------------------------------------
823{-# INLINE_NORMAL hoist #-}
824hoist :: Monad n => (forall x. m x -> n x) -> Stream m a -> Stream n a
825hoist f (Stream step state) = (Stream step' state)
827 {-# INLINE_LATE step' #-}
829 r <- f $ step (adaptState gst) st
831 Yield x s -> Yield x s
835{-# INLINE generally #-}
836generally :: Monad m => Stream Identity a -> Stream m a
837generally = hoist (return . runIdentity)
839{-# INLINE_NORMAL liftInner #-}
840liftInner :: (Monad m, MonadTrans t, Monad (t m))
841 => Stream m a -> Stream (t m) a
842liftInner (Stream step state) = Stream step' state
844 {-# INLINE_LATE step' #-}
846 r <- lift $ step (adaptState gst) st
848 Yield x s -> Yield x s
852{-# INLINE_NORMAL runReaderT #-}
853runReaderT :: Monad m => s -> Stream (ReaderT s m) a -> Stream m a
854runReaderT sval (Stream step state) = Stream step' state
856 {-# INLINE_LATE step' #-}
858 r <- Reader.runReaderT (step (adaptState gst) st) sval
860 Yield x s -> Yield x s
864{-# INLINE_NORMAL evalStateT #-}
865evalStateT :: Monad m => s -> Stream (StateT s m) a -> Stream m a
866evalStateT sval (Stream step state) = Stream step' (state, sval)
868 {-# INLINE_LATE step' #-}
869 step' gst (st, sv) = do
870 (r, sv') <- State.runStateT (step (adaptState gst) st) sv
872 Yield x s -> Yield x (s, sv')
873 Skip s -> Skip (s, sv')
876{-# INLINE_NORMAL runStateT #-}
877runStateT :: Monad m => s -> Stream (StateT s m) a -> Stream m (s, a)
878runStateT sval (Stream step state) = Stream step' (state, sval)
880 {-# INLINE_LATE step' #-}
881 step' gst (st, sv) = do
882 (r, sv') <- State.runStateT (step (adaptState gst) st) sv
884 Yield x s -> Yield (sv', x) (s, sv')
885 Skip s -> Skip (s, sv')
888------------------------------------------------------------------------------
889-- Elimination by Folds
890------------------------------------------------------------------------------
892------------------------------------------------------------------------------
894------------------------------------------------------------------------------
896{-# INLINE_NORMAL foldr1 #-}
897foldr1 :: Monad m => (a -> a -> a) -> Stream m a -> m (Maybe a)
901 Nothing -> return Nothing
902 Just (h, t) -> fmap Just (foldr f h t)
904------------------------------------------------------------------------------
906------------------------------------------------------------------------------
908{-# INLINE_NORMAL foldlT #-}
909foldlT :: (Monad m, Monad (s m), MonadTrans s)
910 => (s m b -> a -> s m b) -> s m b -> Stream m a -> s m b
911foldlT fstep begin (Stream step state) = go SPEC begin state
914 r <- lift $ step defState st
916 Yield x s -> go SPEC (fstep acc x) s
917 Skip s -> go SPEC acc s
920-- Note, this is going to have horrible performance, because of the nature of
921-- the stream type (i.e. direct stream vs CPS). Its only for reference, it is
922-- likely be practically unusable.
923{-# INLINE_NORMAL foldlS #-}
925 => (Stream m b -> a -> Stream m b) -> Stream m b -> Stream m a -> Stream m b
926foldlS fstep begin (Stream step state) = Stream step' (Left (state, begin))
928 step' gst (Left (st, acc)) = do
929 r <- step (adaptState gst) st
931 Yield x s -> Skip (Left (s, fstep acc x))
932 Skip s -> Skip (Left (s, acc))
933 Stop -> Skip (Right acc)
935 step' gst (Right (Stream stp stt)) = do
936 r <- stp (adaptState gst) stt
938 Yield x s -> Yield x (Right (Stream stp s))
939 Skip s -> Skip (Right (Stream stp s))
942------------------------------------------------------------------------------
944------------------------------------------------------------------------------
946-- | Run a streaming composition, discard the results.
947{-# INLINE_LATE drain #-}
948drain :: Monad m => Stream m a -> m ()
949-- drain = foldrM (\_ xs -> xs) (return ())
950drain (Stream step state) = go SPEC state
953 r <- step defState st
955 Yield _ s -> go SPEC s
959{-# INLINE_NORMAL null #-}
960null :: Monad m => Stream m a -> m Bool
961null m = foldrM (\_ _ -> return False) (return True) m
963{-# INLINE_NORMAL head #-}
964head :: Monad m => Stream m a -> m (Maybe a)
965head m = foldrM (\x _ -> return (Just x)) (return Nothing) m
967{-# INLINE_NORMAL headElse #-}
968headElse :: Monad m => a -> Stream m a -> m a
969headElse a m = foldrM (\x _ -> return x) (return a) m
971-- Does not fuse, has the same performance as the StreamK version.
972{-# INLINE_NORMAL tail #-}
973tail :: Monad m => Stream m a -> m (Maybe (Stream m a))
974tail (UnStream step state) = go state
977 r <- step defState st
979 Yield _ s -> return (Just $ Stream step s)
981 Stop -> return Nothing
983-- XXX will it fuse? need custom impl?
984{-# INLINE_NORMAL last #-}
985last :: Monad m => Stream m a -> m (Maybe a)
986last = foldl' (\_ y -> Just y) Nothing
988{-# INLINE_NORMAL elem #-}
989elem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
990-- elem e m = foldrM (\x xs -> if x == e then return True else xs) (return False) m
991elem e (Stream step state) = go state
994 r <- step defState st
997 | x == e -> return True
1000 Stop -> return False
1002{-# INLINE_NORMAL notElem #-}
1003notElem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
1004notElem e s = fmap not (elem e s)
1006{-# INLINE_NORMAL all #-}
1007all :: Monad m => (a -> Bool) -> Stream m a -> m Bool
1008-- all p m = foldrM (\x xs -> if p x then xs else return False) (return True) m
1009all p (Stream step state) = go state
1012 r <- step defState st
1016 | otherwise -> return False
1020{-# INLINE_NORMAL any #-}
1021any :: Monad m => (a -> Bool) -> Stream m a -> m Bool
1022-- any p m = foldrM (\x xs -> if p x then return True else xs) (return False) m
1023any p (Stream step state) = go state
1026 r <- step defState st
1029 | p x -> return True
1032 Stop -> return False
1034{-# INLINE_NORMAL maximum #-}
1035maximum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
1036maximum (Stream step state) = go Nothing state
1039 r <- step defState st
1041 Yield x s -> go (Just x) s
1042 Skip s -> go Nothing s
1043 Stop -> return Nothing
1044 go (Just acc) st = do
1045 r <- step defState st
1048 | acc <= x -> go (Just x) s
1049 | otherwise -> go (Just acc) s
1050 Skip s -> go (Just acc) s
1051 Stop -> return (Just acc)
1053{-# INLINE_NORMAL maximumBy #-}
1054maximumBy :: Monad m => (a -> a -> Ordering) -> Stream m a -> m (Maybe a)
1055maximumBy cmp (Stream step state) = go Nothing state
1058 r <- step defState st
1060 Yield x s -> go (Just x) s
1061 Skip s -> go Nothing s
1062 Stop -> return Nothing
1063 go (Just acc) st = do
1064 r <- step defState st
1066 Yield x s -> case cmp acc x of
1067 GT -> go (Just acc) s
1069 Skip s -> go (Just acc) s
1070 Stop -> return (Just acc)
1072{-# INLINE_NORMAL minimum #-}
1073minimum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
1074minimum (Stream step state) = go Nothing state
1077 r <- step defState st
1079 Yield x s -> go (Just x) s
1080 Skip s -> go Nothing s
1081 Stop -> return Nothing
1082 go (Just acc) st = do
1083 r <- step defState st
1086 | acc <= x -> go (Just acc) s
1087 | otherwise -> go (Just x) s
1088 Skip s -> go (Just acc) s
1089 Stop -> return (Just acc)
1091{-# INLINE_NORMAL minimumBy #-}
1092minimumBy :: Monad m => (a -> a -> Ordering) -> Stream m a -> m (Maybe a)
1093minimumBy cmp (Stream step state) = go Nothing state
1096 r <- step defState st
1098 Yield x s -> go (Just x) s
1099 Skip s -> go Nothing s
1100 Stop -> return Nothing
1101 go (Just acc) st = do
1102 r <- step defState st
1104 Yield x s -> case cmp acc x of
1106 _ -> go (Just acc) s
1107 Skip s -> go (Just acc) s
1108 Stop -> return (Just acc)
1110{-# INLINE_NORMAL (!!) #-}
1111(!!) :: (Monad m) => Stream m a -> Int -> m (Maybe a)
1112(Stream step state) !! i = go i state
1115 r <- step defState st
1117 Yield x s | n < 0 -> return Nothing
1118 | n == 0 -> return $ Just x
1119 | otherwise -> go (n - 1) s
1121 Stop -> return Nothing
1123{-# INLINE_NORMAL lookup #-}
1124lookup :: (Monad m, Eq a) => a -> Stream m (a, b) -> m (Maybe b)
1125lookup e m = foldrM (\(a, b) xs -> if e == a then return (Just b) else xs)
1128{-# INLINE_NORMAL findM #-}
1129findM :: Monad m => (a -> m Bool) -> Stream m a -> m (Maybe a)
1130findM p m = foldrM (\x xs -> p x >>= \r -> if r then return (Just x) else xs)
1134find :: Monad m => (a -> Bool) -> Stream m a -> m (Maybe a)
1135find p = findM (return . p)
1137{-# INLINE_NORMAL findIndices #-}
1138findIndices :: Monad m => (a -> Bool) -> Stream m a -> Stream m Int
1139findIndices p (Stream step state) = Stream step' (state, 0)
1141 {-# INLINE_LATE step' #-}
1142 step' gst (st, i) = i `seq` do
1143 r <- step (adaptState gst) st
1145 Yield x s -> if p x then Yield i (s, i+1) else Skip (s, i+1)
1146 Skip s -> Skip (s, i)
1149{-# INLINE toListRev #-}
1150toListRev :: Monad m => Stream m a -> m [a]
1151toListRev = foldl' (flip (:)) []
1153-- We can implement reverse as:
1155-- > reverse = foldlS (flip cons) nil
1157-- However, this implementation is unusable because of the horrible performance
1158-- of cons. So we just convert it to a list first and then stream from the
1161-- XXX Maybe we can use an Array instead of a list here?
1162{-# INLINE_NORMAL reverse #-}
1163reverse :: Monad m => Stream m a -> Stream m a
1164reverse m = Stream step Nothing
1166 {-# INLINE_LATE step #-}
1169 return $ Skip (Just xs)
1170 step _ (Just (x:xs)) = return $ Yield x (Just xs)
1171 step _ (Just []) = return Stop
1173-- Much faster reverse for Storables
1174{-# INLINE_NORMAL reverse' #-}
1175reverse' :: forall m a. (MonadIO m, Storable a) => Stream m a -> Stream m a
1177-- This commented implementation copies the whole stream into one single array
1178-- and then streams from that array, this is 3-4x faster than the chunked code
1179-- that follows. Though this could be problematic due to unbounded large
1180-- allocations. We need to figure out why the chunked code is slower and if we
1181-- can optimize the chunked code to work as fast as this one. It may be a
1183import Foreign.ForeignPtr (touchForeignPtr)
1184import Foreign.ForeignPtr.Unsafe (unsafeForeignPtrToPtr)
1185import Foreign.Ptr (Ptr, plusPtr)
1186reverse' m = Stream step Nothing
1188 {-# INLINE_LATE step #-}
1190 arr <- A.fromStreamD m
1191 let p = aEnd arr `plusPtr` negate (sizeOf (undefined :: a))
1192 return $ Skip $ Just (aStart arr, p)
1194 step _ (Just (start, p)) | p < unsafeForeignPtrToPtr start = return Stop
1196 step _ (Just (start, p)) = do
1197 let !x = A.unsafeInlineIO $ do
1199 touchForeignPtr start
1201 next = p `plusPtr` negate (sizeOf (undefined :: a))
1202 return $ Yield x (Just (start, next))
1209 $ A.fromStreamDArraysOf A.defaultChunkSize m
1212------------------------------------------------------------------------------
1213-- Grouping/Splitting
1214------------------------------------------------------------------------------
1216{-# INLINE_NORMAL splitSuffixBy' #-}
1217splitSuffixBy' :: Monad m
1218 => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1219splitSuffixBy' predicate f (Stream step state) =
1220 Stream (stepOuter f) (Just state)
1224 {-# INLINE_LATE stepOuter #-}
1225 stepOuter (Fold fstep initial done) gst (Just st) = do
1226 res <- step (adaptState gst) st
1232 then done acc' >>= \val -> return $ Yield val (Just s)
1235 Skip s -> return $ Skip $ Just s
1241 res <- step (adaptState gst) stt
1246 then done acc' >>= \val -> return $ Yield val (Just s)
1248 Skip s -> go SPEC s acc
1249 Stop -> done acc >>= \val -> return $ Yield val Nothing
1251 stepOuter _ _ Nothing = return Stop
1253{-# INLINE_NORMAL groupsBy #-}
1259groupsBy cmp f (Stream step state) = Stream (stepOuter f) (Just state, Nothing)
1263 {-# INLINE_LATE stepOuter #-}
1264 stepOuter (Fold fstep initial done) gst (Just st, Nothing) = do
1265 res <- step (adaptState gst) st
1272 Skip s -> return $ Skip $ (Just s, Nothing)
1277 go !_ prev stt !acc = do
1278 res <- step (adaptState gst) stt
1285 else done acc >>= \r -> return $ Yield r (Just s, Just x)
1286 Skip s -> go SPEC prev s acc
1287 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1289 stepOuter (Fold fstep initial done) gst (Just st, Just prev) = do
1291 acc' <- fstep acc prev
1296 -- XXX code duplicated from the previous equation
1298 res <- step (adaptState gst) stt
1305 else done acc >>= \r -> return $ Yield r (Just s, Just x)
1306 Skip s -> go SPEC s acc
1307 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1309 stepOuter _ _ (Nothing,_) = return Stop
1311{-# INLINE_NORMAL groupsRollingBy #-}
1312groupsRollingBy :: Monad m
1317groupsRollingBy cmp f (Stream step state) =
1318 Stream (stepOuter f) (Just state, Nothing)
1321 {-# INLINE_LATE stepOuter #-}
1322 stepOuter (Fold fstep initial done) gst (Just st, Nothing) = do
1323 res <- step (adaptState gst) st
1330 Skip s -> return $ Skip $ (Just s, Nothing)
1334 go !_ prev stt !acc = do
1335 res <- step (adaptState gst) stt
1343 done acc >>= \r -> return $ Yield r (Just s, Just x)
1344 Skip s -> go SPEC prev s acc
1345 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1347 stepOuter (Fold fstep initial done) gst (Just st, Just prev') = do
1349 acc' <- fstep acc prev'
1350 go SPEC prev' st acc'
1353 go !_ prevv stt !acc = do
1354 res <- step (adaptState gst) stt
1361 else done acc >>= \r -> return $ Yield r (Just s, Just x)
1362 Skip s -> go SPEC prevv s acc
1363 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1365 stepOuter _ _ (Nothing, _) = return Stop
1367{-# INLINE_NORMAL splitBy #-}
1368splitBy :: Monad m => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1369splitBy predicate f (Stream step state) = Stream (step' f) (Just state)
1373 {-# INLINE_LATE step' #-}
1374 step' (Fold fstep initial done) gst (Just st) = initial >>= go SPEC st
1379 res <- step (adaptState gst) stt
1383 then done acc >>= \r -> return $ Yield r (Just s)
1387 Skip s -> go SPEC s acc
1388 Stop -> done acc >>= \r -> return $ Yield r Nothing
1390 step' _ _ Nothing = return Stop
1392-- XXX requires -funfolding-use-threshold=150 in lines-unlines benchmark
1393{-# INLINE_NORMAL splitSuffixBy #-}
1394splitSuffixBy :: Monad m
1395 => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1396splitSuffixBy predicate f (Stream step state) = Stream (step' f) (Just state)
1400 {-# INLINE_LATE step' #-}
1401 step' (Fold fstep initial done) gst (Just st) = do
1402 res <- step (adaptState gst) st
1407 then done acc >>= \val -> return $ Yield val (Just s)
1412 Skip s -> return $ Skip $ Just s
1418 res <- step (adaptState gst) stt
1422 then done acc >>= \r -> return $ Yield r (Just s)
1426 Skip s -> go SPEC s acc
1427 Stop -> done acc >>= \r -> return $ Yield r Nothing
1429 step' _ _ Nothing = return Stop
1431{-# INLINE_NORMAL wordsBy #-}
1432wordsBy :: Monad m => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1433wordsBy predicate f (Stream step state) = Stream (stepOuter f) (Just state)
1437 {-# INLINE_LATE stepOuter #-}
1438 stepOuter (Fold fstep initial done) gst (Just st) = do
1439 res <- step (adaptState gst) st
1443 then return $ Skip (Just s)
1449 Skip s -> return $ Skip $ Just s
1455 res <- step (adaptState gst) stt
1459 then done acc >>= \r -> return $ Yield r (Just s)
1463 Skip s -> go SPEC s acc
1464 Stop -> done acc >>= \r -> return $ Yield r Nothing
1466 stepOuter _ _ Nothing = return Stop
1468-- String search algorithms:
1469-- http://www-igm.univ-mlv.fr/~lecroq/string/index.html
1472-- TODO can we unify the splitting operations using a splitting configuration
1473-- like in the split package.
1475data SplitStyle = Infix | Suffix | Prefix deriving (Eq, Show)
1477data SplitOptions = SplitOptions
1478 { style :: SplitStyle
1479 , withSep :: Bool -- ^ keep the separators in output
1480 -- , compact :: Bool -- ^ treat multiple consecutive separators as one
1481 -- , trimHead :: Bool -- ^ drop blank at head
1482 -- , trimTail :: Bool -- ^ drop blank at tail
1486data SplitOnState s a =
1491 | GO_KARP_RABIN s !(RB.Ring a) !(Ptr a)
1494{-# INLINE_NORMAL splitOn #-}
1496 :: forall m a b. (MonadIO m, Storable a, Enum a, Eq a)
1501splitOn patArr@Array{..} (Fold fstep initial done) (Stream step state) =
1502 Stream stepOuter GO_START
1506 patLen = A.length patArr
1507 maxIndex = patLen - 1
1508 elemBits = sizeOf (undefined :: a) * 8
1510 {-# INLINE_LATE stepOuter #-}
1511 stepOuter _ GO_START =
1513 then return $ Skip $ GO_EMPTY_PAT state
1516 r <- liftIO $ (A.unsafeIndexIO patArr 0)
1517 return $ Skip $ GO_SINGLE_PAT state r
1518 else if sizeOf (undefined :: a) * patLen
1519 <= sizeOf (undefined :: Word)
1520 then return $ Skip $ GO_SHORT_PAT state
1522 (rb, rhead) <- liftIO $ RB.new patLen
1523 return $ Skip $ GO_KARP_RABIN state rb rhead
1525 stepOuter gst (GO_SINGLE_PAT stt pat) = initial >>= go SPEC stt
1530 res <- step (adaptState gst) st
1536 return $ Yield r (GO_SINGLE_PAT s pat)
1537 else fstep acc x >>= go SPEC s
1538 Skip s -> go SPEC s acc
1539 Stop -> done acc >>= \r -> return $ Yield r GO_DONE
1541 stepOuter gst (GO_SHORT_PAT stt) = initial >>= go0 SPEC 0 (0 :: Word) stt
1546 mask = (1 `shiftL` (elemBits * patLen)) - 1
1548 addToWord wrd a = (wrd `shiftL` elemBits) .|. fromIntegral (fromEnum a)
1551 patWord = mask .&. A.foldl' addToWord 0 patArr
1553 go0 !_ !idx wrd st !acc = do
1554 res <- step (adaptState gst) st
1557 let wrd' = addToWord wrd x
1560 if wrd' .&. mask == patWord
1563 return $ Yield r (GO_SHORT_PAT s)
1564 else go1 SPEC wrd' s acc
1565 else go0 SPEC (idx + 1) wrd' s acc
1566 Skip s -> go0 SPEC idx wrd s acc
1569 then go2 wrd idx acc
1571 done acc' >>= \r -> return $ Yield r GO_DONE
1574 go1 !_ wrd st !acc = do
1575 res <- step (adaptState gst) st
1578 let wrd' = addToWord wrd x
1579 old = (mask .&. wrd) `shiftR` (elemBits * (patLen - 1))
1580 acc' <- fstep acc (toEnum $ fromIntegral old)
1581 if wrd' .&. mask == patWord
1582 then done acc' >>= \r -> return $ Yield r (GO_SHORT_PAT s)
1583 else go1 SPEC wrd' s acc'
1584 Skip s -> go1 SPEC wrd s acc
1586 acc' <- go2 wrd patLen acc
1587 done acc' >>= \r -> return $ Yield r GO_DONE
1589 go2 !wrd !n !acc | n > 0 = do
1590 let old = (mask .&. wrd) `shiftR` (elemBits * (n - 1))
1591 fstep acc (toEnum $ fromIntegral old) >>= go2 wrd (n - 1)
1592 go2 _ _ acc = return acc
1594 stepOuter gst (GO_KARP_RABIN stt rb rhead) = do
1595 initial >>= go0 SPEC 0 rhead stt
1599 k = 2891336453 :: Word32
1601 addCksum cksum a = cksum * k + fromIntegral (fromEnum a)
1602 deltaCksum cksum old new =
1603 addCksum cksum new - coeff * fromIntegral (fromEnum old)
1605 -- XXX shall we use a random starting hash or 1 instead of 0?
1606 patHash = A.foldl' addCksum 0 patArr
1609 go0 !_ !idx !rh st !acc = do
1610 res <- step (adaptState gst) st
1613 rh' <- liftIO $ RB.unsafeInsert rb rh x
1616 let fold = RB.unsafeFoldRing (RB.ringBound rb)
1617 let !ringHash = fold addCksum 0 rb
1618 if ringHash == patHash
1619 then go2 SPEC ringHash rh' s acc
1620 else go1 SPEC ringHash rh' s acc
1621 else go0 SPEC (idx + 1) rh' s acc
1622 Skip s -> go0 SPEC idx rh s acc
1624 !acc' <- if idx /= 0
1625 then RB.unsafeFoldRingM rh fstep acc rb
1627 done acc' >>= \r -> return $ Yield r GO_DONE
1629 -- XXX Theoretically this code can do 4 times faster if GHC generates
1630 -- optimal code. If we use just "(cksum' == patHash)" condition it goes
1631 -- 4x faster, as soon as we add the "RB.unsafeEqArray rb v" condition
1632 -- the generated code changes drastically and becomes 4x slower. Need
1633 -- to investigate what is going on with GHC.
1635 go1 !_ !cksum !rh st !acc = do
1636 res <- step (adaptState gst) st
1639 old <- liftIO $ peek rh
1640 let cksum' = deltaCksum cksum old x
1641 acc' <- fstep acc old
1643 if (cksum' == patHash)
1645 rh' <- liftIO (RB.unsafeInsert rb rh x)
1646 go2 SPEC cksum' rh' s acc'
1648 rh' <- liftIO (RB.unsafeInsert rb rh x)
1649 go1 SPEC cksum' rh' s acc'
1650 Skip s -> go1 SPEC cksum rh s acc
1652 acc' <- RB.unsafeFoldRingFullM rh fstep acc rb
1653 done acc' >>= \r -> return $ Yield r GO_DONE
1655 go2 !_ !cksum' !rh' s !acc' = do
1656 if RB.unsafeEqArray rb rh' patArr
1659 return $ Yield r (GO_KARP_RABIN s rb rhead)
1660 else go1 SPEC cksum' rh' s acc'
1662 stepOuter gst (GO_EMPTY_PAT st) = do
1663 res <- step (adaptState gst) st
1668 done acc' >>= \r -> return $ Yield r (GO_EMPTY_PAT s)
1669 Skip s -> return $ Skip (GO_EMPTY_PAT s)
1672 stepOuter _ GO_DONE = return Stop
1674{-# INLINE_NORMAL splitSuffixOn #-}
1676 :: forall m a b. (MonadIO m, Storable a, Enum a, Eq a)
1682splitSuffixOn withSep patArr@Array{..} (Fold fstep initial done)
1683 (Stream step state) =
1684 Stream stepOuter GO_START
1688 patLen = A.length patArr
1689 maxIndex = patLen - 1
1690 elemBits = sizeOf (undefined :: a) * 8
1692 {-# INLINE_LATE stepOuter #-}
1693 stepOuter _ GO_START =
1695 then return $ Skip $ GO_EMPTY_PAT state
1698 r <- liftIO $ (A.unsafeIndexIO patArr 0)
1699 return $ Skip $ GO_SINGLE_PAT state r
1700 else if sizeOf (undefined :: a) * patLen
1701 <= sizeOf (undefined :: Word)
1702 then return $ Skip $ GO_SHORT_PAT state
1704 (rb, rhead) <- liftIO $ RB.new patLen
1705 return $ Skip $ GO_KARP_RABIN state rb rhead
1707 stepOuter gst (GO_SINGLE_PAT stt pat) = do
1708 -- This first part is the only difference between splitOn and
1710 -- If the last element is a separator do not issue a blank segment.
1711 res <- step (adaptState gst) stt
1717 acc' <- if withSep then fstep acc x else return acc
1718 done acc' >>= \r -> return $ Yield r (GO_SINGLE_PAT s pat)
1719 else fstep acc x >>= go SPEC s
1720 Skip s -> return $ Skip $ (GO_SINGLE_PAT s pat)
1725 -- This is identical for splitOn and splitSuffixOn
1727 res <- step (adaptState gst) st
1732 acc' <- if withSep then fstep acc x else return acc
1734 return $ Yield r (GO_SINGLE_PAT s pat)
1735 else fstep acc x >>= go SPEC s
1736 Skip s -> go SPEC s acc
1737 Stop -> done acc >>= \r -> return $ Yield r GO_DONE
1739 stepOuter gst (GO_SHORT_PAT stt) = do
1741 -- Call "initial" only if the stream yields an element, otherwise we
1742 -- may call "initial" but never yield anything. initial may produce a
1743 -- side effect, therefore we will end up doing and discard a side
1748 res <- step (adaptState gst) stt
1752 let wrd' = addToWord wrd x
1753 acc' <- if withSep then fstep acc x else return acc
1756 if wrd' .&. mask == patWord
1757 then done acc' >>= \r -> return $ Yield r (GO_SHORT_PAT s)
1758 else go0 SPEC (idx + 1) wrd' s acc'
1759 else go0 SPEC (idx + 1) wrd' s acc'
1760 Skip s -> return $ Skip (GO_SHORT_PAT s)
1766 mask = (1 `shiftL` (elemBits * patLen)) - 1
1768 addToWord wrd a = (wrd `shiftL` elemBits) .|. fromIntegral (fromEnum a)
1771 patWord = mask .&. A.foldl' addToWord 0 patArr
1773 go0 !_ !idx wrd st !acc = do
1774 res <- step (adaptState gst) st
1777 let wrd' = addToWord wrd x
1778 acc' <- if withSep then fstep acc x else return acc
1781 if wrd' .&. mask == patWord
1784 return $ Yield r (GO_SHORT_PAT s)
1785 else go1 SPEC wrd' s acc'
1786 else go0 SPEC (idx + 1) wrd' s acc'
1787 Skip s -> go0 SPEC idx wrd s acc
1789 if (idx == maxIndex) && (wrd .&. mask == patWord)
1792 acc' <- if idx /= 0 && not withSep
1793 then go2 wrd idx acc
1795 done acc' >>= \r -> return $ Yield r GO_DONE
1798 go1 !_ wrd st !acc = do
1799 res <- step (adaptState gst) st
1802 let wrd' = addToWord wrd x
1803 old = (mask .&. wrd) `shiftR` (elemBits * (patLen - 1))
1806 else fstep acc (toEnum $ fromIntegral old)
1807 if wrd' .&. mask == patWord
1808 then done acc' >>= \r -> return $ Yield r (GO_SHORT_PAT s)
1809 else go1 SPEC wrd' s acc'
1810 Skip s -> go1 SPEC wrd s acc
1812 -- If the last sequence is a separator do not issue a blank
1814 if wrd .&. mask == patWord
1819 else go2 wrd patLen acc
1820 done acc' >>= \r -> return $ Yield r GO_DONE
1822 go2 !wrd !n !acc | n > 0 = do
1823 let old = (mask .&. wrd) `shiftR` (elemBits * (n - 1))
1824 fstep acc (toEnum $ fromIntegral old) >>= go2 wrd (n - 1)
1825 go2 _ _ acc = return acc
1827 stepOuter gst (GO_KARP_RABIN stt rb rhead) = do
1829 res <- step (adaptState gst) stt
1833 acc' <- if withSep then fstep acc x else return acc
1834 rh' <- liftIO (RB.unsafeInsert rb rhead x)
1837 let fold = RB.unsafeFoldRing (RB.ringBound rb)
1838 let !ringHash = fold addCksum 0 rb
1839 if ringHash == patHash
1840 then go2 SPEC ringHash rh' s acc'
1841 else go0 SPEC (idx + 1) rh' s acc'
1842 else go0 SPEC (idx + 1) rh' s acc'
1843 Skip s -> return $ Skip (GO_KARP_RABIN s rb rhead)
1848 k = 2891336453 :: Word32
1850 addCksum cksum a = cksum * k + fromIntegral (fromEnum a)
1851 deltaCksum cksum old new =
1852 addCksum cksum new - coeff * fromIntegral (fromEnum old)
1854 -- XXX shall we use a random starting hash or 1 instead of 0?
1855 patHash = A.foldl' addCksum 0 patArr
1858 go0 !_ !idx !rh st !acc = do
1859 res <- step (adaptState gst) st
1862 acc' <- if withSep then fstep acc x else return acc
1863 rh' <- liftIO (RB.unsafeInsert rb rh x)
1866 let fold = RB.unsafeFoldRing (RB.ringBound rb)
1867 let !ringHash = fold addCksum 0 rb
1868 if ringHash == patHash
1869 then go2 SPEC ringHash rh' s acc'
1870 else go1 SPEC ringHash rh' s acc'
1871 else go0 SPEC (idx + 1) rh' s acc'
1872 Skip s -> go0 SPEC idx rh s acc
1874 -- do not issue a blank segment when we end at pattern
1875 if (idx == maxIndex) && RB.unsafeEqArray rb rh patArr
1878 !acc' <- if idx /= 0 && not withSep
1879 then RB.unsafeFoldRingM rh fstep acc rb
1881 done acc' >>= \r -> return $ Yield r GO_DONE
1883 -- XXX Theoretically this code can do 4 times faster if GHC generates
1884 -- optimal code. If we use just "(cksum' == patHash)" condition it goes
1885 -- 4x faster, as soon as we add the "RB.unsafeEqArray rb v" condition
1886 -- the generated code changes drastically and becomes 4x slower. Need
1887 -- to investigate what is going on with GHC.
1889 go1 !_ !cksum !rh st !acc = do
1890 res <- step (adaptState gst) st
1893 old <- liftIO $ peek rh
1894 let cksum' = deltaCksum cksum old x
1899 if (cksum' == patHash)
1901 rh' <- liftIO (RB.unsafeInsert rb rh x)
1902 go2 SPEC cksum' rh' s acc'
1904 rh' <- liftIO (RB.unsafeInsert rb rh x)
1905 go1 SPEC cksum' rh' s acc'
1906 Skip s -> go1 SPEC cksum rh s acc
1908 if RB.unsafeEqArray rb rh patArr
1913 else RB.unsafeFoldRingFullM rh fstep acc rb
1914 done acc' >>= \r -> return $ Yield r GO_DONE
1916 go2 !_ !cksum' !rh' s !acc' = do
1917 if RB.unsafeEqArray rb rh' patArr
1920 return $ Yield r (GO_KARP_RABIN s rb rhead)
1921 else go1 SPEC cksum' rh' s acc'
1923 stepOuter gst (GO_EMPTY_PAT st) = do
1924 res <- step (adaptState gst) st
1929 done acc' >>= \r -> return $ Yield r (GO_EMPTY_PAT s)
1930 Skip s -> return $ Skip (GO_EMPTY_PAT s)
1933 stepOuter _ GO_DONE = return Stop
1935data SplitState s arr
1937 | SplitBuffering s arr
1938 | SplitSplitting s arr
1939 | SplitYielding arr (SplitState s arr)
1942-- XXX An alternative approach would be to use a partial fold (Fold m a b) to
1943-- split using a splitBy like combinator. The Fold would consume upto the
1944-- separator and return any leftover which can then be fed to the next fold.
1946-- We can revisit this once we have partial folds/parsers.
1948-- | Performs infix separator style splitting.
1949{-# INLINE_NORMAL splitInnerBy #-}
1952 => (f a -> m (f a, Maybe (f a))) -- splitter
1953 -> (f a -> f a -> m (f a)) -- joiner
1956splitInnerBy splitter joiner (Stream step1 state1) =
1957 (Stream step (SplitInitial state1))
1961 {-# INLINE_LATE step #-}
1962 step gst (SplitInitial st) = do
1966 (x1, mx2) <- splitter x
1967 return $ case mx2 of
1968 Nothing -> Skip (SplitBuffering s x1)
1969 Just x2 -> Skip (SplitYielding x1 (SplitSplitting s x2))
1970 Skip s -> return $ Skip (SplitInitial s)
1971 Stop -> return $ Stop
1973 step gst (SplitBuffering st buf) = do
1977 (x1, mx2) <- splitter x
1978 buf' <- joiner buf x1
1979 return $ case mx2 of
1980 Nothing -> Skip (SplitBuffering s buf')
1981 Just x2 -> Skip (SplitYielding buf' (SplitSplitting s x2))
1982 Skip s -> return $ Skip (SplitBuffering s buf)
1983 Stop -> return $ Skip (SplitYielding buf SplitFinishing)
1985 step _ (SplitSplitting st buf) = do
1986 (x1, mx2) <- splitter buf
1987 return $ case mx2 of
1988 Nothing -> Skip $ SplitBuffering st x1
1989 Just x2 -> Skip $ SplitYielding x1 (SplitSplitting st x2)
1991 step _ (SplitYielding x next) = return $ Yield x next
1992 step _ SplitFinishing = return $ Stop
1994-- | Performs infix separator style splitting.
1995{-# INLINE_NORMAL splitInnerBySuffix #-}
1997 :: (Monad m, Eq (f a), Monoid (f a))
1998 => (f a -> m (f a, Maybe (f a))) -- splitter
1999 -> (f a -> f a -> m (f a)) -- joiner
2002splitInnerBySuffix splitter joiner (Stream step1 state1) =
2003 (Stream step (SplitInitial state1))
2007 {-# INLINE_LATE step #-}
2008 step gst (SplitInitial st) = do
2012 (x1, mx2) <- splitter x
2013 return $ case mx2 of
2014 Nothing -> Skip (SplitBuffering s x1)
2015 Just x2 -> Skip (SplitYielding x1 (SplitSplitting s x2))
2016 Skip s -> return $ Skip (SplitInitial s)
2017 Stop -> return $ Stop
2019 step gst (SplitBuffering st buf) = do
2023 (x1, mx2) <- splitter x
2024 buf' <- joiner buf x1
2025 return $ case mx2 of
2026 Nothing -> Skip (SplitBuffering s buf')
2027 Just x2 -> Skip (SplitYielding buf' (SplitSplitting s x2))
2028 Skip s -> return $ Skip (SplitBuffering s buf)
2032 else Skip (SplitYielding buf SplitFinishing)
2034 step _ (SplitSplitting st buf) = do
2035 (x1, mx2) <- splitter buf
2036 return $ case mx2 of
2037 Nothing -> Skip $ SplitBuffering st x1
2038 Just x2 -> Skip $ SplitYielding x1 (SplitSplitting st x2)
2040 step _ (SplitYielding x next) = return $ Yield x next
2041 step _ SplitFinishing = return $ Stop
2043------------------------------------------------------------------------------
2045------------------------------------------------------------------------------
2047{-# INLINE_NORMAL isPrefixOf #-}
2048isPrefixOf :: (Eq a, Monad m) => Stream m a -> Stream m a -> m Bool
2049isPrefixOf (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
2051 go (sa, sb, Nothing) = do
2052 r <- stepa defState sa
2054 Yield x sa' -> go (sa', sb, Just x)
2055 Skip sa' -> go (sa', sb, Nothing)
2058 go (sa, sb, Just x) = do
2059 r <- stepb defState sb
2063 then go (sa, sb', Nothing)
2065 Skip sb' -> go (sa, sb', Just x)
2066 Stop -> return False
2068{-# INLINE_NORMAL isSubsequenceOf #-}
2069isSubsequenceOf :: (Eq a, Monad m) => Stream m a -> Stream m a -> m Bool
2070isSubsequenceOf (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
2072 go (sa, sb, Nothing) = do
2073 r <- stepa defState sa
2075 Yield x sa' -> go (sa', sb, Just x)
2076 Skip sa' -> go (sa', sb, Nothing)
2079 go (sa, sb, Just x) = do
2080 r <- stepb defState sb
2084 then go (sa, sb', Nothing)
2085 else go (sa, sb', Just x)
2086 Skip sb' -> go (sa, sb', Just x)
2087 Stop -> return False
2089{-# INLINE_NORMAL stripPrefix #-}
2092 => Stream m a -> Stream m a -> m (Maybe (Stream m a))
2093stripPrefix (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
2095 go (sa, sb, Nothing) = do
2096 r <- stepa defState sa
2098 Yield x sa' -> go (sa', sb, Just x)
2099 Skip sa' -> go (sa', sb, Nothing)
2100 Stop -> return $ Just (Stream stepb sb)
2102 go (sa, sb, Just x) = do
2103 r <- stepb defState sb
2107 then go (sa, sb', Nothing)
2109 Skip sb' -> go (sa, sb', Just x)
2110 Stop -> return Nothing
2112------------------------------------------------------------------------------
2114------------------------------------------------------------------------------
2116-- | Execute a monadic action for each element of the 'Stream'
2117{-# INLINE_NORMAL mapM_ #-}
2118mapM_ :: Monad m => (a -> m b) -> Stream m a -> m ()
2119mapM_ m = drain . mapM m
2121-------------------------------------------------------------------------------
2122-- Stream transformations using Unfolds
2123-------------------------------------------------------------------------------
2125-- Define a unique structure to use in inspection testing
2126data ConcatMapUState o i =
2128 | ConcatMapUInner o i
2130-- | @concatMapU unfold stream@ uses @unfold@ to map the input stream elements
2131-- to streams and then flattens the generated streams into a single output
2134-- This is like 'concatMap' but uses an unfold with an explicit state to
2135-- generate the stream instead of a 'Stream' type generator. This allows better
2136-- optimization via fusion. This can be many times more efficient than
2139{-# INLINE_NORMAL concatMapU #-}
2140concatMapU :: Monad m => Unfold m a b -> Stream m a -> Stream m b
2141concatMapU (Unfold istep inject) (Stream ostep ost) =
2142 Stream step (ConcatMapUOuter ost)
2144 {-# INLINE_LATE step #-}
2145 step gst (ConcatMapUOuter o) = do
2146 r <- ostep (adaptState gst) o
2150 i `seq` return (Skip (ConcatMapUInner o' i))
2151 Skip o' -> return $ Skip (ConcatMapUOuter o')
2152 Stop -> return $ Stop
2154 step _ (ConcatMapUInner o i) = do
2157 Yield x i' -> Yield x (ConcatMapUInner o i')
2158 Skip i' -> Skip (ConcatMapUInner o i')
2159 Stop -> Skip (ConcatMapUOuter o)
2161data ConcatUnfoldInterleaveState o i =
2162 ConcatUnfoldInterleaveOuter o [i]
2163 | ConcatUnfoldInterleaveInner o [i]
2164 | ConcatUnfoldInterleaveInnerL [i] [i]
2165 | ConcatUnfoldInterleaveInnerR [i] [i]
2167-- XXX use arrays to store state instead of lists.
2168-- XXX In general we can use different scheduling strategies e.g. how to
2169-- schedule the outer vs inner loop or assigning weights to different streams
2170-- or outer and inner loops.
2172-- After a yield, switch to the next stream. Do not switch streams on Skip.
2173-- Yield from outer stream switches to the inner stream.
2175-- There are two choices here, (1) exhaust the outer stream first and then
2176-- start yielding from the inner streams, this is much simpler to implement,
2177-- (2) yield at least one element from an inner stream before going back to
2178-- outer stream and opening the next stream from it.
2180-- Ideally, we need some scheduling bias to inner streams vs outer stream.
2181-- Maybe we can configure the behavior.
2183{-# INLINE_NORMAL concatUnfoldInterleave #-}
2184concatUnfoldInterleave :: Monad m => Unfold m a b -> Stream m a -> Stream m b
2185concatUnfoldInterleave (Unfold istep inject) (Stream ostep ost) =
2186 Stream step (ConcatUnfoldInterleaveOuter ost [])
2188 {-# INLINE_LATE step #-}
2189 step gst (ConcatUnfoldInterleaveOuter o ls) = do
2190 r <- ostep (adaptState gst) o
2194 i `seq` return (Skip (ConcatUnfoldInterleaveInner o' (i : ls)))
2195 Skip o' -> return $ Skip (ConcatUnfoldInterleaveOuter o' ls)
2196 Stop -> return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2198 step _ (ConcatUnfoldInterleaveInner _ []) = undefined
2199 step _ (ConcatUnfoldInterleaveInner o (st:ls)) = do
2202 Yield x s -> Yield x (ConcatUnfoldInterleaveOuter o (s:ls))
2203 Skip s -> Skip (ConcatUnfoldInterleaveInner o (s:ls))
2204 Stop -> Skip (ConcatUnfoldInterleaveOuter o ls)
2206 step _ (ConcatUnfoldInterleaveInnerL [] []) = return Stop
2207 step _ (ConcatUnfoldInterleaveInnerL [] rs) =
2208 return $ Skip (ConcatUnfoldInterleaveInnerR [] rs)
2210 step _ (ConcatUnfoldInterleaveInnerL (st:ls) rs) = do
2213 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerL ls (s:rs))
2214 Skip s -> Skip (ConcatUnfoldInterleaveInnerL (s:ls) rs)
2215 Stop -> Skip (ConcatUnfoldInterleaveInnerL ls rs)
2217 step _ (ConcatUnfoldInterleaveInnerR [] []) = return Stop
2218 step _ (ConcatUnfoldInterleaveInnerR ls []) =
2219 return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2221 step _ (ConcatUnfoldInterleaveInnerR ls (st:rs)) = do
2224 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerR (s:ls) rs)
2225 Skip s -> Skip (ConcatUnfoldInterleaveInnerR ls (s:rs))
2226 Stop -> Skip (ConcatUnfoldInterleaveInnerR ls rs)
2228-- XXX In general we can use different scheduling strategies e.g. how to
2229-- schedule the outer vs inner loop or assigning weights to different streams
2230-- or outer and inner loops.
2232-- This could be inefficient if the tasks are too small.
2234-- Compared to concatUnfoldInterleave this one switches streams on Skips.
2236{-# INLINE_NORMAL concatUnfoldRoundrobin #-}
2237concatUnfoldRoundrobin :: Monad m => Unfold m a b -> Stream m a -> Stream m b
2238concatUnfoldRoundrobin (Unfold istep inject) (Stream ostep ost) =
2239 Stream step (ConcatUnfoldInterleaveOuter ost [])
2241 {-# INLINE_LATE step #-}
2242 step gst (ConcatUnfoldInterleaveOuter o ls) = do
2243 r <- ostep (adaptState gst) o
2247 i `seq` return (Skip (ConcatUnfoldInterleaveInner o' (i : ls)))
2248 Skip o' -> return $ Skip (ConcatUnfoldInterleaveInner o' ls)
2249 Stop -> return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2251 step _ (ConcatUnfoldInterleaveInner o []) =
2252 return $ Skip (ConcatUnfoldInterleaveOuter o [])
2254 step _ (ConcatUnfoldInterleaveInner o (st:ls)) = do
2257 Yield x s -> Yield x (ConcatUnfoldInterleaveOuter o (s:ls))
2258 Skip s -> Skip (ConcatUnfoldInterleaveOuter o (s:ls))
2259 Stop -> Skip (ConcatUnfoldInterleaveOuter o ls)
2261 step _ (ConcatUnfoldInterleaveInnerL [] []) = return Stop
2262 step _ (ConcatUnfoldInterleaveInnerL [] rs) =
2263 return $ Skip (ConcatUnfoldInterleaveInnerR [] rs)
2265 step _ (ConcatUnfoldInterleaveInnerL (st:ls) rs) = do
2268 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerL ls (s:rs))
2269 Skip s -> Skip (ConcatUnfoldInterleaveInnerL ls (s:rs))
2270 Stop -> Skip (ConcatUnfoldInterleaveInnerL ls rs)
2272 step _ (ConcatUnfoldInterleaveInnerR [] []) = return Stop
2273 step _ (ConcatUnfoldInterleaveInnerR ls []) =
2274 return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2276 step _ (ConcatUnfoldInterleaveInnerR ls (st:rs)) = do
2279 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerR (s:ls) rs)
2280 Skip s -> Skip (ConcatUnfoldInterleaveInnerR (s:ls) rs)
2281 Stop -> Skip (ConcatUnfoldInterleaveInnerR ls rs)
2283data AppendState s1 s2 = AppendFirst s1 | AppendSecond s2
2285-- Note that this could be much faster compared to the CPS stream. However, as
2286-- the number of streams being composed increases this may become expensive.
2287-- Need to see where the breaking point is between the two.
2289{-# INLINE_NORMAL append #-}
2290append :: Monad m => Stream m a -> Stream m a -> Stream m a
2291append (Stream step1 state1) (Stream step2 state2) =
2292 Stream step (AppendFirst state1)
2296 {-# INLINE_LATE step #-}
2297 step gst (AppendFirst st) = do
2300 Yield a s -> Yield a (AppendFirst s)
2301 Skip s -> Skip (AppendFirst s)
2302 Stop -> Skip (AppendSecond state2)
2304 step gst (AppendSecond st) = do
2307 Yield a s -> Yield a (AppendSecond s)
2308 Skip s -> Skip (AppendSecond s)
2311data InterleaveState s1 s2 = InterleaveFirst s1 s2 | InterleaveSecond s1 s2
2312 | InterleaveSecondOnly s2 | InterleaveFirstOnly s1
2314{-# INLINE_NORMAL interleave #-}
2315interleave :: Monad m => Stream m a -> Stream m a -> Stream m a
2316interleave (Stream step1 state1) (Stream step2 state2) =
2317 Stream step (InterleaveFirst state1 state2)
2321 {-# INLINE_LATE step #-}
2322 step gst (InterleaveFirst st1 st2) = do
2325 Yield a s -> Yield a (InterleaveSecond s st2)
2326 Skip s -> Skip (InterleaveFirst s st2)
2327 Stop -> Skip (InterleaveSecondOnly st2)
2329 step gst (InterleaveSecond st1 st2) = do
2332 Yield a s -> Yield a (InterleaveFirst st1 s)
2333 Skip s -> Skip (InterleaveSecond st1 s)
2334 Stop -> Skip (InterleaveFirstOnly st1)
2336 step gst (InterleaveFirstOnly st1) = do
2339 Yield a s -> Yield a (InterleaveFirstOnly s)
2340 Skip s -> Skip (InterleaveFirstOnly s)
2343 step gst (InterleaveSecondOnly st2) = do
2346 Yield a s -> Yield a (InterleaveSecondOnly s)
2347 Skip s -> Skip (InterleaveSecondOnly s)
2350{-# INLINE_NORMAL interleaveMin #-}
2351interleaveMin :: Monad m => Stream m a -> Stream m a -> Stream m a
2352interleaveMin (Stream step1 state1) (Stream step2 state2) =
2353 Stream step (InterleaveFirst state1 state2)
2357 {-# INLINE_LATE step #-}
2358 step gst (InterleaveFirst st1 st2) = do
2361 Yield a s -> Yield a (InterleaveSecond s st2)
2362 Skip s -> Skip (InterleaveFirst s st2)
2365 step gst (InterleaveSecond st1 st2) = do
2368 Yield a s -> Yield a (InterleaveFirst st1 s)
2369 Skip s -> Skip (InterleaveSecond st1 s)
2372 step _ (InterleaveFirstOnly _) = undefined
2373 step _ (InterleaveSecondOnly _) = undefined
2375{-# INLINE_NORMAL interleaveSuffix #-}
2376interleaveSuffix :: Monad m => Stream m a -> Stream m a -> Stream m a
2377interleaveSuffix (Stream step1 state1) (Stream step2 state2) =
2378 Stream step (InterleaveFirst state1 state2)
2382 {-# INLINE_LATE step #-}
2383 step gst (InterleaveFirst st1 st2) = do
2386 Yield a s -> Yield a (InterleaveSecond s st2)
2387 Skip s -> Skip (InterleaveFirst s st2)
2390 step gst (InterleaveSecond st1 st2) = do
2393 Yield a s -> Yield a (InterleaveFirst st1 s)
2394 Skip s -> Skip (InterleaveSecond st1 s)
2395 Stop -> Skip (InterleaveFirstOnly st1)
2397 step gst (InterleaveFirstOnly st1) = do
2400 Yield a s -> Yield a (InterleaveFirstOnly s)
2401 Skip s -> Skip (InterleaveFirstOnly s)
2404 step _ (InterleaveSecondOnly _) = undefined
2406data InterleaveInfixState s1 s2 a
2407 = InterleaveInfixFirst s1 s2
2408 | InterleaveInfixSecondBuf s1 s2
2409 | InterleaveInfixSecondYield s1 s2 a
2410 | InterleaveInfixFirstYield s1 s2 a
2411 | InterleaveInfixFirstOnly s1
2413{-# INLINE_NORMAL interleaveInfix #-}
2414interleaveInfix :: Monad m => Stream m a -> Stream m a -> Stream m a
2415interleaveInfix (Stream step1 state1) (Stream step2 state2) =
2416 Stream step (InterleaveInfixFirst state1 state2)
2420 {-# INLINE_LATE step #-}
2421 step gst (InterleaveInfixFirst st1 st2) = do
2424 Yield a s -> Yield a (InterleaveInfixSecondBuf s st2)
2425 Skip s -> Skip (InterleaveInfixFirst s st2)
2428 step gst (InterleaveInfixSecondBuf st1 st2) = do
2431 Yield a s -> Skip (InterleaveInfixSecondYield st1 s a)
2432 Skip s -> Skip (InterleaveInfixSecondBuf st1 s)
2433 Stop -> Skip (InterleaveInfixFirstOnly st1)
2435 step gst (InterleaveInfixSecondYield st1 st2 x) = do
2438 Yield a s -> Yield x (InterleaveInfixFirstYield s st2 a)
2439 Skip s -> Skip (InterleaveInfixSecondYield s st2 x)
2442 step _ (InterleaveInfixFirstYield st1 st2 x) = do
2443 return $ Yield x (InterleaveInfixSecondBuf st1 st2)
2445 step gst (InterleaveInfixFirstOnly st1) = do
2448 Yield a s -> Yield a (InterleaveInfixFirstOnly s)
2449 Skip s -> Skip (InterleaveInfixFirstOnly s)
2452{-# INLINE_NORMAL roundRobin #-}
2453roundRobin :: Monad m => Stream m a -> Stream m a -> Stream m a
2454roundRobin (Stream step1 state1) (Stream step2 state2) =
2455 Stream step (InterleaveFirst state1 state2)
2459 {-# INLINE_LATE step #-}
2460 step gst (InterleaveFirst st1 st2) = do
2463 Yield a s -> Yield a (InterleaveSecond s st2)
2464 Skip s -> Skip (InterleaveSecond s st2)
2465 Stop -> Skip (InterleaveSecondOnly st2)
2467 step gst (InterleaveSecond st1 st2) = do
2470 Yield a s -> Yield a (InterleaveFirst st1 s)
2471 Skip s -> Skip (InterleaveFirst st1 s)
2472 Stop -> Skip (InterleaveFirstOnly st1)
2474 step gst (InterleaveSecondOnly st2) = do
2477 Yield a s -> Yield a (InterleaveSecondOnly s)
2478 Skip s -> Skip (InterleaveSecondOnly s)
2481 step gst (InterleaveFirstOnly st1) = do
2484 Yield a s -> Yield a (InterleaveFirstOnly s)
2485 Skip s -> Skip (InterleaveFirstOnly s)
2488data ICUState s1 s2 i1 i2 =
2493 | ICUFirstInner s1 s2 i1
2494 | ICUSecondInner s1 s2 i2
2495 | ICUFirstOnlyInner s1 i1
2496 | ICUSecondOnlyInner s2 i2
2498-- | Interleave streams (full streams, not the elements) unfolded from two
2499-- input streams and concat. Stop when the first stream stops. If the second
2500-- stream ends before the first one then first stream still keeps running alone
2501-- without any interleaving with the second stream.
2503-- [a1, a2, ... an] [b1, b2 ...]
2504-- => [streamA1, streamA2, ... streamAn] [streamB1, streamB2, ...]
2505-- => [streamA1, streamB1, streamA2...StreamAn, streamBn]
2506-- => [a11, a12, ...a1j, b11, b12, ...b1k, a21, a22, ...]
2508{-# INLINE_NORMAL gintercalateSuffix #-}
2511 => Unfold m a c -> Stream m a -> Unfold m b c -> Stream m b -> Stream m c
2513 (Unfold istep1 inject1) (Stream step1 state1)
2514 (Unfold istep2 inject2) (Stream step2 state2) =
2515 Stream step (ICUFirst state1 state2)
2519 {-# INLINE_LATE step #-}
2520 step gst (ICUFirst s1 s2) = do
2521 r <- step1 (adaptState gst) s1
2525 i `seq` return (Skip (ICUFirstInner s s2 i))
2526 Skip s -> return $ Skip (ICUFirst s s2)
2529 step gst (ICUFirstOnly s1) = do
2530 r <- step1 (adaptState gst) s1
2534 i `seq` return (Skip (ICUFirstOnlyInner s i))
2535 Skip s -> return $ Skip (ICUFirstOnly s)
2538 step _ (ICUFirstInner s1 s2 i1) = do
2541 Yield x i' -> Yield x (ICUFirstInner s1 s2 i')
2542 Skip i' -> Skip (ICUFirstInner s1 s2 i')
2543 Stop -> Skip (ICUSecond s1 s2)
2545 step _ (ICUFirstOnlyInner s1 i1) = do
2548 Yield x i' -> Yield x (ICUFirstOnlyInner s1 i')
2549 Skip i' -> Skip (ICUFirstOnlyInner s1 i')
2550 Stop -> Skip (ICUFirstOnly s1)
2552 step gst (ICUSecond s1 s2) = do
2553 r <- step2 (adaptState gst) s2
2557 i `seq` return (Skip (ICUSecondInner s1 s i))
2558 Skip s -> return $ Skip (ICUSecond s1 s)
2559 Stop -> return $ Skip (ICUFirstOnly s1)
2561 step _ (ICUSecondInner s1 s2 i2) = do
2564 Yield x i' -> Yield x (ICUSecondInner s1 s2 i')
2565 Skip i' -> Skip (ICUSecondInner s1 s2 i')
2566 Stop -> Skip (ICUFirst s1 s2)
2568 step _ (ICUSecondOnly _s2) = undefined
2569 step _ (ICUSecondOnlyInner _s2 _i2) = undefined
2571data InterposeSuffixState s1 i1 =
2572 InterposeSuffixFirst s1
2573 -- | InterposeSuffixFirstYield s1 i1
2574 | InterposeSuffixFirstInner s1 i1
2575 | InterposeSuffixSecond s1
2577-- Note that if an unfolded layer turns out to be nil we still emit the
2578-- separator effect. An alternate behavior could be to emit the separator
2579-- effect only if at least one element has been yielded by the unfolding.
2580-- However, that becomes a bit complicated, so we have chosen the former
2582{-# INLINE_NORMAL interposeSuffix #-}
2585 => m c -> Unfold m b c -> Stream m b -> Stream m c
2588 (Unfold istep1 inject1) (Stream step1 state1) =
2589 Stream step (InterposeSuffixFirst state1)
2593 {-# INLINE_LATE step #-}
2594 step gst (InterposeSuffixFirst s1) = do
2595 r <- step1 (adaptState gst) s1
2599 i `seq` return (Skip (InterposeSuffixFirstInner s i))
2600 -- i `seq` return (Skip (InterposeSuffixFirstYield s i))
2601 Skip s -> return $ Skip (InterposeSuffixFirst s)
2605 step _ (InterposeSuffixFirstYield s1 i1) = do
2608 Yield x i' -> Yield x (InterposeSuffixFirstInner s1 i')
2609 Skip i' -> Skip (InterposeSuffixFirstYield s1 i')
2610 Stop -> Skip (InterposeSuffixFirst s1)
2613 step _ (InterposeSuffixFirstInner s1 i1) = do
2616 Yield x i' -> Yield x (InterposeSuffixFirstInner s1 i')
2617 Skip i' -> Skip (InterposeSuffixFirstInner s1 i')
2618 Stop -> Skip (InterposeSuffixSecond s1)
2620 step _ (InterposeSuffixSecond s1) = do
2622 return $ Yield r (InterposeSuffixFirst s1)
2624data ICALState s1 s2 i1 i2 a =
2626 -- | ICALFirstYield s1 s2 i1
2627 | ICALFirstInner s1 s2 i1
2629 | ICALFirstOnlyInner s1 i1
2630 | ICALSecondInject s1 s2
2631 | ICALFirstInject s1 s2 i2
2632 -- | ICALFirstBuf s1 s2 i1 i2
2633 | ICALSecondInner s1 s2 i1 i2
2634 -- -- | ICALSecondInner s1 s2 i1 i2 a
2635 -- -- | ICALFirstResume s1 s2 i1 i2 a
2637-- | Interleave streams (full streams, not the elements) unfolded from two
2638-- input streams and concat. Stop when the first stream stops. If the second
2639-- stream ends before the first one then first stream still keeps running alone
2640-- without any interleaving with the second stream.
2642-- [a1, a2, ... an] [b1, b2 ...]
2643-- => [streamA1, streamA2, ... streamAn] [streamB1, streamB2, ...]
2644-- => [streamA1, streamB1, streamA2...StreamAn, streamBn]
2645-- => [a11, a12, ...a1j, b11, b12, ...b1k, a21, a22, ...]
2647{-# INLINE_NORMAL gintercalate #-}
2650 => Unfold m a c -> Stream m a -> Unfold m b c -> Stream m b -> Stream m c
2652 (Unfold istep1 inject1) (Stream step1 state1)
2653 (Unfold istep2 inject2) (Stream step2 state2) =
2654 Stream step (ICALFirst state1 state2)
2658 {-# INLINE_LATE step #-}
2659 step gst (ICALFirst s1 s2) = do
2660 r <- step1 (adaptState gst) s1
2664 i `seq` return (Skip (ICALFirstInner s s2 i))
2665 -- i `seq` return (Skip (ICALFirstYield s s2 i))
2666 Skip s -> return $ Skip (ICALFirst s s2)
2670 step _ (ICALFirstYield s1 s2 i1) = do
2673 Yield x i' -> Yield x (ICALFirstInner s1 s2 i')
2674 Skip i' -> Skip (ICALFirstYield s1 s2 i')
2675 Stop -> Skip (ICALFirst s1 s2)
2678 step _ (ICALFirstInner s1 s2 i1) = do
2681 Yield x i' -> Yield x (ICALFirstInner s1 s2 i')
2682 Skip i' -> Skip (ICALFirstInner s1 s2 i')
2683 Stop -> Skip (ICALSecondInject s1 s2)
2685 step gst (ICALFirstOnly s1) = do
2686 r <- step1 (adaptState gst) s1
2690 i `seq` return (Skip (ICALFirstOnlyInner s i))
2691 Skip s -> return $ Skip (ICALFirstOnly s)
2694 step _ (ICALFirstOnlyInner s1 i1) = do
2697 Yield x i' -> Yield x (ICALFirstOnlyInner s1 i')
2698 Skip i' -> Skip (ICALFirstOnlyInner s1 i')
2699 Stop -> Skip (ICALFirstOnly s1)
2701 -- We inject the second stream even before checking if the first stream
2702 -- would yield any more elements. There is no clear choice whether we
2703 -- should do this before or after that. Doing it after may make the state
2704 -- machine a bit simpler though.
2705 step gst (ICALSecondInject s1 s2) = do
2706 r <- step2 (adaptState gst) s2
2710 i `seq` return (Skip (ICALFirstInject s1 s i))
2711 Skip s -> return $ Skip (ICALSecondInject s1 s)
2712 Stop -> return $ Skip (ICALFirstOnly s1)
2714 step gst (ICALFirstInject s1 s2 i2) = do
2715 r <- step1 (adaptState gst) s1
2719 i `seq` return (Skip (ICALSecondInner s s2 i i2))
2720 -- i `seq` return (Skip (ICALFirstBuf s s2 i i2))
2721 Skip s -> return $ Skip (ICALFirstInject s s2 i2)
2725 step _ (ICALFirstBuf s1 s2 i1 i2) = do
2728 Yield x i' -> Skip (ICALSecondInner s1 s2 i' i2 x)
2729 Skip i' -> Skip (ICALFirstBuf s1 s2 i' i2)
2732 step _ (ICALSecondInner s1 s2 i1 i2 v) = do
2735 Yield x i' -> Yield x (ICALSecondInner s1 s2 i1 i' v)
2736 Skip i' -> Skip (ICALSecondInner s1 s2 i1 i' v)
2737 Stop -> Skip (ICALFirstResume s1 s2 i1 i2 v)
2740 step _ (ICALSecondInner s1 s2 i1 i2) = do
2743 Yield x i' -> Yield x (ICALSecondInner s1 s2 i1 i')
2744 Skip i' -> Skip (ICALSecondInner s1 s2 i1 i')
2745 Stop -> Skip (ICALFirstInner s1 s2 i1)
2746 -- Stop -> Skip (ICALFirstResume s1 s2 i1 i2)
2749 step _ (ICALFirstResume s1 s2 i1 i2 x) = do
2750 return $ Yield x (ICALFirstInner s1 s2 i1 i2)
2753data InterposeState s1 i1 a =
2755 -- | InterposeFirstYield s1 i1
2756 | InterposeFirstInner s1 i1
2757 | InterposeFirstInject s1
2758 -- | InterposeFirstBuf s1 i1
2759 | InterposeSecondYield s1 i1
2760 -- -- | InterposeSecondYield s1 i1 a
2761 -- -- | InterposeFirstResume s1 i1 a
2763-- Note that this only interposes the pure values, we may run many effects to
2764-- generate those values as some effects may not generate anything (Skip).
2765{-# INLINE_NORMAL interpose #-}
2766interpose :: Monad m => m c -> Unfold m b c -> Stream m b -> Stream m c
2769 (Unfold istep1 inject1) (Stream step1 state1) =
2770 Stream step (InterposeFirst state1)
2774 {-# INLINE_LATE step #-}
2775 step gst (InterposeFirst s1) = do
2776 r <- step1 (adaptState gst) s1
2780 i `seq` return (Skip (InterposeFirstInner s i))
2781 -- i `seq` return (Skip (InterposeFirstYield s i))
2782 Skip s -> return $ Skip (InterposeFirst s)
2786 step _ (InterposeFirstYield s1 i1) = do
2789 Yield x i' -> Yield x (InterposeFirstInner s1 i')
2790 Skip i' -> Skip (InterposeFirstYield s1 i')
2791 Stop -> Skip (InterposeFirst s1)
2794 step _ (InterposeFirstInner s1 i1) = do
2797 Yield x i' -> Yield x (InterposeFirstInner s1 i')
2798 Skip i' -> Skip (InterposeFirstInner s1 i')
2799 Stop -> Skip (InterposeFirstInject s1)
2801 step gst (InterposeFirstInject s1) = do
2802 r <- step1 (adaptState gst) s1
2806 -- i `seq` return (Skip (InterposeFirstBuf s i))
2807 i `seq` return (Skip (InterposeSecondYield s i))
2808 Skip s -> return $ Skip (InterposeFirstInject s)
2812 step _ (InterposeFirstBuf s1 i1) = do
2815 Yield x i' -> Skip (InterposeSecondYield s1 i' x)
2816 Skip i' -> Skip (InterposeFirstBuf s1 i')
2821 step _ (InterposeSecondYield s1 i1 v) = do
2823 return $ Yield r (InterposeFirstResume s1 i1 v)
2825 step _ (InterposeSecondYield s1 i1) = do
2827 return $ Yield r (InterposeFirstInner s1 i1)
2830 step _ (InterposeFirstResume s1 i1 v) = do
2831 return $ Yield v (InterposeFirstInner s1 i1)
2834------------------------------------------------------------------------------
2836------------------------------------------------------------------------------
2838data GbracketState s1 s2 v
2840 | GBracketNormal s1 v
2841 | GBracketException s2
2843-- | The most general bracketing and exception combinator. All other
2844-- combinators can be expressed in terms of this combinator. This can also be
2845-- used for cases which are not covered by the standard combinators.
2849{-# INLINE_NORMAL gbracket #-}
2853 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
2854 -> (c -> m d) -- ^ after, on normal stop
2855 -> (c -> e -> Stream m b) -- ^ on exception
2856 -> (c -> Stream m b) -- ^ stream generator
2858gbracket bef exc aft fexc fnormal =
2859 Stream step GBracketInit
2863 {-# INLINE_LATE step #-}
2864 step _ GBracketInit = do
2866 return $ Skip $ GBracketNormal (fnormal r) r
2868 step gst (GBracketNormal (UnStream step1 st) v) = do
2869 res <- exc $ step1 gst st
2871 Right r -> case r of
2873 return $ Yield x (GBracketNormal (Stream step1 s) v)
2874 Skip s -> return $ Skip (GBracketNormal (Stream step1 s) v)
2875 Stop -> aft v >> return Stop
2876 Left e -> return $ Skip (GBracketException (fexc v e))
2877 step gst (GBracketException (UnStream step1 st)) = do
2880 Yield x s -> return $ Yield x (GBracketException (Stream step1 s))
2881 Skip s -> return $ Skip (GBracketException (Stream step1 s))
2884-- | Create an IORef holding a finalizer that is called automatically when the
2885-- IORef is garbage collected. The IORef can be written to with a 'Nothing'
2886-- value to deactivate the finalizer.
2887newFinalizedIORef :: (MonadIO m, MonadBaseControl IO m)
2888 => m a -> m (IORef (Maybe (IO ())))
2889newFinalizedIORef finalizer = do
2890 mrun <- captureMonadState
2891 ref <- liftIO $ newIORef $ Just $ liftIO $ void $ do
2892 _ <- runInIO mrun finalizer
2895 res <- readIORef ref
2897 Nothing -> return ()
2899 _ <- liftIO $ mkWeakIORef ref finalizer1
2902data GbracketIOState s1 s2 v wref
2904 | GBracketIONormal s1 v wref
2905 | GBracketIOException s2
2907-- | Like gbracket but also uses a finalizer to make sure when the stream is
2908-- garbage collected we run the finalizing action. This requires a MonadIO and
2909-- MonadBaseControl IO constraint.
2911-- | The most general bracketing and exception combinator. All other
2912-- combinators can be expressed in terms of this combinator. This can also be
2913-- used for cases which are not covered by the standard combinators.
2917{-# INLINE_NORMAL gbracketIO #-}
2919 :: (MonadIO m, MonadBaseControl IO m)
2921 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
2922 -> (c -> m d) -- ^ after, on normal stop or GC
2923 -> (c -> e -> Stream m b) -- ^ on exception
2924 -> (c -> Stream m b) -- ^ stream generator
2926gbracketIO bef exc aft fexc fnormal =
2927 Stream step GBracketIOInit
2931 -- If the stream is never evaluated the "aft" action will never be
2932 -- called. For that to occur we will need the user of this API to pass a
2933 -- weak pointer to us.
2934 {-# INLINE_LATE step #-}
2935 step _ GBracketIOInit = do
2937 ref <- newFinalizedIORef (aft r)
2938 return $ Skip $ GBracketIONormal (fnormal r) r ref
2940 step gst (GBracketIONormal (UnStream step1 st) v ref) = do
2941 res <- exc $ step1 gst st
2943 Right r -> case r of
2945 return $ Yield x (GBracketIONormal (Stream step1 s) v ref)
2947 return $ Skip (GBracketIONormal (Stream step1 s) v ref)
2949 liftIO (writeIORef ref Nothing)
2950 aft v >> return Stop
2952 liftIO (writeIORef ref Nothing)
2953 return $ Skip (GBracketIOException (fexc v e))
2954 step gst (GBracketIOException (UnStream step1 st)) = do
2958 return $ Yield x (GBracketIOException (Stream step1 s))
2959 Skip s -> return $ Skip (GBracketIOException (Stream step1 s))
2962-- | Run a side effect before the stream yields its first element.
2963{-# INLINE_NORMAL before #-}
2964before :: Monad m => m b -> Stream m a -> Stream m a
2965before action (Stream step state) = Stream step' Nothing
2969 {-# INLINE_LATE step' #-}
2970 step' _ Nothing = action >> return (Skip (Just state))
2972 step' gst (Just st) = do
2975 Yield x s -> return $ Yield x (Just s)
2976 Skip s -> return $ Skip (Just s)
2979-- | Run a side effect whenever the stream stops normally.
2980{-# INLINE_NORMAL after #-}
2981after :: Monad m => m b -> Stream m a -> Stream m a
2982after action (Stream step state) = Stream step' state
2986 {-# INLINE_LATE step' #-}
2990 Yield x s -> return $ Yield x s
2991 Skip s -> return $ Skip s
2992 Stop -> action >> return Stop
2994{-# INLINE_NORMAL afterIO #-}
2995afterIO :: (MonadIO m, MonadBaseControl IO m)
2996 => m b -> Stream m a -> Stream m a
2997afterIO action (Stream step state) = Stream step' Nothing
3001 {-# INLINE_LATE step' #-}
3002 step' _ Nothing = do
3003 ref <- newFinalizedIORef action
3004 return $ Skip $ Just (state, ref)
3005 step' gst (Just (st, ref)) = do
3008 Yield x s -> return $ Yield x (Just (s, ref))
3009 Skip s -> return $ Skip (Just (s, ref))
3011 Just f <- readIORef ref
3012 writeIORef ref Nothing
3015-- XXX These combinators are expensive due to the call to
3016-- onException/handle/try on each step. Therefore, when possible, they should
3017-- be called in an outer loop where we perform less iterations. For example, we
3018-- cannot call them on each iteration in a char stream, instead we can call
3019-- them when doing an IO on an array.
3021-- XXX For high performance error checks in busy streams we may need another
3022-- Error constructor in step.
3024-- | Run a side effect whenever the stream aborts due to an exception. The
3025-- exception is not caught, simply rethrown.
3026{-# INLINE_NORMAL onException #-}
3027onException :: MonadCatch m => m b -> Stream m a -> Stream m a
3028onException action str =
3029 gbracket (return ()) MC.try return
3030 (\_ (e :: MC.SomeException) -> nilM (action >> MC.throwM e))
3033{-# INLINE_NORMAL _onException #-}
3034_onException :: MonadCatch m => m b -> Stream m a -> Stream m a
3035_onException action (Stream step state) = Stream step' state
3039 {-# INLINE_LATE step' #-}
3041 res <- step gst st `MC.onException` action
3043 Yield x s -> return $ Yield x s
3044 Skip s -> return $ Skip s
3047-- XXX bracket is like concatMap, it generates a stream and then flattens it.
3048-- Like concatMap it has 10x worse performance compared to linear fused
3051-- | Run the first action before the stream starts and remember its output,
3052-- generate a stream using the output, run the second action providing the
3053-- remembered value as an argument whenever the stream ends normally or due to
3055{-# INLINE_NORMAL bracket #-}
3056bracket :: MonadCatch m => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
3057bracket bef aft bet =
3058 gbracket bef MC.try aft
3059 (\a (e :: SomeException) -> nilM (aft a >> MC.throwM e)) bet
3061{-# INLINE_NORMAL bracketIO #-}
3062bracketIO :: (MonadAsync m, MonadCatch m)
3063 => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
3064bracketIO bef aft bet =
3065 gbracketIO bef MC.try aft
3066 (\a (e :: SomeException) -> nilM (aft a >> MC.throwM e)) bet
3068data BracketState s v = BracketInit | BracketRun s v
3070{-# INLINE_NORMAL _bracket #-}
3071_bracket :: MonadCatch m => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
3072_bracket bef aft bet = Stream step' BracketInit
3076 {-# INLINE_LATE step' #-}
3077 step' _ BracketInit = bef >>= \x -> return (Skip (BracketRun (bet x) x))
3079 -- NOTE: It is important to use UnStream instead of the Stream pattern
3080 -- here, otherwise we get huge perf degradation, see note in concatMap.
3081 step' gst (BracketRun (UnStream step state) v) = do
3082 -- res <- step gst state `MC.onException` aft v
3083 res <- MC.try $ step gst state
3085 Left (e :: SomeException) -> aft v >> MC.throwM e >> return Stop
3086 Right r -> case r of
3087 Yield x s -> return $ Yield x (BracketRun (Stream step s) v)
3088 Skip s -> return $ Skip (BracketRun (Stream step s) v)
3089 Stop -> aft v >> return Stop
3091-- | Run a side effect whenever the stream stops normally or aborts due to an
3093{-# INLINE finally #-}
3094finally :: MonadCatch m => m b -> Stream m a -> Stream m a
3095-- finally action xs = after action $ onException action xs
3096finally action xs = bracket (return ()) (\_ -> action) (const xs)
3098{-# INLINE finallyIO #-}
3099finallyIO :: (MonadAsync m, MonadCatch m) => m b -> Stream m a -> Stream m a
3100finallyIO action xs = bracketIO (return ()) (\_ -> action) (const xs)
3102-- | When evaluating a stream if an exception occurs, stream evaluation aborts
3103-- and the specified exception handler is run with the exception as argument.
3104{-# INLINE_NORMAL handle #-}
3105handle :: (MonadCatch m, Exception e)
3106 => (e -> Stream m a) -> Stream m a -> Stream m a
3108 gbracket (return ()) MC.try return (\_ e -> f e) (\_ -> str)
3110{-# INLINE_NORMAL _handle #-}
3111_handle :: (MonadCatch m, Exception e)
3112 => (e -> Stream m a) -> Stream m a -> Stream m a
3113_handle f (Stream step state) = Stream step' (Left state)
3117 {-# INLINE_LATE step' #-}
3118 step' gst (Left st) = do
3119 res <- MC.try $ step gst st
3121 Left e -> return $ Skip $ Right (f e)
3122 Right r -> case r of
3123 Yield x s -> return $ Yield x (Left s)
3124 Skip s -> return $ Skip (Left s)
3127 step' gst (Right (UnStream step1 st)) = do
3130 Yield x s -> return $ Yield x (Right (Stream step1 s))
3131 Skip s -> return $ Skip (Right (Stream step1 s))
3134-------------------------------------------------------------------------------
3135-- General transformation
3136-------------------------------------------------------------------------------
3138{-# INLINE_NORMAL transform #-}
3139transform :: Monad m => Pipe m a b -> Stream m a -> Stream m b
3140transform (Pipe pstep1 pstep2 pstate) (Stream step state) =
3141 Stream step' (Consume pstate, state)
3145 {-# INLINE_LATE step' #-}
3147 step' gst (Consume pst, st) = pst `seq` do
3148 r <- step (adaptState gst) st
3153 Pipe.Yield b pst' -> return $ Yield b (pst', s)
3154 Pipe.Continue pst' -> return $ Skip (pst', s)
3155 Skip s -> return $ Skip (Consume pst, s)
3158 step' _ (Produce pst, st) = pst `seq` do
3161 Pipe.Yield b pst' -> return $ Yield b (pst', st)
3162 Pipe.Continue pst' -> return $ Skip (pst', st)
3164------------------------------------------------------------------------------
3165-- Transformation by Folding (Scans)
3166------------------------------------------------------------------------------
3168------------------------------------------------------------------------------
3170------------------------------------------------------------------------------
3172-- XXX Is a prescan useful, discarding the last step does not sound useful? I
3173-- am not sure about the utility of this function, so this is implemented but
3174-- not exposed. We can expose it if someone provides good reasons why this is
3177-- XXX We have to execute the stream one step ahead to know that we are at the
3178-- last step. The vector implementation of prescan executes the last fold step
3179-- but does not yield the result. This means we have executed the effect but
3180-- discarded value. This does not sound right. In this implementation we are
3181-- not executing the last fold step.
3182{-# INLINE_NORMAL prescanlM' #-}
3183prescanlM' :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
3184prescanlM' f mz (Stream step state) = Stream step' (state, mz)
3186 {-# INLINE_LATE step' #-}
3187 step' gst (st, prev) = do
3188 r <- step (adaptState gst) st
3192 return $ Yield acc (s, f acc x)
3193 Skip s -> return $ Skip (s, prev)
3196{-# INLINE prescanl' #-}
3197prescanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
3198prescanl' f z = prescanlM' (\a b -> return (f a b)) (return z)
3200------------------------------------------------------------------------------
3201-- Monolithic postscans (postscan followed by a map)
3202------------------------------------------------------------------------------
3204-- The performance of a modular postscan followed by a map seems to be
3205-- equivalent to this monolithic scan followed by map therefore we may not need
3206-- this implementation. We just have it for performance comparison and in case
3207-- modular version does not perform well in some situation.
3209{-# INLINE_NORMAL postscanlMx' #-}
3210postscanlMx' :: Monad m
3211 => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> Stream m b
3212postscanlMx' fstep begin done (Stream step state) = do
3213 Stream step' (state, begin)
3215 {-# INLINE_LATE step' #-}
3216 step' gst (st, acc) = do
3217 r <- step (adaptState gst) st
3223 v `seq` y `seq` return (Yield v (s, return y))
3224 Skip s -> return $ Skip (s, acc)
3227{-# INLINE_NORMAL postscanlx' #-}
3228postscanlx' :: Monad m
3229 => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> Stream m b
3230postscanlx' fstep begin done s =
3231 postscanlMx' (\b a -> return (fstep b a)) (return begin) (return . done) s
3233-- XXX do we need consM strict to evaluate the begin value?
3234{-# INLINE scanlMx' #-}
3236 => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> Stream m b
3237scanlMx' fstep begin done s =
3238 (begin >>= \x -> x `seq` done x) `consM` postscanlMx' fstep begin done s
3240{-# INLINE scanlx' #-}
3242 => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> Stream m b
3243scanlx' fstep begin done s =
3244 scanlMx' (\b a -> return (fstep b a)) (return begin) (return . done) s
3246------------------------------------------------------------------------------
3248------------------------------------------------------------------------------
3250{-# INLINE_NORMAL postscanlM' #-}
3251postscanlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
3252postscanlM' fstep begin (Stream step state) =
3253 begin `seq` Stream step' (state, begin)
3255 {-# INLINE_LATE step' #-}
3256 step' gst (st, acc) = acc `seq` do
3257 r <- step (adaptState gst) st
3261 y `seq` return (Yield y (s, y))
3262 Skip s -> return $ Skip (s, acc)
3265{-# INLINE_NORMAL postscanl' #-}
3266postscanl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
3267postscanl' f = postscanlM' (\a b -> return (f a b))
3269{-# INLINE_NORMAL postscanlM #-}
3270postscanlM :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
3271postscanlM fstep begin (Stream step state) = Stream step' (state, begin)
3273 {-# INLINE_LATE step' #-}
3274 step' gst (st, acc) = do
3275 r <- step (adaptState gst) st
3279 return (Yield y (s, y))
3280 Skip s -> return $ Skip (s, acc)
3283{-# INLINE_NORMAL postscanl #-}
3284postscanl :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
3285postscanl f = postscanlM (\a b -> return (f a b))
3287{-# INLINE_NORMAL scanlM' #-}
3288scanlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
3289scanlM' fstep begin s = begin `seq` (begin `cons` postscanlM' fstep begin s)
3291{-# INLINE scanl' #-}
3292scanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
3293scanl' f = scanlM' (\a b -> return (f a b))
3295{-# INLINE_NORMAL scanlM #-}
3296scanlM :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
3297scanlM fstep begin s = begin `cons` postscanlM fstep begin s
3300scanl :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
3301scanl f = scanlM (\a b -> return (f a b))
3303{-# INLINE_NORMAL scanl1M #-}
3304scanl1M :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
3305scanl1M fstep (Stream step state) = Stream step' (state, Nothing)
3307 {-# INLINE_LATE step' #-}
3308 step' gst (st, Nothing) = do
3311 Yield x s -> return $ Yield x (s, Just x)
3312 Skip s -> return $ Skip (s, Nothing)
3315 step' gst (st, Just acc) = do
3320 return $ Yield z (s, Just z)
3321 Skip s -> return $ Skip (s, Just acc)
3324{-# INLINE scanl1 #-}
3325scanl1 :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
3326scanl1 f = scanl1M (\x y -> return (f x y))
3328{-# INLINE_NORMAL scanl1M' #-}
3329scanl1M' :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
3330scanl1M' fstep (Stream step state) = Stream step' (state, Nothing)
3332 {-# INLINE_LATE step' #-}
3333 step' gst (st, Nothing) = do
3336 Yield x s -> x `seq` return $ Yield x (s, Just x)
3337 Skip s -> return $ Skip (s, Nothing)
3340 step' gst (st, Just acc) = acc `seq` do
3345 z `seq` return $ Yield z (s, Just z)
3346 Skip s -> return $ Skip (s, Just acc)
3349{-# INLINE scanl1' #-}
3350scanl1' :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
3351scanl1' f = scanl1M' (\x y -> return (f x y))
3353------------------------------------------------------------------------------
3355------------------------------------------------------------------------------
3357data RollingMapState s a = RollingMapInit s | RollingMapGo s a
3359{-# INLINE rollingMapM #-}
3360rollingMapM :: Monad m => (a -> a -> m b) -> Stream m a -> Stream m b
3361rollingMapM f (Stream step1 state1) = Stream step (RollingMapInit state1)
3363 step gst (RollingMapInit st) = do
3364 r <- step1 (adaptState gst) st
3366 Yield x s -> Skip $ RollingMapGo s x
3367 Skip s -> Skip $ RollingMapInit s
3370 step gst (RollingMapGo s1 x1) = do
3371 r <- step1 (adaptState gst) s1
3375 return $ Yield res $ RollingMapGo s x
3376 Skip s -> return $ Skip $ RollingMapGo s x1
3377 Stop -> return $ Stop
3379{-# INLINE rollingMap #-}
3380rollingMap :: Monad m => (a -> a -> b) -> Stream m a -> Stream m b
3381rollingMap f = rollingMapM (\x y -> return $ f x y)
3383------------------------------------------------------------------------------
3384-- Tapping/Distributing
3385------------------------------------------------------------------------------
3388tap :: Monad m => Fold m a b -> Stream m a -> Stream m a
3389tap (Fold fstep initial extract) (Stream step state) = Stream step' Nothing
3393 step' _ Nothing = do
3395 return $ Skip (Just (r, state))
3397 step' gst (Just (acc, st)) = acc `seq` do
3402 return $ Yield x (Just (acc', s))
3403 Skip s -> return $ Skip (Just (acc, s))
3408{-# INLINE_NORMAL tapOffsetEvery #-}
3409tapOffsetEvery :: Monad m
3410 => Int -> Int -> Fold m a b -> Stream m a -> Stream m a
3411tapOffsetEvery offset n (Fold fstep initial extract) (Stream step state) =
3412 Stream step' Nothing
3416 {-# INLINE_LATE step' #-}
3417 step' _ Nothing = do
3419 return $ Skip (Just (r, state, offset `mod` n))
3421 step' gst (Just (acc, st, count)) | count <= 0 = do
3425 !acc' <- fstep acc x
3426 return $ Yield x (Just (acc', s, n - 1))
3427 Skip s -> return $ Skip (Just (acc, s, count))
3432 step' gst (Just (acc, st, count)) = do
3435 Yield x s -> return $ Yield x (Just (acc, s, count - 1))
3436 Skip s -> return $ Skip (Just (acc, s, count))
3441{-# INLINE_NORMAL pollCounts #-}
3444 => (Stream m Int -> Stream m Int)
3448pollCounts transf fld (Stream step state) = Stream step' Nothing
3451 {-# INLINE_LATE step' #-}
3452 step' _ Nothing = do
3453 -- As long as we are using an "Int" for counts lockfree reads from
3454 -- Var should work correctly on both 32-bit and 64-bit machines.
3455 -- However, an Int on a 32-bit machine may overflow quickly.
3456 countVar <- liftIO $ newVar (0 :: Int)
3458 $ void $ runFold fld
3459 $ transf $ fromPrimVar countVar
3460 return $ Skip (Just (countVar, tid, state))
3462 step' gst (Just (countVar, tid, st)) = do
3466 liftIO $ modifyVar' countVar (+ 1)
3467 return $ Yield x (Just (countVar, tid, s))
3468 Skip s -> return $ Skip (Just (countVar, tid, s))
3470 liftIO $ killThread tid
3473{-# INLINE_NORMAL tapRate #-}
3475 (MonadAsync m, MonadCatch m)
3480tapRate samplingRate action (Stream step state) = Stream step' Nothing
3482 {-# NOINLINE loop #-}
3483 loop countVar prev = do
3486 (do liftIO $ threadDelay (round $ samplingRate * 1000000)
3487 i <- liftIO $ readVar countVar
3488 let !diff = i - prev
3491 (\(e :: AsyncException) -> do
3492 i <- liftIO $ readVar countVar
3493 let !diff = i - prev
3495 throwM (MC.toException e))
3498 {-# INLINE_LATE step' #-}
3499 step' _ Nothing = do
3500 countVar <- liftIO $ newVar 0
3501 tid <- fork $ loop countVar 0
3502 ref <- liftIO $ newIORef ()
3503 _ <- liftIO $ mkWeakIORef ref (killThread tid)
3504 return $ Skip (Just (countVar, tid, state, ref))
3506 step' gst (Just (countVar, tid, st, ref)) = do
3510 liftIO $ modifyVar' countVar (+ 1)
3511 return $ Yield x (Just (countVar, tid, s, ref))
3512 Skip s -> return $ Skip (Just (countVar, tid, s, ref))
3514 liftIO $ killThread tid
3518-------------------------------------------------------------------------------
3520-------------------------------------------------------------------------------
3522{-# INLINE_NORMAL takeWhileM #-}
3523takeWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
3524takeWhileM f (Stream step state) = Stream step' state
3526 {-# INLINE_LATE step' #-}
3532 return $ if b then Yield x s else Stop
3533 Skip s -> return $ Skip s
3536{-# INLINE takeWhile #-}
3537takeWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
3538takeWhile f = takeWhileM (return . f)
3540{-# INLINE_NORMAL drop #-}
3541drop :: Monad m => Int -> Stream m a -> Stream m a
3542drop n (Stream step state) = Stream step' (state, Just n)
3544 {-# INLINE_LATE step' #-}
3545 step' gst (st, Just i)
3550 Yield _ s -> Skip (s, Just (i - 1))
3551 Skip s -> Skip (s, Just i)
3553 | otherwise = return $ Skip (st, Nothing)
3555 step' gst (st, Nothing) = do
3559 Yield x s -> Yield x (s, Nothing)
3560 Skip s -> Skip (s, Nothing)
3563data DropWhileState s a
3565 | DropWhileYield a s
3568{-# INLINE_NORMAL dropWhileM #-}
3569dropWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
3570dropWhileM f (Stream step state) = Stream step' (DropWhileDrop state)
3572 {-# INLINE_LATE step' #-}
3573 step' gst (DropWhileDrop st) = do
3579 then return $ Skip (DropWhileDrop s)
3580 else return $ Skip (DropWhileYield x s)
3581 Skip s -> return $ Skip (DropWhileDrop s)
3584 step' gst (DropWhileNext st) = do
3587 Yield x s -> return $ Skip (DropWhileYield x s)
3588 Skip s -> return $ Skip (DropWhileNext s)
3591 step' _ (DropWhileYield x st) = return $ Yield x (DropWhileNext st)
3593{-# INLINE dropWhile #-}
3594dropWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
3595dropWhile f = dropWhileM (return . f)
3597{-# INLINE_NORMAL filterM #-}
3598filterM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
3599filterM f (Stream step state) = Stream step' state
3601 {-# INLINE_LATE step' #-}
3610 Skip s -> return $ Skip s
3613{-# INLINE filter #-}
3614filter :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
3615filter f = filterM (return . f)
3617{-# INLINE_NORMAL uniq #-}
3618uniq :: (Eq a, Monad m) => Stream m a -> Stream m a
3619uniq (Stream step state) = Stream step' (Nothing, state)
3621 {-# INLINE_LATE step' #-}
3622 step' gst (Nothing, st) = do
3625 Yield x s -> return $ Yield x (Just x, s)
3626 Skip s -> return $ Skip (Nothing, s)
3628 step' gst (Just x, st) = do
3631 Yield y s | x == y -> return $ Skip (Just x, s)
3632 | otherwise -> return $ Yield y (Just y, s)
3633 Skip s -> return $ Skip (Just x, s)
3636------------------------------------------------------------------------------
3637-- Transformation by Mapping
3638------------------------------------------------------------------------------
3640{-# INLINE_NORMAL sequence #-}
3641sequence :: Monad m => Stream m (m a) -> Stream m a
3642sequence (Stream step state) = Stream step' state
3644 {-# INLINE_LATE step' #-}
3646 r <- step (adaptState gst) st
3648 Yield x s -> x >>= \a -> return (Yield a s)
3649 Skip s -> return $ Skip s
3652------------------------------------------------------------------------------
3654------------------------------------------------------------------------------
3656data LoopState x s = FirstYield s
3657 | InterspersingYield s
3660{-# INLINE_NORMAL intersperseM #-}
3661intersperseM :: Monad m => m a -> Stream m a -> Stream m a
3662intersperseM m (Stream step state) = Stream step' (FirstYield state)
3664 {-# INLINE_LATE step' #-}
3665 step' gst (FirstYield st) = do
3669 Yield x s -> Skip (YieldAndCarry x s)
3670 Skip s -> Skip (FirstYield s)
3673 step' gst (InterspersingYield st) = do
3678 return $ Yield a (YieldAndCarry x s)
3679 Skip s -> return $ Skip $ InterspersingYield s
3682 step' _ (YieldAndCarry x st) = return $ Yield x (InterspersingYield st)
3687 | SuffixYield a (SuffixState s a)
3689{-# INLINE_NORMAL intersperseSuffix #-}
3690intersperseSuffix :: forall m a. Monad m => m a -> Stream m a -> Stream m a
3691intersperseSuffix action (Stream step state) = Stream step' (SuffixElem state)
3693 {-# INLINE_LATE step' #-}
3694 step' gst (SuffixElem st) = do
3697 Yield x s -> Skip (SuffixYield x (SuffixSuffix s))
3698 Skip s -> Skip (SuffixElem s)
3701 step' _ (SuffixSuffix st) = do
3702 action >>= \r -> return $ Skip (SuffixYield r (SuffixElem st))
3704 step' _ (SuffixYield x next) = return $ Yield x next
3706data SuffixSpanState s a
3707 = SuffixSpanElem s Int
3708 | SuffixSpanSuffix s
3709 | SuffixSpanYield a (SuffixSpanState s a)
3713-- | intersperse after every n items
3714{-# INLINE_NORMAL intersperseSuffixBySpan #-}
3715intersperseSuffixBySpan :: forall m a. Monad m
3716 => Int -> m a -> Stream m a -> Stream m a
3717intersperseSuffixBySpan n action (Stream step state) =
3718 Stream step' (SuffixSpanElem state n)
3720 {-# INLINE_LATE step' #-}
3721 step' gst (SuffixSpanElem st i) | i > 0 = do
3724 Yield x s -> Skip (SuffixSpanYield x (SuffixSpanElem s (i - 1)))
3725 Skip s -> Skip (SuffixSpanElem s i)
3726 Stop -> if i == n then Stop else Skip SuffixSpanLast
3727 step' _ (SuffixSpanElem st _) = return $ Skip (SuffixSpanSuffix st)
3729 step' _ (SuffixSpanSuffix st) = do
3730 action >>= \r -> return $ Skip (SuffixSpanYield r (SuffixSpanElem st n))
3732 step' _ (SuffixSpanLast) = do
3733 action >>= \r -> return $ Skip (SuffixSpanYield r SuffixSpanStop)
3735 step' _ (SuffixSpanYield x next) = return $ Yield x next
3737 step' _ (SuffixSpanStop) = return Stop
3739{-# INLINE intersperse #-}
3740intersperse :: Monad m => a -> Stream m a -> Stream m a
3741intersperse a = intersperseM (return a)
3743{-# INLINE_NORMAL insertBy #-}
3744insertBy :: Monad m => (a -> a -> Ordering) -> a -> Stream m a -> Stream m a
3745insertBy cmp a (Stream step state) = Stream step' (state, False, Nothing)
3747 {-# INLINE_LATE step' #-}
3748 step' gst (st, False, _) = do
3751 Yield x s -> case cmp a x of
3752 GT -> return $ Yield x (s, False, Nothing)
3753 _ -> return $ Yield a (s, True, Just x)
3754 Skip s -> return $ Skip (s, False, Nothing)
3755 Stop -> return $ Yield a (st, True, Nothing)
3757 step' _ (_, True, Nothing) = return Stop
3759 step' gst (st, True, Just prev) = do
3762 Yield x s -> return $ Yield prev (s, True, Just x)
3763 Skip s -> return $ Skip (s, True, Just prev)
3764 Stop -> return $ Yield prev (st, True, Nothing)
3766------------------------------------------------------------------------------
3768------------------------------------------------------------------------------
3770{-# INLINE_NORMAL deleteBy #-}
3771deleteBy :: Monad m => (a -> a -> Bool) -> a -> Stream m a -> Stream m a
3772deleteBy eq x (Stream step state) = Stream step' (state, False)
3774 {-# INLINE_LATE step' #-}
3775 step' gst (st, False) = do
3778 Yield y s -> return $
3779 if eq x y then Skip (s, True) else Yield y (s, False)
3780 Skip s -> return $ Skip (s, False)
3783 step' gst (st, True) = do
3786 Yield y s -> return $ Yield y (s, True)
3787 Skip s -> return $ Skip (s, True)
3790------------------------------------------------------------------------------
3791-- Transformation by Map and Filter
3792------------------------------------------------------------------------------
3794-- XXX Will this always fuse properly?
3795{-# INLINE_NORMAL mapMaybe #-}
3796mapMaybe :: Monad m => (a -> Maybe b) -> Stream m a -> Stream m b
3797mapMaybe f = fmap fromJust . filter isJust . map f
3799{-# INLINE_NORMAL mapMaybeM #-}
3800mapMaybeM :: Monad m => (a -> m (Maybe b)) -> Stream m a -> Stream m b
3801mapMaybeM f = fmap fromJust . filter isJust . mapM f
3803------------------------------------------------------------------------------
3805------------------------------------------------------------------------------
3807{-# INLINE_NORMAL indexed #-}
3808indexed :: Monad m => Stream m a -> Stream m (Int, a)
3809indexed (Stream step state) = Stream step' (state, 0)
3811 {-# INLINE_LATE step' #-}
3812 step' gst (st, i) = i `seq` do
3813 r <- step (adaptState gst) st
3815 Yield x s -> return $ Yield (i, x) (s, i+1)
3816 Skip s -> return $ Skip (s, i)
3819{-# INLINE_NORMAL indexedR #-}
3820indexedR :: Monad m => Int -> Stream m a -> Stream m (Int, a)
3821indexedR m (Stream step state) = Stream step' (state, m)
3823 {-# INLINE_LATE step' #-}
3824 step' gst (st, i) = i `seq` do
3825 r <- step (adaptState gst) st
3827 Yield x s -> let i' = i - 1
3828 in return $ Yield (i, x) (s, i')
3829 Skip s -> return $ Skip (s, i)
3832{-# INLINE_NORMAL zipWithM #-}
3834 => (a -> b -> m c) -> Stream m a -> Stream m b -> Stream m c
3835zipWithM f (Stream stepa ta) (Stream stepb tb) = Stream step (ta, tb, Nothing)
3837 {-# INLINE_LATE step #-}
3838 step gst (sa, sb, Nothing) = do
3839 r <- stepa (adaptState gst) sa
3842 Yield x sa' -> Skip (sa', sb, Just x)
3843 Skip sa' -> Skip (sa', sb, Nothing)
3846 step gst (sa, sb, Just x) = do
3847 r <- stepb (adaptState gst) sb
3851 return $ Yield z (sa, sb', Nothing)
3852 Skip sb' -> return $ Skip (sa, sb', Just x)
3855#if __GLASGOW_HASKELL__ >= 801
3856{-# RULES "zipWithM xs xs"
3857 forall f xs. zipWithM @Identity f xs xs = mapM (\x -> f x x) xs #-}
3860{-# INLINE zipWith #-}
3861zipWith :: Monad m => (a -> b -> c) -> Stream m a -> Stream m b -> Stream m c
3862zipWith f = zipWithM (\a b -> return (f a b))
3864------------------------------------------------------------------------------
3866------------------------------------------------------------------------------
3868{-# INLINE_NORMAL mergeByM #-}
3871 => (a -> a -> m Ordering) -> Stream m a -> Stream m a -> Stream m a
3872mergeByM cmp (Stream stepa ta) (Stream stepb tb) =
3873 Stream step (Just ta, Just tb, Nothing, Nothing)
3875 {-# INLINE_LATE step #-}
3877 -- one of the values is missing, and the corresponding stream is running
3878 step gst (Just sa, sb, Nothing, b) = do
3881 Yield a sa' -> Skip (Just sa', sb, Just a, b)
3882 Skip sa' -> Skip (Just sa', sb, Nothing, b)
3883 Stop -> Skip (Nothing, sb, Nothing, b)
3885 step gst (sa, Just sb, a, Nothing) = do
3888 Yield b sb' -> Skip (sa, Just sb', a, Just b)
3889 Skip sb' -> Skip (sa, Just sb', a, Nothing)
3890 Stop -> Skip (sa, Nothing, a, Nothing)
3892 -- both the values are available
3893 step _ (sa, sb, Just a, Just b) = do
3895 return $ case res of
3896 GT -> Yield b (sa, sb, Just a, Nothing)
3897 _ -> Yield a (sa, sb, Nothing, Just b)
3899 -- one of the values is missing, corresponding stream is done
3900 step _ (Nothing, sb, Nothing, Just b) =
3901 return $ Yield b (Nothing, sb, Nothing, Nothing)
3903 step _ (sa, Nothing, Just a, Nothing) =
3904 return $ Yield a (sa, Nothing, Nothing, Nothing)
3906 step _ (Nothing, Nothing, Nothing, Nothing) = return Stop
3908{-# INLINE mergeBy #-}
3911 => (a -> a -> Ordering) -> Stream m a -> Stream m a -> Stream m a
3912mergeBy cmp = mergeByM (\a b -> return $ cmp a b)
3914------------------------------------------------------------------------------
3915-- Transformation comprehensions
3916------------------------------------------------------------------------------
3918{-# INLINE_NORMAL the #-}
3919the :: (Eq a, Monad m) => Stream m a -> m (Maybe a)
3920the (Stream step state) = go state
3923 r <- step defState st
3925 Yield x s -> go' x s
3927 Stop -> return Nothing
3929 r <- step defState st
3931 Yield x s | x == n -> go' n s
3932 | otherwise -> return Nothing
3934 Stop -> return (Just n)
3936{-# INLINE runFold #-}
3937runFold :: (Monad m) => Fold m a b -> Stream m a -> m b
3938runFold (Fold step begin done) = foldlMx' step begin done
3940-------------------------------------------------------------------------------
3941-- Concurrent application and fold
3942-------------------------------------------------------------------------------
3944-- XXX These functions should be moved to Stream/Parallel.hs
3946-- Using StreamD the worker stream producing code can fuse with the code to
3947-- queue output to the SVar giving some perf boost.
3949-- Note that StreamD can only be used in limited situations, specifically, we
3950-- cannot implement joinStreamVarPar using this.
3952-- XXX make sure that the SVar passed is a Parallel style SVar.
3954-- | Fold the supplied stream to the SVar asynchronously using Parallel
3955-- concurrency style.
3956-- {-# INLINE_NORMAL toSVarParallel #-}
3957{-# INLINE toSVarParallel #-}
3958toSVarParallel :: MonadAsync m
3959 => State t m a -> SVar t m a -> Stream m a -> m ()
3960toSVarParallel st sv xs =
3961 if svarInspectMode sv
3965 case getYieldLimit st of
3966 Nothing -> doFork (work Nothing)
3968 (handleChildException sv)
3969 Just _ -> doFork (workLim Nothing)
3971 (handleChildException sv)
3976 {-# NOINLINE work #-}
3977 work info = (runFold (FL.toParallelSVar sv info) xs)
3979 {-# NOINLINE workLim #-}
3980 workLim info = runFold (FL.toParallelSVarLimited sv info) xs
3982 {-# NOINLINE forkWithDiag #-}
3984 -- We do not use workerCount in case of ParallelVar but still there is
3985 -- no harm in maintaining it correctly.
3986 liftIO $ atomicModifyIORefCAS_ (workerCount sv) $ \n -> n + 1
3988 -- This allocation matters when significant number of workers are being
3989 -- sent. We allocate it only when needed. The overhead increases by 4x.
3991 case yieldRateInfo sv of
3992 Nothing -> return Nothing
3993 Just _ -> liftIO $ do
3994 cntRef <- newIORef 0
3995 t <- getTime Monotonic
3996 lat <- newIORef (0, t)
3997 return $ Just WorkerInfo
3998 { workerYieldMax = 0
3999 , workerYieldCount = cntRef
4000 , workerLatencyStart = lat
4003 case getYieldLimit st of
4004 Nothing -> doFork (work winfo)
4006 (handleChildException sv)
4007 Just _ -> doFork (workLim winfo)
4009 (handleChildException sv)
4012{-# INLINE_NORMAL mkParallelD #-}
4013mkParallelD :: MonadAsync m => Stream m a -> Stream m a
4014mkParallelD m = Stream step Nothing
4017 step gst Nothing = do
4018 sv <- newParallelVar StopNone gst
4019 toSVarParallel gst sv m
4020 -- XXX use unfold instead?
4021 return $ Skip $ Just $ fromSVar sv
4023 step gst (Just (UnStream step1 st)) = do
4026 Yield a s -> Yield a (Just $ Stream step1 s)
4027 Skip s -> Skip (Just $ Stream step1 s)
4030-- Compare with mkAsync. mkAsync uses an Async style SVar whereas this uses a
4031-- parallel style SVar for evaluation. Currently, parallel style cannot use
4032-- rate control whereas Async style can use rate control. In async style SVar
4033-- the worker thread terminates when the buffer is full whereas in Parallel
4036-- | Make the stream producer and consumer run concurrently by introducing a
4037-- buffer between them. The producer thread evaluates the input stream until
4038-- the buffer fills, it blocks if the buffer is full until there is space in
4039-- the buffer. The consumer consumes the stream lazily from the buffer.
4043{-# INLINE_NORMAL mkParallel #-}
4044mkParallel :: (K.IsStream t, MonadAsync m) => t m a -> t m a
4045mkParallel = fromStreamD . mkParallelD . toStreamD
4047-------------------------------------------------------------------------------
4049-------------------------------------------------------------------------------
4051-- | Create an SVar with a fold consumer that will fold any elements sent to it
4052-- using the supplied fold function.
4053{-# INLINE newFoldSVar #-}
4054newFoldSVar :: MonadAsync m => State t m a -> Fold m a b -> m (SVar t m a)
4055newFoldSVar stt f = do
4056 -- Buffer size for the SVar is derived from the current state
4057 sv <- newParallelVar StopAny (adaptState stt)
4058 -- Add the producer thread-id to the SVar.
4059 liftIO myThreadId >>= modifyThread sv
4060 void $ doFork (work sv) (svarMrun sv) (handleFoldException sv)
4065 {-# NOINLINE work #-}
4066 work sv = void $ runFold f $ fromProducer sv
4068data TapState sv st = TapInit | Tapping sv st | TapDone st
4070{-# INLINE_NORMAL tapAsync #-}
4071tapAsync :: MonadAsync m => Fold m a b -> Stream m a -> Stream m a
4072tapAsync f (Stream step1 state1) = Stream step TapInit
4076 -- In general, a Stop event would come equipped with the result
4077 -- of the fold. It is not used here but it would be useful in
4078 -- applicative and distribute.
4079 done <- fromConsumer svr
4080 when (not done) $ do
4081 liftIO $ withDiagMVar svr "teeToSVar: waiting to drain"
4082 $ takeMVar (outputDoorBellFromConsumer svr)
4086 liftIO $ sendStop svr Nothing
4087 -- drain/wait until a stop event arrives from the fold.
4090 {-# INLINE_LATE step #-}
4091 step gst TapInit = do
4092 sv <- newFoldSVar gst f
4093 return $ Skip (Tapping sv state1)
4095 step gst (Tapping sv st) = do
4099 done <- pushToFold sv a
4102 -- XXX we do not need to wait synchronously here
4104 return $ Yield a (TapDone s)
4105 else return $ Yield a (Tapping sv s)
4106 Skip s -> return $ Skip (Tapping sv s)
4111 step gst (TapDone st) = do
4114 Yield a s -> Yield a (TapDone s)
4115 Skip s -> Skip (TapDone s)
4118-- XXX Exported from Array again as this fold is specific to Array
4119-- | Take last 'n' elements from the stream and discard the rest.
4121lastN :: (Storable a, MonadIO m) => Int -> Fold m a (Array a)
4122lastN n = Fold step initial done
4124 step (Tuple3' rb rh i) a = do
4125 rh1 <- liftIO $ RB.unsafeInsert rb rh a
4126 return $ Tuple3' rb rh1 (i + 1)
4127 initial = fmap (\(a, b) -> Tuple3' a b (0 :: Int)) $ liftIO $ RB.new n
4128 done (Tuple3' rb rh i) = do
4129 arr <- liftIO $ A.newArray n
4130 foldFunc i rh snoc' arr rb
4131 snoc' b a = liftIO $ A.unsafeSnoc b a
4133 | i < n = RB.unsafeFoldRingM
4134 | otherwise = RB.unsafeFoldRingFullM
4136------------------------------------------------------------------------------
4138------------------------------------------------------------------------------
4140-- XXX using getTime in the loop can be pretty expensive especially for
4141-- computations where iterations are lightweight. We have the following
4144-- 1) Run a timeout thread updating a flag asynchronously and check that
4145-- flag here, that way we can have a cheap termination check.
4147-- 2) Use COARSE clock to get time with lower resolution but more efficiently.
4149-- 3) Use rdtscp/rdtsc to get time directly from the processor, compute the
4150-- termination value of rdtsc in the beginning and then in each iteration just
4151-- get rdtsc and check if we should terminate.
4155 | TakeByTimeCheck st s
4156 | TakeByTimeYield st s
4158{-# INLINE_NORMAL takeByTime #-}
4159takeByTime :: (MonadIO m, TimeUnit64 t) => t -> Stream m a -> Stream m a
4160takeByTime duration (Stream step1 state1) = Stream step (TakeByTimeInit state1)
4163 lim = toRelTime64 duration
4165 {-# INLINE_LATE step #-}
4166 step _ (TakeByTimeInit _) | lim == 0 = return Stop
4167 step _ (TakeByTimeInit st) = do
4168 t0 <- liftIO $ getTime Monotonic
4169 return $ Skip (TakeByTimeYield st t0)
4170 step _ (TakeByTimeCheck st t0) = do
4171 t <- liftIO $ getTime Monotonic
4173 if diffAbsTime64 t t0 > lim
4175 else Skip (TakeByTimeYield st t0)
4176 step gst (TakeByTimeYield st t0) = do
4179 Yield x s -> Yield x (TakeByTimeCheck s t0)
4180 Skip s -> Skip (TakeByTimeCheck s t0)
4183data DropByTime st s x
4185 | DropByTimeGen st s
4186 | DropByTimeCheck st s x
4187 | DropByTimeYield st
4189{-# INLINE_NORMAL dropByTime #-}
4190dropByTime :: (MonadIO m, TimeUnit64 t) => t -> Stream m a -> Stream m a
4191dropByTime duration (Stream step1 state1) = Stream step (DropByTimeInit state1)
4194 lim = toRelTime64 duration
4196 {-# INLINE_LATE step #-}
4197 step _ (DropByTimeInit st) = do
4198 t0 <- liftIO $ getTime Monotonic
4199 return $ Skip (DropByTimeGen st t0)
4200 step gst (DropByTimeGen st t0) = do
4203 Yield x s -> Skip (DropByTimeCheck s t0 x)
4204 Skip s -> Skip (DropByTimeGen s t0)
4206 step _ (DropByTimeCheck st t0 x) = do
4207 t <- liftIO $ getTime Monotonic
4208 if diffAbsTime64 t t0 <= lim
4209 then return $ Skip $ DropByTimeGen st t0
4210 else return $ Yield x $ DropByTimeYield st
4211 step gst (DropByTimeYield st) = do
4214 Yield x s -> Yield x (DropByTimeYield s)
4215 Skip s -> Skip (DropByTimeYield s)
4218-- XXX we should move this to stream generation section of this file. Also, the
4219-- take/drop combinators above should be moved to filtering section.
4220{-# INLINE_NORMAL currentTime #-}
4221currentTime :: MonadAsync m => Double -> Stream m AbsTime
4222currentTime g = Stream step Nothing
4226 g' = g * 10 ^ (6 :: Int)
4228 -- XXX should have a minimum granularity to avoid high CPU usage?
4229 {-# INLINE delayTime #-}
4231 if g' >= fromIntegral (maxBound :: Int)
4235 updateTimeVar timeVar = do
4236 threadDelay $ delayTime
4237 MicroSecond64 t <- fromAbsTime <$> getTime Monotonic
4238 modifyVar' timeVar (const t)
4240 {-# INLINE_LATE step #-}
4242 -- XXX note that this is safe only on a 64-bit machine. On a 32-bit
4243 -- machine a 64-bit 'Var' cannot be read consistently without a lock
4244 -- while another thread is writing to it.
4245 timeVar <- liftIO $ newVar (0 :: Int64)
4246 tid <- forkManaged $ liftIO $ forever (updateTimeVar timeVar)
4247 return $ Skip $ Just (timeVar, tid)
4249 step _ s@(Just (timeVar, _)) = do
4250 a <- liftIO $ readVar timeVar
4251 -- XXX we can perhaps use an AbsTime64 using a 64 bit Int for
4252 -- efficiency. or maybe we can use a representation using Double for
4253 -- floating precision time
4254 return $ Yield (toAbsTime (MicroSecond64 a)) 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