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-- Copyright : (c) Roman Leshchinskiy 2008-2010
25-- Copyright : (c) The University of Glasgow, 2009
26-- Copyright : (c) Bjoern Hoehrmann 2008-2009
29-- Maintainer : streamly@composewell.com
30-- Stability : experimental
33-- Direct style re-implementation of CPS style stream in StreamK module. The
34-- symbol or suffix 'D' in this module denotes the "Direct" style. GHC is able
35-- to INLINE and fuse direct style better, providing better performance than
39-- import qualified Streamly.Internal.Data.Stream.StreamD as D
42-- Some of the functions in this file have been adapted from the vector
43-- library, https://hackage.haskell.org/package/vector.
45module Streamly.Internal.Data.Stream.StreamD
50#if __GLASGOW_HASKELL__ >= 800
51 , Stream (Stream, UnStream)
71 -- ** Specialized Generation
72 -- | Generate a monadic stream from a seed.
85 , enumerateFromStepIntegral
86 , enumerateFromIntegral
87 , enumerateFromThenIntegral
88 , enumerateFromToIntegral
89 , enumerateFromThenToIntegral
91 , enumerateFromStepNum
94 , enumerateFromToFractional
95 , enumerateFromThenToFractional
98 -- | Transform an input structure into a stream.
99 -- | Direct style stream does not support @fromFoldable@.
129 -- ** Specialized Folds
156 -- ** Flattening nested streams
159 , ConcatMapUState (..)
161 , ConcatUnfoldInterleaveState (..)
162 , concatUnfoldInterleave
163 , concatUnfoldRoundrobin
166 , InterleaveState(..)
171 , roundRobin -- interleaveFair?/ParallelFair
204 -- | Transform a stream into another type.
221 -- ** By folding (scans)
266 , intersperseSuffixBySpan
292 -- * Transformation comprehensions
304 -- * Concurrent Application
309 -- XXX Move this later
310 -- XXX Exported from Array again
315import Control.Concurrent (killThread, myThreadId, takeMVar, threadDelay)
316import Control.Exception
317 (Exception, SomeException, AsyncException, fromException)
318import Control.Monad (void, when)
319import Control.Monad.Catch (MonadCatch, throwM)
320import Control.Monad.IO.Class (MonadIO(..))
321import Control.Monad.Reader (ReaderT)
322import Control.Monad.State.Strict (StateT)
323import Control.Monad.Trans (MonadTrans(lift))
324import Data.Bits (shiftR, shiftL, (.|.), (.&.))
325import Data.Functor.Identity (Identity(..))
326import Data.IORef (newIORef, readIORef, mkWeakIORef, writeIORef)
327import Data.Maybe (fromJust, isJust, isNothing)
328import Data.Word (Word32)
329import Foreign.Ptr (Ptr)
330import Foreign.Storable (Storable(..))
331import GHC.Types (SPEC(..))
332import System.Mem (performMajorGC)
334 hiding (map, mapM, mapM_, repeat, foldr, last, take, filter,
335 takeWhile, drop, dropWhile, all, any, maximum, minimum, elem,
336 notElem, null, head, tail, zipWith, lookup, foldr1, sequence,
337 (!!), scanl, scanl1, concatMap, replicate, enumFromTo, concat,
340import qualified Control.Monad.Catch as MC
341import qualified Control.Monad.Reader as Reader
342import qualified Control.Monad.State.Strict as State
343import qualified Prelude
345import Streamly.Internal.Mutable.Prim.Var
346 (Prim, Var, readVar, newVar, modifyVar')
347import Streamly.Internal.Data.Atomics (atomicModifyIORefCAS_)
348import Streamly.Internal.Memory.Array.Types (Array(..))
349import Streamly.Internal.Data.Fold.Types (Fold(..))
350import Streamly.Internal.Data.Pipe.Types (Pipe(..), PipeState(..))
351import Streamly.Internal.Data.Time.Clock (Clock(Monotonic), getTime)
352import Streamly.Internal.Data.Unfold.Types (Unfold(..))
353import Streamly.Internal.Data.Strict (Tuple3'(..))
355import Streamly.Internal.Data.Stream.StreamD.Type
356import Streamly.Internal.Data.SVar
357import Streamly.Internal.Data.Stream.SVar (fromConsumer, pushToFold)
359import qualified Streamly.Internal.Data.Pipe.Types as Pipe
360import qualified Streamly.Internal.Memory.Array.Types as A
361import qualified Streamly.Internal.Data.Fold as FL
362import qualified Streamly.Memory.Ring as RB
363import qualified Streamly.Internal.Data.Stream.StreamK as K
365------------------------------------------------------------------------------
367------------------------------------------------------------------------------
369-- | An empty 'Stream'.
370{-# INLINE_NORMAL nil #-}
371nil :: Monad m => Stream m a
372nil = Stream (\_ _ -> return Stop) ()
374-- | An empty 'Stream' with a side effect.
375{-# INLINE_NORMAL nilM #-}
376nilM :: Monad m => m b -> Stream m a
377nilM m = Stream (\_ _ -> m >> return Stop) ()
379{-# INLINE_NORMAL consM #-}
380consM :: Monad m => m a -> Stream m a -> Stream m a
381consM m (Stream step state) = Stream step1 Nothing
383 {-# INLINE_LATE step1 #-}
384 step1 _ Nothing = m >>= \x -> return $ Yield x (Just state)
385 step1 gst (Just st) = do
389 Yield a s -> Yield a (Just s)
390 Skip s -> Skip (Just s)
393-- XXX implement in terms of consM?
394-- cons x = consM (return x)
396-- | Can fuse but has O(n^2) complexity.
397{-# INLINE_NORMAL cons #-}
398cons :: Monad m => a -> Stream m a -> Stream m a
399cons x (Stream step state) = Stream step1 Nothing
401 {-# INLINE_LATE step1 #-}
402 step1 _ Nothing = return $ Yield x (Just state)
403 step1 gst (Just st) = do
407 Yield a s -> Yield a (Just s)
408 Skip s -> Skip (Just s)
411-------------------------------------------------------------------------------
413-------------------------------------------------------------------------------
415-- Does not fuse, has the same performance as the StreamK version.
416{-# INLINE_NORMAL uncons #-}
417uncons :: Monad m => Stream m a -> m (Maybe (a, Stream m a))
418uncons (UnStream step state) = go state
421 r <- step defState st
423 Yield x s -> return $ Just (x, Stream step s)
425 Stop -> return Nothing
427------------------------------------------------------------------------------
428-- Generation by unfold
429------------------------------------------------------------------------------
431{-# INLINE_NORMAL unfoldrM #-}
432unfoldrM :: Monad m => (s -> m (Maybe (a, s))) -> s -> Stream m a
433unfoldrM next state = Stream step state
435 {-# INLINE_LATE step #-}
439 Just (x, s) -> Yield x s
442{-# INLINE_LATE unfoldr #-}
443unfoldr :: Monad m => (s -> Maybe (a, s)) -> s -> Stream m a
444unfoldr f = unfoldrM (return . f)
446-- | Convert an 'Unfold' into a 'Stream' by supplying it a seed.
448{-# INLINE_NORMAL unfold #-}
449unfold :: Monad m => Unfold m a b -> a -> Stream m b
450unfold (Unfold ustep inject) seed = Stream step Nothing
452 {-# INLINE_LATE step #-}
453 step _ Nothing = inject seed >>= return . Skip . Just
454 step _ (Just st) = do
457 Yield x s -> Yield x (Just s)
458 Skip s -> Skip (Just s)
461------------------------------------------------------------------------------
462-- Specialized Generation
463------------------------------------------------------------------------------
465{-# INLINE_NORMAL repeatM #-}
466repeatM :: Monad m => m a -> Stream m a
467repeatM x = Stream (\_ _ -> x >>= \r -> return $ Yield r ()) ()
469{-# INLINE_NORMAL repeat #-}
470repeat :: Monad m => a -> Stream m a
471repeat x = Stream (\_ _ -> return $ Yield x ()) ()
473{-# INLINE_NORMAL iterateM #-}
474iterateM :: Monad m => (a -> m a) -> m a -> Stream m a
475iterateM step = Stream (\_ st -> st >>= \x -> return $ Yield x (step x))
477{-# INLINE_NORMAL iterate #-}
478iterate :: Monad m => (a -> a) -> a -> Stream m a
479iterate step st = iterateM (return . step) (return st)
481{-# INLINE_NORMAL replicateM #-}
482replicateM :: forall m a. Monad m => Int -> m a -> Stream m a
483replicateM n p = Stream step n
485 {-# INLINE_LATE step #-}
487 | i <= 0 = return Stop
490 return $ Yield x (i - 1)
492{-# INLINE_NORMAL replicate #-}
493replicate :: Monad m => Int -> a -> Stream m a
494replicate n x = replicateM n (return x)
496-- This would not work properly for floats, therefore we put an Integral
498-- | Can be used to enumerate unbounded integrals. This does not check for
499-- overflow or underflow for bounded integrals.
500{-# INLINE_NORMAL enumerateFromStepIntegral #-}
501enumerateFromStepIntegral :: (Integral a, Monad m) => a -> a -> Stream m a
502enumerateFromStepIntegral from stride =
503 from `seq` stride `seq` Stream step from
505 {-# INLINE_LATE step #-}
506 step _ !x = return $ Yield x $! (x + stride)
508-- We are assuming that "to" is constrained by the type to be within
510{-# INLINE enumerateFromToIntegral #-}
511enumerateFromToIntegral :: (Monad m, Integral a) => a -> a -> Stream m a
512enumerateFromToIntegral from to =
513 takeWhile (<= to) $ enumerateFromStepIntegral from 1
515{-# INLINE enumerateFromIntegral #-}
516enumerateFromIntegral :: (Monad m, Integral a, Bounded a) => a -> Stream m a
517enumerateFromIntegral from = enumerateFromToIntegral from maxBound
519data EnumState a = EnumInit | EnumYield a a a | EnumStop
521{-# INLINE_NORMAL enumerateFromThenToIntegralUp #-}
522enumerateFromThenToIntegralUp
523 :: (Monad m, Integral a)
524 => a -> a -> a -> Stream m a
525enumerateFromThenToIntegralUp from next to = Stream step EnumInit
527 {-# INLINE_LATE step #-}
533 else Yield from EnumStop
534 else -- from <= next <= to
535 let stride = next - from
536 in Skip $ EnumYield from stride (to - stride)
538 step _ (EnumYield x stride toMinus) =
541 then Yield x EnumStop
542 else Yield x $ EnumYield (x + stride) stride toMinus
544 step _ EnumStop = return Stop
546{-# INLINE_NORMAL enumerateFromThenToIntegralDn #-}
547enumerateFromThenToIntegralDn
548 :: (Monad m, Integral a)
549 => a -> a -> a -> Stream m a
550enumerateFromThenToIntegralDn from next to = Stream step EnumInit
552 {-# INLINE_LATE step #-}
554 return $ if to > next
557 else Yield from EnumStop
558 else -- from >= next >= to
559 let stride = next - from
560 in Skip $ EnumYield from stride (to - stride)
562 step _ (EnumYield x stride toMinus) =
565 then Yield x EnumStop
566 else Yield x $ EnumYield (x + stride) stride toMinus
568 step _ EnumStop = return Stop
570{-# INLINE_NORMAL enumerateFromThenToIntegral #-}
571enumerateFromThenToIntegral
572 :: (Monad m, Integral a)
573 => a -> a -> a -> Stream m a
574enumerateFromThenToIntegral from next to
575 | next >= from = enumerateFromThenToIntegralUp from next to
576 | otherwise = enumerateFromThenToIntegralDn from next to
578{-# INLINE_NORMAL enumerateFromThenIntegral #-}
579enumerateFromThenIntegral
580 :: (Monad m, Integral a, Bounded a)
581 => a -> a -> Stream m a
582enumerateFromThenIntegral from next =
584 then enumerateFromThenToIntegralUp from next maxBound
585 else enumerateFromThenToIntegralDn from next minBound
587-- For floating point numbers if the increment is less than the precision then
588-- it just gets lost. Therefore we cannot always increment it correctly by just
590-- 9007199254740992 + 1 + 1 :: Double => 9.007199254740992e15
591-- 9007199254740992 + 2 :: Double => 9.007199254740994e15
593-- Instead we accumulate the increment counter and compute the increment
594-- every time before adding it to the starting number.
596-- This works for Integrals as well as floating point numbers, but
597-- enumerateFromStepIntegral is faster for integrals.
598{-# INLINE_NORMAL enumerateFromStepNum #-}
599enumerateFromStepNum :: (Monad m, Num a) => a -> a -> Stream m a
600enumerateFromStepNum from stride = Stream step 0
602 {-# INLINE_LATE step #-}
603 step _ !i = return $ (Yield $! (from + i * stride)) $! (i + 1)
605{-# INLINE_NORMAL numFrom #-}
606numFrom :: (Monad m, Num a) => a -> Stream m a
607numFrom from = enumerateFromStepNum from 1
609{-# INLINE_NORMAL numFromThen #-}
610numFromThen :: (Monad m, Num a) => a -> a -> Stream m a
611numFromThen from next = enumerateFromStepNum from (next - from)
613-- We cannot write a general function for Num. The only way to write code
614-- portable between the two is to use a 'Real' constraint and convert between
615-- Fractional and Integral using fromRational which is horribly slow.
616{-# INLINE_NORMAL enumerateFromToFractional #-}
617enumerateFromToFractional
618 :: (Monad m, Fractional a, Ord a)
619 => a -> a -> Stream m a
620enumerateFromToFractional from to =
621 takeWhile (<= to + 1 / 2) $ enumerateFromStepNum from 1
623{-# INLINE_NORMAL enumerateFromThenToFractional #-}
624enumerateFromThenToFractional
625 :: (Monad m, Fractional a, Ord a)
626 => a -> a -> a -> Stream m a
627enumerateFromThenToFractional from next to =
628 takeWhile predicate $ numFromThen from next
630 mid = (next - from) / 2
631 predicate | next >= from = (<= to + mid)
632 | otherwise = (>= to + mid)
634-------------------------------------------------------------------------------
635-- Generation by Conversion
636-------------------------------------------------------------------------------
638{-# INLINE_NORMAL fromIndicesM #-}
639fromIndicesM :: Monad m => (Int -> m a) -> Stream m a
640fromIndicesM gen = Stream step 0
642 {-# INLINE_LATE step #-}
645 return $ Yield x (i + 1)
647{-# INLINE fromIndices #-}
648fromIndices :: Monad m => (Int -> a) -> Stream m a
649fromIndices gen = fromIndicesM (return . gen)
651{-# INLINE_NORMAL generateM #-}
652generateM :: Monad m => Int -> (Int -> m a) -> Stream m a
653generateM n gen = n `seq` Stream step 0
655 {-# INLINE_LATE step #-}
656 step _ i | i < n = do
658 return $ Yield x (i + 1)
659 | otherwise = return Stop
661{-# INLINE generate #-}
662generate :: Monad m => Int -> (Int -> a) -> Stream m a
663generate n gen = generateM n (return . gen)
665-- XXX we need the MonadAsync constraint because of a rewrite rule.
666-- | Convert a list of monadic actions to a 'Stream'
667{-# INLINE_LATE fromListM #-}
668fromListM :: MonadAsync m => [m a] -> Stream m a
669fromListM = Stream step
671 {-# INLINE_LATE step #-}
672 step _ (m:ms) = m >>= \x -> return $ Yield x ms
673 step _ [] = return Stop
675{-# INLINE toStreamD #-}
676toStreamD :: (K.IsStream t, Monad m) => t m a -> Stream m a
677toStreamD = fromStreamK . K.toStream
679{-# INLINE_NORMAL fromPrimVar #-}
680fromPrimVar :: (MonadIO m, Prim a) => Var IO a -> Stream m a
681fromPrimVar var = Stream step ()
683 {-# INLINE_LATE step #-}
684 step _ () = liftIO (readVar var) >>= \x -> return $ Yield x ()
686-------------------------------------------------------------------------------
687-- Generation from SVar
688-------------------------------------------------------------------------------
690data FromSVarState t m a =
692 | FromSVarRead (SVar t m a)
693 | FromSVarLoop (SVar t m a) [ChildEvent a]
694 | FromSVarDone (SVar t m a)
696{-# INLINE_NORMAL fromSVar #-}
697fromSVar :: (MonadAsync m) => SVar t m a -> Stream m a
698fromSVar svar = Stream step FromSVarInit
701 {-# INLINE_LATE step #-}
702 step _ FromSVarInit = do
703 ref <- liftIO $ newIORef ()
704 _ <- liftIO $ mkWeakIORef ref hook
705 -- when this copy of svar gets garbage collected "ref" will get
706 -- garbage collected and our GC hook will be called.
707 let sv = svar{svarRef = Just ref}
708 return $ Skip (FromSVarRead sv)
712 {-# NOINLINE hook #-}
714 when (svarInspectMode svar) $ do
715 r <- liftIO $ readIORef (svarStopTime (svarStats svar))
717 printSVar svar "SVar Garbage Collected"
719 -- If there are any SVars referenced by this SVar a GC will prompt
720 -- them to be cleaned up quickly.
721 when (svarInspectMode svar) performMajorGC
723 step _ (FromSVarRead sv) = do
724 list <- readOutputQ sv
725 -- Reversing the output is important to guarantee that we process the
726 -- outputs in the same order as they were generated by the constituent
728 return $ Skip $ FromSVarLoop sv (Prelude.reverse list)
730 step _ (FromSVarLoop sv []) = do
731 done <- postProcess sv
732 return $ Skip $ if done
733 then (FromSVarDone sv)
734 else (FromSVarRead sv)
736 step _ (FromSVarLoop sv (ev : es)) = do
738 ChildYield a -> return $ Yield a (FromSVarLoop sv es)
739 ChildStop tid e -> do
743 stop <- shouldStop tid
746 liftIO (cleanupSVar sv)
747 return $ Skip (FromSVarDone sv)
748 else return $ Skip (FromSVarLoop sv es)
750 case fromException ex of
752 return $ Skip (FromSVarLoop sv es)
753 Nothing -> liftIO (cleanupSVar sv) >> throwM ex
757 case svarStopStyle sv of
758 StopNone -> return False
759 StopAny -> return True
761 sid <- liftIO $ readIORef (svarStopBy sv)
762 return $ if tid == sid then True else False
764 step _ (FromSVarDone sv) = do
765 when (svarInspectMode sv) $ do
766 t <- liftIO $ getTime Monotonic
767 liftIO $ writeIORef (svarStopTime (svarStats sv)) (Just t)
768 liftIO $ printSVar sv "SVar Done"
771-------------------------------------------------------------------------------
772-- Process events received by a fold consumer from a stream producer
773-------------------------------------------------------------------------------
775{-# INLINE_NORMAL fromProducer #-}
776fromProducer :: (MonadAsync m) => SVar t m a -> Stream m a
777fromProducer svar = Stream step (FromSVarRead svar)
780 {-# INLINE_LATE step #-}
781 step _ (FromSVarRead sv) = do
782 list <- readOutputQ sv
783 -- Reversing the output is important to guarantee that we process the
784 -- outputs in the same order as they were generated by the constituent
786 return $ Skip $ FromSVarLoop sv (Prelude.reverse list)
788 step _ (FromSVarLoop sv []) = return $ Skip $ FromSVarRead sv
789 step _ (FromSVarLoop sv (ev : es)) = do
791 ChildYield a -> return $ Yield a (FromSVarLoop sv es)
792 ChildStop tid e -> do
796 sendStopToProducer sv
797 return $ Skip (FromSVarDone sv)
798 Just _ -> error "Bug: fromProducer: received exception"
800 step _ (FromSVarDone sv) = do
801 when (svarInspectMode sv) $ do
802 t <- liftIO $ getTime Monotonic
803 liftIO $ writeIORef (svarStopTime (svarStats sv)) (Just t)
804 liftIO $ printSVar sv "SVar Done"
807 step _ FromSVarInit = undefined
809-------------------------------------------------------------------------------
810-- Hoisting the inner monad
811-------------------------------------------------------------------------------
813{-# INLINE_NORMAL hoist #-}
814hoist :: Monad n => (forall x. m x -> n x) -> Stream m a -> Stream n a
815hoist f (Stream step state) = (Stream step' state)
817 {-# INLINE_LATE step' #-}
819 r <- f $ step (adaptState gst) st
821 Yield x s -> Yield x s
825{-# INLINE generally #-}
826generally :: Monad m => Stream Identity a -> Stream m a
827generally = hoist (return . runIdentity)
829{-# INLINE_NORMAL liftInner #-}
830liftInner :: (Monad m, MonadTrans t, Monad (t m))
831 => Stream m a -> Stream (t m) a
832liftInner (Stream step state) = Stream step' state
834 {-# INLINE_LATE step' #-}
836 r <- lift $ step (adaptState gst) st
838 Yield x s -> Yield x s
842{-# INLINE_NORMAL runReaderT #-}
843runReaderT :: Monad m => s -> Stream (ReaderT s m) a -> Stream m a
844runReaderT sval (Stream step state) = Stream step' state
846 {-# INLINE_LATE step' #-}
848 r <- Reader.runReaderT (step (adaptState gst) st) sval
850 Yield x s -> Yield x s
854{-# INLINE_NORMAL evalStateT #-}
855evalStateT :: Monad m => s -> Stream (StateT s m) a -> Stream m a
856evalStateT sval (Stream step state) = Stream step' (state, sval)
858 {-# INLINE_LATE step' #-}
859 step' gst (st, sv) = do
860 (r, sv') <- State.runStateT (step (adaptState gst) st) sv
862 Yield x s -> Yield x (s, sv')
863 Skip s -> Skip (s, sv')
866{-# INLINE_NORMAL runStateT #-}
867runStateT :: Monad m => s -> Stream (StateT s m) a -> Stream m (s, a)
868runStateT sval (Stream step state) = Stream step' (state, sval)
870 {-# INLINE_LATE step' #-}
871 step' gst (st, sv) = do
872 (r, sv') <- State.runStateT (step (adaptState gst) st) sv
874 Yield x s -> Yield (sv', x) (s, sv')
875 Skip s -> Skip (s, sv')
878------------------------------------------------------------------------------
879-- Elimination by Folds
880------------------------------------------------------------------------------
882------------------------------------------------------------------------------
884------------------------------------------------------------------------------
886{-# INLINE_NORMAL foldr1 #-}
887foldr1 :: Monad m => (a -> a -> a) -> Stream m a -> m (Maybe a)
891 Nothing -> return Nothing
892 Just (h, t) -> fmap Just (foldr f h t)
894------------------------------------------------------------------------------
896------------------------------------------------------------------------------
898{-# INLINE_NORMAL foldlT #-}
899foldlT :: (Monad m, Monad (s m), MonadTrans s)
900 => (s m b -> a -> s m b) -> s m b -> Stream m a -> s m b
901foldlT fstep begin (Stream step state) = go SPEC begin state
904 r <- lift $ step defState st
906 Yield x s -> go SPEC (fstep acc x) s
907 Skip s -> go SPEC acc s
910-- Note, this is going to have horrible performance, because of the nature of
911-- the stream type (i.e. direct stream vs CPS). Its only for reference, it is
912-- likely be practically unusable.
913{-# INLINE_NORMAL foldlS #-}
915 => (Stream m b -> a -> Stream m b) -> Stream m b -> Stream m a -> Stream m b
916foldlS fstep begin (Stream step state) = Stream step' (Left (state, begin))
918 step' gst (Left (st, acc)) = do
919 r <- step (adaptState gst) st
921 Yield x s -> Skip (Left (s, fstep acc x))
922 Skip s -> Skip (Left (s, acc))
923 Stop -> Skip (Right acc)
925 step' gst (Right (Stream stp stt)) = do
926 r <- stp (adaptState gst) stt
928 Yield x s -> Yield x (Right (Stream stp s))
929 Skip s -> Skip (Right (Stream stp s))
932------------------------------------------------------------------------------
934------------------------------------------------------------------------------
936-- | Run a streaming composition, discard the results.
937{-# INLINE_LATE drain #-}
938drain :: Monad m => Stream m a -> m ()
939-- drain = foldrM (\_ xs -> xs) (return ())
940drain (Stream step state) = go SPEC state
943 r <- step defState st
945 Yield _ s -> go SPEC s
949{-# INLINE_NORMAL null #-}
950null :: Monad m => Stream m a -> m Bool
951null m = foldrM (\_ _ -> return False) (return True) m
953{-# INLINE_NORMAL head #-}
954head :: Monad m => Stream m a -> m (Maybe a)
955head m = foldrM (\x _ -> return (Just x)) (return Nothing) m
957{-# INLINE_NORMAL headElse #-}
958headElse :: Monad m => a -> Stream m a -> m a
959headElse a m = foldrM (\x _ -> return x) (return a) m
961-- Does not fuse, has the same performance as the StreamK version.
962{-# INLINE_NORMAL tail #-}
963tail :: Monad m => Stream m a -> m (Maybe (Stream m a))
964tail (UnStream step state) = go state
967 r <- step defState st
969 Yield _ s -> return (Just $ Stream step s)
971 Stop -> return Nothing
973-- XXX will it fuse? need custom impl?
974{-# INLINE_NORMAL last #-}
975last :: Monad m => Stream m a -> m (Maybe a)
976last = foldl' (\_ y -> Just y) Nothing
978{-# INLINE_NORMAL elem #-}
979elem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
980-- elem e m = foldrM (\x xs -> if x == e then return True else xs) (return False) m
981elem e (Stream step state) = go state
984 r <- step defState st
987 | x == e -> return True
992{-# INLINE_NORMAL notElem #-}
993notElem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
994notElem e s = fmap not (elem e s)
996{-# INLINE_NORMAL all #-}
997all :: Monad m => (a -> Bool) -> Stream m a -> m Bool
998-- all p m = foldrM (\x xs -> if p x then xs else return False) (return True) m
999all p (Stream step state) = go state
1002 r <- step defState st
1006 | otherwise -> return False
1010{-# INLINE_NORMAL any #-}
1011any :: Monad m => (a -> Bool) -> Stream m a -> m Bool
1012-- any p m = foldrM (\x xs -> if p x then return True else xs) (return False) m
1013any p (Stream step state) = go state
1016 r <- step defState st
1019 | p x -> return True
1022 Stop -> return False
1024{-# INLINE_NORMAL maximum #-}
1025maximum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
1026maximum (Stream step state) = go Nothing state
1029 r <- step defState st
1031 Yield x s -> go (Just x) s
1032 Skip s -> go Nothing s
1033 Stop -> return Nothing
1034 go (Just acc) st = do
1035 r <- step defState st
1038 | acc <= x -> go (Just x) s
1039 | otherwise -> go (Just acc) s
1040 Skip s -> go (Just acc) s
1041 Stop -> return (Just acc)
1043{-# INLINE_NORMAL maximumBy #-}
1044maximumBy :: Monad m => (a -> a -> Ordering) -> Stream m a -> m (Maybe a)
1045maximumBy cmp (Stream step state) = go Nothing state
1048 r <- step defState st
1050 Yield x s -> go (Just x) s
1051 Skip s -> go Nothing s
1052 Stop -> return Nothing
1053 go (Just acc) st = do
1054 r <- step defState st
1056 Yield x s -> case cmp acc x of
1057 GT -> go (Just acc) s
1059 Skip s -> go (Just acc) s
1060 Stop -> return (Just acc)
1062{-# INLINE_NORMAL minimum #-}
1063minimum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
1064minimum (Stream step state) = go Nothing state
1067 r <- step defState st
1069 Yield x s -> go (Just x) s
1070 Skip s -> go Nothing s
1071 Stop -> return Nothing
1072 go (Just acc) st = do
1073 r <- step defState st
1076 | acc <= x -> go (Just acc) s
1077 | otherwise -> go (Just x) s
1078 Skip s -> go (Just acc) s
1079 Stop -> return (Just acc)
1081{-# INLINE_NORMAL minimumBy #-}
1082minimumBy :: Monad m => (a -> a -> Ordering) -> Stream m a -> m (Maybe a)
1083minimumBy cmp (Stream step state) = go Nothing state
1086 r <- step defState st
1088 Yield x s -> go (Just x) s
1089 Skip s -> go Nothing s
1090 Stop -> return Nothing
1091 go (Just acc) st = do
1092 r <- step defState st
1094 Yield x s -> case cmp acc x of
1096 _ -> go (Just acc) s
1097 Skip s -> go (Just acc) s
1098 Stop -> return (Just acc)
1100{-# INLINE_NORMAL (!!) #-}
1101(!!) :: (Monad m) => Stream m a -> Int -> m (Maybe a)
1102(Stream step state) !! i = go i state
1105 r <- step defState st
1107 Yield x s | n < 0 -> return Nothing
1108 | n == 0 -> return $ Just x
1109 | otherwise -> go (n - 1) s
1111 Stop -> return Nothing
1113{-# INLINE_NORMAL lookup #-}
1114lookup :: (Monad m, Eq a) => a -> Stream m (a, b) -> m (Maybe b)
1115lookup e m = foldrM (\(a, b) xs -> if e == a then return (Just b) else xs)
1118{-# INLINE_NORMAL findM #-}
1119findM :: Monad m => (a -> m Bool) -> Stream m a -> m (Maybe a)
1120findM p m = foldrM (\x xs -> p x >>= \r -> if r then return (Just x) else xs)
1124find :: Monad m => (a -> Bool) -> Stream m a -> m (Maybe a)
1125find p = findM (return . p)
1127{-# INLINE_NORMAL findIndices #-}
1128findIndices :: Monad m => (a -> Bool) -> Stream m a -> Stream m Int
1129findIndices p (Stream step state) = Stream step' (state, 0)
1131 {-# INLINE_LATE step' #-}
1132 step' gst (st, i) = i `seq` do
1133 r <- step (adaptState gst) st
1135 Yield x s -> if p x then Yield i (s, i+1) else Skip (s, i+1)
1136 Skip s -> Skip (s, i)
1139{-# INLINE toListRev #-}
1140toListRev :: Monad m => Stream m a -> m [a]
1141toListRev = foldl' (flip (:)) []
1143-- We can implement reverse as:
1145-- > reverse = foldlS (flip cons) nil
1147-- However, this implementation is unusable because of the horrible performance
1148-- of cons. So we just convert it to a list first and then stream from the
1151-- XXX Maybe we can use an Array instead of a list here?
1152{-# INLINE_NORMAL reverse #-}
1153reverse :: Monad m => Stream m a -> Stream m a
1154reverse m = Stream step Nothing
1156 {-# INLINE_LATE step #-}
1159 return $ Skip (Just xs)
1160 step _ (Just (x:xs)) = return $ Yield x (Just xs)
1161 step _ (Just []) = return Stop
1163-- Much faster reverse for Storables
1164{-# INLINE_NORMAL reverse' #-}
1165reverse' :: forall m a. (MonadIO m, Storable a) => Stream m a -> Stream m a
1167-- This commented implementation copies the whole stream into one single array
1168-- and then streams from that array, this is 3-4x faster than the chunked code
1169-- that follows. Though this could be problematic due to unbounded large
1170-- allocations. We need to figure out why the chunked code is slower and if we
1171-- can optimize the chunked code to work as fast as this one. It may be a
1173import Foreign.ForeignPtr (touchForeignPtr)
1174import Foreign.ForeignPtr.Unsafe (unsafeForeignPtrToPtr)
1175import Foreign.Ptr (Ptr, plusPtr)
1176reverse' m = Stream step Nothing
1178 {-# INLINE_LATE step #-}
1180 arr <- A.fromStreamD m
1181 let p = aEnd arr `plusPtr` negate (sizeOf (undefined :: a))
1182 return $ Skip $ Just (aStart arr, p)
1184 step _ (Just (start, p)) | p < unsafeForeignPtrToPtr start = return Stop
1186 step _ (Just (start, p)) = do
1187 let !x = A.unsafeInlineIO $ do
1189 touchForeignPtr start
1191 next = p `plusPtr` negate (sizeOf (undefined :: a))
1192 return $ Yield x (Just (start, next))
1199 $ A.fromStreamDArraysOf A.defaultChunkSize m
1202------------------------------------------------------------------------------
1203-- Grouping/Splitting
1204------------------------------------------------------------------------------
1206{-# INLINE_NORMAL splitSuffixBy' #-}
1207splitSuffixBy' :: Monad m
1208 => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1209splitSuffixBy' predicate f (Stream step state) =
1210 Stream (stepOuter f) (Just state)
1214 {-# INLINE_LATE stepOuter #-}
1215 stepOuter (Fold fstep initial done) gst (Just st) = do
1216 res <- step (adaptState gst) st
1222 then done acc' >>= \val -> return $ Yield val (Just s)
1225 Skip s -> return $ Skip $ Just s
1231 res <- step (adaptState gst) stt
1236 then done acc' >>= \val -> return $ Yield val (Just s)
1238 Skip s -> go SPEC s acc
1239 Stop -> done acc >>= \val -> return $ Yield val Nothing
1241 stepOuter _ _ Nothing = return Stop
1243{-# INLINE_NORMAL groupsBy #-}
1249groupsBy cmp f (Stream step state) = Stream (stepOuter f) (Just state, Nothing)
1253 {-# INLINE_LATE stepOuter #-}
1254 stepOuter (Fold fstep initial done) gst (Just st, Nothing) = do
1255 res <- step (adaptState gst) st
1262 Skip s -> return $ Skip $ (Just s, Nothing)
1267 go !_ prev stt !acc = do
1268 res <- step (adaptState gst) stt
1275 else done acc >>= \r -> return $ Yield r (Just s, Just x)
1276 Skip s -> go SPEC prev s acc
1277 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1279 stepOuter (Fold fstep initial done) gst (Just st, Just prev) = do
1281 acc' <- fstep acc prev
1286 -- XXX code duplicated from the previous equation
1288 res <- step (adaptState gst) stt
1295 else done acc >>= \r -> return $ Yield r (Just s, Just x)
1296 Skip s -> go SPEC s acc
1297 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1299 stepOuter _ _ (Nothing,_) = return Stop
1301{-# INLINE_NORMAL groupsRollingBy #-}
1302groupsRollingBy :: Monad m
1307groupsRollingBy cmp f (Stream step state) =
1308 Stream (stepOuter f) (Just state, Nothing)
1311 {-# INLINE_LATE stepOuter #-}
1312 stepOuter (Fold fstep initial done) gst (Just st, Nothing) = do
1313 res <- step (adaptState gst) st
1320 Skip s -> return $ Skip $ (Just s, Nothing)
1324 go !_ prev stt !acc = do
1325 res <- step (adaptState gst) stt
1333 done acc >>= \r -> return $ Yield r (Just s, Just x)
1334 Skip s -> go SPEC prev s acc
1335 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1337 stepOuter (Fold fstep initial done) gst (Just st, Just prev') = do
1339 acc' <- fstep acc prev'
1340 go SPEC prev' st acc'
1343 go !_ prevv stt !acc = do
1344 res <- step (adaptState gst) stt
1351 else done acc >>= \r -> return $ Yield r (Just s, Just x)
1352 Skip s -> go SPEC prevv s acc
1353 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1355 stepOuter _ _ (Nothing, _) = return Stop
1357{-# INLINE_NORMAL splitBy #-}
1358splitBy :: Monad m => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1359splitBy predicate f (Stream step state) = Stream (step' f) (Just state)
1363 {-# INLINE_LATE step' #-}
1364 step' (Fold fstep initial done) gst (Just st) = initial >>= go SPEC st
1369 res <- step (adaptState gst) stt
1373 then done acc >>= \r -> return $ Yield r (Just s)
1377 Skip s -> go SPEC s acc
1378 Stop -> done acc >>= \r -> return $ Yield r Nothing
1380 step' _ _ Nothing = return Stop
1382-- XXX requires -funfolding-use-threshold=150 in lines-unlines benchmark
1383{-# INLINE_NORMAL splitSuffixBy #-}
1384splitSuffixBy :: Monad m
1385 => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1386splitSuffixBy predicate f (Stream step state) = Stream (step' f) (Just state)
1390 {-# INLINE_LATE step' #-}
1391 step' (Fold fstep initial done) gst (Just st) = do
1392 res <- step (adaptState gst) st
1397 then done acc >>= \val -> return $ Yield val (Just s)
1402 Skip s -> return $ Skip $ Just s
1408 res <- step (adaptState gst) stt
1412 then done acc >>= \r -> return $ Yield r (Just s)
1416 Skip s -> go SPEC s acc
1417 Stop -> done acc >>= \r -> return $ Yield r Nothing
1419 step' _ _ Nothing = return Stop
1421{-# INLINE_NORMAL wordsBy #-}
1422wordsBy :: Monad m => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1423wordsBy predicate f (Stream step state) = Stream (stepOuter f) (Just state)
1427 {-# INLINE_LATE stepOuter #-}
1428 stepOuter (Fold fstep initial done) gst (Just st) = do
1429 res <- step (adaptState gst) st
1433 then return $ Skip (Just s)
1439 Skip s -> return $ Skip $ Just s
1445 res <- step (adaptState gst) stt
1449 then done acc >>= \r -> return $ Yield r (Just s)
1453 Skip s -> go SPEC s acc
1454 Stop -> done acc >>= \r -> return $ Yield r Nothing
1456 stepOuter _ _ Nothing = return Stop
1458-- String search algorithms:
1459-- http://www-igm.univ-mlv.fr/~lecroq/string/index.html
1462-- TODO can we unify the splitting operations using a splitting configuration
1463-- like in the split package.
1465data SplitStyle = Infix | Suffix | Prefix deriving (Eq, Show)
1467data SplitOptions = SplitOptions
1468 { style :: SplitStyle
1469 , withSep :: Bool -- ^ keep the separators in output
1470 -- , compact :: Bool -- ^ treat multiple consecutive separators as one
1471 -- , trimHead :: Bool -- ^ drop blank at head
1472 -- , trimTail :: Bool -- ^ drop blank at tail
1476data SplitOnState s a =
1481 | GO_KARP_RABIN s !(RB.Ring a) !(Ptr a)
1484{-# INLINE_NORMAL splitOn #-}
1486 :: forall m a b. (MonadIO m, Storable a, Enum a, Eq a)
1491splitOn patArr@Array{..} (Fold fstep initial done) (Stream step state) =
1492 Stream stepOuter GO_START
1496 patLen = A.length patArr
1497 maxIndex = patLen - 1
1498 elemBits = sizeOf (undefined :: a) * 8
1500 {-# INLINE_LATE stepOuter #-}
1501 stepOuter _ GO_START =
1503 then return $ Skip $ GO_EMPTY_PAT state
1506 r <- liftIO $ (A.unsafeIndexIO patArr 0)
1507 return $ Skip $ GO_SINGLE_PAT state r
1508 else if sizeOf (undefined :: a) * patLen
1509 <= sizeOf (undefined :: Word)
1510 then return $ Skip $ GO_SHORT_PAT state
1512 (rb, rhead) <- liftIO $ RB.new patLen
1513 return $ Skip $ GO_KARP_RABIN state rb rhead
1515 stepOuter gst (GO_SINGLE_PAT stt pat) = initial >>= go SPEC stt
1520 res <- step (adaptState gst) st
1526 return $ Yield r (GO_SINGLE_PAT s pat)
1527 else fstep acc x >>= go SPEC s
1528 Skip s -> go SPEC s acc
1529 Stop -> done acc >>= \r -> return $ Yield r GO_DONE
1531 stepOuter gst (GO_SHORT_PAT stt) = initial >>= go0 SPEC 0 (0 :: Word) stt
1536 mask = (1 `shiftL` (elemBits * patLen)) - 1
1538 addToWord wrd a = (wrd `shiftL` elemBits) .|. fromIntegral (fromEnum a)
1541 patWord = mask .&. A.foldl' addToWord 0 patArr
1543 go0 !_ !idx wrd st !acc = do
1544 res <- step (adaptState gst) st
1547 let wrd' = addToWord wrd x
1550 if wrd' .&. mask == patWord
1553 return $ Yield r (GO_SHORT_PAT s)
1554 else go1 SPEC wrd' s acc
1555 else go0 SPEC (idx + 1) wrd' s acc
1556 Skip s -> go0 SPEC idx wrd s acc
1559 then go2 wrd idx acc
1561 done acc' >>= \r -> return $ Yield r GO_DONE
1564 go1 !_ wrd st !acc = do
1565 res <- step (adaptState gst) st
1568 let wrd' = addToWord wrd x
1569 old = (mask .&. wrd) `shiftR` (elemBits * (patLen - 1))
1570 acc' <- fstep acc (toEnum $ fromIntegral old)
1571 if wrd' .&. mask == patWord
1572 then done acc' >>= \r -> return $ Yield r (GO_SHORT_PAT s)
1573 else go1 SPEC wrd' s acc'
1574 Skip s -> go1 SPEC wrd s acc
1576 acc' <- go2 wrd patLen acc
1577 done acc' >>= \r -> return $ Yield r GO_DONE
1579 go2 !wrd !n !acc | n > 0 = do
1580 let old = (mask .&. wrd) `shiftR` (elemBits * (n - 1))
1581 fstep acc (toEnum $ fromIntegral old) >>= go2 wrd (n - 1)
1582 go2 _ _ acc = return acc
1584 stepOuter gst (GO_KARP_RABIN stt rb rhead) = do
1585 initial >>= go0 SPEC 0 rhead stt
1589 k = 2891336453 :: Word32
1591 addCksum cksum a = cksum * k + fromIntegral (fromEnum a)
1592 deltaCksum cksum old new =
1593 addCksum cksum new - coeff * fromIntegral (fromEnum old)
1595 -- XXX shall we use a random starting hash or 1 instead of 0?
1596 patHash = A.foldl' addCksum 0 patArr
1599 go0 !_ !idx !rh st !acc = do
1600 res <- step (adaptState gst) st
1603 rh' <- liftIO $ RB.unsafeInsert rb rh x
1606 let fold = RB.unsafeFoldRing (RB.ringBound rb)
1607 let !ringHash = fold addCksum 0 rb
1608 if ringHash == patHash
1609 then go2 SPEC ringHash rh' s acc
1610 else go1 SPEC ringHash rh' s acc
1611 else go0 SPEC (idx + 1) rh' s acc
1612 Skip s -> go0 SPEC idx rh s acc
1614 !acc' <- if idx /= 0
1615 then RB.unsafeFoldRingM rh fstep acc rb
1617 done acc' >>= \r -> return $ Yield r GO_DONE
1619 -- XXX Theoretically this code can do 4 times faster if GHC generates
1620 -- optimal code. If we use just "(cksum' == patHash)" condition it goes
1621 -- 4x faster, as soon as we add the "RB.unsafeEqArray rb v" condition
1622 -- the generated code changes drastically and becomes 4x slower. Need
1623 -- to investigate what is going on with GHC.
1625 go1 !_ !cksum !rh st !acc = do
1626 res <- step (adaptState gst) st
1629 old <- liftIO $ peek rh
1630 let cksum' = deltaCksum cksum old x
1631 acc' <- fstep acc old
1633 if (cksum' == patHash)
1635 rh' <- liftIO (RB.unsafeInsert rb rh x)
1636 go2 SPEC cksum' rh' s acc'
1638 rh' <- liftIO (RB.unsafeInsert rb rh x)
1639 go1 SPEC cksum' rh' s acc'
1640 Skip s -> go1 SPEC cksum rh s acc
1642 acc' <- RB.unsafeFoldRingFullM rh fstep acc rb
1643 done acc' >>= \r -> return $ Yield r GO_DONE
1645 go2 !_ !cksum' !rh' s !acc' = do
1646 if RB.unsafeEqArray rb rh' patArr
1649 return $ Yield r (GO_KARP_RABIN s rb rhead)
1650 else go1 SPEC cksum' rh' s acc'
1652 stepOuter gst (GO_EMPTY_PAT st) = do
1653 res <- step (adaptState gst) st
1658 done acc' >>= \r -> return $ Yield r (GO_EMPTY_PAT s)
1659 Skip s -> return $ Skip (GO_EMPTY_PAT s)
1662 stepOuter _ GO_DONE = return Stop
1664{-# INLINE_NORMAL splitSuffixOn #-}
1666 :: forall m a b. (MonadIO m, Storable a, Enum a, Eq a)
1672splitSuffixOn withSep patArr@Array{..} (Fold fstep initial done)
1673 (Stream step state) =
1674 Stream stepOuter GO_START
1678 patLen = A.length patArr
1679 maxIndex = patLen - 1
1680 elemBits = sizeOf (undefined :: a) * 8
1682 {-# INLINE_LATE stepOuter #-}
1683 stepOuter _ GO_START =
1685 then return $ Skip $ GO_EMPTY_PAT state
1688 r <- liftIO $ (A.unsafeIndexIO patArr 0)
1689 return $ Skip $ GO_SINGLE_PAT state r
1690 else if sizeOf (undefined :: a) * patLen
1691 <= sizeOf (undefined :: Word)
1692 then return $ Skip $ GO_SHORT_PAT state
1694 (rb, rhead) <- liftIO $ RB.new patLen
1695 return $ Skip $ GO_KARP_RABIN state rb rhead
1697 stepOuter gst (GO_SINGLE_PAT stt pat) = do
1698 -- This first part is the only difference between splitOn and
1700 -- If the last element is a separator do not issue a blank segment.
1701 res <- step (adaptState gst) stt
1707 acc' <- if withSep then fstep acc x else return acc
1708 done acc' >>= \r -> return $ Yield r (GO_SINGLE_PAT s pat)
1709 else fstep acc x >>= go SPEC s
1710 Skip s -> return $ Skip $ (GO_SINGLE_PAT s pat)
1715 -- This is identical for splitOn and splitSuffixOn
1717 res <- step (adaptState gst) st
1722 acc' <- if withSep then fstep acc x else return acc
1724 return $ Yield r (GO_SINGLE_PAT s pat)
1725 else fstep acc x >>= go SPEC s
1726 Skip s -> go SPEC s acc
1727 Stop -> done acc >>= \r -> return $ Yield r GO_DONE
1729 stepOuter gst (GO_SHORT_PAT stt) = do
1731 -- Call "initial" only if the stream yields an element, otherwise we
1732 -- may call "initial" but never yield anything. initial may produce a
1733 -- side effect, therefore we will end up doing and discard a side
1738 res <- step (adaptState gst) stt
1742 let wrd' = addToWord wrd x
1743 acc' <- if withSep then fstep acc x else return acc
1746 if wrd' .&. mask == patWord
1747 then done acc' >>= \r -> return $ Yield r (GO_SHORT_PAT s)
1748 else go0 SPEC (idx + 1) wrd' s acc'
1749 else go0 SPEC (idx + 1) wrd' s acc'
1750 Skip s -> return $ Skip (GO_SHORT_PAT s)
1756 mask = (1 `shiftL` (elemBits * patLen)) - 1
1758 addToWord wrd a = (wrd `shiftL` elemBits) .|. fromIntegral (fromEnum a)
1761 patWord = mask .&. A.foldl' addToWord 0 patArr
1763 go0 !_ !idx wrd st !acc = do
1764 res <- step (adaptState gst) st
1767 let wrd' = addToWord wrd x
1768 acc' <- if withSep then fstep acc x else return acc
1771 if wrd' .&. mask == patWord
1774 return $ Yield r (GO_SHORT_PAT s)
1775 else go1 SPEC wrd' s acc'
1776 else go0 SPEC (idx + 1) wrd' s acc'
1777 Skip s -> go0 SPEC idx wrd s acc
1779 if (idx == maxIndex) && (wrd .&. mask == patWord)
1782 acc' <- if idx /= 0 && not withSep
1783 then go2 wrd idx acc
1785 done acc' >>= \r -> return $ Yield r GO_DONE
1788 go1 !_ wrd st !acc = do
1789 res <- step (adaptState gst) st
1792 let wrd' = addToWord wrd x
1793 old = (mask .&. wrd) `shiftR` (elemBits * (patLen - 1))
1796 else fstep acc (toEnum $ fromIntegral old)
1797 if wrd' .&. mask == patWord
1798 then done acc' >>= \r -> return $ Yield r (GO_SHORT_PAT s)
1799 else go1 SPEC wrd' s acc'
1800 Skip s -> go1 SPEC wrd s acc
1802 -- If the last sequence is a separator do not issue a blank
1804 if wrd .&. mask == patWord
1809 else go2 wrd patLen acc
1810 done acc' >>= \r -> return $ Yield r GO_DONE
1812 go2 !wrd !n !acc | n > 0 = do
1813 let old = (mask .&. wrd) `shiftR` (elemBits * (n - 1))
1814 fstep acc (toEnum $ fromIntegral old) >>= go2 wrd (n - 1)
1815 go2 _ _ acc = return acc
1817 stepOuter gst (GO_KARP_RABIN stt rb rhead) = do
1819 res <- step (adaptState gst) stt
1823 acc' <- if withSep then fstep acc x else return acc
1824 rh' <- liftIO (RB.unsafeInsert rb rhead x)
1827 let fold = RB.unsafeFoldRing (RB.ringBound rb)
1828 let !ringHash = fold addCksum 0 rb
1829 if ringHash == patHash
1830 then go2 SPEC ringHash rh' s acc'
1831 else go0 SPEC (idx + 1) rh' s acc'
1832 else go0 SPEC (idx + 1) rh' s acc'
1833 Skip s -> return $ Skip (GO_KARP_RABIN s rb rhead)
1838 k = 2891336453 :: Word32
1840 addCksum cksum a = cksum * k + fromIntegral (fromEnum a)
1841 deltaCksum cksum old new =
1842 addCksum cksum new - coeff * fromIntegral (fromEnum old)
1844 -- XXX shall we use a random starting hash or 1 instead of 0?
1845 patHash = A.foldl' addCksum 0 patArr
1848 go0 !_ !idx !rh st !acc = do
1849 res <- step (adaptState gst) st
1852 acc' <- if withSep then fstep acc x else return acc
1853 rh' <- liftIO (RB.unsafeInsert rb rh x)
1856 let fold = RB.unsafeFoldRing (RB.ringBound rb)
1857 let !ringHash = fold addCksum 0 rb
1858 if ringHash == patHash
1859 then go2 SPEC ringHash rh' s acc'
1860 else go1 SPEC ringHash rh' s acc'
1861 else go0 SPEC (idx + 1) rh' s acc'
1862 Skip s -> go0 SPEC idx rh s acc
1864 -- do not issue a blank segment when we end at pattern
1865 if (idx == maxIndex) && RB.unsafeEqArray rb rh patArr
1868 !acc' <- if idx /= 0 && not withSep
1869 then RB.unsafeFoldRingM rh fstep acc rb
1871 done acc' >>= \r -> return $ Yield r GO_DONE
1873 -- XXX Theoretically this code can do 4 times faster if GHC generates
1874 -- optimal code. If we use just "(cksum' == patHash)" condition it goes
1875 -- 4x faster, as soon as we add the "RB.unsafeEqArray rb v" condition
1876 -- the generated code changes drastically and becomes 4x slower. Need
1877 -- to investigate what is going on with GHC.
1879 go1 !_ !cksum !rh st !acc = do
1880 res <- step (adaptState gst) st
1883 old <- liftIO $ peek rh
1884 let cksum' = deltaCksum cksum old x
1889 if (cksum' == patHash)
1891 rh' <- liftIO (RB.unsafeInsert rb rh x)
1892 go2 SPEC cksum' rh' s acc'
1894 rh' <- liftIO (RB.unsafeInsert rb rh x)
1895 go1 SPEC cksum' rh' s acc'
1896 Skip s -> go1 SPEC cksum rh s acc
1898 if RB.unsafeEqArray rb rh patArr
1903 else RB.unsafeFoldRingFullM rh fstep acc rb
1904 done acc' >>= \r -> return $ Yield r GO_DONE
1906 go2 !_ !cksum' !rh' s !acc' = do
1907 if RB.unsafeEqArray rb rh' patArr
1910 return $ Yield r (GO_KARP_RABIN s rb rhead)
1911 else go1 SPEC cksum' rh' s acc'
1913 stepOuter gst (GO_EMPTY_PAT st) = do
1914 res <- step (adaptState gst) st
1919 done acc' >>= \r -> return $ Yield r (GO_EMPTY_PAT s)
1920 Skip s -> return $ Skip (GO_EMPTY_PAT s)
1923 stepOuter _ GO_DONE = return Stop
1925data SplitState s arr
1927 | SplitBuffering s arr
1928 | SplitSplitting s arr
1929 | SplitYielding arr (SplitState s arr)
1932-- XXX An alternative approach would be to use a partial fold (Fold m a b) to
1933-- split using a splitBy like combinator. The Fold would consume upto the
1934-- separator and return any leftover which can then be fed to the next fold.
1936-- We can revisit this once we have partial folds/parsers.
1938-- | Performs infix separator style splitting.
1939{-# INLINE_NORMAL splitInnerBy #-}
1942 => (f a -> m (f a, Maybe (f a))) -- splitter
1943 -> (f a -> f a -> m (f a)) -- joiner
1946splitInnerBy splitter joiner (Stream step1 state1) =
1947 (Stream step (SplitInitial state1))
1951 {-# INLINE_LATE step #-}
1952 step gst (SplitInitial st) = do
1956 (x1, mx2) <- splitter x
1957 return $ case mx2 of
1958 Nothing -> Skip (SplitBuffering s x1)
1959 Just x2 -> Skip (SplitYielding x1 (SplitSplitting s x2))
1960 Skip s -> return $ Skip (SplitInitial s)
1961 Stop -> return $ Stop
1963 step gst (SplitBuffering st buf) = do
1967 (x1, mx2) <- splitter x
1968 buf' <- joiner buf x1
1969 return $ case mx2 of
1970 Nothing -> Skip (SplitBuffering s buf')
1971 Just x2 -> Skip (SplitYielding buf' (SplitSplitting s x2))
1972 Skip s -> return $ Skip (SplitBuffering s buf)
1973 Stop -> return $ Skip (SplitYielding buf SplitFinishing)
1975 step _ (SplitSplitting st buf) = do
1976 (x1, mx2) <- splitter buf
1977 return $ case mx2 of
1978 Nothing -> Skip $ SplitBuffering st x1
1979 Just x2 -> Skip $ SplitYielding x1 (SplitSplitting st x2)
1981 step _ (SplitYielding x next) = return $ Yield x next
1982 step _ SplitFinishing = return $ Stop
1984-- | Performs infix separator style splitting.
1985{-# INLINE_NORMAL splitInnerBySuffix #-}
1987 :: (Monad m, Eq (f a), Monoid (f a))
1988 => (f a -> m (f a, Maybe (f a))) -- splitter
1989 -> (f a -> f a -> m (f a)) -- joiner
1992splitInnerBySuffix splitter joiner (Stream step1 state1) =
1993 (Stream step (SplitInitial state1))
1997 {-# INLINE_LATE step #-}
1998 step gst (SplitInitial st) = do
2002 (x1, mx2) <- splitter x
2003 return $ case mx2 of
2004 Nothing -> Skip (SplitBuffering s x1)
2005 Just x2 -> Skip (SplitYielding x1 (SplitSplitting s x2))
2006 Skip s -> return $ Skip (SplitInitial s)
2007 Stop -> return $ Stop
2009 step gst (SplitBuffering st buf) = do
2013 (x1, mx2) <- splitter x
2014 buf' <- joiner buf x1
2015 return $ case mx2 of
2016 Nothing -> Skip (SplitBuffering s buf')
2017 Just x2 -> Skip (SplitYielding buf' (SplitSplitting s x2))
2018 Skip s -> return $ Skip (SplitBuffering s buf)
2022 else Skip (SplitYielding buf SplitFinishing)
2024 step _ (SplitSplitting st buf) = do
2025 (x1, mx2) <- splitter buf
2026 return $ case mx2 of
2027 Nothing -> Skip $ SplitBuffering st x1
2028 Just x2 -> Skip $ SplitYielding x1 (SplitSplitting st x2)
2030 step _ (SplitYielding x next) = return $ Yield x next
2031 step _ SplitFinishing = return $ Stop
2033------------------------------------------------------------------------------
2035------------------------------------------------------------------------------
2037{-# INLINE_NORMAL isPrefixOf #-}
2038isPrefixOf :: (Eq a, Monad m) => Stream m a -> Stream m a -> m Bool
2039isPrefixOf (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
2041 go (sa, sb, Nothing) = do
2042 r <- stepa defState sa
2044 Yield x sa' -> go (sa', sb, Just x)
2045 Skip sa' -> go (sa', sb, Nothing)
2048 go (sa, sb, Just x) = do
2049 r <- stepb defState sb
2053 then go (sa, sb', Nothing)
2055 Skip sb' -> go (sa, sb', Just x)
2056 Stop -> return False
2058{-# INLINE_NORMAL isSubsequenceOf #-}
2059isSubsequenceOf :: (Eq a, Monad m) => Stream m a -> Stream m a -> m Bool
2060isSubsequenceOf (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
2062 go (sa, sb, Nothing) = do
2063 r <- stepa defState sa
2065 Yield x sa' -> go (sa', sb, Just x)
2066 Skip sa' -> go (sa', sb, Nothing)
2069 go (sa, sb, Just x) = do
2070 r <- stepb defState sb
2074 then go (sa, sb', Nothing)
2075 else go (sa, sb', Just x)
2076 Skip sb' -> go (sa, sb', Just x)
2077 Stop -> return False
2079{-# INLINE_NORMAL stripPrefix #-}
2082 => Stream m a -> Stream m a -> m (Maybe (Stream m a))
2083stripPrefix (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
2085 go (sa, sb, Nothing) = do
2086 r <- stepa defState sa
2088 Yield x sa' -> go (sa', sb, Just x)
2089 Skip sa' -> go (sa', sb, Nothing)
2090 Stop -> return $ Just (Stream stepb sb)
2092 go (sa, sb, Just x) = do
2093 r <- stepb defState sb
2097 then go (sa, sb', Nothing)
2099 Skip sb' -> go (sa, sb', Just x)
2100 Stop -> return Nothing
2102------------------------------------------------------------------------------
2104------------------------------------------------------------------------------
2106-- | Execute a monadic action for each element of the 'Stream'
2107{-# INLINE_NORMAL mapM_ #-}
2108mapM_ :: Monad m => (a -> m b) -> Stream m a -> m ()
2109mapM_ m = drain . mapM m
2111-------------------------------------------------------------------------------
2112-- Stream transformations using Unfolds
2113-------------------------------------------------------------------------------
2115-- Define a unique structure to use in inspection testing
2116data ConcatMapUState o i =
2118 | ConcatMapUInner o i
2120-- | @concatMapU unfold stream@ uses @unfold@ to map the input stream elements
2121-- to streams and then flattens the generated streams into a single output
2124-- This is like 'concatMap' but uses an unfold with an explicit state to
2125-- generate the stream instead of a 'Stream' type generator. This allows better
2126-- optimization via fusion. This can be many times more efficient than
2129{-# INLINE_NORMAL concatMapU #-}
2130concatMapU :: Monad m => Unfold m a b -> Stream m a -> Stream m b
2131concatMapU (Unfold istep inject) (Stream ostep ost) =
2132 Stream step (ConcatMapUOuter ost)
2134 {-# INLINE_LATE step #-}
2135 step gst (ConcatMapUOuter o) = do
2136 r <- ostep (adaptState gst) o
2140 i `seq` return (Skip (ConcatMapUInner o' i))
2141 Skip o' -> return $ Skip (ConcatMapUOuter o')
2142 Stop -> return $ Stop
2144 step _ (ConcatMapUInner o i) = do
2147 Yield x i' -> Yield x (ConcatMapUInner o i')
2148 Skip i' -> Skip (ConcatMapUInner o i')
2149 Stop -> Skip (ConcatMapUOuter o)
2151data ConcatUnfoldInterleaveState o i =
2152 ConcatUnfoldInterleaveOuter o [i]
2153 | ConcatUnfoldInterleaveInner o [i]
2154 | ConcatUnfoldInterleaveInnerL [i] [i]
2155 | ConcatUnfoldInterleaveInnerR [i] [i]
2157-- XXX use arrays to store state instead of lists.
2158-- XXX In general we can use different scheduling strategies e.g. how to
2159-- schedule the outer vs inner loop or assigning weights to different streams
2160-- or outer and inner loops.
2162-- After a yield, switch to the next stream. Do not switch streams on Skip.
2163-- Yield from outer stream switches to the inner stream.
2165-- There are two choices here, (1) exhaust the outer stream first and then
2166-- start yielding from the inner streams, this is much simpler to implement,
2167-- (2) yield at least one element from an inner stream before going back to
2168-- outer stream and opening the next stream from it.
2170-- Ideally, we need some scheduling bias to inner streams vs outer stream.
2171-- Maybe we can configure the behavior.
2173{-# INLINE_NORMAL concatUnfoldInterleave #-}
2174concatUnfoldInterleave :: Monad m => Unfold m a b -> Stream m a -> Stream m b
2175concatUnfoldInterleave (Unfold istep inject) (Stream ostep ost) =
2176 Stream step (ConcatUnfoldInterleaveOuter ost [])
2178 {-# INLINE_LATE step #-}
2179 step gst (ConcatUnfoldInterleaveOuter o ls) = do
2180 r <- ostep (adaptState gst) o
2184 i `seq` return (Skip (ConcatUnfoldInterleaveInner o' (i : ls)))
2185 Skip o' -> return $ Skip (ConcatUnfoldInterleaveOuter o' ls)
2186 Stop -> return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2188 step _ (ConcatUnfoldInterleaveInner _ []) = undefined
2189 step _ (ConcatUnfoldInterleaveInner o (st:ls)) = do
2192 Yield x s -> Yield x (ConcatUnfoldInterleaveOuter o (s:ls))
2193 Skip s -> Skip (ConcatUnfoldInterleaveInner o (s:ls))
2194 Stop -> Skip (ConcatUnfoldInterleaveOuter o ls)
2196 step _ (ConcatUnfoldInterleaveInnerL [] []) = return Stop
2197 step _ (ConcatUnfoldInterleaveInnerL [] rs) =
2198 return $ Skip (ConcatUnfoldInterleaveInnerR [] rs)
2200 step _ (ConcatUnfoldInterleaveInnerL (st:ls) rs) = do
2203 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerL ls (s:rs))
2204 Skip s -> Skip (ConcatUnfoldInterleaveInnerL (s:ls) rs)
2205 Stop -> Skip (ConcatUnfoldInterleaveInnerL ls rs)
2207 step _ (ConcatUnfoldInterleaveInnerR [] []) = return Stop
2208 step _ (ConcatUnfoldInterleaveInnerR ls []) =
2209 return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2211 step _ (ConcatUnfoldInterleaveInnerR ls (st:rs)) = do
2214 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerR (s:ls) rs)
2215 Skip s -> Skip (ConcatUnfoldInterleaveInnerR ls (s:rs))
2216 Stop -> Skip (ConcatUnfoldInterleaveInnerR ls rs)
2218-- XXX In general we can use different scheduling strategies e.g. how to
2219-- schedule the outer vs inner loop or assigning weights to different streams
2220-- or outer and inner loops.
2222-- This could be inefficient if the tasks are too small.
2224-- Compared to concatUnfoldInterleave this one switches streams on Skips.
2226{-# INLINE_NORMAL concatUnfoldRoundrobin #-}
2227concatUnfoldRoundrobin :: Monad m => Unfold m a b -> Stream m a -> Stream m b
2228concatUnfoldRoundrobin (Unfold istep inject) (Stream ostep ost) =
2229 Stream step (ConcatUnfoldInterleaveOuter ost [])
2231 {-# INLINE_LATE step #-}
2232 step gst (ConcatUnfoldInterleaveOuter o ls) = do
2233 r <- ostep (adaptState gst) o
2237 i `seq` return (Skip (ConcatUnfoldInterleaveInner o' (i : ls)))
2238 Skip o' -> return $ Skip (ConcatUnfoldInterleaveInner o' ls)
2239 Stop -> return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2241 step _ (ConcatUnfoldInterleaveInner o []) =
2242 return $ Skip (ConcatUnfoldInterleaveOuter o [])
2244 step _ (ConcatUnfoldInterleaveInner o (st:ls)) = do
2247 Yield x s -> Yield x (ConcatUnfoldInterleaveOuter o (s:ls))
2248 Skip s -> Skip (ConcatUnfoldInterleaveOuter o (s:ls))
2249 Stop -> Skip (ConcatUnfoldInterleaveOuter o ls)
2251 step _ (ConcatUnfoldInterleaveInnerL [] []) = return Stop
2252 step _ (ConcatUnfoldInterleaveInnerL [] rs) =
2253 return $ Skip (ConcatUnfoldInterleaveInnerR [] rs)
2255 step _ (ConcatUnfoldInterleaveInnerL (st:ls) rs) = do
2258 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerL ls (s:rs))
2259 Skip s -> Skip (ConcatUnfoldInterleaveInnerL ls (s:rs))
2260 Stop -> Skip (ConcatUnfoldInterleaveInnerL ls rs)
2262 step _ (ConcatUnfoldInterleaveInnerR [] []) = return Stop
2263 step _ (ConcatUnfoldInterleaveInnerR ls []) =
2264 return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2266 step _ (ConcatUnfoldInterleaveInnerR ls (st:rs)) = do
2269 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerR (s:ls) rs)
2270 Skip s -> Skip (ConcatUnfoldInterleaveInnerR (s:ls) rs)
2271 Stop -> Skip (ConcatUnfoldInterleaveInnerR ls rs)
2273data AppendState s1 s2 = AppendFirst s1 | AppendSecond s2
2275-- Note that this could be much faster compared to the CPS stream. However, as
2276-- the number of streams being composed increases this may become expensive.
2277-- Need to see where the breaking point is between the two.
2279{-# INLINE_NORMAL append #-}
2280append :: Monad m => Stream m a -> Stream m a -> Stream m a
2281append (Stream step1 state1) (Stream step2 state2) =
2282 Stream step (AppendFirst state1)
2286 {-# INLINE_LATE step #-}
2287 step gst (AppendFirst st) = do
2290 Yield a s -> Yield a (AppendFirst s)
2291 Skip s -> Skip (AppendFirst s)
2292 Stop -> Skip (AppendSecond state2)
2294 step gst (AppendSecond st) = do
2297 Yield a s -> Yield a (AppendSecond s)
2298 Skip s -> Skip (AppendSecond s)
2301data InterleaveState s1 s2 = InterleaveFirst s1 s2 | InterleaveSecond s1 s2
2302 | InterleaveSecondOnly s2 | InterleaveFirstOnly s1
2304{-# INLINE_NORMAL interleave #-}
2305interleave :: Monad m => Stream m a -> Stream m a -> Stream m a
2306interleave (Stream step1 state1) (Stream step2 state2) =
2307 Stream step (InterleaveFirst state1 state2)
2311 {-# INLINE_LATE step #-}
2312 step gst (InterleaveFirst st1 st2) = do
2315 Yield a s -> Yield a (InterleaveSecond s st2)
2316 Skip s -> Skip (InterleaveFirst s st2)
2317 Stop -> Skip (InterleaveSecondOnly st2)
2319 step gst (InterleaveSecond st1 st2) = do
2322 Yield a s -> Yield a (InterleaveFirst st1 s)
2323 Skip s -> Skip (InterleaveSecond st1 s)
2324 Stop -> Skip (InterleaveFirstOnly st1)
2326 step gst (InterleaveFirstOnly st1) = do
2329 Yield a s -> Yield a (InterleaveFirstOnly s)
2330 Skip s -> Skip (InterleaveFirstOnly s)
2333 step gst (InterleaveSecondOnly st2) = do
2336 Yield a s -> Yield a (InterleaveSecondOnly s)
2337 Skip s -> Skip (InterleaveSecondOnly s)
2340{-# INLINE_NORMAL interleaveMin #-}
2341interleaveMin :: Monad m => Stream m a -> Stream m a -> Stream m a
2342interleaveMin (Stream step1 state1) (Stream step2 state2) =
2343 Stream step (InterleaveFirst state1 state2)
2347 {-# INLINE_LATE step #-}
2348 step gst (InterleaveFirst st1 st2) = do
2351 Yield a s -> Yield a (InterleaveSecond s st2)
2352 Skip s -> Skip (InterleaveFirst s st2)
2355 step gst (InterleaveSecond st1 st2) = do
2358 Yield a s -> Yield a (InterleaveFirst st1 s)
2359 Skip s -> Skip (InterleaveSecond st1 s)
2362 step _ (InterleaveFirstOnly _) = undefined
2363 step _ (InterleaveSecondOnly _) = undefined
2365{-# INLINE_NORMAL interleaveSuffix #-}
2366interleaveSuffix :: Monad m => Stream m a -> Stream m a -> Stream m a
2367interleaveSuffix (Stream step1 state1) (Stream step2 state2) =
2368 Stream step (InterleaveFirst state1 state2)
2372 {-# INLINE_LATE step #-}
2373 step gst (InterleaveFirst st1 st2) = do
2376 Yield a s -> Yield a (InterleaveSecond s st2)
2377 Skip s -> Skip (InterleaveFirst s st2)
2380 step gst (InterleaveSecond st1 st2) = do
2383 Yield a s -> Yield a (InterleaveFirst st1 s)
2384 Skip s -> Skip (InterleaveSecond st1 s)
2385 Stop -> Skip (InterleaveFirstOnly st1)
2387 step gst (InterleaveFirstOnly st1) = do
2390 Yield a s -> Yield a (InterleaveFirstOnly s)
2391 Skip s -> Skip (InterleaveFirstOnly s)
2394 step _ (InterleaveSecondOnly _) = undefined
2396data InterleaveInfixState s1 s2 a
2397 = InterleaveInfixFirst s1 s2
2398 | InterleaveInfixSecondBuf s1 s2
2399 | InterleaveInfixSecondYield s1 s2 a
2400 | InterleaveInfixFirstYield s1 s2 a
2401 | InterleaveInfixFirstOnly s1
2403{-# INLINE_NORMAL interleaveInfix #-}
2404interleaveInfix :: Monad m => Stream m a -> Stream m a -> Stream m a
2405interleaveInfix (Stream step1 state1) (Stream step2 state2) =
2406 Stream step (InterleaveInfixFirst state1 state2)
2410 {-# INLINE_LATE step #-}
2411 step gst (InterleaveInfixFirst st1 st2) = do
2414 Yield a s -> Yield a (InterleaveInfixSecondBuf s st2)
2415 Skip s -> Skip (InterleaveInfixFirst s st2)
2418 step gst (InterleaveInfixSecondBuf st1 st2) = do
2421 Yield a s -> Skip (InterleaveInfixSecondYield st1 s a)
2422 Skip s -> Skip (InterleaveInfixSecondBuf st1 s)
2423 Stop -> Skip (InterleaveInfixFirstOnly st1)
2425 step gst (InterleaveInfixSecondYield st1 st2 x) = do
2428 Yield a s -> Yield x (InterleaveInfixFirstYield s st2 a)
2429 Skip s -> Skip (InterleaveInfixSecondYield s st2 x)
2432 step _ (InterleaveInfixFirstYield st1 st2 x) = do
2433 return $ Yield x (InterleaveInfixSecondBuf st1 st2)
2435 step gst (InterleaveInfixFirstOnly st1) = do
2438 Yield a s -> Yield a (InterleaveInfixFirstOnly s)
2439 Skip s -> Skip (InterleaveInfixFirstOnly s)
2442{-# INLINE_NORMAL roundRobin #-}
2443roundRobin :: Monad m => Stream m a -> Stream m a -> Stream m a
2444roundRobin (Stream step1 state1) (Stream step2 state2) =
2445 Stream step (InterleaveFirst state1 state2)
2449 {-# INLINE_LATE step #-}
2450 step gst (InterleaveFirst st1 st2) = do
2453 Yield a s -> Yield a (InterleaveSecond s st2)
2454 Skip s -> Skip (InterleaveSecond s st2)
2455 Stop -> Skip (InterleaveSecondOnly st2)
2457 step gst (InterleaveSecond st1 st2) = do
2460 Yield a s -> Yield a (InterleaveFirst st1 s)
2461 Skip s -> Skip (InterleaveFirst st1 s)
2462 Stop -> Skip (InterleaveFirstOnly st1)
2464 step gst (InterleaveSecondOnly st2) = do
2467 Yield a s -> Yield a (InterleaveSecondOnly s)
2468 Skip s -> Skip (InterleaveSecondOnly s)
2471 step gst (InterleaveFirstOnly st1) = do
2474 Yield a s -> Yield a (InterleaveFirstOnly s)
2475 Skip s -> Skip (InterleaveFirstOnly s)
2478data ICUState s1 s2 i1 i2 =
2483 | ICUFirstInner s1 s2 i1
2484 | ICUSecondInner s1 s2 i2
2485 | ICUFirstOnlyInner s1 i1
2486 | ICUSecondOnlyInner s2 i2
2488-- | Interleave streams (full streams, not the elements) unfolded from two
2489-- input streams and concat. Stop when the first stream stops. If the second
2490-- stream ends before the first one then first stream still keeps running alone
2491-- without any interleaving with the second stream.
2493-- [a1, a2, ... an] [b1, b2 ...]
2494-- => [streamA1, streamA2, ... streamAn] [streamB1, streamB2, ...]
2495-- => [streamA1, streamB1, streamA2...StreamAn, streamBn]
2496-- => [a11, a12, ...a1j, b11, b12, ...b1k, a21, a22, ...]
2498{-# INLINE_NORMAL gintercalateSuffix #-}
2501 => Unfold m a c -> Stream m a -> Unfold m b c -> Stream m b -> Stream m c
2503 (Unfold istep1 inject1) (Stream step1 state1)
2504 (Unfold istep2 inject2) (Stream step2 state2) =
2505 Stream step (ICUFirst state1 state2)
2509 {-# INLINE_LATE step #-}
2510 step gst (ICUFirst s1 s2) = do
2511 r <- step1 (adaptState gst) s1
2515 i `seq` return (Skip (ICUFirstInner s s2 i))
2516 Skip s -> return $ Skip (ICUFirst s s2)
2519 step gst (ICUFirstOnly s1) = do
2520 r <- step1 (adaptState gst) s1
2524 i `seq` return (Skip (ICUFirstOnlyInner s i))
2525 Skip s -> return $ Skip (ICUFirstOnly s)
2528 step _ (ICUFirstInner s1 s2 i1) = do
2531 Yield x i' -> Yield x (ICUFirstInner s1 s2 i')
2532 Skip i' -> Skip (ICUFirstInner s1 s2 i')
2533 Stop -> Skip (ICUSecond s1 s2)
2535 step _ (ICUFirstOnlyInner s1 i1) = do
2538 Yield x i' -> Yield x (ICUFirstOnlyInner s1 i')
2539 Skip i' -> Skip (ICUFirstOnlyInner s1 i')
2540 Stop -> Skip (ICUFirstOnly s1)
2542 step gst (ICUSecond s1 s2) = do
2543 r <- step2 (adaptState gst) s2
2547 i `seq` return (Skip (ICUSecondInner s1 s i))
2548 Skip s -> return $ Skip (ICUSecond s1 s)
2549 Stop -> return $ Skip (ICUFirstOnly s1)
2551 step _ (ICUSecondInner s1 s2 i2) = do
2554 Yield x i' -> Yield x (ICUSecondInner s1 s2 i')
2555 Skip i' -> Skip (ICUSecondInner s1 s2 i')
2556 Stop -> Skip (ICUFirst s1 s2)
2558 step _ (ICUSecondOnly _s2) = undefined
2559 step _ (ICUSecondOnlyInner _s2 _i2) = undefined
2561data InterposeSuffixState s1 i1 =
2562 InterposeSuffixFirst s1
2563 -- | InterposeSuffixFirstYield s1 i1
2564 | InterposeSuffixFirstInner s1 i1
2565 | InterposeSuffixSecond s1
2567-- Note that if an unfolded layer turns out to be nil we still emit the
2568-- separator effect. An alternate behavior could be to emit the separator
2569-- effect only if at least one element has been yielded by the unfolding.
2570-- However, that becomes a bit complicated, so we have chosen the former
2572{-# INLINE_NORMAL interposeSuffix #-}
2575 => m c -> Unfold m b c -> Stream m b -> Stream m c
2578 (Unfold istep1 inject1) (Stream step1 state1) =
2579 Stream step (InterposeSuffixFirst state1)
2583 {-# INLINE_LATE step #-}
2584 step gst (InterposeSuffixFirst s1) = do
2585 r <- step1 (adaptState gst) s1
2589 i `seq` return (Skip (InterposeSuffixFirstInner s i))
2590 -- i `seq` return (Skip (InterposeSuffixFirstYield s i))
2591 Skip s -> return $ Skip (InterposeSuffixFirst s)
2595 step _ (InterposeSuffixFirstYield s1 i1) = do
2598 Yield x i' -> Yield x (InterposeSuffixFirstInner s1 i')
2599 Skip i' -> Skip (InterposeSuffixFirstYield s1 i')
2600 Stop -> Skip (InterposeSuffixFirst s1)
2603 step _ (InterposeSuffixFirstInner s1 i1) = do
2606 Yield x i' -> Yield x (InterposeSuffixFirstInner s1 i')
2607 Skip i' -> Skip (InterposeSuffixFirstInner s1 i')
2608 Stop -> Skip (InterposeSuffixSecond s1)
2610 step _ (InterposeSuffixSecond s1) = do
2612 return $ Yield r (InterposeSuffixFirst s1)
2614data ICALState s1 s2 i1 i2 a =
2616 -- | ICALFirstYield s1 s2 i1
2617 | ICALFirstInner s1 s2 i1
2619 | ICALFirstOnlyInner s1 i1
2620 | ICALSecondInject s1 s2
2621 | ICALFirstInject s1 s2 i2
2622 -- | ICALFirstBuf s1 s2 i1 i2
2623 | ICALSecondInner s1 s2 i1 i2
2624 -- -- | ICALSecondInner s1 s2 i1 i2 a
2625 -- -- | ICALFirstResume s1 s2 i1 i2 a
2627-- | Interleave streams (full streams, not the elements) unfolded from two
2628-- input streams and concat. Stop when the first stream stops. If the second
2629-- stream ends before the first one then first stream still keeps running alone
2630-- without any interleaving with the second stream.
2632-- [a1, a2, ... an] [b1, b2 ...]
2633-- => [streamA1, streamA2, ... streamAn] [streamB1, streamB2, ...]
2634-- => [streamA1, streamB1, streamA2...StreamAn, streamBn]
2635-- => [a11, a12, ...a1j, b11, b12, ...b1k, a21, a22, ...]
2637{-# INLINE_NORMAL gintercalate #-}
2640 => Unfold m a c -> Stream m a -> Unfold m b c -> Stream m b -> Stream m c
2642 (Unfold istep1 inject1) (Stream step1 state1)
2643 (Unfold istep2 inject2) (Stream step2 state2) =
2644 Stream step (ICALFirst state1 state2)
2648 {-# INLINE_LATE step #-}
2649 step gst (ICALFirst s1 s2) = do
2650 r <- step1 (adaptState gst) s1
2654 i `seq` return (Skip (ICALFirstInner s s2 i))
2655 -- i `seq` return (Skip (ICALFirstYield s s2 i))
2656 Skip s -> return $ Skip (ICALFirst s s2)
2660 step _ (ICALFirstYield s1 s2 i1) = do
2663 Yield x i' -> Yield x (ICALFirstInner s1 s2 i')
2664 Skip i' -> Skip (ICALFirstYield s1 s2 i')
2665 Stop -> Skip (ICALFirst s1 s2)
2668 step _ (ICALFirstInner s1 s2 i1) = do
2671 Yield x i' -> Yield x (ICALFirstInner s1 s2 i')
2672 Skip i' -> Skip (ICALFirstInner s1 s2 i')
2673 Stop -> Skip (ICALSecondInject s1 s2)
2675 step gst (ICALFirstOnly s1) = do
2676 r <- step1 (adaptState gst) s1
2680 i `seq` return (Skip (ICALFirstOnlyInner s i))
2681 Skip s -> return $ Skip (ICALFirstOnly s)
2684 step _ (ICALFirstOnlyInner s1 i1) = do
2687 Yield x i' -> Yield x (ICALFirstOnlyInner s1 i')
2688 Skip i' -> Skip (ICALFirstOnlyInner s1 i')
2689 Stop -> Skip (ICALFirstOnly s1)
2691 -- We inject the second stream even before checking if the first stream
2692 -- would yield any more elements. There is no clear choice whether we
2693 -- should do this before or after that. Doing it after may make the state
2694 -- machine a bit simpler though.
2695 step gst (ICALSecondInject s1 s2) = do
2696 r <- step2 (adaptState gst) s2
2700 i `seq` return (Skip (ICALFirstInject s1 s i))
2701 Skip s -> return $ Skip (ICALSecondInject s1 s)
2702 Stop -> return $ Skip (ICALFirstOnly s1)
2704 step gst (ICALFirstInject s1 s2 i2) = do
2705 r <- step1 (adaptState gst) s1
2709 i `seq` return (Skip (ICALSecondInner s s2 i i2))
2710 -- i `seq` return (Skip (ICALFirstBuf s s2 i i2))
2711 Skip s -> return $ Skip (ICALFirstInject s s2 i2)
2715 step _ (ICALFirstBuf s1 s2 i1 i2) = do
2718 Yield x i' -> Skip (ICALSecondInner s1 s2 i' i2 x)
2719 Skip i' -> Skip (ICALFirstBuf s1 s2 i' i2)
2722 step _ (ICALSecondInner s1 s2 i1 i2 v) = do
2725 Yield x i' -> Yield x (ICALSecondInner s1 s2 i1 i' v)
2726 Skip i' -> Skip (ICALSecondInner s1 s2 i1 i' v)
2727 Stop -> Skip (ICALFirstResume s1 s2 i1 i2 v)
2730 step _ (ICALSecondInner s1 s2 i1 i2) = do
2733 Yield x i' -> Yield x (ICALSecondInner s1 s2 i1 i')
2734 Skip i' -> Skip (ICALSecondInner s1 s2 i1 i')
2735 Stop -> Skip (ICALFirstInner s1 s2 i1)
2736 -- Stop -> Skip (ICALFirstResume s1 s2 i1 i2)
2739 step _ (ICALFirstResume s1 s2 i1 i2 x) = do
2740 return $ Yield x (ICALFirstInner s1 s2 i1 i2)
2743data InterposeState s1 i1 a =
2745 -- | InterposeFirstYield s1 i1
2746 | InterposeFirstInner s1 i1
2747 | InterposeFirstInject s1
2748 -- | InterposeFirstBuf s1 i1
2749 | InterposeSecondYield s1 i1
2750 -- -- | InterposeSecondYield s1 i1 a
2751 -- -- | InterposeFirstResume s1 i1 a
2753-- Note that this only interposes the pure values, we may run many effects to
2754-- generate those values as some effects may not generate anything (Skip).
2755{-# INLINE_NORMAL interpose #-}
2756interpose :: Monad m => m c -> Unfold m b c -> Stream m b -> Stream m c
2759 (Unfold istep1 inject1) (Stream step1 state1) =
2760 Stream step (InterposeFirst state1)
2764 {-# INLINE_LATE step #-}
2765 step gst (InterposeFirst s1) = do
2766 r <- step1 (adaptState gst) s1
2770 i `seq` return (Skip (InterposeFirstInner s i))
2771 -- i `seq` return (Skip (InterposeFirstYield s i))
2772 Skip s -> return $ Skip (InterposeFirst s)
2776 step _ (InterposeFirstYield s1 i1) = do
2779 Yield x i' -> Yield x (InterposeFirstInner s1 i')
2780 Skip i' -> Skip (InterposeFirstYield s1 i')
2781 Stop -> Skip (InterposeFirst s1)
2784 step _ (InterposeFirstInner s1 i1) = do
2787 Yield x i' -> Yield x (InterposeFirstInner s1 i')
2788 Skip i' -> Skip (InterposeFirstInner s1 i')
2789 Stop -> Skip (InterposeFirstInject s1)
2791 step gst (InterposeFirstInject s1) = do
2792 r <- step1 (adaptState gst) s1
2796 -- i `seq` return (Skip (InterposeFirstBuf s i))
2797 i `seq` return (Skip (InterposeSecondYield s i))
2798 Skip s -> return $ Skip (InterposeFirstInject s)
2802 step _ (InterposeFirstBuf s1 i1) = do
2805 Yield x i' -> Skip (InterposeSecondYield s1 i' x)
2806 Skip i' -> Skip (InterposeFirstBuf s1 i')
2811 step _ (InterposeSecondYield s1 i1 v) = do
2813 return $ Yield r (InterposeFirstResume s1 i1 v)
2815 step _ (InterposeSecondYield s1 i1) = do
2817 return $ Yield r (InterposeFirstInner s1 i1)
2820 step _ (InterposeFirstResume s1 i1 v) = do
2821 return $ Yield v (InterposeFirstInner s1 i1)
2824------------------------------------------------------------------------------
2826------------------------------------------------------------------------------
2828data GbracketState s1 s2 v
2830 | GBracketNormal s1 v
2831 | GBracketException s2
2833-- | The most general bracketing and exception combinator. All other
2834-- combinators can be expressed in terms of this combinator. This can also be
2835-- used for cases which are not covered by the standard combinators.
2839{-# INLINE_NORMAL gbracket #-}
2843 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
2844 -> (c -> m d) -- ^ after, on normal stop
2845 -> (c -> e -> Stream m b) -- ^ on exception
2846 -> (c -> Stream m b) -- ^ stream generator
2848gbracket bef exc aft fexc fnormal =
2849 Stream step GBracketInit
2853 {-# INLINE_LATE step #-}
2854 step _ GBracketInit = do
2856 return $ Skip $ GBracketNormal (fnormal r) r
2858 step gst (GBracketNormal (UnStream step1 st) v) = do
2859 res <- exc $ step1 gst st
2861 Right r -> case r of
2863 return $ Yield x (GBracketNormal (Stream step1 s) v)
2864 Skip s -> return $ Skip (GBracketNormal (Stream step1 s) v)
2865 Stop -> aft v >> return Stop
2866 Left e -> return $ Skip (GBracketException (fexc v e))
2867 step gst (GBracketException (UnStream step1 st)) = do
2870 Yield x s -> return $ Yield x (GBracketException (Stream step1 s))
2871 Skip s -> return $ Skip (GBracketException (Stream step1 s))
2874-- | Run a side effect before the stream yields its first element.
2875{-# INLINE_NORMAL before #-}
2876before :: Monad m => m b -> Stream m a -> Stream m a
2877before action (Stream step state) = Stream step' Nothing
2881 {-# INLINE_LATE step' #-}
2882 step' _ Nothing = action >> return (Skip (Just state))
2884 step' gst (Just st) = do
2887 Yield x s -> return $ Yield x (Just s)
2888 Skip s -> return $ Skip (Just s)
2891-- | Run a side effect whenever the stream stops normally.
2892{-# INLINE_NORMAL after #-}
2893after :: Monad m => m b -> Stream m a -> Stream m a
2894after action (Stream step state) = Stream step' state
2898 {-# INLINE_LATE step' #-}
2902 Yield x s -> return $ Yield x s
2903 Skip s -> return $ Skip s
2904 Stop -> action >> return Stop
2906-- XXX These combinators are expensive due to the call to
2907-- onException/handle/try on each step. Therefore, when possible, they should
2908-- be called in an outer loop where we perform less iterations. For example, we
2909-- cannot call them on each iteration in a char stream, instead we can call
2910-- them when doing an IO on an array.
2912-- XXX For high performance error checks in busy streams we may need another
2913-- Error constructor in step.
2915-- | Run a side effect whenever the stream aborts due to an exception. The
2916-- exception is not caught, simply rethrown.
2917{-# INLINE_NORMAL onException #-}
2918onException :: MonadCatch m => m b -> Stream m a -> Stream m a
2919onException action str =
2920 gbracket (return ()) MC.try return
2921 (\_ (e :: MC.SomeException) -> nilM (action >> MC.throwM e))
2924{-# INLINE_NORMAL _onException #-}
2925_onException :: MonadCatch m => m b -> Stream m a -> Stream m a
2926_onException action (Stream step state) = Stream step' state
2930 {-# INLINE_LATE step' #-}
2932 res <- step gst st `MC.onException` action
2934 Yield x s -> return $ Yield x s
2935 Skip s -> return $ Skip s
2938-- XXX bracket is like concatMap, it generates a stream and then flattens it.
2939-- Like concatMap it has 10x worse performance compared to linear fused
2942-- | Run the first action before the stream starts and remember its output,
2943-- generate a stream using the output, run the second action providing the
2944-- remembered value as an argument whenever the stream ends normally or due to
2946{-# INLINE_NORMAL bracket #-}
2947bracket :: MonadCatch m => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
2948bracket bef aft bet =
2949 gbracket bef MC.try aft
2950 (\a (e :: SomeException) -> nilM (aft a >> MC.throwM e)) bet
2952data BracketState s v = BracketInit | BracketRun s v
2954{-# INLINE_NORMAL _bracket #-}
2955_bracket :: MonadCatch m => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
2956_bracket bef aft bet = Stream step' BracketInit
2960 {-# INLINE_LATE step' #-}
2961 step' _ BracketInit = bef >>= \x -> return (Skip (BracketRun (bet x) x))
2963 -- NOTE: It is important to use UnStream instead of the Stream pattern
2964 -- here, otherwise we get huge perf degradation, see note in concatMap.
2965 step' gst (BracketRun (UnStream step state) v) = do
2966 -- res <- step gst state `MC.onException` aft v
2967 res <- MC.try $ step gst state
2969 Left (e :: SomeException) -> aft v >> MC.throwM e >> return Stop
2970 Right r -> case r of
2971 Yield x s -> return $ Yield x (BracketRun (Stream step s) v)
2972 Skip s -> return $ Skip (BracketRun (Stream step s) v)
2973 Stop -> aft v >> return Stop
2975-- | Run a side effect whenever the stream stops normally or aborts due to an
2977{-# INLINE finally #-}
2978finally :: MonadCatch m => m b -> Stream m a -> Stream m a
2979-- finally action xs = after action $ onException action xs
2980finally action xs = bracket (return ()) (\_ -> action) (const xs)
2982-- | When evaluating a stream if an exception occurs, stream evaluation aborts
2983-- and the specified exception handler is run with the exception as argument.
2984{-# INLINE_NORMAL handle #-}
2985handle :: (MonadCatch m, Exception e)
2986 => (e -> Stream m a) -> Stream m a -> Stream m a
2988 gbracket (return ()) MC.try return (\_ e -> f e) (\_ -> str)
2990{-# INLINE_NORMAL _handle #-}
2991_handle :: (MonadCatch m, Exception e)
2992 => (e -> Stream m a) -> Stream m a -> Stream m a
2993_handle f (Stream step state) = Stream step' (Left state)
2997 {-# INLINE_LATE step' #-}
2998 step' gst (Left st) = do
2999 res <- MC.try $ step gst st
3001 Left e -> return $ Skip $ Right (f e)
3002 Right r -> case r of
3003 Yield x s -> return $ Yield x (Left s)
3004 Skip s -> return $ Skip (Left s)
3007 step' gst (Right (UnStream step1 st)) = do
3010 Yield x s -> return $ Yield x (Right (Stream step1 s))
3011 Skip s -> return $ Skip (Right (Stream step1 s))
3014-------------------------------------------------------------------------------
3015-- General transformation
3016-------------------------------------------------------------------------------
3018{-# INLINE_NORMAL transform #-}
3019transform :: Monad m => Pipe m a b -> Stream m a -> Stream m b
3020transform (Pipe pstep1 pstep2 pstate) (Stream step state) =
3021 Stream step' (Consume pstate, state)
3025 {-# INLINE_LATE step' #-}
3027 step' gst (Consume pst, st) = pst `seq` do
3028 r <- step (adaptState gst) st
3033 Pipe.Yield b pst' -> return $ Yield b (pst', s)
3034 Pipe.Continue pst' -> return $ Skip (pst', s)
3035 Skip s -> return $ Skip (Consume pst, s)
3038 step' _ (Produce pst, st) = pst `seq` do
3041 Pipe.Yield b pst' -> return $ Yield b (pst', st)
3042 Pipe.Continue pst' -> return $ Skip (pst', st)
3044------------------------------------------------------------------------------
3045-- Transformation by Folding (Scans)
3046------------------------------------------------------------------------------
3048------------------------------------------------------------------------------
3050------------------------------------------------------------------------------
3052-- XXX Is a prescan useful, discarding the last step does not sound useful? I
3053-- am not sure about the utility of this function, so this is implemented but
3054-- not exposed. We can expose it if someone provides good reasons why this is
3057-- XXX We have to execute the stream one step ahead to know that we are at the
3058-- last step. The vector implementation of prescan executes the last fold step
3059-- but does not yield the result. This means we have executed the effect but
3060-- discarded value. This does not sound right. In this implementation we are
3061-- not executing the last fold step.
3062{-# INLINE_NORMAL prescanlM' #-}
3063prescanlM' :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
3064prescanlM' f mz (Stream step state) = Stream step' (state, mz)
3066 {-# INLINE_LATE step' #-}
3067 step' gst (st, prev) = do
3068 r <- step (adaptState gst) st
3072 return $ Yield acc (s, f acc x)
3073 Skip s -> return $ Skip (s, prev)
3076{-# INLINE prescanl' #-}
3077prescanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
3078prescanl' f z = prescanlM' (\a b -> return (f a b)) (return z)
3080------------------------------------------------------------------------------
3081-- Monolithic postscans (postscan followed by a map)
3082------------------------------------------------------------------------------
3084-- The performance of a modular postscan followed by a map seems to be
3085-- equivalent to this monolithic scan followed by map therefore we may not need
3086-- this implementation. We just have it for performance comparison and in case
3087-- modular version does not perform well in some situation.
3089{-# INLINE_NORMAL postscanlMx' #-}
3090postscanlMx' :: Monad m
3091 => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> Stream m b
3092postscanlMx' fstep begin done (Stream step state) = do
3093 Stream step' (state, begin)
3095 {-# INLINE_LATE step' #-}
3096 step' gst (st, acc) = do
3097 r <- step (adaptState gst) st
3103 v `seq` y `seq` return (Yield v (s, return y))
3104 Skip s -> return $ Skip (s, acc)
3107{-# INLINE_NORMAL postscanlx' #-}
3108postscanlx' :: Monad m
3109 => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> Stream m b
3110postscanlx' fstep begin done s =
3111 postscanlMx' (\b a -> return (fstep b a)) (return begin) (return . done) s
3113-- XXX do we need consM strict to evaluate the begin value?
3114{-# INLINE scanlMx' #-}
3116 => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> Stream m b
3117scanlMx' fstep begin done s =
3118 (begin >>= \x -> x `seq` done x) `consM` postscanlMx' fstep begin done s
3120{-# INLINE scanlx' #-}
3122 => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> Stream m b
3123scanlx' fstep begin done s =
3124 scanlMx' (\b a -> return (fstep b a)) (return begin) (return . done) s
3126------------------------------------------------------------------------------
3128------------------------------------------------------------------------------
3130{-# INLINE_NORMAL postscanlM' #-}
3131postscanlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
3132postscanlM' fstep begin (Stream step state) =
3133 begin `seq` Stream step' (state, begin)
3135 {-# INLINE_LATE step' #-}
3136 step' gst (st, acc) = acc `seq` do
3137 r <- step (adaptState gst) st
3141 y `seq` return (Yield y (s, y))
3142 Skip s -> return $ Skip (s, acc)
3145{-# INLINE_NORMAL postscanl' #-}
3146postscanl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
3147postscanl' f = postscanlM' (\a b -> return (f a b))
3149{-# INLINE_NORMAL postscanlM #-}
3150postscanlM :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
3151postscanlM fstep begin (Stream step state) = Stream step' (state, begin)
3153 {-# INLINE_LATE step' #-}
3154 step' gst (st, acc) = do
3155 r <- step (adaptState gst) st
3159 return (Yield y (s, y))
3160 Skip s -> return $ Skip (s, acc)
3163{-# INLINE_NORMAL postscanl #-}
3164postscanl :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
3165postscanl f = postscanlM (\a b -> return (f a b))
3167{-# INLINE_NORMAL scanlM' #-}
3168scanlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
3169scanlM' fstep begin s = begin `seq` (begin `cons` postscanlM' fstep begin s)
3171{-# INLINE scanl' #-}
3172scanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
3173scanl' f = scanlM' (\a b -> return (f a b))
3175{-# INLINE_NORMAL scanlM #-}
3176scanlM :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
3177scanlM fstep begin s = begin `cons` postscanlM fstep begin s
3180scanl :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
3181scanl f = scanlM (\a b -> return (f a b))
3183{-# INLINE_NORMAL scanl1M #-}
3184scanl1M :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
3185scanl1M fstep (Stream step state) = Stream step' (state, Nothing)
3187 {-# INLINE_LATE step' #-}
3188 step' gst (st, Nothing) = do
3191 Yield x s -> return $ Yield x (s, Just x)
3192 Skip s -> return $ Skip (s, Nothing)
3195 step' gst (st, Just acc) = do
3200 return $ Yield z (s, Just z)
3201 Skip s -> return $ Skip (s, Just acc)
3204{-# INLINE scanl1 #-}
3205scanl1 :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
3206scanl1 f = scanl1M (\x y -> return (f x y))
3208{-# INLINE_NORMAL scanl1M' #-}
3209scanl1M' :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
3210scanl1M' fstep (Stream step state) = Stream step' (state, Nothing)
3212 {-# INLINE_LATE step' #-}
3213 step' gst (st, Nothing) = do
3216 Yield x s -> x `seq` return $ Yield x (s, Just x)
3217 Skip s -> return $ Skip (s, Nothing)
3220 step' gst (st, Just acc) = acc `seq` do
3225 z `seq` return $ Yield z (s, Just z)
3226 Skip s -> return $ Skip (s, Just acc)
3229{-# INLINE scanl1' #-}
3230scanl1' :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
3231scanl1' f = scanl1M' (\x y -> return (f x y))
3233------------------------------------------------------------------------------
3235------------------------------------------------------------------------------
3237data RollingMapState s a = RollingMapInit s | RollingMapGo s a
3239{-# INLINE rollingMapM #-}
3240rollingMapM :: Monad m => (a -> a -> m b) -> Stream m a -> Stream m b
3241rollingMapM f (Stream step1 state1) = Stream step (RollingMapInit state1)
3243 step gst (RollingMapInit st) = do
3244 r <- step1 (adaptState gst) st
3246 Yield x s -> Skip $ RollingMapGo s x
3247 Skip s -> Skip $ RollingMapInit s
3250 step gst (RollingMapGo s1 x1) = do
3251 r <- step1 (adaptState gst) s1
3255 return $ Yield res $ RollingMapGo s x
3256 Skip s -> return $ Skip $ RollingMapGo s x1
3257 Stop -> return $ Stop
3259{-# INLINE rollingMap #-}
3260rollingMap :: Monad m => (a -> a -> b) -> Stream m a -> Stream m b
3261rollingMap f = rollingMapM (\x y -> return $ f x y)
3263------------------------------------------------------------------------------
3264-- Tapping/Distributing
3265------------------------------------------------------------------------------
3268tap :: Monad m => Fold m a b -> Stream m a -> Stream m a
3269tap (Fold fstep initial extract) (Stream step state) = Stream step' Nothing
3273 step' _ Nothing = do
3275 return $ Skip (Just (r, state))
3277 step' gst (Just (acc, st)) = acc `seq` do
3282 return $ Yield x (Just (acc', s))
3283 Skip s -> return $ Skip (Just (acc, s))
3288{-# INLINE_NORMAL tapOffsetEvery #-}
3289tapOffsetEvery :: Monad m
3290 => Int -> Int -> Fold m a b -> Stream m a -> Stream m a
3291tapOffsetEvery offset n (Fold fstep initial extract) (Stream step state) =
3292 Stream step' Nothing
3296 {-# INLINE_LATE step' #-}
3297 step' _ Nothing = do
3299 return $ Skip (Just (r, state, offset `mod` n))
3301 step' gst (Just (acc, st, count)) | count <= 0 = do
3305 !acc' <- fstep acc x
3306 return $ Yield x (Just (acc', s, n - 1))
3307 Skip s -> return $ Skip (Just (acc, s, count))
3312 step' gst (Just (acc, st, count)) = do
3315 Yield x s -> return $ Yield x (Just (acc, s, count - 1))
3316 Skip s -> return $ Skip (Just (acc, s, count))
3321{-# INLINE_NORMAL pollCounts #-}
3324 => (Stream m Int -> Stream m Int)
3328pollCounts transf fld (Stream step state) = Stream step' Nothing
3331 {-# INLINE_LATE step' #-}
3332 step' _ Nothing = do
3333 countVar <- liftIO $ newVar (0 :: Int)
3335 $ void $ runFold fld
3336 $ transf $ fromPrimVar countVar
3337 return $ Skip (Just (countVar, tid, state))
3339 step' gst (Just (countVar, tid, st)) = do
3343 liftIO $ modifyVar' countVar (+ 1)
3344 return $ Yield x (Just (countVar, tid, s))
3345 Skip s -> return $ Skip (Just (countVar, tid, s))
3347 liftIO $ killThread tid
3350{-# INLINE_NORMAL tapRate #-}
3352 (MonadAsync m, MonadCatch m)
3357tapRate samplingRate action (Stream step state) = Stream step' Nothing
3359 {-# NOINLINE loop #-}
3360 loop countVar prev = do
3363 (do liftIO $ threadDelay (round $ samplingRate * 1000000)
3364 i <- liftIO $ readVar countVar
3365 let !diff = i - prev
3368 (\(e :: AsyncException) -> do
3369 i <- liftIO $ readVar countVar
3370 let !diff = i - prev
3372 throwM (MC.toException e))
3375 {-# INLINE_LATE step' #-}
3376 step' _ Nothing = do
3377 countVar <- liftIO $ newVar 0
3378 tid <- fork $ loop countVar 0
3379 ref <- liftIO $ newIORef ()
3380 _ <- liftIO $ mkWeakIORef ref (killThread tid)
3381 return $ Skip (Just (countVar, tid, state, ref))
3383 step' gst (Just (countVar, tid, st, ref)) = do
3387 liftIO $ modifyVar' countVar (+ 1)
3388 return $ Yield x (Just (countVar, tid, s, ref))
3389 Skip s -> return $ Skip (Just (countVar, tid, s, ref))
3391 liftIO $ killThread tid
3395-------------------------------------------------------------------------------
3397-------------------------------------------------------------------------------
3399{-# INLINE_NORMAL takeWhileM #-}
3400takeWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
3401takeWhileM f (Stream step state) = Stream step' state
3403 {-# INLINE_LATE step' #-}
3409 return $ if b then Yield x s else Stop
3410 Skip s -> return $ Skip s
3413{-# INLINE takeWhile #-}
3414takeWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
3415takeWhile f = takeWhileM (return . f)
3417{-# INLINE_NORMAL drop #-}
3418drop :: Monad m => Int -> Stream m a -> Stream m a
3419drop n (Stream step state) = Stream step' (state, Just n)
3421 {-# INLINE_LATE step' #-}
3422 step' gst (st, Just i)
3427 Yield _ s -> Skip (s, Just (i - 1))
3428 Skip s -> Skip (s, Just i)
3430 | otherwise = return $ Skip (st, Nothing)
3432 step' gst (st, Nothing) = do
3436 Yield x s -> Yield x (s, Nothing)
3437 Skip s -> Skip (s, Nothing)
3440data DropWhileState s a
3442 | DropWhileYield a s
3445{-# INLINE_NORMAL dropWhileM #-}
3446dropWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
3447dropWhileM f (Stream step state) = Stream step' (DropWhileDrop state)
3449 {-# INLINE_LATE step' #-}
3450 step' gst (DropWhileDrop st) = do
3456 then return $ Skip (DropWhileDrop s)
3457 else return $ Skip (DropWhileYield x s)
3458 Skip s -> return $ Skip (DropWhileDrop s)
3461 step' gst (DropWhileNext st) = do
3464 Yield x s -> return $ Skip (DropWhileYield x s)
3465 Skip s -> return $ Skip (DropWhileNext s)
3468 step' _ (DropWhileYield x st) = return $ Yield x (DropWhileNext st)
3470{-# INLINE dropWhile #-}
3471dropWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
3472dropWhile f = dropWhileM (return . f)
3474{-# INLINE_NORMAL filterM #-}
3475filterM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
3476filterM f (Stream step state) = Stream step' state
3478 {-# INLINE_LATE step' #-}
3487 Skip s -> return $ Skip s
3490{-# INLINE filter #-}
3491filter :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
3492filter f = filterM (return . f)
3494{-# INLINE_NORMAL uniq #-}
3495uniq :: (Eq a, Monad m) => Stream m a -> Stream m a
3496uniq (Stream step state) = Stream step' (Nothing, state)
3498 {-# INLINE_LATE step' #-}
3499 step' gst (Nothing, st) = do
3502 Yield x s -> return $ Yield x (Just x, s)
3503 Skip s -> return $ Skip (Nothing, s)
3505 step' gst (Just x, st) = do
3508 Yield y s | x == y -> return $ Skip (Just x, s)
3509 | otherwise -> return $ Yield y (Just y, s)
3510 Skip s -> return $ Skip (Just x, s)
3513------------------------------------------------------------------------------
3514-- Transformation by Mapping
3515------------------------------------------------------------------------------
3517{-# INLINE_NORMAL sequence #-}
3518sequence :: Monad m => Stream m (m a) -> Stream m a
3519sequence (Stream step state) = Stream step' state
3521 {-# INLINE_LATE step' #-}
3523 r <- step (adaptState gst) st
3525 Yield x s -> x >>= \a -> return (Yield a s)
3526 Skip s -> return $ Skip s
3529------------------------------------------------------------------------------
3531------------------------------------------------------------------------------
3533data LoopState x s = FirstYield s
3534 | InterspersingYield s
3537{-# INLINE_NORMAL intersperseM #-}
3538intersperseM :: Monad m => m a -> Stream m a -> Stream m a
3539intersperseM m (Stream step state) = Stream step' (FirstYield state)
3541 {-# INLINE_LATE step' #-}
3542 step' gst (FirstYield st) = do
3546 Yield x s -> Skip (YieldAndCarry x s)
3547 Skip s -> Skip (FirstYield s)
3550 step' gst (InterspersingYield st) = do
3555 return $ Yield a (YieldAndCarry x s)
3556 Skip s -> return $ Skip $ InterspersingYield s
3559 step' _ (YieldAndCarry x st) = return $ Yield x (InterspersingYield st)
3564 | SuffixYield a (SuffixState s a)
3566{-# INLINE_NORMAL intersperseSuffix #-}
3567intersperseSuffix :: forall m a. Monad m => m a -> Stream m a -> Stream m a
3568intersperseSuffix action (Stream step state) = Stream step' (SuffixElem state)
3570 {-# INLINE_LATE step' #-}
3571 step' gst (SuffixElem st) = do
3574 Yield x s -> Skip (SuffixYield x (SuffixSuffix s))
3575 Skip s -> Skip (SuffixElem s)
3578 step' _ (SuffixSuffix st) = do
3579 action >>= \r -> return $ Skip (SuffixYield r (SuffixElem st))
3581 step' _ (SuffixYield x next) = return $ Yield x next
3583data SuffixSpanState s a
3584 = SuffixSpanElem s Int
3585 | SuffixSpanSuffix s
3586 | SuffixSpanYield a (SuffixSpanState s a)
3590-- | intersperse after every n items
3591{-# INLINE_NORMAL intersperseSuffixBySpan #-}
3592intersperseSuffixBySpan :: forall m a. Monad m
3593 => Int -> m a -> Stream m a -> Stream m a
3594intersperseSuffixBySpan n action (Stream step state) =
3595 Stream step' (SuffixSpanElem state n)
3597 {-# INLINE_LATE step' #-}
3598 step' gst (SuffixSpanElem st i) | i > 0 = do
3601 Yield x s -> Skip (SuffixSpanYield x (SuffixSpanElem s (i - 1)))
3602 Skip s -> Skip (SuffixSpanElem s i)
3603 Stop -> if i == n then Stop else Skip SuffixSpanLast
3604 step' _ (SuffixSpanElem st _) = return $ Skip (SuffixSpanSuffix st)
3606 step' _ (SuffixSpanSuffix st) = do
3607 action >>= \r -> return $ Skip (SuffixSpanYield r (SuffixSpanElem st n))
3609 step' _ (SuffixSpanLast) = do
3610 action >>= \r -> return $ Skip (SuffixSpanYield r SuffixSpanStop)
3612 step' _ (SuffixSpanYield x next) = return $ Yield x next
3614 step' _ (SuffixSpanStop) = return Stop
3616{-# INLINE intersperse #-}
3617intersperse :: Monad m => a -> Stream m a -> Stream m a
3618intersperse a = intersperseM (return a)
3620{-# INLINE_NORMAL insertBy #-}
3621insertBy :: Monad m => (a -> a -> Ordering) -> a -> Stream m a -> Stream m a
3622insertBy cmp a (Stream step state) = Stream step' (state, False, Nothing)
3624 {-# INLINE_LATE step' #-}
3625 step' gst (st, False, _) = do
3628 Yield x s -> case cmp a x of
3629 GT -> return $ Yield x (s, False, Nothing)
3630 _ -> return $ Yield a (s, True, Just x)
3631 Skip s -> return $ Skip (s, False, Nothing)
3632 Stop -> return $ Yield a (st, True, Nothing)
3634 step' _ (_, True, Nothing) = return Stop
3636 step' gst (st, True, Just prev) = do
3639 Yield x s -> return $ Yield prev (s, True, Just x)
3640 Skip s -> return $ Skip (s, True, Just prev)
3641 Stop -> return $ Yield prev (st, True, Nothing)
3643------------------------------------------------------------------------------
3645------------------------------------------------------------------------------
3647{-# INLINE_NORMAL deleteBy #-}
3648deleteBy :: Monad m => (a -> a -> Bool) -> a -> Stream m a -> Stream m a
3649deleteBy eq x (Stream step state) = Stream step' (state, False)
3651 {-# INLINE_LATE step' #-}
3652 step' gst (st, False) = do
3655 Yield y s -> return $
3656 if eq x y then Skip (s, True) else Yield y (s, False)
3657 Skip s -> return $ Skip (s, False)
3660 step' gst (st, True) = do
3663 Yield y s -> return $ Yield y (s, True)
3664 Skip s -> return $ Skip (s, True)
3667------------------------------------------------------------------------------
3668-- Transformation by Map and Filter
3669------------------------------------------------------------------------------
3671-- XXX Will this always fuse properly?
3672{-# INLINE_NORMAL mapMaybe #-}
3673mapMaybe :: Monad m => (a -> Maybe b) -> Stream m a -> Stream m b
3674mapMaybe f = fmap fromJust . filter isJust . map f
3676{-# INLINE_NORMAL mapMaybeM #-}
3677mapMaybeM :: Monad m => (a -> m (Maybe b)) -> Stream m a -> Stream m b
3678mapMaybeM f = fmap fromJust . filter isJust . mapM f
3680------------------------------------------------------------------------------
3682------------------------------------------------------------------------------
3684{-# INLINE_NORMAL indexed #-}
3685indexed :: Monad m => Stream m a -> Stream m (Int, a)
3686indexed (Stream step state) = Stream step' (state, 0)
3688 {-# INLINE_LATE step' #-}
3689 step' gst (st, i) = i `seq` do
3690 r <- step (adaptState gst) st
3692 Yield x s -> return $ Yield (i, x) (s, i+1)
3693 Skip s -> return $ Skip (s, i)
3696{-# INLINE_NORMAL indexedR #-}
3697indexedR :: Monad m => Int -> Stream m a -> Stream m (Int, a)
3698indexedR m (Stream step state) = Stream step' (state, m)
3700 {-# INLINE_LATE step' #-}
3701 step' gst (st, i) = i `seq` do
3702 r <- step (adaptState gst) st
3704 Yield x s -> let i' = i - 1
3705 in return $ Yield (i, x) (s, i')
3706 Skip s -> return $ Skip (s, i)
3709{-# INLINE_NORMAL zipWithM #-}
3711 => (a -> b -> m c) -> Stream m a -> Stream m b -> Stream m c
3712zipWithM f (Stream stepa ta) (Stream stepb tb) = Stream step (ta, tb, Nothing)
3714 {-# INLINE_LATE step #-}
3715 step gst (sa, sb, Nothing) = do
3716 r <- stepa (adaptState gst) sa
3719 Yield x sa' -> Skip (sa', sb, Just x)
3720 Skip sa' -> Skip (sa', sb, Nothing)
3723 step gst (sa, sb, Just x) = do
3724 r <- stepb (adaptState gst) sb
3728 return $ Yield z (sa, sb', Nothing)
3729 Skip sb' -> return $ Skip (sa, sb', Just x)
3732#if __GLASGOW_HASKELL__ >= 801
3733{-# RULES "zipWithM xs xs"
3734 forall f xs. zipWithM @Identity f xs xs = mapM (\x -> f x x) xs #-}
3737{-# INLINE zipWith #-}
3738zipWith :: Monad m => (a -> b -> c) -> Stream m a -> Stream m b -> Stream m c
3739zipWith f = zipWithM (\a b -> return (f a b))
3741-- | Like 'zipWithM' but zips concurrently i.e. both the streams being zipped
3742-- are generated concurrently.
3744{-# INLINE_NORMAL zipAsyncWithM #-}
3745zipAsyncWithM :: MonadAsync m
3746 => (a -> b -> m c) -> Stream m a -> Stream m b -> Stream m c
3747zipAsyncWithM f m1 m2 = Stream step Nothing
3751 {-# INLINE_LATE step #-}
3752 step gst Nothing = do
3753 ma <- mkParallel (adaptState gst) m1
3754 mb <- mkParallel (adaptState gst) m2
3755 return $ Skip (Just $ zipWithM f ma mb)
3757 step gst (Just (UnStream step1 st)) = do
3760 Yield a s -> Yield a (Just $ Stream step1 s)
3761 Skip s -> Skip (Just $ Stream step1 s)
3764-- | Like 'zipWith' but zips concurrently i.e. both the streams being zipped
3765-- are generated concurrently.
3767{-# INLINE zipAsyncWith #-}
3768zipAsyncWith :: MonadAsync m
3769 => (a -> b -> c) -> Stream m a -> Stream m b -> Stream m c
3770zipAsyncWith f = zipAsyncWithM (\a b -> return (f a b))
3772------------------------------------------------------------------------------
3774------------------------------------------------------------------------------
3776{-# INLINE_NORMAL mergeByM #-}
3779 => (a -> a -> m Ordering) -> Stream m a -> Stream m a -> Stream m a
3780mergeByM cmp (Stream stepa ta) (Stream stepb tb) =
3781 Stream step (Just ta, Just tb, Nothing, Nothing)
3783 {-# INLINE_LATE step #-}
3785 -- one of the values is missing, and the corresponding stream is running
3786 step gst (Just sa, sb, Nothing, b) = do
3789 Yield a sa' -> Skip (Just sa', sb, Just a, b)
3790 Skip sa' -> Skip (Just sa', sb, Nothing, b)
3791 Stop -> Skip (Nothing, sb, Nothing, b)
3793 step gst (sa, Just sb, a, Nothing) = do
3796 Yield b sb' -> Skip (sa, Just sb', a, Just b)
3797 Skip sb' -> Skip (sa, Just sb', a, Nothing)
3798 Stop -> Skip (sa, Nothing, a, Nothing)
3800 -- both the values are available
3801 step _ (sa, sb, Just a, Just b) = do
3803 return $ case res of
3804 GT -> Yield b (sa, sb, Just a, Nothing)
3805 _ -> Yield a (sa, sb, Nothing, Just b)
3807 -- one of the values is missing, corresponding stream is done
3808 step _ (Nothing, sb, Nothing, Just b) =
3809 return $ Yield b (Nothing, sb, Nothing, Nothing)
3811 step _ (sa, Nothing, Just a, Nothing) =
3812 return $ Yield a (sa, Nothing, Nothing, Nothing)
3814 step _ (Nothing, Nothing, Nothing, Nothing) = return Stop
3816{-# INLINE mergeBy #-}
3819 => (a -> a -> Ordering) -> Stream m a -> Stream m a -> Stream m a
3820mergeBy cmp = mergeByM (\a b -> return $ cmp a b)
3822------------------------------------------------------------------------------
3823-- Transformation comprehensions
3824------------------------------------------------------------------------------
3826{-# INLINE_NORMAL the #-}
3827the :: (Eq a, Monad m) => Stream m a -> m (Maybe a)
3828the (Stream step state) = go state
3831 r <- step defState st
3833 Yield x s -> go' x s
3835 Stop -> return Nothing
3837 r <- step defState st
3839 Yield x s | x == n -> go' n s
3840 | otherwise -> return Nothing
3842 Stop -> return (Just n)
3844{-# INLINE runFold #-}
3845runFold :: (Monad m) => Fold m a b -> Stream m a -> m b
3846runFold (Fold step begin done) = foldlMx' step begin done
3848-------------------------------------------------------------------------------
3849-- Concurrent application and fold
3850-------------------------------------------------------------------------------
3852-- Using StreamD the worker stream producing code can fuse with the code to
3853-- queue output to the SVar giving some perf boost.
3855-- Note that StreamD can only be used in limited situations, specifically, we
3856-- cannot implement joinStreamVarPar using this.
3858-- XXX make sure that the SVar passed is a Parallel style SVar.
3860-- | Fold the supplied stream to the SVar asynchronously using Parallel
3861-- concurrency style.
3862-- {-# INLINE_NORMAL toSVarParallel #-}
3863{-# INLINE toSVarParallel #-}
3864toSVarParallel :: MonadAsync m
3865 => State t m a -> SVar t m a -> Stream m a -> m ()
3866toSVarParallel st sv xs =
3867 if svarInspectMode sv
3871 case getYieldLimit st of
3872 Nothing -> doFork (work Nothing)
3874 (handleChildException sv)
3875 Just _ -> doFork (workLim Nothing)
3877 (handleChildException sv)
3882 {-# NOINLINE work #-}
3883 work info = (runFold (FL.toParallelSVar sv info) xs)
3885 {-# NOINLINE workLim #-}
3886 workLim info = runFold (FL.toParallelSVarLimited sv info) xs
3888 {-# NOINLINE forkWithDiag #-}
3890 -- We do not use workerCount in case of ParallelVar but still there is
3891 -- no harm in maintaining it correctly.
3892 liftIO $ atomicModifyIORefCAS_ (workerCount sv) $ \n -> n + 1
3894 -- This allocation matters when significant number of workers are being
3895 -- sent. We allocate it only when needed. The overhead increases by 4x.
3897 case yieldRateInfo sv of
3898 Nothing -> return Nothing
3899 Just _ -> liftIO $ do
3900 cntRef <- newIORef 0
3901 t <- getTime Monotonic
3902 lat <- newIORef (0, t)
3903 return $ Just WorkerInfo
3904 { workerYieldMax = 0
3905 , workerYieldCount = cntRef
3906 , workerLatencyStart = lat
3909 case getYieldLimit st of
3910 Nothing -> doFork (work winfo)
3912 (handleChildException sv)
3913 Just _ -> doFork (workLim winfo)
3915 (handleChildException sv)
3918{-# INLINE_NORMAL mkParallel #-}
3919mkParallel :: MonadAsync m => State t m a -> Stream m a -> m (Stream m a)
3921 sv <- newParallelVar StopNone defState
3922 toSVarParallel st sv m
3923 return $ fromSVar sv
3925{-# INLINE_NORMAL applyParallel #-}
3926applyParallel :: MonadAsync m
3927 => (Stream m a -> Stream m b) -> Stream m a -> Stream m b
3928applyParallel f m = Stream step Nothing
3931 {-# INLINE_LATE step #-}
3932 step gst Nothing = do
3933 sv <- newParallelVar StopNone (adaptState gst)
3934 toSVarParallel (adaptState gst) sv m
3935 return $ Skip $ Just $ f (fromSVar sv)
3937 step gst (Just (UnStream step1 st)) = do
3940 Yield a s -> Yield a (Just $ Stream step1 s)
3941 Skip s -> Skip (Just $ Stream step1 s)
3944{-# INLINE foldParallel #-}
3945foldParallel :: (K.IsStream t, MonadAsync m) => (t m a -> m b) -> t m a -> m b
3946foldParallel f m = do
3947 sv <- newParallelVar StopNone defState
3948 toSVarParallel defState sv (toStreamD m)
3949 f $ fromStreamD $ fromSVar sv
3951-------------------------------------------------------------------------------
3953-------------------------------------------------------------------------------
3955-- | Create an SVar with a fold consumer that will fold any elements sent to it
3956-- using the supplied fold function.
3957{-# INLINE newFoldSVar #-}
3958newFoldSVar :: MonadAsync m => State t m a -> Fold m a b -> m (SVar t m a)
3959newFoldSVar stt f = do
3960 -- Buffer size for the SVar is derived from the current state
3961 sv <- newParallelVar StopAny (adaptState stt)
3962 -- Add the producer thread-id to the SVar.
3963 liftIO myThreadId >>= modifyThread sv
3964 void $ doFork (work sv) (svarMrun sv) (handleFoldException sv)
3969 {-# NOINLINE work #-}
3970 work sv = void $ runFold f $ fromProducer sv
3972data TapState sv st = TapInit | Tapping sv st | TapDone st
3974{-# INLINE_NORMAL tapAsync #-}
3975tapAsync :: MonadAsync m => Fold m a b -> Stream m a -> Stream m a
3976tapAsync f (Stream step1 state1) = Stream step TapInit
3980 -- In general, a Stop event would come equipped with the result
3981 -- of the fold. It is not used here but it would be useful in
3982 -- applicative and distribute.
3983 done <- fromConsumer svr
3984 when (not done) $ do
3985 liftIO $ withDiagMVar svr "teeToSVar: waiting to drain"
3986 $ takeMVar (outputDoorBellFromConsumer svr)
3990 liftIO $ sendStop svr Nothing
3991 -- drain/wait until a stop event arrives from the fold.
3994 {-# INLINE_LATE step #-}
3995 step gst TapInit = do
3996 sv <- newFoldSVar gst f
3997 return $ Skip (Tapping sv state1)
3999 step gst (Tapping sv st) = do
4003 done <- pushToFold sv a
4006 -- XXX we do not need to wait synchronously here
4008 return $ Yield a (TapDone s)
4009 else return $ Yield a (Tapping sv s)
4010 Skip s -> return $ Skip (Tapping sv s)
4015 step gst (TapDone st) = do
4018 Yield a s -> Yield a (TapDone s)
4019 Skip s -> Skip (TapDone s)
4022-- XXX Is there room for improvement?
4023-- | Take last 'n' elements from the stream and discard the rest.
4024{-# INLINABLE lastN #-}
4025lastN :: (Storable a, MonadIO m) => Int -> Fold m a (Array a)
4026lastN n = Fold step initial done
4028 step (Tuple3' rb rh i) a = do
4029 rh1 <- liftIO $ RB.unsafeInsert rb rh a
4030 return $ Tuple3' rb rh1 (i + 1)
4031 initial = fmap (\(a, b) -> Tuple3' a b 0) $ liftIO $ RB.new n
4032 done (Tuple3' rb rh i) = do
4033 arr <- liftIO $ A.newArray n
4034 foldFunc i rh snoc' arr rb
4035 snoc' b a = liftIO $ A.unsafeSnoc b a
4037 | i < n = RB.unsafeFoldRingM
4038 | otherwise = RB.unsafeFoldRingFullM
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