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
97 -- | Transform an input structure into a stream.
98 -- | Direct style stream does not support @fromFoldable@.
128 -- ** Specialized Folds
155 -- ** Flattening nested streams
158 , ConcatMapUState (..)
160 , ConcatUnfoldInterleaveState (..)
161 , concatUnfoldInterleave
162 , concatUnfoldRoundrobin
165 , InterleaveState(..)
170 , roundRobin -- interleaveFair?/ParallelFair
203 -- | Transform a stream into another type.
220 -- ** By folding (scans)
265 , intersperseSuffixBySpan
289 -- * Transformation comprehensions
305 -- * Concurrent Application
317import Control.Concurrent (killThread, myThreadId, takeMVar, threadDelay)
318import Control.Exception
319 (Exception, SomeException, AsyncException, fromException)
320import Control.Monad (void, when)
321import Control.Monad.Catch (MonadCatch, throwM)
322import Control.Monad.IO.Class (MonadIO(..))
323import Control.Monad.Reader (ReaderT)
324import Control.Monad.State.Strict (StateT)
325import Control.Monad.Trans (MonadTrans(lift))
326import Control.Monad.Trans.Control (MonadBaseControl)
327import Data.Bits (shiftR, shiftL, (.|.), (.&.))
328import Data.Functor.Identity (Identity(..))
329import Data.IORef (newIORef, readIORef, mkWeakIORef, writeIORef, IORef)
330import Data.Maybe (fromJust, isJust, isNothing)
331import Data.Word (Word32)
332import Foreign.Ptr (Ptr)
333import Foreign.Storable (Storable(..))
334import GHC.Types (SPEC(..))
335import System.Mem (performMajorGC)
337 hiding (map, mapM, mapM_, repeat, foldr, last, take, filter,
338 takeWhile, drop, dropWhile, all, any, maximum, minimum, elem,
339 notElem, null, head, tail, zipWith, lookup, foldr1, sequence,
340 (!!), scanl, scanl1, concatMap, replicate, enumFromTo, concat,
343import qualified Control.Monad.Catch as MC
344import qualified Control.Monad.Reader as Reader
345import qualified Control.Monad.State.Strict as State
346import qualified Prelude
348import Streamly.Internal.Mutable.Prim.Var
349 (Prim, Var, readVar, newVar, modifyVar')
350import Streamly.Internal.Data.Time.Units
351 (TimeUnit64, toRelTime64, diffAbsTime64)
353import Streamly.Internal.Data.Atomics (atomicModifyIORefCAS_)
354import Streamly.Internal.Memory.Array.Types (Array(..))
355import Streamly.Internal.Data.Fold.Types (Fold(..))
356import Streamly.Internal.Data.Pipe.Types (Pipe(..), PipeState(..))
357import Streamly.Internal.Data.Time.Clock (Clock(Monotonic), getTime)
358import Streamly.Internal.Data.Unfold.Types (Unfold(..))
359import Streamly.Internal.Data.Strict (Tuple3'(..))
361import Streamly.Internal.Data.Stream.StreamD.Type
362import Streamly.Internal.Data.SVar
363import Streamly.Internal.Data.Stream.SVar (fromConsumer, pushToFold)
365import qualified Streamly.Internal.Data.Pipe.Types as Pipe
366import qualified Streamly.Internal.Memory.Array.Types as A
367import qualified Streamly.Internal.Data.Fold as FL
368import qualified Streamly.Memory.Ring as RB
369import qualified Streamly.Internal.Data.Stream.StreamK as K
371------------------------------------------------------------------------------
373------------------------------------------------------------------------------
375-- | An empty 'Stream'.
376{-# INLINE_NORMAL nil #-}
377nil :: Monad m => Stream m a
378nil = Stream (\_ _ -> return Stop) ()
380-- | An empty 'Stream' with a side effect.
381{-# INLINE_NORMAL nilM #-}
382nilM :: Monad m => m b -> Stream m a
383nilM m = Stream (\_ _ -> m >> return Stop) ()
385{-# INLINE_NORMAL consM #-}
386consM :: Monad m => m a -> Stream m a -> Stream m a
387consM m (Stream step state) = Stream step1 Nothing
389 {-# INLINE_LATE step1 #-}
390 step1 _ Nothing = m >>= \x -> return $ Yield x (Just state)
391 step1 gst (Just st) = do
395 Yield a s -> Yield a (Just s)
396 Skip s -> Skip (Just s)
399-- XXX implement in terms of consM?
400-- cons x = consM (return x)
402-- | Can fuse but has O(n^2) complexity.
403{-# INLINE_NORMAL cons #-}
404cons :: Monad m => a -> Stream m a -> Stream m a
405cons x (Stream step state) = Stream step1 Nothing
407 {-# INLINE_LATE step1 #-}
408 step1 _ Nothing = return $ Yield x (Just state)
409 step1 gst (Just st) = do
413 Yield a s -> Yield a (Just s)
414 Skip s -> Skip (Just s)
417-------------------------------------------------------------------------------
419-------------------------------------------------------------------------------
421-- Does not fuse, has the same performance as the StreamK version.
422{-# INLINE_NORMAL uncons #-}
423uncons :: Monad m => Stream m a -> m (Maybe (a, Stream m a))
424uncons (UnStream step state) = go state
427 r <- step defState st
429 Yield x s -> return $ Just (x, Stream step s)
431 Stop -> return Nothing
433------------------------------------------------------------------------------
434-- Generation by unfold
435------------------------------------------------------------------------------
437{-# INLINE_NORMAL unfoldrM #-}
438unfoldrM :: Monad m => (s -> m (Maybe (a, s))) -> s -> Stream m a
439unfoldrM next state = Stream step state
441 {-# INLINE_LATE step #-}
445 Just (x, s) -> Yield x s
448{-# INLINE_LATE unfoldr #-}
449unfoldr :: Monad m => (s -> Maybe (a, s)) -> s -> Stream m a
450unfoldr f = unfoldrM (return . f)
452-- | Convert an 'Unfold' into a 'Stream' by supplying it a seed.
454{-# INLINE_NORMAL unfold #-}
455unfold :: Monad m => Unfold m a b -> a -> Stream m b
456unfold (Unfold ustep inject) seed = Stream step Nothing
458 {-# INLINE_LATE step #-}
459 step _ Nothing = inject seed >>= return . Skip . Just
460 step _ (Just st) = do
463 Yield x s -> Yield x (Just s)
464 Skip s -> Skip (Just s)
467------------------------------------------------------------------------------
468-- Specialized Generation
469------------------------------------------------------------------------------
471{-# INLINE_NORMAL repeatM #-}
472repeatM :: Monad m => m a -> Stream m a
473repeatM x = Stream (\_ _ -> x >>= \r -> return $ Yield r ()) ()
475{-# INLINE_NORMAL repeat #-}
476repeat :: Monad m => a -> Stream m a
477repeat x = Stream (\_ _ -> return $ Yield x ()) ()
479{-# INLINE_NORMAL iterateM #-}
480iterateM :: Monad m => (a -> m a) -> m a -> Stream m a
481iterateM step = Stream (\_ st -> st >>= \x -> return $ Yield x (step x))
483{-# INLINE_NORMAL iterate #-}
484iterate :: Monad m => (a -> a) -> a -> Stream m a
485iterate step st = iterateM (return . step) (return st)
487{-# INLINE_NORMAL replicateM #-}
488replicateM :: forall m a. Monad m => Int -> m a -> Stream m a
489replicateM n p = Stream step n
491 {-# INLINE_LATE step #-}
493 | i <= 0 = return Stop
496 return $ Yield x (i - 1)
498{-# INLINE_NORMAL replicate #-}
499replicate :: Monad m => Int -> a -> Stream m a
500replicate n x = replicateM n (return x)
502-- This would not work properly for floats, therefore we put an Integral
504-- | Can be used to enumerate unbounded integrals. This does not check for
505-- overflow or underflow for bounded integrals.
506{-# INLINE_NORMAL enumerateFromStepIntegral #-}
507enumerateFromStepIntegral :: (Integral a, Monad m) => a -> a -> Stream m a
508enumerateFromStepIntegral from stride =
509 from `seq` stride `seq` Stream step from
511 {-# INLINE_LATE step #-}
512 step _ !x = return $ Yield x $! (x + stride)
514-- We are assuming that "to" is constrained by the type to be within
516{-# INLINE enumerateFromToIntegral #-}
517enumerateFromToIntegral :: (Monad m, Integral a) => a -> a -> Stream m a
518enumerateFromToIntegral from to =
519 takeWhile (<= to) $ enumerateFromStepIntegral from 1
521{-# INLINE enumerateFromIntegral #-}
522enumerateFromIntegral :: (Monad m, Integral a, Bounded a) => a -> Stream m a
523enumerateFromIntegral from = enumerateFromToIntegral from maxBound
525data EnumState a = EnumInit | EnumYield a a a | EnumStop
527{-# INLINE_NORMAL enumerateFromThenToIntegralUp #-}
528enumerateFromThenToIntegralUp
529 :: (Monad m, Integral a)
530 => a -> a -> a -> Stream m a
531enumerateFromThenToIntegralUp from next to = Stream step EnumInit
533 {-# INLINE_LATE step #-}
539 else Yield from EnumStop
540 else -- from <= next <= to
541 let stride = next - from
542 in Skip $ EnumYield from stride (to - stride)
544 step _ (EnumYield x stride toMinus) =
547 then Yield x EnumStop
548 else Yield x $ EnumYield (x + stride) stride toMinus
550 step _ EnumStop = return Stop
552{-# INLINE_NORMAL enumerateFromThenToIntegralDn #-}
553enumerateFromThenToIntegralDn
554 :: (Monad m, Integral a)
555 => a -> a -> a -> Stream m a
556enumerateFromThenToIntegralDn from next to = Stream step EnumInit
558 {-# INLINE_LATE step #-}
560 return $ if to > next
563 else Yield from EnumStop
564 else -- from >= next >= to
565 let stride = next - from
566 in Skip $ EnumYield from stride (to - stride)
568 step _ (EnumYield x stride toMinus) =
571 then Yield x EnumStop
572 else Yield x $ EnumYield (x + stride) stride toMinus
574 step _ EnumStop = return Stop
576{-# INLINE_NORMAL enumerateFromThenToIntegral #-}
577enumerateFromThenToIntegral
578 :: (Monad m, Integral a)
579 => a -> a -> a -> Stream m a
580enumerateFromThenToIntegral from next to
581 | next >= from = enumerateFromThenToIntegralUp from next to
582 | otherwise = enumerateFromThenToIntegralDn from next to
584{-# INLINE_NORMAL enumerateFromThenIntegral #-}
585enumerateFromThenIntegral
586 :: (Monad m, Integral a, Bounded a)
587 => a -> a -> Stream m a
588enumerateFromThenIntegral from next =
590 then enumerateFromThenToIntegralUp from next maxBound
591 else enumerateFromThenToIntegralDn from next minBound
593-- For floating point numbers if the increment is less than the precision then
594-- it just gets lost. Therefore we cannot always increment it correctly by just
596-- 9007199254740992 + 1 + 1 :: Double => 9.007199254740992e15
597-- 9007199254740992 + 2 :: Double => 9.007199254740994e15
599-- Instead we accumulate the increment counter and compute the increment
600-- every time before adding it to the starting number.
602-- This works for Integrals as well as floating point numbers, but
603-- enumerateFromStepIntegral is faster for integrals.
604{-# INLINE_NORMAL enumerateFromStepNum #-}
605enumerateFromStepNum :: (Monad m, Num a) => a -> a -> Stream m a
606enumerateFromStepNum from stride = Stream step 0
608 {-# INLINE_LATE step #-}
609 step _ !i = return $ (Yield $! (from + i * stride)) $! (i + 1)
611{-# INLINE_NORMAL numFrom #-}
612numFrom :: (Monad m, Num a) => a -> Stream m a
613numFrom from = enumerateFromStepNum from 1
615{-# INLINE_NORMAL numFromThen #-}
616numFromThen :: (Monad m, Num a) => a -> a -> Stream m a
617numFromThen from next = enumerateFromStepNum from (next - from)
619-- We cannot write a general function for Num. The only way to write code
620-- portable between the two is to use a 'Real' constraint and convert between
621-- Fractional and Integral using fromRational which is horribly slow.
622{-# INLINE_NORMAL enumerateFromToFractional #-}
623enumerateFromToFractional
624 :: (Monad m, Fractional a, Ord a)
625 => a -> a -> Stream m a
626enumerateFromToFractional from to =
627 takeWhile (<= to + 1 / 2) $ enumerateFromStepNum from 1
629{-# INLINE_NORMAL enumerateFromThenToFractional #-}
630enumerateFromThenToFractional
631 :: (Monad m, Fractional a, Ord a)
632 => a -> a -> a -> Stream m a
633enumerateFromThenToFractional from next to =
634 takeWhile predicate $ numFromThen from next
636 mid = (next - from) / 2
637 predicate | next >= from = (<= to + mid)
638 | otherwise = (>= to + mid)
640-------------------------------------------------------------------------------
641-- Generation by Conversion
642-------------------------------------------------------------------------------
644{-# INLINE_NORMAL fromIndicesM #-}
645fromIndicesM :: Monad m => (Int -> m a) -> Stream m a
646fromIndicesM gen = Stream step 0
648 {-# INLINE_LATE step #-}
651 return $ Yield x (i + 1)
653{-# INLINE fromIndices #-}
654fromIndices :: Monad m => (Int -> a) -> Stream m a
655fromIndices gen = fromIndicesM (return . gen)
657{-# INLINE_NORMAL generateM #-}
658generateM :: Monad m => Int -> (Int -> m a) -> Stream m a
659generateM n gen = n `seq` Stream step 0
661 {-# INLINE_LATE step #-}
662 step _ i | i < n = do
664 return $ Yield x (i + 1)
665 | otherwise = return Stop
667{-# INLINE generate #-}
668generate :: Monad m => Int -> (Int -> a) -> Stream m a
669generate n gen = generateM n (return . gen)
671-- XXX we need the MonadAsync constraint because of a rewrite rule.
672-- | Convert a list of monadic actions to a 'Stream'
673{-# INLINE_LATE fromListM #-}
674fromListM :: MonadAsync m => [m a] -> Stream m a
675fromListM = Stream step
677 {-# INLINE_LATE step #-}
678 step _ (m:ms) = m >>= \x -> return $ Yield x ms
679 step _ [] = return Stop
681{-# INLINE toStreamD #-}
682toStreamD :: (K.IsStream t, Monad m) => t m a -> Stream m a
683toStreamD = fromStreamK . K.toStream
685{-# INLINE_NORMAL fromPrimVar #-}
686fromPrimVar :: (MonadIO m, Prim a) => Var IO a -> Stream m a
687fromPrimVar var = Stream step ()
689 {-# INLINE_LATE step #-}
690 step _ () = liftIO (readVar var) >>= \x -> return $ Yield x ()
692-------------------------------------------------------------------------------
693-- Generation from SVar
694-------------------------------------------------------------------------------
696data FromSVarState t m a =
698 | FromSVarRead (SVar t m a)
699 | FromSVarLoop (SVar t m a) [ChildEvent a]
700 | FromSVarDone (SVar t m a)
702{-# INLINE_NORMAL fromSVar #-}
703fromSVar :: (MonadAsync m) => SVar t m a -> Stream m a
704fromSVar svar = Stream step FromSVarInit
707 {-# INLINE_LATE step #-}
708 step _ FromSVarInit = do
709 ref <- liftIO $ newIORef ()
710 _ <- liftIO $ mkWeakIORef ref hook
711 -- when this copy of svar gets garbage collected "ref" will get
712 -- garbage collected and our GC hook will be called.
713 let sv = svar{svarRef = Just ref}
714 return $ Skip (FromSVarRead sv)
718 {-# NOINLINE hook #-}
720 when (svarInspectMode svar) $ do
721 r <- liftIO $ readIORef (svarStopTime (svarStats svar))
723 printSVar svar "SVar Garbage Collected"
725 -- If there are any SVars referenced by this SVar a GC will prompt
726 -- them to be cleaned up quickly.
727 when (svarInspectMode svar) performMajorGC
729 step _ (FromSVarRead sv) = do
730 list <- readOutputQ sv
731 -- Reversing the output is important to guarantee that we process the
732 -- outputs in the same order as they were generated by the constituent
734 return $ Skip $ FromSVarLoop sv (Prelude.reverse list)
736 step _ (FromSVarLoop sv []) = do
737 done <- postProcess sv
738 return $ Skip $ if done
739 then (FromSVarDone sv)
740 else (FromSVarRead sv)
742 step _ (FromSVarLoop sv (ev : es)) = do
744 ChildYield a -> return $ Yield a (FromSVarLoop sv es)
745 ChildStop tid e -> do
749 stop <- shouldStop tid
752 liftIO (cleanupSVar sv)
753 return $ Skip (FromSVarDone sv)
754 else return $ Skip (FromSVarLoop sv es)
756 case fromException ex of
758 return $ Skip (FromSVarLoop sv es)
759 Nothing -> liftIO (cleanupSVar sv) >> throwM ex
763 case svarStopStyle sv of
764 StopNone -> return False
765 StopAny -> return True
767 sid <- liftIO $ readIORef (svarStopBy sv)
768 return $ if tid == sid then True else False
770 step _ (FromSVarDone sv) = do
771 when (svarInspectMode sv) $ do
772 t <- liftIO $ getTime Monotonic
773 liftIO $ writeIORef (svarStopTime (svarStats sv)) (Just t)
774 liftIO $ printSVar sv "SVar Done"
777-------------------------------------------------------------------------------
778-- Process events received by a fold consumer from a stream producer
779-------------------------------------------------------------------------------
781{-# INLINE_NORMAL fromProducer #-}
782fromProducer :: (MonadAsync m) => SVar t m a -> Stream m a
783fromProducer svar = Stream step (FromSVarRead svar)
786 {-# INLINE_LATE step #-}
787 step _ (FromSVarRead sv) = do
788 list <- readOutputQ sv
789 -- Reversing the output is important to guarantee that we process the
790 -- outputs in the same order as they were generated by the constituent
792 return $ Skip $ FromSVarLoop sv (Prelude.reverse list)
794 step _ (FromSVarLoop sv []) = return $ Skip $ FromSVarRead sv
795 step _ (FromSVarLoop sv (ev : es)) = do
797 ChildYield a -> return $ Yield a (FromSVarLoop sv es)
798 ChildStop tid e -> do
802 sendStopToProducer sv
803 return $ Skip (FromSVarDone sv)
804 Just _ -> error "Bug: fromProducer: received exception"
806 step _ (FromSVarDone sv) = do
807 when (svarInspectMode sv) $ do
808 t <- liftIO $ getTime Monotonic
809 liftIO $ writeIORef (svarStopTime (svarStats sv)) (Just t)
810 liftIO $ printSVar sv "SVar Done"
813 step _ FromSVarInit = undefined
815-------------------------------------------------------------------------------
816-- Hoisting the inner monad
817-------------------------------------------------------------------------------
819{-# INLINE_NORMAL hoist #-}
820hoist :: Monad n => (forall x. m x -> n x) -> Stream m a -> Stream n a
821hoist f (Stream step state) = (Stream step' state)
823 {-# INLINE_LATE step' #-}
825 r <- f $ step (adaptState gst) st
827 Yield x s -> Yield x s
831{-# INLINE generally #-}
832generally :: Monad m => Stream Identity a -> Stream m a
833generally = hoist (return . runIdentity)
835{-# INLINE_NORMAL liftInner #-}
836liftInner :: (Monad m, MonadTrans t, Monad (t m))
837 => Stream m a -> Stream (t m) a
838liftInner (Stream step state) = Stream step' state
840 {-# INLINE_LATE step' #-}
842 r <- lift $ step (adaptState gst) st
844 Yield x s -> Yield x s
848{-# INLINE_NORMAL runReaderT #-}
849runReaderT :: Monad m => s -> Stream (ReaderT s m) a -> Stream m a
850runReaderT sval (Stream step state) = Stream step' state
852 {-# INLINE_LATE step' #-}
854 r <- Reader.runReaderT (step (adaptState gst) st) sval
856 Yield x s -> Yield x s
860{-# INLINE_NORMAL evalStateT #-}
861evalStateT :: Monad m => s -> Stream (StateT s m) a -> Stream m a
862evalStateT sval (Stream step state) = Stream step' (state, sval)
864 {-# INLINE_LATE step' #-}
865 step' gst (st, sv) = do
866 (r, sv') <- State.runStateT (step (adaptState gst) st) sv
868 Yield x s -> Yield x (s, sv')
869 Skip s -> Skip (s, sv')
872{-# INLINE_NORMAL runStateT #-}
873runStateT :: Monad m => s -> Stream (StateT s m) a -> Stream m (s, a)
874runStateT sval (Stream step state) = Stream step' (state, sval)
876 {-# INLINE_LATE step' #-}
877 step' gst (st, sv) = do
878 (r, sv') <- State.runStateT (step (adaptState gst) st) sv
880 Yield x s -> Yield (sv', x) (s, sv')
881 Skip s -> Skip (s, sv')
884------------------------------------------------------------------------------
885-- Elimination by Folds
886------------------------------------------------------------------------------
888------------------------------------------------------------------------------
890------------------------------------------------------------------------------
892{-# INLINE_NORMAL foldr1 #-}
893foldr1 :: Monad m => (a -> a -> a) -> Stream m a -> m (Maybe a)
897 Nothing -> return Nothing
898 Just (h, t) -> fmap Just (foldr f h t)
900------------------------------------------------------------------------------
902------------------------------------------------------------------------------
904{-# INLINE_NORMAL foldlT #-}
905foldlT :: (Monad m, Monad (s m), MonadTrans s)
906 => (s m b -> a -> s m b) -> s m b -> Stream m a -> s m b
907foldlT fstep begin (Stream step state) = go SPEC begin state
910 r <- lift $ step defState st
912 Yield x s -> go SPEC (fstep acc x) s
913 Skip s -> go SPEC acc s
916-- Note, this is going to have horrible performance, because of the nature of
917-- the stream type (i.e. direct stream vs CPS). Its only for reference, it is
918-- likely be practically unusable.
919{-# INLINE_NORMAL foldlS #-}
921 => (Stream m b -> a -> Stream m b) -> Stream m b -> Stream m a -> Stream m b
922foldlS fstep begin (Stream step state) = Stream step' (Left (state, begin))
924 step' gst (Left (st, acc)) = do
925 r <- step (adaptState gst) st
927 Yield x s -> Skip (Left (s, fstep acc x))
928 Skip s -> Skip (Left (s, acc))
929 Stop -> Skip (Right acc)
931 step' gst (Right (Stream stp stt)) = do
932 r <- stp (adaptState gst) stt
934 Yield x s -> Yield x (Right (Stream stp s))
935 Skip s -> Skip (Right (Stream stp s))
938------------------------------------------------------------------------------
940------------------------------------------------------------------------------
942-- | Run a streaming composition, discard the results.
943{-# INLINE_LATE drain #-}
944drain :: Monad m => Stream m a -> m ()
945-- drain = foldrM (\_ xs -> xs) (return ())
946drain (Stream step state) = go SPEC state
949 r <- step defState st
951 Yield _ s -> go SPEC s
955{-# INLINE_NORMAL null #-}
956null :: Monad m => Stream m a -> m Bool
957null m = foldrM (\_ _ -> return False) (return True) m
959{-# INLINE_NORMAL head #-}
960head :: Monad m => Stream m a -> m (Maybe a)
961head m = foldrM (\x _ -> return (Just x)) (return Nothing) m
963{-# INLINE_NORMAL headElse #-}
964headElse :: Monad m => a -> Stream m a -> m a
965headElse a m = foldrM (\x _ -> return x) (return a) m
967-- Does not fuse, has the same performance as the StreamK version.
968{-# INLINE_NORMAL tail #-}
969tail :: Monad m => Stream m a -> m (Maybe (Stream m a))
970tail (UnStream step state) = go state
973 r <- step defState st
975 Yield _ s -> return (Just $ Stream step s)
977 Stop -> return Nothing
979-- XXX will it fuse? need custom impl?
980{-# INLINE_NORMAL last #-}
981last :: Monad m => Stream m a -> m (Maybe a)
982last = foldl' (\_ y -> Just y) Nothing
984{-# INLINE_NORMAL elem #-}
985elem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
986-- elem e m = foldrM (\x xs -> if x == e then return True else xs) (return False) m
987elem e (Stream step state) = go state
990 r <- step defState st
993 | x == e -> return True
998{-# INLINE_NORMAL notElem #-}
999notElem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
1000notElem e s = fmap not (elem e s)
1002{-# INLINE_NORMAL all #-}
1003all :: Monad m => (a -> Bool) -> Stream m a -> m Bool
1004-- all p m = foldrM (\x xs -> if p x then xs else return False) (return True) m
1005all p (Stream step state) = go state
1008 r <- step defState st
1012 | otherwise -> return False
1016{-# INLINE_NORMAL any #-}
1017any :: Monad m => (a -> Bool) -> Stream m a -> m Bool
1018-- any p m = foldrM (\x xs -> if p x then return True else xs) (return False) m
1019any p (Stream step state) = go state
1022 r <- step defState st
1025 | p x -> return True
1028 Stop -> return False
1030{-# INLINE_NORMAL maximum #-}
1031maximum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
1032maximum (Stream step state) = go Nothing state
1035 r <- step defState st
1037 Yield x s -> go (Just x) s
1038 Skip s -> go Nothing s
1039 Stop -> return Nothing
1040 go (Just acc) st = do
1041 r <- step defState st
1044 | acc <= x -> go (Just x) s
1045 | otherwise -> go (Just acc) s
1046 Skip s -> go (Just acc) s
1047 Stop -> return (Just acc)
1049{-# INLINE_NORMAL maximumBy #-}
1050maximumBy :: Monad m => (a -> a -> Ordering) -> Stream m a -> m (Maybe a)
1051maximumBy cmp (Stream step state) = go Nothing state
1054 r <- step defState st
1056 Yield x s -> go (Just x) s
1057 Skip s -> go Nothing s
1058 Stop -> return Nothing
1059 go (Just acc) st = do
1060 r <- step defState st
1062 Yield x s -> case cmp acc x of
1063 GT -> go (Just acc) s
1065 Skip s -> go (Just acc) s
1066 Stop -> return (Just acc)
1068{-# INLINE_NORMAL minimum #-}
1069minimum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
1070minimum (Stream step state) = go Nothing state
1073 r <- step defState st
1075 Yield x s -> go (Just x) s
1076 Skip s -> go Nothing s
1077 Stop -> return Nothing
1078 go (Just acc) st = do
1079 r <- step defState st
1082 | acc <= x -> go (Just acc) s
1083 | otherwise -> go (Just x) s
1084 Skip s -> go (Just acc) s
1085 Stop -> return (Just acc)
1087{-# INLINE_NORMAL minimumBy #-}
1088minimumBy :: Monad m => (a -> a -> Ordering) -> Stream m a -> m (Maybe a)
1089minimumBy cmp (Stream step state) = go Nothing state
1092 r <- step defState st
1094 Yield x s -> go (Just x) s
1095 Skip s -> go Nothing s
1096 Stop -> return Nothing
1097 go (Just acc) st = do
1098 r <- step defState st
1100 Yield x s -> case cmp acc x of
1102 _ -> go (Just acc) s
1103 Skip s -> go (Just acc) s
1104 Stop -> return (Just acc)
1106{-# INLINE_NORMAL (!!) #-}
1107(!!) :: (Monad m) => Stream m a -> Int -> m (Maybe a)
1108(Stream step state) !! i = go i state
1111 r <- step defState st
1113 Yield x s | n < 0 -> return Nothing
1114 | n == 0 -> return $ Just x
1115 | otherwise -> go (n - 1) s
1117 Stop -> return Nothing
1119{-# INLINE_NORMAL lookup #-}
1120lookup :: (Monad m, Eq a) => a -> Stream m (a, b) -> m (Maybe b)
1121lookup e m = foldrM (\(a, b) xs -> if e == a then return (Just b) else xs)
1124{-# INLINE_NORMAL findM #-}
1125findM :: Monad m => (a -> m Bool) -> Stream m a -> m (Maybe a)
1126findM p m = foldrM (\x xs -> p x >>= \r -> if r then return (Just x) else xs)
1130find :: Monad m => (a -> Bool) -> Stream m a -> m (Maybe a)
1131find p = findM (return . p)
1133{-# INLINE_NORMAL findIndices #-}
1134findIndices :: Monad m => (a -> Bool) -> Stream m a -> Stream m Int
1135findIndices p (Stream step state) = Stream step' (state, 0)
1137 {-# INLINE_LATE step' #-}
1138 step' gst (st, i) = i `seq` do
1139 r <- step (adaptState gst) st
1141 Yield x s -> if p x then Yield i (s, i+1) else Skip (s, i+1)
1142 Skip s -> Skip (s, i)
1145{-# INLINE toListRev #-}
1146toListRev :: Monad m => Stream m a -> m [a]
1147toListRev = foldl' (flip (:)) []
1149-- We can implement reverse as:
1151-- > reverse = foldlS (flip cons) nil
1153-- However, this implementation is unusable because of the horrible performance
1154-- of cons. So we just convert it to a list first and then stream from the
1157-- XXX Maybe we can use an Array instead of a list here?
1158{-# INLINE_NORMAL reverse #-}
1159reverse :: Monad m => Stream m a -> Stream m a
1160reverse m = Stream step Nothing
1162 {-# INLINE_LATE step #-}
1165 return $ Skip (Just xs)
1166 step _ (Just (x:xs)) = return $ Yield x (Just xs)
1167 step _ (Just []) = return Stop
1169-- Much faster reverse for Storables
1170{-# INLINE_NORMAL reverse' #-}
1171reverse' :: forall m a. (MonadIO m, Storable a) => Stream m a -> Stream m a
1173-- This commented implementation copies the whole stream into one single array
1174-- and then streams from that array, this is 3-4x faster than the chunked code
1175-- that follows. Though this could be problematic due to unbounded large
1176-- allocations. We need to figure out why the chunked code is slower and if we
1177-- can optimize the chunked code to work as fast as this one. It may be a
1179import Foreign.ForeignPtr (touchForeignPtr)
1180import Foreign.ForeignPtr.Unsafe (unsafeForeignPtrToPtr)
1181import Foreign.Ptr (Ptr, plusPtr)
1182reverse' m = Stream step Nothing
1184 {-# INLINE_LATE step #-}
1186 arr <- A.fromStreamD m
1187 let p = aEnd arr `plusPtr` negate (sizeOf (undefined :: a))
1188 return $ Skip $ Just (aStart arr, p)
1190 step _ (Just (start, p)) | p < unsafeForeignPtrToPtr start = return Stop
1192 step _ (Just (start, p)) = do
1193 let !x = A.unsafeInlineIO $ do
1195 touchForeignPtr start
1197 next = p `plusPtr` negate (sizeOf (undefined :: a))
1198 return $ Yield x (Just (start, next))
1205 $ A.fromStreamDArraysOf A.defaultChunkSize m
1208------------------------------------------------------------------------------
1209-- Grouping/Splitting
1210------------------------------------------------------------------------------
1212{-# INLINE_NORMAL splitSuffixBy' #-}
1213splitSuffixBy' :: Monad m
1214 => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1215splitSuffixBy' predicate f (Stream step state) =
1216 Stream (stepOuter f) (Just state)
1220 {-# INLINE_LATE stepOuter #-}
1221 stepOuter (Fold fstep initial done) gst (Just st) = do
1222 res <- step (adaptState gst) st
1228 then done acc' >>= \val -> return $ Yield val (Just s)
1231 Skip s -> return $ Skip $ Just s
1237 res <- step (adaptState gst) stt
1242 then done acc' >>= \val -> return $ Yield val (Just s)
1244 Skip s -> go SPEC s acc
1245 Stop -> done acc >>= \val -> return $ Yield val Nothing
1247 stepOuter _ _ Nothing = return Stop
1249{-# INLINE_NORMAL groupsBy #-}
1255groupsBy cmp f (Stream step state) = Stream (stepOuter f) (Just state, Nothing)
1259 {-# INLINE_LATE stepOuter #-}
1260 stepOuter (Fold fstep initial done) gst (Just st, Nothing) = do
1261 res <- step (adaptState gst) st
1268 Skip s -> return $ Skip $ (Just s, Nothing)
1273 go !_ prev stt !acc = do
1274 res <- step (adaptState gst) stt
1281 else done acc >>= \r -> return $ Yield r (Just s, Just x)
1282 Skip s -> go SPEC prev s acc
1283 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1285 stepOuter (Fold fstep initial done) gst (Just st, Just prev) = do
1287 acc' <- fstep acc prev
1292 -- XXX code duplicated from the previous equation
1294 res <- step (adaptState gst) stt
1301 else done acc >>= \r -> return $ Yield r (Just s, Just x)
1302 Skip s -> go SPEC s acc
1303 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1305 stepOuter _ _ (Nothing,_) = return Stop
1307{-# INLINE_NORMAL groupsRollingBy #-}
1308groupsRollingBy :: Monad m
1313groupsRollingBy cmp f (Stream step state) =
1314 Stream (stepOuter f) (Just state, Nothing)
1317 {-# INLINE_LATE stepOuter #-}
1318 stepOuter (Fold fstep initial done) gst (Just st, Nothing) = do
1319 res <- step (adaptState gst) st
1326 Skip s -> return $ Skip $ (Just s, Nothing)
1330 go !_ prev stt !acc = do
1331 res <- step (adaptState gst) stt
1339 done acc >>= \r -> return $ Yield r (Just s, Just x)
1340 Skip s -> go SPEC prev s acc
1341 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1343 stepOuter (Fold fstep initial done) gst (Just st, Just prev') = do
1345 acc' <- fstep acc prev'
1346 go SPEC prev' st acc'
1349 go !_ prevv stt !acc = do
1350 res <- step (adaptState gst) stt
1357 else done acc >>= \r -> return $ Yield r (Just s, Just x)
1358 Skip s -> go SPEC prevv s acc
1359 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1361 stepOuter _ _ (Nothing, _) = return Stop
1363{-# INLINE_NORMAL splitBy #-}
1364splitBy :: Monad m => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1365splitBy predicate f (Stream step state) = Stream (step' f) (Just state)
1369 {-# INLINE_LATE step' #-}
1370 step' (Fold fstep initial done) gst (Just st) = initial >>= go SPEC st
1375 res <- step (adaptState gst) stt
1379 then done acc >>= \r -> return $ Yield r (Just s)
1383 Skip s -> go SPEC s acc
1384 Stop -> done acc >>= \r -> return $ Yield r Nothing
1386 step' _ _ Nothing = return Stop
1388-- XXX requires -funfolding-use-threshold=150 in lines-unlines benchmark
1389{-# INLINE_NORMAL splitSuffixBy #-}
1390splitSuffixBy :: Monad m
1391 => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1392splitSuffixBy predicate f (Stream step state) = Stream (step' f) (Just state)
1396 {-# INLINE_LATE step' #-}
1397 step' (Fold fstep initial done) gst (Just st) = do
1398 res <- step (adaptState gst) st
1403 then done acc >>= \val -> return $ Yield val (Just s)
1408 Skip s -> return $ Skip $ Just s
1414 res <- step (adaptState gst) stt
1418 then done acc >>= \r -> return $ Yield r (Just s)
1422 Skip s -> go SPEC s acc
1423 Stop -> done acc >>= \r -> return $ Yield r Nothing
1425 step' _ _ Nothing = return Stop
1427{-# INLINE_NORMAL wordsBy #-}
1428wordsBy :: Monad m => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1429wordsBy predicate f (Stream step state) = Stream (stepOuter f) (Just state)
1433 {-# INLINE_LATE stepOuter #-}
1434 stepOuter (Fold fstep initial done) gst (Just st) = do
1435 res <- step (adaptState gst) st
1439 then return $ Skip (Just s)
1445 Skip s -> return $ Skip $ Just s
1451 res <- step (adaptState gst) stt
1455 then done acc >>= \r -> return $ Yield r (Just s)
1459 Skip s -> go SPEC s acc
1460 Stop -> done acc >>= \r -> return $ Yield r Nothing
1462 stepOuter _ _ Nothing = return Stop
1464-- String search algorithms:
1465-- http://www-igm.univ-mlv.fr/~lecroq/string/index.html
1468-- TODO can we unify the splitting operations using a splitting configuration
1469-- like in the split package.
1471data SplitStyle = Infix | Suffix | Prefix deriving (Eq, Show)
1473data SplitOptions = SplitOptions
1474 { style :: SplitStyle
1475 , withSep :: Bool -- ^ keep the separators in output
1476 -- , compact :: Bool -- ^ treat multiple consecutive separators as one
1477 -- , trimHead :: Bool -- ^ drop blank at head
1478 -- , trimTail :: Bool -- ^ drop blank at tail
1482data SplitOnState s a =
1487 | GO_KARP_RABIN s !(RB.Ring a) !(Ptr a)
1490{-# INLINE_NORMAL splitOn #-}
1492 :: forall m a b. (MonadIO m, Storable a, Enum a, Eq a)
1497splitOn patArr@Array{..} (Fold fstep initial done) (Stream step state) =
1498 Stream stepOuter GO_START
1502 patLen = A.length patArr
1503 maxIndex = patLen - 1
1504 elemBits = sizeOf (undefined :: a) * 8
1506 {-# INLINE_LATE stepOuter #-}
1507 stepOuter _ GO_START =
1509 then return $ Skip $ GO_EMPTY_PAT state
1512 r <- liftIO $ (A.unsafeIndexIO patArr 0)
1513 return $ Skip $ GO_SINGLE_PAT state r
1514 else if sizeOf (undefined :: a) * patLen
1515 <= sizeOf (undefined :: Word)
1516 then return $ Skip $ GO_SHORT_PAT state
1518 (rb, rhead) <- liftIO $ RB.new patLen
1519 return $ Skip $ GO_KARP_RABIN state rb rhead
1521 stepOuter gst (GO_SINGLE_PAT stt pat) = initial >>= go SPEC stt
1526 res <- step (adaptState gst) st
1532 return $ Yield r (GO_SINGLE_PAT s pat)
1533 else fstep acc x >>= go SPEC s
1534 Skip s -> go SPEC s acc
1535 Stop -> done acc >>= \r -> return $ Yield r GO_DONE
1537 stepOuter gst (GO_SHORT_PAT stt) = initial >>= go0 SPEC 0 (0 :: Word) stt
1542 mask = (1 `shiftL` (elemBits * patLen)) - 1
1544 addToWord wrd a = (wrd `shiftL` elemBits) .|. fromIntegral (fromEnum a)
1547 patWord = mask .&. A.foldl' addToWord 0 patArr
1549 go0 !_ !idx wrd st !acc = do
1550 res <- step (adaptState gst) st
1553 let wrd' = addToWord wrd x
1556 if wrd' .&. mask == patWord
1559 return $ Yield r (GO_SHORT_PAT s)
1560 else go1 SPEC wrd' s acc
1561 else go0 SPEC (idx + 1) wrd' s acc
1562 Skip s -> go0 SPEC idx wrd s acc
1565 then go2 wrd idx acc
1567 done acc' >>= \r -> return $ Yield r GO_DONE
1570 go1 !_ wrd st !acc = do
1571 res <- step (adaptState gst) st
1574 let wrd' = addToWord wrd x
1575 old = (mask .&. wrd) `shiftR` (elemBits * (patLen - 1))
1576 acc' <- fstep acc (toEnum $ fromIntegral old)
1577 if wrd' .&. mask == patWord
1578 then done acc' >>= \r -> return $ Yield r (GO_SHORT_PAT s)
1579 else go1 SPEC wrd' s acc'
1580 Skip s -> go1 SPEC wrd s acc
1582 acc' <- go2 wrd patLen acc
1583 done acc' >>= \r -> return $ Yield r GO_DONE
1585 go2 !wrd !n !acc | n > 0 = do
1586 let old = (mask .&. wrd) `shiftR` (elemBits * (n - 1))
1587 fstep acc (toEnum $ fromIntegral old) >>= go2 wrd (n - 1)
1588 go2 _ _ acc = return acc
1590 stepOuter gst (GO_KARP_RABIN stt rb rhead) = do
1591 initial >>= go0 SPEC 0 rhead stt
1595 k = 2891336453 :: Word32
1597 addCksum cksum a = cksum * k + fromIntegral (fromEnum a)
1598 deltaCksum cksum old new =
1599 addCksum cksum new - coeff * fromIntegral (fromEnum old)
1601 -- XXX shall we use a random starting hash or 1 instead of 0?
1602 patHash = A.foldl' addCksum 0 patArr
1605 go0 !_ !idx !rh st !acc = do
1606 res <- step (adaptState gst) st
1609 rh' <- liftIO $ RB.unsafeInsert rb rh x
1612 let fold = RB.unsafeFoldRing (RB.ringBound rb)
1613 let !ringHash = fold addCksum 0 rb
1614 if ringHash == patHash
1615 then go2 SPEC ringHash rh' s acc
1616 else go1 SPEC ringHash rh' s acc
1617 else go0 SPEC (idx + 1) rh' s acc
1618 Skip s -> go0 SPEC idx rh s acc
1620 !acc' <- if idx /= 0
1621 then RB.unsafeFoldRingM rh fstep acc rb
1623 done acc' >>= \r -> return $ Yield r GO_DONE
1625 -- XXX Theoretically this code can do 4 times faster if GHC generates
1626 -- optimal code. If we use just "(cksum' == patHash)" condition it goes
1627 -- 4x faster, as soon as we add the "RB.unsafeEqArray rb v" condition
1628 -- the generated code changes drastically and becomes 4x slower. Need
1629 -- to investigate what is going on with GHC.
1631 go1 !_ !cksum !rh st !acc = do
1632 res <- step (adaptState gst) st
1635 old <- liftIO $ peek rh
1636 let cksum' = deltaCksum cksum old x
1637 acc' <- fstep acc old
1639 if (cksum' == patHash)
1641 rh' <- liftIO (RB.unsafeInsert rb rh x)
1642 go2 SPEC cksum' rh' s acc'
1644 rh' <- liftIO (RB.unsafeInsert rb rh x)
1645 go1 SPEC cksum' rh' s acc'
1646 Skip s -> go1 SPEC cksum rh s acc
1648 acc' <- RB.unsafeFoldRingFullM rh fstep acc rb
1649 done acc' >>= \r -> return $ Yield r GO_DONE
1651 go2 !_ !cksum' !rh' s !acc' = do
1652 if RB.unsafeEqArray rb rh' patArr
1655 return $ Yield r (GO_KARP_RABIN s rb rhead)
1656 else go1 SPEC cksum' rh' s acc'
1658 stepOuter gst (GO_EMPTY_PAT st) = do
1659 res <- step (adaptState gst) st
1664 done acc' >>= \r -> return $ Yield r (GO_EMPTY_PAT s)
1665 Skip s -> return $ Skip (GO_EMPTY_PAT s)
1668 stepOuter _ GO_DONE = return Stop
1670{-# INLINE_NORMAL splitSuffixOn #-}
1672 :: forall m a b. (MonadIO m, Storable a, Enum a, Eq a)
1678splitSuffixOn withSep patArr@Array{..} (Fold fstep initial done)
1679 (Stream step state) =
1680 Stream stepOuter GO_START
1684 patLen = A.length patArr
1685 maxIndex = patLen - 1
1686 elemBits = sizeOf (undefined :: a) * 8
1688 {-# INLINE_LATE stepOuter #-}
1689 stepOuter _ GO_START =
1691 then return $ Skip $ GO_EMPTY_PAT state
1694 r <- liftIO $ (A.unsafeIndexIO patArr 0)
1695 return $ Skip $ GO_SINGLE_PAT state r
1696 else if sizeOf (undefined :: a) * patLen
1697 <= sizeOf (undefined :: Word)
1698 then return $ Skip $ GO_SHORT_PAT state
1700 (rb, rhead) <- liftIO $ RB.new patLen
1701 return $ Skip $ GO_KARP_RABIN state rb rhead
1703 stepOuter gst (GO_SINGLE_PAT stt pat) = do
1704 -- This first part is the only difference between splitOn and
1706 -- If the last element is a separator do not issue a blank segment.
1707 res <- step (adaptState gst) stt
1713 acc' <- if withSep then fstep acc x else return acc
1714 done acc' >>= \r -> return $ Yield r (GO_SINGLE_PAT s pat)
1715 else fstep acc x >>= go SPEC s
1716 Skip s -> return $ Skip $ (GO_SINGLE_PAT s pat)
1721 -- This is identical for splitOn and splitSuffixOn
1723 res <- step (adaptState gst) st
1728 acc' <- if withSep then fstep acc x else return acc
1730 return $ Yield r (GO_SINGLE_PAT s pat)
1731 else fstep acc x >>= go SPEC s
1732 Skip s -> go SPEC s acc
1733 Stop -> done acc >>= \r -> return $ Yield r GO_DONE
1735 stepOuter gst (GO_SHORT_PAT stt) = do
1737 -- Call "initial" only if the stream yields an element, otherwise we
1738 -- may call "initial" but never yield anything. initial may produce a
1739 -- side effect, therefore we will end up doing and discard a side
1744 res <- step (adaptState gst) stt
1748 let wrd' = addToWord wrd x
1749 acc' <- if withSep then fstep acc x else return acc
1752 if wrd' .&. mask == patWord
1753 then done acc' >>= \r -> return $ Yield r (GO_SHORT_PAT s)
1754 else go0 SPEC (idx + 1) wrd' s acc'
1755 else go0 SPEC (idx + 1) wrd' s acc'
1756 Skip s -> return $ Skip (GO_SHORT_PAT s)
1762 mask = (1 `shiftL` (elemBits * patLen)) - 1
1764 addToWord wrd a = (wrd `shiftL` elemBits) .|. fromIntegral (fromEnum a)
1767 patWord = mask .&. A.foldl' addToWord 0 patArr
1769 go0 !_ !idx wrd st !acc = do
1770 res <- step (adaptState gst) st
1773 let wrd' = addToWord wrd x
1774 acc' <- if withSep then fstep acc x else return acc
1777 if wrd' .&. mask == patWord
1780 return $ Yield r (GO_SHORT_PAT s)
1781 else go1 SPEC wrd' s acc'
1782 else go0 SPEC (idx + 1) wrd' s acc'
1783 Skip s -> go0 SPEC idx wrd s acc
1785 if (idx == maxIndex) && (wrd .&. mask == patWord)
1788 acc' <- if idx /= 0 && not withSep
1789 then go2 wrd idx acc
1791 done acc' >>= \r -> return $ Yield r GO_DONE
1794 go1 !_ wrd st !acc = do
1795 res <- step (adaptState gst) st
1798 let wrd' = addToWord wrd x
1799 old = (mask .&. wrd) `shiftR` (elemBits * (patLen - 1))
1802 else fstep acc (toEnum $ fromIntegral old)
1803 if wrd' .&. mask == patWord
1804 then done acc' >>= \r -> return $ Yield r (GO_SHORT_PAT s)
1805 else go1 SPEC wrd' s acc'
1806 Skip s -> go1 SPEC wrd s acc
1808 -- If the last sequence is a separator do not issue a blank
1810 if wrd .&. mask == patWord
1815 else go2 wrd patLen acc
1816 done acc' >>= \r -> return $ Yield r GO_DONE
1818 go2 !wrd !n !acc | n > 0 = do
1819 let old = (mask .&. wrd) `shiftR` (elemBits * (n - 1))
1820 fstep acc (toEnum $ fromIntegral old) >>= go2 wrd (n - 1)
1821 go2 _ _ acc = return acc
1823 stepOuter gst (GO_KARP_RABIN stt rb rhead) = do
1825 res <- step (adaptState gst) stt
1829 acc' <- if withSep then fstep acc x else return acc
1830 rh' <- liftIO (RB.unsafeInsert rb rhead x)
1833 let fold = RB.unsafeFoldRing (RB.ringBound rb)
1834 let !ringHash = fold addCksum 0 rb
1835 if ringHash == patHash
1836 then go2 SPEC ringHash rh' s acc'
1837 else go0 SPEC (idx + 1) rh' s acc'
1838 else go0 SPEC (idx + 1) rh' s acc'
1839 Skip s -> return $ Skip (GO_KARP_RABIN s rb rhead)
1844 k = 2891336453 :: Word32
1846 addCksum cksum a = cksum * k + fromIntegral (fromEnum a)
1847 deltaCksum cksum old new =
1848 addCksum cksum new - coeff * fromIntegral (fromEnum old)
1850 -- XXX shall we use a random starting hash or 1 instead of 0?
1851 patHash = A.foldl' addCksum 0 patArr
1854 go0 !_ !idx !rh st !acc = do
1855 res <- step (adaptState gst) st
1858 acc' <- if withSep then fstep acc x else return acc
1859 rh' <- liftIO (RB.unsafeInsert rb rh x)
1862 let fold = RB.unsafeFoldRing (RB.ringBound rb)
1863 let !ringHash = fold addCksum 0 rb
1864 if ringHash == patHash
1865 then go2 SPEC ringHash rh' s acc'
1866 else go1 SPEC ringHash rh' s acc'
1867 else go0 SPEC (idx + 1) rh' s acc'
1868 Skip s -> go0 SPEC idx rh s acc
1870 -- do not issue a blank segment when we end at pattern
1871 if (idx == maxIndex) && RB.unsafeEqArray rb rh patArr
1874 !acc' <- if idx /= 0 && not withSep
1875 then RB.unsafeFoldRingM rh fstep acc rb
1877 done acc' >>= \r -> return $ Yield r GO_DONE
1879 -- XXX Theoretically this code can do 4 times faster if GHC generates
1880 -- optimal code. If we use just "(cksum' == patHash)" condition it goes
1881 -- 4x faster, as soon as we add the "RB.unsafeEqArray rb v" condition
1882 -- the generated code changes drastically and becomes 4x slower. Need
1883 -- to investigate what is going on with GHC.
1885 go1 !_ !cksum !rh st !acc = do
1886 res <- step (adaptState gst) st
1889 old <- liftIO $ peek rh
1890 let cksum' = deltaCksum cksum old x
1895 if (cksum' == patHash)
1897 rh' <- liftIO (RB.unsafeInsert rb rh x)
1898 go2 SPEC cksum' rh' s acc'
1900 rh' <- liftIO (RB.unsafeInsert rb rh x)
1901 go1 SPEC cksum' rh' s acc'
1902 Skip s -> go1 SPEC cksum rh s acc
1904 if RB.unsafeEqArray rb rh patArr
1909 else RB.unsafeFoldRingFullM rh fstep acc rb
1910 done acc' >>= \r -> return $ Yield r GO_DONE
1912 go2 !_ !cksum' !rh' s !acc' = do
1913 if RB.unsafeEqArray rb rh' patArr
1916 return $ Yield r (GO_KARP_RABIN s rb rhead)
1917 else go1 SPEC cksum' rh' s acc'
1919 stepOuter gst (GO_EMPTY_PAT st) = do
1920 res <- step (adaptState gst) st
1925 done acc' >>= \r -> return $ Yield r (GO_EMPTY_PAT s)
1926 Skip s -> return $ Skip (GO_EMPTY_PAT s)
1929 stepOuter _ GO_DONE = return Stop
1931data SplitState s arr
1933 | SplitBuffering s arr
1934 | SplitSplitting s arr
1935 | SplitYielding arr (SplitState s arr)
1938-- XXX An alternative approach would be to use a partial fold (Fold m a b) to
1939-- split using a splitBy like combinator. The Fold would consume upto the
1940-- separator and return any leftover which can then be fed to the next fold.
1942-- We can revisit this once we have partial folds/parsers.
1944-- | Performs infix separator style splitting.
1945{-# INLINE_NORMAL splitInnerBy #-}
1948 => (f a -> m (f a, Maybe (f a))) -- splitter
1949 -> (f a -> f a -> m (f a)) -- joiner
1952splitInnerBy splitter joiner (Stream step1 state1) =
1953 (Stream step (SplitInitial state1))
1957 {-# INLINE_LATE step #-}
1958 step gst (SplitInitial st) = do
1962 (x1, mx2) <- splitter x
1963 return $ case mx2 of
1964 Nothing -> Skip (SplitBuffering s x1)
1965 Just x2 -> Skip (SplitYielding x1 (SplitSplitting s x2))
1966 Skip s -> return $ Skip (SplitInitial s)
1967 Stop -> return $ Stop
1969 step gst (SplitBuffering st buf) = do
1973 (x1, mx2) <- splitter x
1974 buf' <- joiner buf x1
1975 return $ case mx2 of
1976 Nothing -> Skip (SplitBuffering s buf')
1977 Just x2 -> Skip (SplitYielding buf' (SplitSplitting s x2))
1978 Skip s -> return $ Skip (SplitBuffering s buf)
1979 Stop -> return $ Skip (SplitYielding buf SplitFinishing)
1981 step _ (SplitSplitting st buf) = do
1982 (x1, mx2) <- splitter buf
1983 return $ case mx2 of
1984 Nothing -> Skip $ SplitBuffering st x1
1985 Just x2 -> Skip $ SplitYielding x1 (SplitSplitting st x2)
1987 step _ (SplitYielding x next) = return $ Yield x next
1988 step _ SplitFinishing = return $ Stop
1990-- | Performs infix separator style splitting.
1991{-# INLINE_NORMAL splitInnerBySuffix #-}
1993 :: (Monad m, Eq (f a), Monoid (f a))
1994 => (f a -> m (f a, Maybe (f a))) -- splitter
1995 -> (f a -> f a -> m (f a)) -- joiner
1998splitInnerBySuffix splitter joiner (Stream step1 state1) =
1999 (Stream step (SplitInitial state1))
2003 {-# INLINE_LATE step #-}
2004 step gst (SplitInitial st) = do
2008 (x1, mx2) <- splitter x
2009 return $ case mx2 of
2010 Nothing -> Skip (SplitBuffering s x1)
2011 Just x2 -> Skip (SplitYielding x1 (SplitSplitting s x2))
2012 Skip s -> return $ Skip (SplitInitial s)
2013 Stop -> return $ Stop
2015 step gst (SplitBuffering st buf) = do
2019 (x1, mx2) <- splitter x
2020 buf' <- joiner buf x1
2021 return $ case mx2 of
2022 Nothing -> Skip (SplitBuffering s buf')
2023 Just x2 -> Skip (SplitYielding buf' (SplitSplitting s x2))
2024 Skip s -> return $ Skip (SplitBuffering s buf)
2028 else Skip (SplitYielding buf SplitFinishing)
2030 step _ (SplitSplitting st buf) = do
2031 (x1, mx2) <- splitter buf
2032 return $ case mx2 of
2033 Nothing -> Skip $ SplitBuffering st x1
2034 Just x2 -> Skip $ SplitYielding x1 (SplitSplitting st x2)
2036 step _ (SplitYielding x next) = return $ Yield x next
2037 step _ SplitFinishing = return $ Stop
2039------------------------------------------------------------------------------
2041------------------------------------------------------------------------------
2043{-# INLINE_NORMAL isPrefixOf #-}
2044isPrefixOf :: (Eq a, Monad m) => Stream m a -> Stream m a -> m Bool
2045isPrefixOf (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
2047 go (sa, sb, Nothing) = do
2048 r <- stepa defState sa
2050 Yield x sa' -> go (sa', sb, Just x)
2051 Skip sa' -> go (sa', sb, Nothing)
2054 go (sa, sb, Just x) = do
2055 r <- stepb defState sb
2059 then go (sa, sb', Nothing)
2061 Skip sb' -> go (sa, sb', Just x)
2062 Stop -> return False
2064{-# INLINE_NORMAL isSubsequenceOf #-}
2065isSubsequenceOf :: (Eq a, Monad m) => Stream m a -> Stream m a -> m Bool
2066isSubsequenceOf (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
2068 go (sa, sb, Nothing) = do
2069 r <- stepa defState sa
2071 Yield x sa' -> go (sa', sb, Just x)
2072 Skip sa' -> go (sa', sb, Nothing)
2075 go (sa, sb, Just x) = do
2076 r <- stepb defState sb
2080 then go (sa, sb', Nothing)
2081 else go (sa, sb', Just x)
2082 Skip sb' -> go (sa, sb', Just x)
2083 Stop -> return False
2085{-# INLINE_NORMAL stripPrefix #-}
2088 => Stream m a -> Stream m a -> m (Maybe (Stream m a))
2089stripPrefix (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
2091 go (sa, sb, Nothing) = do
2092 r <- stepa defState sa
2094 Yield x sa' -> go (sa', sb, Just x)
2095 Skip sa' -> go (sa', sb, Nothing)
2096 Stop -> return $ Just (Stream stepb sb)
2098 go (sa, sb, Just x) = do
2099 r <- stepb defState sb
2103 then go (sa, sb', Nothing)
2105 Skip sb' -> go (sa, sb', Just x)
2106 Stop -> return Nothing
2108------------------------------------------------------------------------------
2110------------------------------------------------------------------------------
2112-- | Execute a monadic action for each element of the 'Stream'
2113{-# INLINE_NORMAL mapM_ #-}
2114mapM_ :: Monad m => (a -> m b) -> Stream m a -> m ()
2115mapM_ m = drain . mapM m
2117-------------------------------------------------------------------------------
2118-- Stream transformations using Unfolds
2119-------------------------------------------------------------------------------
2121-- Define a unique structure to use in inspection testing
2122data ConcatMapUState o i =
2124 | ConcatMapUInner o i
2126-- | @concatMapU unfold stream@ uses @unfold@ to map the input stream elements
2127-- to streams and then flattens the generated streams into a single output
2130-- This is like 'concatMap' but uses an unfold with an explicit state to
2131-- generate the stream instead of a 'Stream' type generator. This allows better
2132-- optimization via fusion. This can be many times more efficient than
2135{-# INLINE_NORMAL concatMapU #-}
2136concatMapU :: Monad m => Unfold m a b -> Stream m a -> Stream m b
2137concatMapU (Unfold istep inject) (Stream ostep ost) =
2138 Stream step (ConcatMapUOuter ost)
2140 {-# INLINE_LATE step #-}
2141 step gst (ConcatMapUOuter o) = do
2142 r <- ostep (adaptState gst) o
2146 i `seq` return (Skip (ConcatMapUInner o' i))
2147 Skip o' -> return $ Skip (ConcatMapUOuter o')
2148 Stop -> return $ Stop
2150 step _ (ConcatMapUInner o i) = do
2153 Yield x i' -> Yield x (ConcatMapUInner o i')
2154 Skip i' -> Skip (ConcatMapUInner o i')
2155 Stop -> Skip (ConcatMapUOuter o)
2157data ConcatUnfoldInterleaveState o i =
2158 ConcatUnfoldInterleaveOuter o [i]
2159 | ConcatUnfoldInterleaveInner o [i]
2160 | ConcatUnfoldInterleaveInnerL [i] [i]
2161 | ConcatUnfoldInterleaveInnerR [i] [i]
2163-- XXX use arrays to store state instead of lists.
2164-- XXX In general we can use different scheduling strategies e.g. how to
2165-- schedule the outer vs inner loop or assigning weights to different streams
2166-- or outer and inner loops.
2168-- After a yield, switch to the next stream. Do not switch streams on Skip.
2169-- Yield from outer stream switches to the inner stream.
2171-- There are two choices here, (1) exhaust the outer stream first and then
2172-- start yielding from the inner streams, this is much simpler to implement,
2173-- (2) yield at least one element from an inner stream before going back to
2174-- outer stream and opening the next stream from it.
2176-- Ideally, we need some scheduling bias to inner streams vs outer stream.
2177-- Maybe we can configure the behavior.
2179{-# INLINE_NORMAL concatUnfoldInterleave #-}
2180concatUnfoldInterleave :: Monad m => Unfold m a b -> Stream m a -> Stream m b
2181concatUnfoldInterleave (Unfold istep inject) (Stream ostep ost) =
2182 Stream step (ConcatUnfoldInterleaveOuter ost [])
2184 {-# INLINE_LATE step #-}
2185 step gst (ConcatUnfoldInterleaveOuter o ls) = do
2186 r <- ostep (adaptState gst) o
2190 i `seq` return (Skip (ConcatUnfoldInterleaveInner o' (i : ls)))
2191 Skip o' -> return $ Skip (ConcatUnfoldInterleaveOuter o' ls)
2192 Stop -> return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2194 step _ (ConcatUnfoldInterleaveInner _ []) = undefined
2195 step _ (ConcatUnfoldInterleaveInner o (st:ls)) = do
2198 Yield x s -> Yield x (ConcatUnfoldInterleaveOuter o (s:ls))
2199 Skip s -> Skip (ConcatUnfoldInterleaveInner o (s:ls))
2200 Stop -> Skip (ConcatUnfoldInterleaveOuter o ls)
2202 step _ (ConcatUnfoldInterleaveInnerL [] []) = return Stop
2203 step _ (ConcatUnfoldInterleaveInnerL [] rs) =
2204 return $ Skip (ConcatUnfoldInterleaveInnerR [] rs)
2206 step _ (ConcatUnfoldInterleaveInnerL (st:ls) rs) = do
2209 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerL ls (s:rs))
2210 Skip s -> Skip (ConcatUnfoldInterleaveInnerL (s:ls) rs)
2211 Stop -> Skip (ConcatUnfoldInterleaveInnerL ls rs)
2213 step _ (ConcatUnfoldInterleaveInnerR [] []) = return Stop
2214 step _ (ConcatUnfoldInterleaveInnerR ls []) =
2215 return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2217 step _ (ConcatUnfoldInterleaveInnerR ls (st:rs)) = do
2220 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerR (s:ls) rs)
2221 Skip s -> Skip (ConcatUnfoldInterleaveInnerR ls (s:rs))
2222 Stop -> Skip (ConcatUnfoldInterleaveInnerR ls rs)
2224-- XXX In general we can use different scheduling strategies e.g. how to
2225-- schedule the outer vs inner loop or assigning weights to different streams
2226-- or outer and inner loops.
2228-- This could be inefficient if the tasks are too small.
2230-- Compared to concatUnfoldInterleave this one switches streams on Skips.
2232{-# INLINE_NORMAL concatUnfoldRoundrobin #-}
2233concatUnfoldRoundrobin :: Monad m => Unfold m a b -> Stream m a -> Stream m b
2234concatUnfoldRoundrobin (Unfold istep inject) (Stream ostep ost) =
2235 Stream step (ConcatUnfoldInterleaveOuter ost [])
2237 {-# INLINE_LATE step #-}
2238 step gst (ConcatUnfoldInterleaveOuter o ls) = do
2239 r <- ostep (adaptState gst) o
2243 i `seq` return (Skip (ConcatUnfoldInterleaveInner o' (i : ls)))
2244 Skip o' -> return $ Skip (ConcatUnfoldInterleaveInner o' ls)
2245 Stop -> return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2247 step _ (ConcatUnfoldInterleaveInner o []) =
2248 return $ Skip (ConcatUnfoldInterleaveOuter o [])
2250 step _ (ConcatUnfoldInterleaveInner o (st:ls)) = do
2253 Yield x s -> Yield x (ConcatUnfoldInterleaveOuter o (s:ls))
2254 Skip s -> Skip (ConcatUnfoldInterleaveOuter o (s:ls))
2255 Stop -> Skip (ConcatUnfoldInterleaveOuter o ls)
2257 step _ (ConcatUnfoldInterleaveInnerL [] []) = return Stop
2258 step _ (ConcatUnfoldInterleaveInnerL [] rs) =
2259 return $ Skip (ConcatUnfoldInterleaveInnerR [] rs)
2261 step _ (ConcatUnfoldInterleaveInnerL (st:ls) rs) = do
2264 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerL ls (s:rs))
2265 Skip s -> Skip (ConcatUnfoldInterleaveInnerL ls (s:rs))
2266 Stop -> Skip (ConcatUnfoldInterleaveInnerL ls rs)
2268 step _ (ConcatUnfoldInterleaveInnerR [] []) = return Stop
2269 step _ (ConcatUnfoldInterleaveInnerR ls []) =
2270 return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2272 step _ (ConcatUnfoldInterleaveInnerR ls (st:rs)) = do
2275 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerR (s:ls) rs)
2276 Skip s -> Skip (ConcatUnfoldInterleaveInnerR (s:ls) rs)
2277 Stop -> Skip (ConcatUnfoldInterleaveInnerR ls rs)
2279data AppendState s1 s2 = AppendFirst s1 | AppendSecond s2
2281-- Note that this could be much faster compared to the CPS stream. However, as
2282-- the number of streams being composed increases this may become expensive.
2283-- Need to see where the breaking point is between the two.
2285{-# INLINE_NORMAL append #-}
2286append :: Monad m => Stream m a -> Stream m a -> Stream m a
2287append (Stream step1 state1) (Stream step2 state2) =
2288 Stream step (AppendFirst state1)
2292 {-# INLINE_LATE step #-}
2293 step gst (AppendFirst st) = do
2296 Yield a s -> Yield a (AppendFirst s)
2297 Skip s -> Skip (AppendFirst s)
2298 Stop -> Skip (AppendSecond state2)
2300 step gst (AppendSecond st) = do
2303 Yield a s -> Yield a (AppendSecond s)
2304 Skip s -> Skip (AppendSecond s)
2307data InterleaveState s1 s2 = InterleaveFirst s1 s2 | InterleaveSecond s1 s2
2308 | InterleaveSecondOnly s2 | InterleaveFirstOnly s1
2310{-# INLINE_NORMAL interleave #-}
2311interleave :: Monad m => Stream m a -> Stream m a -> Stream m a
2312interleave (Stream step1 state1) (Stream step2 state2) =
2313 Stream step (InterleaveFirst state1 state2)
2317 {-# INLINE_LATE step #-}
2318 step gst (InterleaveFirst st1 st2) = do
2321 Yield a s -> Yield a (InterleaveSecond s st2)
2322 Skip s -> Skip (InterleaveFirst s st2)
2323 Stop -> Skip (InterleaveSecondOnly st2)
2325 step gst (InterleaveSecond st1 st2) = do
2328 Yield a s -> Yield a (InterleaveFirst st1 s)
2329 Skip s -> Skip (InterleaveSecond st1 s)
2330 Stop -> Skip (InterleaveFirstOnly st1)
2332 step gst (InterleaveFirstOnly st1) = do
2335 Yield a s -> Yield a (InterleaveFirstOnly s)
2336 Skip s -> Skip (InterleaveFirstOnly s)
2339 step gst (InterleaveSecondOnly st2) = do
2342 Yield a s -> Yield a (InterleaveSecondOnly s)
2343 Skip s -> Skip (InterleaveSecondOnly s)
2346{-# INLINE_NORMAL interleaveMin #-}
2347interleaveMin :: Monad m => Stream m a -> Stream m a -> Stream m a
2348interleaveMin (Stream step1 state1) (Stream step2 state2) =
2349 Stream step (InterleaveFirst state1 state2)
2353 {-# INLINE_LATE step #-}
2354 step gst (InterleaveFirst st1 st2) = do
2357 Yield a s -> Yield a (InterleaveSecond s st2)
2358 Skip s -> Skip (InterleaveFirst s st2)
2361 step gst (InterleaveSecond st1 st2) = do
2364 Yield a s -> Yield a (InterleaveFirst st1 s)
2365 Skip s -> Skip (InterleaveSecond st1 s)
2368 step _ (InterleaveFirstOnly _) = undefined
2369 step _ (InterleaveSecondOnly _) = undefined
2371{-# INLINE_NORMAL interleaveSuffix #-}
2372interleaveSuffix :: Monad m => Stream m a -> Stream m a -> Stream m a
2373interleaveSuffix (Stream step1 state1) (Stream step2 state2) =
2374 Stream step (InterleaveFirst state1 state2)
2378 {-# INLINE_LATE step #-}
2379 step gst (InterleaveFirst st1 st2) = do
2382 Yield a s -> Yield a (InterleaveSecond s st2)
2383 Skip s -> Skip (InterleaveFirst s st2)
2386 step gst (InterleaveSecond st1 st2) = do
2389 Yield a s -> Yield a (InterleaveFirst st1 s)
2390 Skip s -> Skip (InterleaveSecond st1 s)
2391 Stop -> Skip (InterleaveFirstOnly st1)
2393 step gst (InterleaveFirstOnly st1) = do
2396 Yield a s -> Yield a (InterleaveFirstOnly s)
2397 Skip s -> Skip (InterleaveFirstOnly s)
2400 step _ (InterleaveSecondOnly _) = undefined
2402data InterleaveInfixState s1 s2 a
2403 = InterleaveInfixFirst s1 s2
2404 | InterleaveInfixSecondBuf s1 s2
2405 | InterleaveInfixSecondYield s1 s2 a
2406 | InterleaveInfixFirstYield s1 s2 a
2407 | InterleaveInfixFirstOnly s1
2409{-# INLINE_NORMAL interleaveInfix #-}
2410interleaveInfix :: Monad m => Stream m a -> Stream m a -> Stream m a
2411interleaveInfix (Stream step1 state1) (Stream step2 state2) =
2412 Stream step (InterleaveInfixFirst state1 state2)
2416 {-# INLINE_LATE step #-}
2417 step gst (InterleaveInfixFirst st1 st2) = do
2420 Yield a s -> Yield a (InterleaveInfixSecondBuf s st2)
2421 Skip s -> Skip (InterleaveInfixFirst s st2)
2424 step gst (InterleaveInfixSecondBuf st1 st2) = do
2427 Yield a s -> Skip (InterleaveInfixSecondYield st1 s a)
2428 Skip s -> Skip (InterleaveInfixSecondBuf st1 s)
2429 Stop -> Skip (InterleaveInfixFirstOnly st1)
2431 step gst (InterleaveInfixSecondYield st1 st2 x) = do
2434 Yield a s -> Yield x (InterleaveInfixFirstYield s st2 a)
2435 Skip s -> Skip (InterleaveInfixSecondYield s st2 x)
2438 step _ (InterleaveInfixFirstYield st1 st2 x) = do
2439 return $ Yield x (InterleaveInfixSecondBuf st1 st2)
2441 step gst (InterleaveInfixFirstOnly st1) = do
2444 Yield a s -> Yield a (InterleaveInfixFirstOnly s)
2445 Skip s -> Skip (InterleaveInfixFirstOnly s)
2448{-# INLINE_NORMAL roundRobin #-}
2449roundRobin :: Monad m => Stream m a -> Stream m a -> Stream m a
2450roundRobin (Stream step1 state1) (Stream step2 state2) =
2451 Stream step (InterleaveFirst state1 state2)
2455 {-# INLINE_LATE step #-}
2456 step gst (InterleaveFirst st1 st2) = do
2459 Yield a s -> Yield a (InterleaveSecond s st2)
2460 Skip s -> Skip (InterleaveSecond s st2)
2461 Stop -> Skip (InterleaveSecondOnly st2)
2463 step gst (InterleaveSecond st1 st2) = do
2466 Yield a s -> Yield a (InterleaveFirst st1 s)
2467 Skip s -> Skip (InterleaveFirst st1 s)
2468 Stop -> Skip (InterleaveFirstOnly st1)
2470 step gst (InterleaveSecondOnly st2) = do
2473 Yield a s -> Yield a (InterleaveSecondOnly s)
2474 Skip s -> Skip (InterleaveSecondOnly s)
2477 step gst (InterleaveFirstOnly st1) = do
2480 Yield a s -> Yield a (InterleaveFirstOnly s)
2481 Skip s -> Skip (InterleaveFirstOnly s)
2484data ICUState s1 s2 i1 i2 =
2489 | ICUFirstInner s1 s2 i1
2490 | ICUSecondInner s1 s2 i2
2491 | ICUFirstOnlyInner s1 i1
2492 | ICUSecondOnlyInner s2 i2
2494-- | Interleave streams (full streams, not the elements) unfolded from two
2495-- input streams and concat. Stop when the first stream stops. If the second
2496-- stream ends before the first one then first stream still keeps running alone
2497-- without any interleaving with the second stream.
2499-- [a1, a2, ... an] [b1, b2 ...]
2500-- => [streamA1, streamA2, ... streamAn] [streamB1, streamB2, ...]
2501-- => [streamA1, streamB1, streamA2...StreamAn, streamBn]
2502-- => [a11, a12, ...a1j, b11, b12, ...b1k, a21, a22, ...]
2504{-# INLINE_NORMAL gintercalateSuffix #-}
2507 => Unfold m a c -> Stream m a -> Unfold m b c -> Stream m b -> Stream m c
2509 (Unfold istep1 inject1) (Stream step1 state1)
2510 (Unfold istep2 inject2) (Stream step2 state2) =
2511 Stream step (ICUFirst state1 state2)
2515 {-# INLINE_LATE step #-}
2516 step gst (ICUFirst s1 s2) = do
2517 r <- step1 (adaptState gst) s1
2521 i `seq` return (Skip (ICUFirstInner s s2 i))
2522 Skip s -> return $ Skip (ICUFirst s s2)
2525 step gst (ICUFirstOnly s1) = do
2526 r <- step1 (adaptState gst) s1
2530 i `seq` return (Skip (ICUFirstOnlyInner s i))
2531 Skip s -> return $ Skip (ICUFirstOnly s)
2534 step _ (ICUFirstInner s1 s2 i1) = do
2537 Yield x i' -> Yield x (ICUFirstInner s1 s2 i')
2538 Skip i' -> Skip (ICUFirstInner s1 s2 i')
2539 Stop -> Skip (ICUSecond s1 s2)
2541 step _ (ICUFirstOnlyInner s1 i1) = do
2544 Yield x i' -> Yield x (ICUFirstOnlyInner s1 i')
2545 Skip i' -> Skip (ICUFirstOnlyInner s1 i')
2546 Stop -> Skip (ICUFirstOnly s1)
2548 step gst (ICUSecond s1 s2) = do
2549 r <- step2 (adaptState gst) s2
2553 i `seq` return (Skip (ICUSecondInner s1 s i))
2554 Skip s -> return $ Skip (ICUSecond s1 s)
2555 Stop -> return $ Skip (ICUFirstOnly s1)
2557 step _ (ICUSecondInner s1 s2 i2) = do
2560 Yield x i' -> Yield x (ICUSecondInner s1 s2 i')
2561 Skip i' -> Skip (ICUSecondInner s1 s2 i')
2562 Stop -> Skip (ICUFirst s1 s2)
2564 step _ (ICUSecondOnly _s2) = undefined
2565 step _ (ICUSecondOnlyInner _s2 _i2) = undefined
2567data InterposeSuffixState s1 i1 =
2568 InterposeSuffixFirst s1
2569 -- | InterposeSuffixFirstYield s1 i1
2570 | InterposeSuffixFirstInner s1 i1
2571 | InterposeSuffixSecond s1
2573-- Note that if an unfolded layer turns out to be nil we still emit the
2574-- separator effect. An alternate behavior could be to emit the separator
2575-- effect only if at least one element has been yielded by the unfolding.
2576-- However, that becomes a bit complicated, so we have chosen the former
2578{-# INLINE_NORMAL interposeSuffix #-}
2581 => m c -> Unfold m b c -> Stream m b -> Stream m c
2584 (Unfold istep1 inject1) (Stream step1 state1) =
2585 Stream step (InterposeSuffixFirst state1)
2589 {-# INLINE_LATE step #-}
2590 step gst (InterposeSuffixFirst s1) = do
2591 r <- step1 (adaptState gst) s1
2595 i `seq` return (Skip (InterposeSuffixFirstInner s i))
2596 -- i `seq` return (Skip (InterposeSuffixFirstYield s i))
2597 Skip s -> return $ Skip (InterposeSuffixFirst s)
2601 step _ (InterposeSuffixFirstYield s1 i1) = do
2604 Yield x i' -> Yield x (InterposeSuffixFirstInner s1 i')
2605 Skip i' -> Skip (InterposeSuffixFirstYield s1 i')
2606 Stop -> Skip (InterposeSuffixFirst s1)
2609 step _ (InterposeSuffixFirstInner s1 i1) = do
2612 Yield x i' -> Yield x (InterposeSuffixFirstInner s1 i')
2613 Skip i' -> Skip (InterposeSuffixFirstInner s1 i')
2614 Stop -> Skip (InterposeSuffixSecond s1)
2616 step _ (InterposeSuffixSecond s1) = do
2618 return $ Yield r (InterposeSuffixFirst s1)
2620data ICALState s1 s2 i1 i2 a =
2622 -- | ICALFirstYield s1 s2 i1
2623 | ICALFirstInner s1 s2 i1
2625 | ICALFirstOnlyInner s1 i1
2626 | ICALSecondInject s1 s2
2627 | ICALFirstInject s1 s2 i2
2628 -- | ICALFirstBuf s1 s2 i1 i2
2629 | ICALSecondInner s1 s2 i1 i2
2630 -- -- | ICALSecondInner s1 s2 i1 i2 a
2631 -- -- | ICALFirstResume s1 s2 i1 i2 a
2633-- | Interleave streams (full streams, not the elements) unfolded from two
2634-- input streams and concat. Stop when the first stream stops. If the second
2635-- stream ends before the first one then first stream still keeps running alone
2636-- without any interleaving with the second stream.
2638-- [a1, a2, ... an] [b1, b2 ...]
2639-- => [streamA1, streamA2, ... streamAn] [streamB1, streamB2, ...]
2640-- => [streamA1, streamB1, streamA2...StreamAn, streamBn]
2641-- => [a11, a12, ...a1j, b11, b12, ...b1k, a21, a22, ...]
2643{-# INLINE_NORMAL gintercalate #-}
2646 => Unfold m a c -> Stream m a -> Unfold m b c -> Stream m b -> Stream m c
2648 (Unfold istep1 inject1) (Stream step1 state1)
2649 (Unfold istep2 inject2) (Stream step2 state2) =
2650 Stream step (ICALFirst state1 state2)
2654 {-# INLINE_LATE step #-}
2655 step gst (ICALFirst s1 s2) = do
2656 r <- step1 (adaptState gst) s1
2660 i `seq` return (Skip (ICALFirstInner s s2 i))
2661 -- i `seq` return (Skip (ICALFirstYield s s2 i))
2662 Skip s -> return $ Skip (ICALFirst s s2)
2666 step _ (ICALFirstYield s1 s2 i1) = do
2669 Yield x i' -> Yield x (ICALFirstInner s1 s2 i')
2670 Skip i' -> Skip (ICALFirstYield s1 s2 i')
2671 Stop -> Skip (ICALFirst s1 s2)
2674 step _ (ICALFirstInner s1 s2 i1) = do
2677 Yield x i' -> Yield x (ICALFirstInner s1 s2 i')
2678 Skip i' -> Skip (ICALFirstInner s1 s2 i')
2679 Stop -> Skip (ICALSecondInject s1 s2)
2681 step gst (ICALFirstOnly s1) = do
2682 r <- step1 (adaptState gst) s1
2686 i `seq` return (Skip (ICALFirstOnlyInner s i))
2687 Skip s -> return $ Skip (ICALFirstOnly s)
2690 step _ (ICALFirstOnlyInner s1 i1) = do
2693 Yield x i' -> Yield x (ICALFirstOnlyInner s1 i')
2694 Skip i' -> Skip (ICALFirstOnlyInner s1 i')
2695 Stop -> Skip (ICALFirstOnly s1)
2697 -- We inject the second stream even before checking if the first stream
2698 -- would yield any more elements. There is no clear choice whether we
2699 -- should do this before or after that. Doing it after may make the state
2700 -- machine a bit simpler though.
2701 step gst (ICALSecondInject s1 s2) = do
2702 r <- step2 (adaptState gst) s2
2706 i `seq` return (Skip (ICALFirstInject s1 s i))
2707 Skip s -> return $ Skip (ICALSecondInject s1 s)
2708 Stop -> return $ Skip (ICALFirstOnly s1)
2710 step gst (ICALFirstInject s1 s2 i2) = do
2711 r <- step1 (adaptState gst) s1
2715 i `seq` return (Skip (ICALSecondInner s s2 i i2))
2716 -- i `seq` return (Skip (ICALFirstBuf s s2 i i2))
2717 Skip s -> return $ Skip (ICALFirstInject s s2 i2)
2721 step _ (ICALFirstBuf s1 s2 i1 i2) = do
2724 Yield x i' -> Skip (ICALSecondInner s1 s2 i' i2 x)
2725 Skip i' -> Skip (ICALFirstBuf s1 s2 i' i2)
2728 step _ (ICALSecondInner s1 s2 i1 i2 v) = do
2731 Yield x i' -> Yield x (ICALSecondInner s1 s2 i1 i' v)
2732 Skip i' -> Skip (ICALSecondInner s1 s2 i1 i' v)
2733 Stop -> Skip (ICALFirstResume s1 s2 i1 i2 v)
2736 step _ (ICALSecondInner s1 s2 i1 i2) = do
2739 Yield x i' -> Yield x (ICALSecondInner s1 s2 i1 i')
2740 Skip i' -> Skip (ICALSecondInner s1 s2 i1 i')
2741 Stop -> Skip (ICALFirstInner s1 s2 i1)
2742 -- Stop -> Skip (ICALFirstResume s1 s2 i1 i2)
2745 step _ (ICALFirstResume s1 s2 i1 i2 x) = do
2746 return $ Yield x (ICALFirstInner s1 s2 i1 i2)
2749data InterposeState s1 i1 a =
2751 -- | InterposeFirstYield s1 i1
2752 | InterposeFirstInner s1 i1
2753 | InterposeFirstInject s1
2754 -- | InterposeFirstBuf s1 i1
2755 | InterposeSecondYield s1 i1
2756 -- -- | InterposeSecondYield s1 i1 a
2757 -- -- | InterposeFirstResume s1 i1 a
2759-- Note that this only interposes the pure values, we may run many effects to
2760-- generate those values as some effects may not generate anything (Skip).
2761{-# INLINE_NORMAL interpose #-}
2762interpose :: Monad m => m c -> Unfold m b c -> Stream m b -> Stream m c
2765 (Unfold istep1 inject1) (Stream step1 state1) =
2766 Stream step (InterposeFirst state1)
2770 {-# INLINE_LATE step #-}
2771 step gst (InterposeFirst s1) = do
2772 r <- step1 (adaptState gst) s1
2776 i `seq` return (Skip (InterposeFirstInner s i))
2777 -- i `seq` return (Skip (InterposeFirstYield s i))
2778 Skip s -> return $ Skip (InterposeFirst s)
2782 step _ (InterposeFirstYield s1 i1) = do
2785 Yield x i' -> Yield x (InterposeFirstInner s1 i')
2786 Skip i' -> Skip (InterposeFirstYield s1 i')
2787 Stop -> Skip (InterposeFirst s1)
2790 step _ (InterposeFirstInner s1 i1) = do
2793 Yield x i' -> Yield x (InterposeFirstInner s1 i')
2794 Skip i' -> Skip (InterposeFirstInner s1 i')
2795 Stop -> Skip (InterposeFirstInject s1)
2797 step gst (InterposeFirstInject s1) = do
2798 r <- step1 (adaptState gst) s1
2802 -- i `seq` return (Skip (InterposeFirstBuf s i))
2803 i `seq` return (Skip (InterposeSecondYield s i))
2804 Skip s -> return $ Skip (InterposeFirstInject s)
2808 step _ (InterposeFirstBuf s1 i1) = do
2811 Yield x i' -> Skip (InterposeSecondYield s1 i' x)
2812 Skip i' -> Skip (InterposeFirstBuf s1 i')
2817 step _ (InterposeSecondYield s1 i1 v) = do
2819 return $ Yield r (InterposeFirstResume s1 i1 v)
2821 step _ (InterposeSecondYield s1 i1) = do
2823 return $ Yield r (InterposeFirstInner s1 i1)
2826 step _ (InterposeFirstResume s1 i1 v) = do
2827 return $ Yield v (InterposeFirstInner s1 i1)
2830------------------------------------------------------------------------------
2832------------------------------------------------------------------------------
2834data GbracketState s1 s2 v
2836 | GBracketNormal s1 v
2837 | GBracketException s2
2839-- | The most general bracketing and exception combinator. All other
2840-- combinators can be expressed in terms of this combinator. This can also be
2841-- used for cases which are not covered by the standard combinators.
2845{-# INLINE_NORMAL gbracket #-}
2849 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
2850 -> (c -> m d) -- ^ after, on normal stop
2851 -> (c -> e -> Stream m b) -- ^ on exception
2852 -> (c -> Stream m b) -- ^ stream generator
2854gbracket bef exc aft fexc fnormal =
2855 Stream step GBracketInit
2859 {-# INLINE_LATE step #-}
2860 step _ GBracketInit = do
2862 return $ Skip $ GBracketNormal (fnormal r) r
2864 step gst (GBracketNormal (UnStream step1 st) v) = do
2865 res <- exc $ step1 gst st
2867 Right r -> case r of
2869 return $ Yield x (GBracketNormal (Stream step1 s) v)
2870 Skip s -> return $ Skip (GBracketNormal (Stream step1 s) v)
2871 Stop -> aft v >> return Stop
2872 Left e -> return $ Skip (GBracketException (fexc v e))
2873 step gst (GBracketException (UnStream step1 st)) = do
2876 Yield x s -> return $ Yield x (GBracketException (Stream step1 s))
2877 Skip s -> return $ Skip (GBracketException (Stream step1 s))
2880-- | Create an IORef holding a finalizer that is called automatically when the
2881-- IORef is garbage collected. The IORef can be written to with a 'Nothing'
2882-- value to deactivate the finalizer.
2883newFinalizedIORef :: (MonadIO m, MonadBaseControl IO m)
2884 => m a -> m (IORef (Maybe (IO ())))
2885newFinalizedIORef finalizer = do
2886 mrun <- captureMonadState
2887 ref <- liftIO $ newIORef $ Just $ liftIO $ void $ do
2888 _ <- runInIO mrun finalizer
2891 res <- readIORef ref
2893 Nothing -> return ()
2895 _ <- liftIO $ mkWeakIORef ref finalizer1
2898data GbracketIOState s1 s2 v wref
2900 | GBracketIONormal s1 v wref
2901 | GBracketIOException s2
2903-- | Like gbracket but also uses a finalizer to make sure when the stream is
2904-- garbage collected we run the finalizing action. This requires a MonadIO and
2905-- MonadBaseControl IO constraint.
2907-- | The most general bracketing and exception combinator. All other
2908-- combinators can be expressed in terms of this combinator. This can also be
2909-- used for cases which are not covered by the standard combinators.
2913{-# INLINE_NORMAL gbracketIO #-}
2915 :: (MonadIO m, MonadBaseControl IO m)
2917 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
2918 -> (c -> m d) -- ^ after, on normal stop or GC
2919 -> (c -> e -> Stream m b) -- ^ on exception
2920 -> (c -> Stream m b) -- ^ stream generator
2922gbracketIO bef exc aft fexc fnormal =
2923 Stream step GBracketIOInit
2927 -- If the stream is never evaluated the "aft" action will never be
2928 -- called. For that to occur we will need the user of this API to pass a
2929 -- weak pointer to us.
2930 {-# INLINE_LATE step #-}
2931 step _ GBracketIOInit = do
2933 ref <- newFinalizedIORef (aft r)
2934 return $ Skip $ GBracketIONormal (fnormal r) r ref
2936 step gst (GBracketIONormal (UnStream step1 st) v ref) = do
2937 res <- exc $ step1 gst st
2939 Right r -> case r of
2941 return $ Yield x (GBracketIONormal (Stream step1 s) v ref)
2943 return $ Skip (GBracketIONormal (Stream step1 s) v ref)
2945 liftIO (writeIORef ref Nothing)
2946 aft v >> return Stop
2948 liftIO (writeIORef ref Nothing)
2949 return $ Skip (GBracketIOException (fexc v e))
2950 step gst (GBracketIOException (UnStream step1 st)) = do
2954 return $ Yield x (GBracketIOException (Stream step1 s))
2955 Skip s -> return $ Skip (GBracketIOException (Stream step1 s))
2958-- | Run a side effect before the stream yields its first element.
2959{-# INLINE_NORMAL before #-}
2960before :: Monad m => m b -> Stream m a -> Stream m a
2961before action (Stream step state) = Stream step' Nothing
2965 {-# INLINE_LATE step' #-}
2966 step' _ Nothing = action >> return (Skip (Just state))
2968 step' gst (Just st) = do
2971 Yield x s -> return $ Yield x (Just s)
2972 Skip s -> return $ Skip (Just s)
2975-- | Run a side effect whenever the stream stops normally.
2976{-# INLINE_NORMAL after #-}
2977after :: Monad m => m b -> Stream m a -> Stream m a
2978after action (Stream step state) = Stream step' state
2982 {-# INLINE_LATE step' #-}
2986 Yield x s -> return $ Yield x s
2987 Skip s -> return $ Skip s
2988 Stop -> action >> return Stop
2990{-# INLINE_NORMAL afterIO #-}
2991afterIO :: (MonadIO m, MonadBaseControl IO m)
2992 => m b -> Stream m a -> Stream m a
2993afterIO action (Stream step state) = Stream step' Nothing
2997 {-# INLINE_LATE step' #-}
2998 step' _ Nothing = do
2999 ref <- newFinalizedIORef action
3000 return $ Skip $ Just (state, ref)
3001 step' gst (Just (st, ref)) = do
3004 Yield x s -> return $ Yield x (Just (s, ref))
3005 Skip s -> return $ Skip (Just (s, ref))
3007 Just f <- readIORef ref
3008 writeIORef ref Nothing
3011-- XXX These combinators are expensive due to the call to
3012-- onException/handle/try on each step. Therefore, when possible, they should
3013-- be called in an outer loop where we perform less iterations. For example, we
3014-- cannot call them on each iteration in a char stream, instead we can call
3015-- them when doing an IO on an array.
3017-- XXX For high performance error checks in busy streams we may need another
3018-- Error constructor in step.
3020-- | Run a side effect whenever the stream aborts due to an exception. The
3021-- exception is not caught, simply rethrown.
3022{-# INLINE_NORMAL onException #-}
3023onException :: MonadCatch m => m b -> Stream m a -> Stream m a
3024onException action str =
3025 gbracket (return ()) MC.try return
3026 (\_ (e :: MC.SomeException) -> nilM (action >> MC.throwM e))
3029{-# INLINE_NORMAL _onException #-}
3030_onException :: MonadCatch m => m b -> Stream m a -> Stream m a
3031_onException action (Stream step state) = Stream step' state
3035 {-# INLINE_LATE step' #-}
3037 res <- step gst st `MC.onException` action
3039 Yield x s -> return $ Yield x s
3040 Skip s -> return $ Skip s
3043-- XXX bracket is like concatMap, it generates a stream and then flattens it.
3044-- Like concatMap it has 10x worse performance compared to linear fused
3047-- | Run the first action before the stream starts and remember its output,
3048-- generate a stream using the output, run the second action providing the
3049-- remembered value as an argument whenever the stream ends normally or due to
3051{-# INLINE_NORMAL bracket #-}
3052bracket :: MonadCatch m => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
3053bracket bef aft bet =
3054 gbracket bef MC.try aft
3055 (\a (e :: SomeException) -> nilM (aft a >> MC.throwM e)) bet
3057{-# INLINE_NORMAL bracketIO #-}
3058bracketIO :: (MonadAsync m, MonadCatch m)
3059 => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
3060bracketIO bef aft bet =
3061 gbracketIO bef MC.try aft
3062 (\a (e :: SomeException) -> nilM (aft a >> MC.throwM e)) bet
3064data BracketState s v = BracketInit | BracketRun s v
3066{-# INLINE_NORMAL _bracket #-}
3067_bracket :: MonadCatch m => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
3068_bracket bef aft bet = Stream step' BracketInit
3072 {-# INLINE_LATE step' #-}
3073 step' _ BracketInit = bef >>= \x -> return (Skip (BracketRun (bet x) x))
3075 -- NOTE: It is important to use UnStream instead of the Stream pattern
3076 -- here, otherwise we get huge perf degradation, see note in concatMap.
3077 step' gst (BracketRun (UnStream step state) v) = do
3078 -- res <- step gst state `MC.onException` aft v
3079 res <- MC.try $ step gst state
3081 Left (e :: SomeException) -> aft v >> MC.throwM e >> return Stop
3082 Right r -> case r of
3083 Yield x s -> return $ Yield x (BracketRun (Stream step s) v)
3084 Skip s -> return $ Skip (BracketRun (Stream step s) v)
3085 Stop -> aft v >> return Stop
3087-- | Run a side effect whenever the stream stops normally or aborts due to an
3089{-# INLINE finally #-}
3090finally :: MonadCatch m => m b -> Stream m a -> Stream m a
3091-- finally action xs = after action $ onException action xs
3092finally action xs = bracket (return ()) (\_ -> action) (const xs)
3094{-# INLINE finallyIO #-}
3095finallyIO :: (MonadAsync m, MonadCatch m) => m b -> Stream m a -> Stream m a
3096finallyIO action xs = bracketIO (return ()) (\_ -> action) (const xs)
3098-- | When evaluating a stream if an exception occurs, stream evaluation aborts
3099-- and the specified exception handler is run with the exception as argument.
3100{-# INLINE_NORMAL handle #-}
3101handle :: (MonadCatch m, Exception e)
3102 => (e -> Stream m a) -> Stream m a -> Stream m a
3104 gbracket (return ()) MC.try return (\_ e -> f e) (\_ -> str)
3106{-# INLINE_NORMAL _handle #-}
3107_handle :: (MonadCatch m, Exception e)
3108 => (e -> Stream m a) -> Stream m a -> Stream m a
3109_handle f (Stream step state) = Stream step' (Left state)
3113 {-# INLINE_LATE step' #-}
3114 step' gst (Left st) = do
3115 res <- MC.try $ step gst st
3117 Left e -> return $ Skip $ Right (f e)
3118 Right r -> case r of
3119 Yield x s -> return $ Yield x (Left s)
3120 Skip s -> return $ Skip (Left s)
3123 step' gst (Right (UnStream step1 st)) = do
3126 Yield x s -> return $ Yield x (Right (Stream step1 s))
3127 Skip s -> return $ Skip (Right (Stream step1 s))
3130-------------------------------------------------------------------------------
3131-- General transformation
3132-------------------------------------------------------------------------------
3134{-# INLINE_NORMAL transform #-}
3135transform :: Monad m => Pipe m a b -> Stream m a -> Stream m b
3136transform (Pipe pstep1 pstep2 pstate) (Stream step state) =
3137 Stream step' (Consume pstate, state)
3141 {-# INLINE_LATE step' #-}
3143 step' gst (Consume pst, st) = pst `seq` do
3144 r <- step (adaptState gst) st
3149 Pipe.Yield b pst' -> return $ Yield b (pst', s)
3150 Pipe.Continue pst' -> return $ Skip (pst', s)
3151 Skip s -> return $ Skip (Consume pst, s)
3154 step' _ (Produce pst, st) = pst `seq` do
3157 Pipe.Yield b pst' -> return $ Yield b (pst', st)
3158 Pipe.Continue pst' -> return $ Skip (pst', st)
3160------------------------------------------------------------------------------
3161-- Transformation by Folding (Scans)
3162------------------------------------------------------------------------------
3164------------------------------------------------------------------------------
3166------------------------------------------------------------------------------
3168-- XXX Is a prescan useful, discarding the last step does not sound useful? I
3169-- am not sure about the utility of this function, so this is implemented but
3170-- not exposed. We can expose it if someone provides good reasons why this is
3173-- XXX We have to execute the stream one step ahead to know that we are at the
3174-- last step. The vector implementation of prescan executes the last fold step
3175-- but does not yield the result. This means we have executed the effect but
3176-- discarded value. This does not sound right. In this implementation we are
3177-- not executing the last fold step.
3178{-# INLINE_NORMAL prescanlM' #-}
3179prescanlM' :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
3180prescanlM' f mz (Stream step state) = Stream step' (state, mz)
3182 {-# INLINE_LATE step' #-}
3183 step' gst (st, prev) = do
3184 r <- step (adaptState gst) st
3188 return $ Yield acc (s, f acc x)
3189 Skip s -> return $ Skip (s, prev)
3192{-# INLINE prescanl' #-}
3193prescanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
3194prescanl' f z = prescanlM' (\a b -> return (f a b)) (return z)
3196------------------------------------------------------------------------------
3197-- Monolithic postscans (postscan followed by a map)
3198------------------------------------------------------------------------------
3200-- The performance of a modular postscan followed by a map seems to be
3201-- equivalent to this monolithic scan followed by map therefore we may not need
3202-- this implementation. We just have it for performance comparison and in case
3203-- modular version does not perform well in some situation.
3205{-# INLINE_NORMAL postscanlMx' #-}
3206postscanlMx' :: Monad m
3207 => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> Stream m b
3208postscanlMx' fstep begin done (Stream step state) = do
3209 Stream step' (state, begin)
3211 {-# INLINE_LATE step' #-}
3212 step' gst (st, acc) = do
3213 r <- step (adaptState gst) st
3219 v `seq` y `seq` return (Yield v (s, return y))
3220 Skip s -> return $ Skip (s, acc)
3223{-# INLINE_NORMAL postscanlx' #-}
3224postscanlx' :: Monad m
3225 => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> Stream m b
3226postscanlx' fstep begin done s =
3227 postscanlMx' (\b a -> return (fstep b a)) (return begin) (return . done) s
3229-- XXX do we need consM strict to evaluate the begin value?
3230{-# INLINE scanlMx' #-}
3232 => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> Stream m b
3233scanlMx' fstep begin done s =
3234 (begin >>= \x -> x `seq` done x) `consM` postscanlMx' fstep begin done s
3236{-# INLINE scanlx' #-}
3238 => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> Stream m b
3239scanlx' fstep begin done s =
3240 scanlMx' (\b a -> return (fstep b a)) (return begin) (return . done) s
3242------------------------------------------------------------------------------
3244------------------------------------------------------------------------------
3246{-# INLINE_NORMAL postscanlM' #-}
3247postscanlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
3248postscanlM' fstep begin (Stream step state) =
3249 begin `seq` Stream step' (state, begin)
3251 {-# INLINE_LATE step' #-}
3252 step' gst (st, acc) = acc `seq` do
3253 r <- step (adaptState gst) st
3257 y `seq` return (Yield y (s, y))
3258 Skip s -> return $ Skip (s, acc)
3261{-# INLINE_NORMAL postscanl' #-}
3262postscanl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
3263postscanl' f = postscanlM' (\a b -> return (f a b))
3265{-# INLINE_NORMAL postscanlM #-}
3266postscanlM :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
3267postscanlM fstep begin (Stream step state) = Stream step' (state, begin)
3269 {-# INLINE_LATE step' #-}
3270 step' gst (st, acc) = do
3271 r <- step (adaptState gst) st
3275 return (Yield y (s, y))
3276 Skip s -> return $ Skip (s, acc)
3279{-# INLINE_NORMAL postscanl #-}
3280postscanl :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
3281postscanl f = postscanlM (\a b -> return (f a b))
3283{-# INLINE_NORMAL scanlM' #-}
3284scanlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
3285scanlM' fstep begin s = begin `seq` (begin `cons` postscanlM' fstep begin s)
3287{-# INLINE scanl' #-}
3288scanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
3289scanl' f = scanlM' (\a b -> return (f a b))
3291{-# INLINE_NORMAL scanlM #-}
3292scanlM :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
3293scanlM fstep begin s = begin `cons` postscanlM fstep begin s
3296scanl :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
3297scanl f = scanlM (\a b -> return (f a b))
3299{-# INLINE_NORMAL scanl1M #-}
3300scanl1M :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
3301scanl1M fstep (Stream step state) = Stream step' (state, Nothing)
3303 {-# INLINE_LATE step' #-}
3304 step' gst (st, Nothing) = do
3307 Yield x s -> return $ Yield x (s, Just x)
3308 Skip s -> return $ Skip (s, Nothing)
3311 step' gst (st, Just acc) = do
3316 return $ Yield z (s, Just z)
3317 Skip s -> return $ Skip (s, Just acc)
3320{-# INLINE scanl1 #-}
3321scanl1 :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
3322scanl1 f = scanl1M (\x y -> return (f x y))
3324{-# INLINE_NORMAL scanl1M' #-}
3325scanl1M' :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
3326scanl1M' fstep (Stream step state) = Stream step' (state, Nothing)
3328 {-# INLINE_LATE step' #-}
3329 step' gst (st, Nothing) = do
3332 Yield x s -> x `seq` return $ Yield x (s, Just x)
3333 Skip s -> return $ Skip (s, Nothing)
3336 step' gst (st, Just acc) = acc `seq` do
3341 z `seq` return $ Yield z (s, Just z)
3342 Skip s -> return $ Skip (s, Just acc)
3345{-# INLINE scanl1' #-}
3346scanl1' :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
3347scanl1' f = scanl1M' (\x y -> return (f x y))
3349------------------------------------------------------------------------------
3351------------------------------------------------------------------------------
3353data RollingMapState s a = RollingMapInit s | RollingMapGo s a
3355{-# INLINE rollingMapM #-}
3356rollingMapM :: Monad m => (a -> a -> m b) -> Stream m a -> Stream m b
3357rollingMapM f (Stream step1 state1) = Stream step (RollingMapInit state1)
3359 step gst (RollingMapInit st) = do
3360 r <- step1 (adaptState gst) st
3362 Yield x s -> Skip $ RollingMapGo s x
3363 Skip s -> Skip $ RollingMapInit s
3366 step gst (RollingMapGo s1 x1) = do
3367 r <- step1 (adaptState gst) s1
3371 return $ Yield res $ RollingMapGo s x
3372 Skip s -> return $ Skip $ RollingMapGo s x1
3373 Stop -> return $ Stop
3375{-# INLINE rollingMap #-}
3376rollingMap :: Monad m => (a -> a -> b) -> Stream m a -> Stream m b
3377rollingMap f = rollingMapM (\x y -> return $ f x y)
3379------------------------------------------------------------------------------
3380-- Tapping/Distributing
3381------------------------------------------------------------------------------
3384tap :: Monad m => Fold m a b -> Stream m a -> Stream m a
3385tap (Fold fstep initial extract) (Stream step state) = Stream step' Nothing
3389 step' _ Nothing = do
3391 return $ Skip (Just (r, state))
3393 step' gst (Just (acc, st)) = acc `seq` do
3398 return $ Yield x (Just (acc', s))
3399 Skip s -> return $ Skip (Just (acc, s))
3404{-# INLINE_NORMAL tapOffsetEvery #-}
3405tapOffsetEvery :: Monad m
3406 => Int -> Int -> Fold m a b -> Stream m a -> Stream m a
3407tapOffsetEvery offset n (Fold fstep initial extract) (Stream step state) =
3408 Stream step' Nothing
3412 {-# INLINE_LATE step' #-}
3413 step' _ Nothing = do
3415 return $ Skip (Just (r, state, offset `mod` n))
3417 step' gst (Just (acc, st, count)) | count <= 0 = do
3421 !acc' <- fstep acc x
3422 return $ Yield x (Just (acc', s, n - 1))
3423 Skip s -> return $ Skip (Just (acc, s, count))
3428 step' gst (Just (acc, st, count)) = do
3431 Yield x s -> return $ Yield x (Just (acc, s, count - 1))
3432 Skip s -> return $ Skip (Just (acc, s, count))
3437{-# INLINE_NORMAL pollCounts #-}
3440 => (Stream m Int -> Stream m Int)
3444pollCounts transf fld (Stream step state) = Stream step' Nothing
3447 {-# INLINE_LATE step' #-}
3448 step' _ Nothing = do
3449 countVar <- liftIO $ newVar (0 :: Int)
3451 $ void $ runFold fld
3452 $ transf $ fromPrimVar countVar
3453 return $ Skip (Just (countVar, tid, state))
3455 step' gst (Just (countVar, tid, st)) = do
3459 liftIO $ modifyVar' countVar (+ 1)
3460 return $ Yield x (Just (countVar, tid, s))
3461 Skip s -> return $ Skip (Just (countVar, tid, s))
3463 liftIO $ killThread tid
3466{-# INLINE_NORMAL tapRate #-}
3468 (MonadAsync m, MonadCatch m)
3473tapRate samplingRate action (Stream step state) = Stream step' Nothing
3475 {-# NOINLINE loop #-}
3476 loop countVar prev = do
3479 (do liftIO $ threadDelay (round $ samplingRate * 1000000)
3480 i <- liftIO $ readVar countVar
3481 let !diff = i - prev
3484 (\(e :: AsyncException) -> do
3485 i <- liftIO $ readVar countVar
3486 let !diff = i - prev
3488 throwM (MC.toException e))
3491 {-# INLINE_LATE step' #-}
3492 step' _ Nothing = do
3493 countVar <- liftIO $ newVar 0
3494 tid <- fork $ loop countVar 0
3495 ref <- liftIO $ newIORef ()
3496 _ <- liftIO $ mkWeakIORef ref (killThread tid)
3497 return $ Skip (Just (countVar, tid, state, ref))
3499 step' gst (Just (countVar, tid, st, ref)) = do
3503 liftIO $ modifyVar' countVar (+ 1)
3504 return $ Yield x (Just (countVar, tid, s, ref))
3505 Skip s -> return $ Skip (Just (countVar, tid, s, ref))
3507 liftIO $ killThread tid
3511-------------------------------------------------------------------------------
3513-------------------------------------------------------------------------------
3515{-# INLINE_NORMAL takeWhileM #-}
3516takeWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
3517takeWhileM f (Stream step state) = Stream step' state
3519 {-# INLINE_LATE step' #-}
3525 return $ if b then Yield x s else Stop
3526 Skip s -> return $ Skip s
3529{-# INLINE takeWhile #-}
3530takeWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
3531takeWhile f = takeWhileM (return . f)
3533{-# INLINE_NORMAL drop #-}
3534drop :: Monad m => Int -> Stream m a -> Stream m a
3535drop n (Stream step state) = Stream step' (state, Just n)
3537 {-# INLINE_LATE step' #-}
3538 step' gst (st, Just i)
3543 Yield _ s -> Skip (s, Just (i - 1))
3544 Skip s -> Skip (s, Just i)
3546 | otherwise = return $ Skip (st, Nothing)
3548 step' gst (st, Nothing) = do
3552 Yield x s -> Yield x (s, Nothing)
3553 Skip s -> Skip (s, Nothing)
3556data DropWhileState s a
3558 | DropWhileYield a s
3561{-# INLINE_NORMAL dropWhileM #-}
3562dropWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
3563dropWhileM f (Stream step state) = Stream step' (DropWhileDrop state)
3565 {-# INLINE_LATE step' #-}
3566 step' gst (DropWhileDrop st) = do
3572 then return $ Skip (DropWhileDrop s)
3573 else return $ Skip (DropWhileYield x s)
3574 Skip s -> return $ Skip (DropWhileDrop s)
3577 step' gst (DropWhileNext st) = do
3580 Yield x s -> return $ Skip (DropWhileYield x s)
3581 Skip s -> return $ Skip (DropWhileNext s)
3584 step' _ (DropWhileYield x st) = return $ Yield x (DropWhileNext st)
3586{-# INLINE dropWhile #-}
3587dropWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
3588dropWhile f = dropWhileM (return . f)
3590{-# INLINE_NORMAL filterM #-}
3591filterM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
3592filterM f (Stream step state) = Stream step' state
3594 {-# INLINE_LATE step' #-}
3603 Skip s -> return $ Skip s
3606{-# INLINE filter #-}
3607filter :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
3608filter f = filterM (return . f)
3610{-# INLINE_NORMAL uniq #-}
3611uniq :: (Eq a, Monad m) => Stream m a -> Stream m a
3612uniq (Stream step state) = Stream step' (Nothing, state)
3614 {-# INLINE_LATE step' #-}
3615 step' gst (Nothing, st) = do
3618 Yield x s -> return $ Yield x (Just x, s)
3619 Skip s -> return $ Skip (Nothing, s)
3621 step' gst (Just x, st) = do
3624 Yield y s | x == y -> return $ Skip (Just x, s)
3625 | otherwise -> return $ Yield y (Just y, s)
3626 Skip s -> return $ Skip (Just x, s)
3629------------------------------------------------------------------------------
3630-- Transformation by Mapping
3631------------------------------------------------------------------------------
3633{-# INLINE_NORMAL sequence #-}
3634sequence :: Monad m => Stream m (m a) -> Stream m a
3635sequence (Stream step state) = Stream step' state
3637 {-# INLINE_LATE step' #-}
3639 r <- step (adaptState gst) st
3641 Yield x s -> x >>= \a -> return (Yield a s)
3642 Skip s -> return $ Skip s
3645------------------------------------------------------------------------------
3647------------------------------------------------------------------------------
3649data LoopState x s = FirstYield s
3650 | InterspersingYield s
3653{-# INLINE_NORMAL intersperseM #-}
3654intersperseM :: Monad m => m a -> Stream m a -> Stream m a
3655intersperseM m (Stream step state) = Stream step' (FirstYield state)
3657 {-# INLINE_LATE step' #-}
3658 step' gst (FirstYield st) = do
3662 Yield x s -> Skip (YieldAndCarry x s)
3663 Skip s -> Skip (FirstYield s)
3666 step' gst (InterspersingYield st) = do
3671 return $ Yield a (YieldAndCarry x s)
3672 Skip s -> return $ Skip $ InterspersingYield s
3675 step' _ (YieldAndCarry x st) = return $ Yield x (InterspersingYield st)
3680 | SuffixYield a (SuffixState s a)
3682{-# INLINE_NORMAL intersperseSuffix #-}
3683intersperseSuffix :: forall m a. Monad m => m a -> Stream m a -> Stream m a
3684intersperseSuffix action (Stream step state) = Stream step' (SuffixElem state)
3686 {-# INLINE_LATE step' #-}
3687 step' gst (SuffixElem st) = do
3690 Yield x s -> Skip (SuffixYield x (SuffixSuffix s))
3691 Skip s -> Skip (SuffixElem s)
3694 step' _ (SuffixSuffix st) = do
3695 action >>= \r -> return $ Skip (SuffixYield r (SuffixElem st))
3697 step' _ (SuffixYield x next) = return $ Yield x next
3699data SuffixSpanState s a
3700 = SuffixSpanElem s Int
3701 | SuffixSpanSuffix s
3702 | SuffixSpanYield a (SuffixSpanState s a)
3706-- | intersperse after every n items
3707{-# INLINE_NORMAL intersperseSuffixBySpan #-}
3708intersperseSuffixBySpan :: forall m a. Monad m
3709 => Int -> m a -> Stream m a -> Stream m a
3710intersperseSuffixBySpan n action (Stream step state) =
3711 Stream step' (SuffixSpanElem state n)
3713 {-# INLINE_LATE step' #-}
3714 step' gst (SuffixSpanElem st i) | i > 0 = do
3717 Yield x s -> Skip (SuffixSpanYield x (SuffixSpanElem s (i - 1)))
3718 Skip s -> Skip (SuffixSpanElem s i)
3719 Stop -> if i == n then Stop else Skip SuffixSpanLast
3720 step' _ (SuffixSpanElem st _) = return $ Skip (SuffixSpanSuffix st)
3722 step' _ (SuffixSpanSuffix st) = do
3723 action >>= \r -> return $ Skip (SuffixSpanYield r (SuffixSpanElem st n))
3725 step' _ (SuffixSpanLast) = do
3726 action >>= \r -> return $ Skip (SuffixSpanYield r SuffixSpanStop)
3728 step' _ (SuffixSpanYield x next) = return $ Yield x next
3730 step' _ (SuffixSpanStop) = return Stop
3732{-# INLINE intersperse #-}
3733intersperse :: Monad m => a -> Stream m a -> Stream m a
3734intersperse a = intersperseM (return a)
3736{-# INLINE_NORMAL insertBy #-}
3737insertBy :: Monad m => (a -> a -> Ordering) -> a -> Stream m a -> Stream m a
3738insertBy cmp a (Stream step state) = Stream step' (state, False, Nothing)
3740 {-# INLINE_LATE step' #-}
3741 step' gst (st, False, _) = do
3744 Yield x s -> case cmp a x of
3745 GT -> return $ Yield x (s, False, Nothing)
3746 _ -> return $ Yield a (s, True, Just x)
3747 Skip s -> return $ Skip (s, False, Nothing)
3748 Stop -> return $ Yield a (st, True, Nothing)
3750 step' _ (_, True, Nothing) = return Stop
3752 step' gst (st, True, Just prev) = do
3755 Yield x s -> return $ Yield prev (s, True, Just x)
3756 Skip s -> return $ Skip (s, True, Just prev)
3757 Stop -> return $ Yield prev (st, True, Nothing)
3759------------------------------------------------------------------------------
3761------------------------------------------------------------------------------
3763{-# INLINE_NORMAL deleteBy #-}
3764deleteBy :: Monad m => (a -> a -> Bool) -> a -> Stream m a -> Stream m a
3765deleteBy eq x (Stream step state) = Stream step' (state, False)
3767 {-# INLINE_LATE step' #-}
3768 step' gst (st, False) = do
3771 Yield y s -> return $
3772 if eq x y then Skip (s, True) else Yield y (s, False)
3773 Skip s -> return $ Skip (s, False)
3776 step' gst (st, True) = do
3779 Yield y s -> return $ Yield y (s, True)
3780 Skip s -> return $ Skip (s, True)
3783------------------------------------------------------------------------------
3784-- Transformation by Map and Filter
3785------------------------------------------------------------------------------
3787-- XXX Will this always fuse properly?
3788{-# INLINE_NORMAL mapMaybe #-}
3789mapMaybe :: Monad m => (a -> Maybe b) -> Stream m a -> Stream m b
3790mapMaybe f = fmap fromJust . filter isJust . map f
3792{-# INLINE_NORMAL mapMaybeM #-}
3793mapMaybeM :: Monad m => (a -> m (Maybe b)) -> Stream m a -> Stream m b
3794mapMaybeM f = fmap fromJust . filter isJust . mapM f
3796------------------------------------------------------------------------------
3798------------------------------------------------------------------------------
3800{-# INLINE_NORMAL indexed #-}
3801indexed :: Monad m => Stream m a -> Stream m (Int, a)
3802indexed (Stream step state) = Stream step' (state, 0)
3804 {-# INLINE_LATE step' #-}
3805 step' gst (st, i) = i `seq` do
3806 r <- step (adaptState gst) st
3808 Yield x s -> return $ Yield (i, x) (s, i+1)
3809 Skip s -> return $ Skip (s, i)
3812{-# INLINE_NORMAL indexedR #-}
3813indexedR :: Monad m => Int -> Stream m a -> Stream m (Int, a)
3814indexedR m (Stream step state) = Stream step' (state, m)
3816 {-# INLINE_LATE step' #-}
3817 step' gst (st, i) = i `seq` do
3818 r <- step (adaptState gst) st
3820 Yield x s -> let i' = i - 1
3821 in return $ Yield (i, x) (s, i')
3822 Skip s -> return $ Skip (s, i)
3825{-# INLINE_NORMAL zipWithM #-}
3827 => (a -> b -> m c) -> Stream m a -> Stream m b -> Stream m c
3828zipWithM f (Stream stepa ta) (Stream stepb tb) = Stream step (ta, tb, Nothing)
3830 {-# INLINE_LATE step #-}
3831 step gst (sa, sb, Nothing) = do
3832 r <- stepa (adaptState gst) sa
3835 Yield x sa' -> Skip (sa', sb, Just x)
3836 Skip sa' -> Skip (sa', sb, Nothing)
3839 step gst (sa, sb, Just x) = do
3840 r <- stepb (adaptState gst) sb
3844 return $ Yield z (sa, sb', Nothing)
3845 Skip sb' -> return $ Skip (sa, sb', Just x)
3848#if __GLASGOW_HASKELL__ >= 801
3849{-# RULES "zipWithM xs xs"
3850 forall f xs. zipWithM @Identity f xs xs = mapM (\x -> f x x) xs #-}
3853{-# INLINE zipWith #-}
3854zipWith :: Monad m => (a -> b -> c) -> Stream m a -> Stream m b -> Stream m c
3855zipWith f = zipWithM (\a b -> return (f a b))
3857------------------------------------------------------------------------------
3859------------------------------------------------------------------------------
3861{-# INLINE_NORMAL mergeByM #-}
3864 => (a -> a -> m Ordering) -> Stream m a -> Stream m a -> Stream m a
3865mergeByM cmp (Stream stepa ta) (Stream stepb tb) =
3866 Stream step (Just ta, Just tb, Nothing, Nothing)
3868 {-# INLINE_LATE step #-}
3870 -- one of the values is missing, and the corresponding stream is running
3871 step gst (Just sa, sb, Nothing, b) = do
3874 Yield a sa' -> Skip (Just sa', sb, Just a, b)
3875 Skip sa' -> Skip (Just sa', sb, Nothing, b)
3876 Stop -> Skip (Nothing, sb, Nothing, b)
3878 step gst (sa, Just sb, a, Nothing) = do
3881 Yield b sb' -> Skip (sa, Just sb', a, Just b)
3882 Skip sb' -> Skip (sa, Just sb', a, Nothing)
3883 Stop -> Skip (sa, Nothing, a, Nothing)
3885 -- both the values are available
3886 step _ (sa, sb, Just a, Just b) = do
3888 return $ case res of
3889 GT -> Yield b (sa, sb, Just a, Nothing)
3890 _ -> Yield a (sa, sb, Nothing, Just b)
3892 -- one of the values is missing, corresponding stream is done
3893 step _ (Nothing, sb, Nothing, Just b) =
3894 return $ Yield b (Nothing, sb, Nothing, Nothing)
3896 step _ (sa, Nothing, Just a, Nothing) =
3897 return $ Yield a (sa, Nothing, Nothing, Nothing)
3899 step _ (Nothing, Nothing, Nothing, Nothing) = return Stop
3901{-# INLINE mergeBy #-}
3904 => (a -> a -> Ordering) -> Stream m a -> Stream m a -> Stream m a
3905mergeBy cmp = mergeByM (\a b -> return $ cmp a b)
3907------------------------------------------------------------------------------
3908-- Transformation comprehensions
3909------------------------------------------------------------------------------
3911{-# INLINE_NORMAL the #-}
3912the :: (Eq a, Monad m) => Stream m a -> m (Maybe a)
3913the (Stream step state) = go state
3916 r <- step defState st
3918 Yield x s -> go' x s
3920 Stop -> return Nothing
3922 r <- step defState st
3924 Yield x s | x == n -> go' n s
3925 | otherwise -> return Nothing
3927 Stop -> return (Just n)
3929{-# INLINE runFold #-}
3930runFold :: (Monad m) => Fold m a b -> Stream m a -> m b
3931runFold (Fold step begin done) = foldlMx' step begin done
3933-------------------------------------------------------------------------------
3934-- Concurrent application and fold
3935-------------------------------------------------------------------------------
3937-- XXX These functions should be moved to Stream/Parallel.hs
3939-- Using StreamD the worker stream producing code can fuse with the code to
3940-- queue output to the SVar giving some perf boost.
3942-- Note that StreamD can only be used in limited situations, specifically, we
3943-- cannot implement joinStreamVarPar using this.
3945-- XXX make sure that the SVar passed is a Parallel style SVar.
3947-- | Fold the supplied stream to the SVar asynchronously using Parallel
3948-- concurrency style.
3949-- {-# INLINE_NORMAL toSVarParallel #-}
3950{-# INLINE toSVarParallel #-}
3951toSVarParallel :: MonadAsync m
3952 => State t m a -> SVar t m a -> Stream m a -> m ()
3953toSVarParallel st sv xs =
3954 if svarInspectMode sv
3958 case getYieldLimit st of
3959 Nothing -> doFork (work Nothing)
3961 (handleChildException sv)
3962 Just _ -> doFork (workLim Nothing)
3964 (handleChildException sv)
3969 {-# NOINLINE work #-}
3970 work info = (runFold (FL.toParallelSVar sv info) xs)
3972 {-# NOINLINE workLim #-}
3973 workLim info = runFold (FL.toParallelSVarLimited sv info) xs
3975 {-# NOINLINE forkWithDiag #-}
3977 -- We do not use workerCount in case of ParallelVar but still there is
3978 -- no harm in maintaining it correctly.
3979 liftIO $ atomicModifyIORefCAS_ (workerCount sv) $ \n -> n + 1
3981 -- This allocation matters when significant number of workers are being
3982 -- sent. We allocate it only when needed. The overhead increases by 4x.
3984 case yieldRateInfo sv of
3985 Nothing -> return Nothing
3986 Just _ -> liftIO $ do
3987 cntRef <- newIORef 0
3988 t <- getTime Monotonic
3989 lat <- newIORef (0, t)
3990 return $ Just WorkerInfo
3991 { workerYieldMax = 0
3992 , workerYieldCount = cntRef
3993 , workerLatencyStart = lat
3996 case getYieldLimit st of
3997 Nothing -> doFork (work winfo)
3999 (handleChildException sv)
4000 Just _ -> doFork (workLim winfo)
4002 (handleChildException sv)
4005{-# INLINE_NORMAL mkParallelD #-}
4006mkParallelD :: MonadAsync m => Stream m a -> Stream m a
4007mkParallelD m = Stream step Nothing
4010 step gst Nothing = do
4011 sv <- newParallelVar StopNone gst
4012 toSVarParallel gst sv m
4013 -- XXX use unfold instead?
4014 return $ Skip $ Just $ fromSVar sv
4016 step gst (Just (UnStream step1 st)) = do
4019 Yield a s -> Yield a (Just $ Stream step1 s)
4020 Skip s -> Skip (Just $ Stream step1 s)
4023-- Compare with mkAsync. mkAsync uses an Async style SVar whereas this uses a
4024-- parallel style SVar for evaluation. Currently, parallel style cannot use
4025-- rate control whereas Async style can use rate control. In async style SVar
4026-- the worker thread terminates when the buffer is full whereas in Parallel
4029-- | Make the stream producer and consumer run concurrently by introducing a
4030-- buffer between them. The producer thread evaluates the input stream until
4031-- the buffer fills, it blocks if the buffer is full until there is space in
4032-- the buffer. The consumer consumes the stream lazily from the buffer.
4036{-# INLINE_NORMAL mkParallel #-}
4037mkParallel :: (K.IsStream t, MonadAsync m) => t m a -> t m a
4038mkParallel = fromStreamD . mkParallelD . toStreamD
4040-------------------------------------------------------------------------------
4042-------------------------------------------------------------------------------
4044-- | Create an SVar with a fold consumer that will fold any elements sent to it
4045-- using the supplied fold function.
4046{-# INLINE newFoldSVar #-}
4047newFoldSVar :: MonadAsync m => State t m a -> Fold m a b -> m (SVar t m a)
4048newFoldSVar stt f = do
4049 -- Buffer size for the SVar is derived from the current state
4050 sv <- newParallelVar StopAny (adaptState stt)
4051 -- Add the producer thread-id to the SVar.
4052 liftIO myThreadId >>= modifyThread sv
4053 void $ doFork (work sv) (svarMrun sv) (handleFoldException sv)
4058 {-# NOINLINE work #-}
4059 work sv = void $ runFold f $ fromProducer sv
4061data TapState sv st = TapInit | Tapping sv st | TapDone st
4063{-# INLINE_NORMAL tapAsync #-}
4064tapAsync :: MonadAsync m => Fold m a b -> Stream m a -> Stream m a
4065tapAsync f (Stream step1 state1) = Stream step TapInit
4069 -- In general, a Stop event would come equipped with the result
4070 -- of the fold. It is not used here but it would be useful in
4071 -- applicative and distribute.
4072 done <- fromConsumer svr
4073 when (not done) $ do
4074 liftIO $ withDiagMVar svr "teeToSVar: waiting to drain"
4075 $ takeMVar (outputDoorBellFromConsumer svr)
4079 liftIO $ sendStop svr Nothing
4080 -- drain/wait until a stop event arrives from the fold.
4083 {-# INLINE_LATE step #-}
4084 step gst TapInit = do
4085 sv <- newFoldSVar gst f
4086 return $ Skip (Tapping sv state1)
4088 step gst (Tapping sv st) = do
4092 done <- pushToFold sv a
4095 -- XXX we do not need to wait synchronously here
4097 return $ Yield a (TapDone s)
4098 else return $ Yield a (Tapping sv s)
4099 Skip s -> return $ Skip (Tapping sv s)
4104 step gst (TapDone st) = do
4107 Yield a s -> Yield a (TapDone s)
4108 Skip s -> Skip (TapDone s)
4111-- XXX Exported from Array again as this fold is specific to Array
4112-- | Take last 'n' elements from the stream and discard the rest.
4114lastN :: (Storable a, MonadIO m) => Int -> Fold m a (Array a)
4115lastN n = Fold step initial done
4117 step (Tuple3' rb rh i) a = do
4118 rh1 <- liftIO $ RB.unsafeInsert rb rh a
4119 return $ Tuple3' rb rh1 (i + 1)
4120 initial = fmap (\(a, b) -> Tuple3' a b (0 :: Int)) $ liftIO $ RB.new n
4121 done (Tuple3' rb rh i) = do
4122 arr <- liftIO $ A.newArray n
4123 foldFunc i rh snoc' arr rb
4124 snoc' b a = liftIO $ A.unsafeSnoc b a
4126 | i < n = RB.unsafeFoldRingM
4127 | otherwise = RB.unsafeFoldRingFullM
4129------------------------------------------------------------------------------
4131------------------------------------------------------------------------------
4133-- XXX using getTime in the loop can be pretty expensive especially for
4134-- computations where iterations are lightweight. We have the following
4137-- 1) Run a timeout thread updating a flag asynchronously and check that
4138-- flag here, that way we can have a cheap termination check.
4140-- 2) Use COARSE clock to get time with lower resolution but more efficiently.
4142-- 3) Use rdtscp/rdtsc to get time directly from the processor, compute the
4143-- termination value of rdtsc in the beginning and then in each iteration just
4144-- get rdtsc and check if we should terminate.
4148 | TakeByTimeCheck st s
4149 | TakeByTimeYield st s
4151{-# INLINE_NORMAL takeByTime #-}
4152takeByTime :: (MonadIO m, TimeUnit64 t) => t -> Stream m a -> Stream m a
4153takeByTime duration (Stream step1 state1) = Stream step (TakeByTimeInit state1)
4156 lim = toRelTime64 duration
4158 {-# INLINE_LATE step #-}
4159 step _ (TakeByTimeInit _) | lim == 0 = return Stop
4160 step _ (TakeByTimeInit st) = do
4161 t0 <- liftIO $ getTime Monotonic
4162 return $ Skip (TakeByTimeYield st t0)
4163 step _ (TakeByTimeCheck st t0) = do
4164 t <- liftIO $ getTime Monotonic
4166 if diffAbsTime64 t t0 > lim
4168 else Skip (TakeByTimeYield st t0)
4169 step gst (TakeByTimeYield st t0) = do
4172 Yield x s -> Yield x (TakeByTimeCheck s t0)
4173 Skip s -> Skip (TakeByTimeCheck s t0)
4176data DropByTime st s x
4178 | DropByTimeGen st s
4179 | DropByTimeCheck st s x
4180 | DropByTimeYield st
4182{-# INLINE_NORMAL dropByTime #-}
4183dropByTime :: (MonadIO m, TimeUnit64 t) => t -> Stream m a -> Stream m a
4184dropByTime duration (Stream step1 state1) = Stream step (DropByTimeInit state1)
4187 lim = toRelTime64 duration
4189 {-# INLINE_LATE step #-}
4190 step _ (DropByTimeInit st) = do
4191 t0 <- liftIO $ getTime Monotonic
4192 return $ Skip (DropByTimeGen st t0)
4193 step gst (DropByTimeGen st t0) = do
4196 Yield x s -> Skip (DropByTimeCheck s t0 x)
4197 Skip s -> Skip (DropByTimeGen s t0)
4199 step _ (DropByTimeCheck st t0 x) = do
4200 t <- liftIO $ getTime Monotonic
4201 if diffAbsTime64 t t0 <= lim
4202 then return $ Skip $ DropByTimeGen st t0
4203 else return $ Yield x $ DropByTimeYield st
4204 step gst (DropByTimeYield st) = do
4207 Yield x s -> Yield x (DropByTimeYield s)
4208 Skip s -> Skip (DropByTimeYield 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