1{-# LANGUAGE BangPatterns #-}
3{-# LANGUAGE ConstraintKinds #-}
4{-# LANGUAGE ExistentialQuantification #-}
5{-# LANGUAGE FlexibleContexts #-}
6{-# LANGUAGE FlexibleInstances #-}
7{-# LANGUAGE MultiParamTypeClasses #-}
8{-# LANGUAGE PatternSynonyms #-}
9{-# LANGUAGE RecordWildCards #-}
10{-# LANGUAGE ScopedTypeVariables #-}
11{-# LANGUAGE ViewPatterns #-}
12{-# LANGUAGE RankNTypes #-}
13{-# LANGUAGE MagicHash #-}
15#if __GLASGOW_HASKELL__ >= 801
16{-# LANGUAGE TypeApplications #-}
22-- Module : Streamly.Internal.Data.Stream.StreamD
23-- Copyright : (c) 2018 Harendra Kumar
24-- (c) Roman Leshchinskiy 2008-2010
25-- (c) The University of Glasgow, 2009
28-- Maintainer : streamly@composewell.com
29-- Stability : experimental
32-- Direct style re-implementation of CPS style stream in StreamK module. The
33-- symbol or suffix 'D' in this module denotes the "Direct" style. GHC is able
34-- to INLINE and fuse direct style better, providing better performance than
38-- import qualified Streamly.Internal.Data.Stream.StreamD as D
41-- Some of the functions in this file have been adapted from the vector
42-- library, https://hackage.haskell.org/package/vector.
44module Streamly.Internal.Data.Stream.StreamD
49#if __GLASGOW_HASKELL__ >= 800
50 , Stream (Stream, UnStream)
70 -- ** Specialized Generation
71 -- | Generate a monadic stream from a seed.
84 , enumerateFromStepIntegral
85 , enumerateFromIntegral
86 , enumerateFromThenIntegral
87 , enumerateFromToIntegral
88 , enumerateFromThenToIntegral
90 , enumerateFromStepNum
93 , enumerateFromToFractional
94 , enumerateFromThenToFractional
97 -- | Transform an input structure into a stream.
98 -- | Direct style stream does not support @fromFoldable@.
128 -- ** Specialized Folds
155 -- ** Flattening nested streams
158 , ConcatMapUState (..)
160 , ConcatUnfoldInterleaveState (..)
161 , concatUnfoldInterleave
162 , concatUnfoldRoundrobin
165 , InterleaveState(..)
170 , roundRobin -- interleaveFair?/ParallelFair
203 -- | Transform a stream into another type.
220 -- ** By folding (scans)
265 , intersperseSuffixBySpan
289 -- * Transformation comprehensions
305 -- * Concurrent Application
318import Control.Concurrent (killThread, myThreadId, takeMVar, threadDelay)
319import Control.Exception
320 (Exception, SomeException, AsyncException, fromException)
321import Control.Monad (void, when)
322import Control.Monad.Catch (MonadCatch, throwM)
323import Control.Monad.IO.Class (MonadIO(..))
324import Control.Monad.Reader (ReaderT)
325import Control.Monad.State.Strict (StateT)
326import Control.Monad.Trans (MonadTrans(lift))
327import Control.Monad.Trans.Control (MonadBaseControl)
328import Data.Bits (shiftR, shiftL, (.|.), (.&.))
329import Data.Functor.Identity (Identity(..))
330import Data.IORef (newIORef, readIORef, mkWeakIORef, writeIORef, IORef)
331import Data.Maybe (fromJust, isJust, isNothing)
332import Data.Word (Word32)
333import Foreign.Ptr (Ptr)
334import Foreign.Storable (Storable(..))
335import GHC.Types (SPEC(..))
336import System.Mem (performMajorGC)
338 hiding (map, mapM, mapM_, repeat, foldr, last, take, filter,
339 takeWhile, drop, dropWhile, all, any, maximum, minimum, elem,
340 notElem, null, head, tail, zipWith, lookup, foldr1, sequence,
341 (!!), scanl, scanl1, concatMap, replicate, enumFromTo, concat,
344import qualified Control.Monad.Catch as MC
345import qualified Control.Monad.Reader as Reader
346import qualified Control.Monad.State.Strict as State
347import qualified Prelude
349import Data.Int (Int64)
350import Streamly.Internal.Mutable.Prim.Var
351 (Prim, Var, readVar, newVar, modifyVar')
352import Streamly.Internal.Data.Time.Units
353 (TimeUnit64, toRelTime64, diffAbsTime64)
355import Streamly.Internal.Data.Atomics (atomicModifyIORefCAS_)
356import Streamly.Internal.Memory.Array.Types (Array(..))
357import Streamly.Internal.Data.Fold.Types (Fold(..))
358import Streamly.Internal.Data.Pipe.Types (Pipe(..), PipeState(..))
359import Streamly.Internal.Data.Time.Clock (Clock(Monotonic), getTime)
360import Streamly.Internal.Data.Time.Units (MicroSecond64(..), fromAbsTime, toAbsTime, AbsTime)
361import Streamly.Internal.Data.Unfold.Types (Unfold(..))
362import Streamly.Internal.Data.Strict (Tuple3'(..))
364import Streamly.Internal.Data.Stream.StreamD.Type
365import Streamly.Internal.Data.SVar
366import Streamly.Internal.Data.Stream.SVar (fromConsumer, pushToFold)
368import qualified Streamly.Internal.Data.Pipe.Types as Pipe
369import qualified Streamly.Internal.Memory.Array.Types as A
370import qualified Streamly.Internal.Data.Fold as FL
371import qualified Streamly.Memory.Ring as RB
372import qualified Streamly.Internal.Data.Stream.StreamK as K
374------------------------------------------------------------------------------
376------------------------------------------------------------------------------
378-- | An empty 'Stream'.
379{-# INLINE_NORMAL nil #-}
380nil :: Monad m => Stream m a
381nil = Stream (\_ _ -> return Stop) ()
383-- | An empty 'Stream' with a side effect.
384{-# INLINE_NORMAL nilM #-}
385nilM :: Monad m => m b -> Stream m a
386nilM m = Stream (\_ _ -> m >> return Stop) ()
388{-# INLINE_NORMAL consM #-}
389consM :: Monad m => m a -> Stream m a -> Stream m a
390consM m (Stream step state) = Stream step1 Nothing
392 {-# INLINE_LATE step1 #-}
393 step1 _ Nothing = m >>= \x -> return $ Yield x (Just state)
394 step1 gst (Just st) = do
398 Yield a s -> Yield a (Just s)
399 Skip s -> Skip (Just s)
402-- XXX implement in terms of consM?
403-- cons x = consM (return x)
405-- | Can fuse but has O(n^2) complexity.
406{-# INLINE_NORMAL cons #-}
407cons :: Monad m => a -> Stream m a -> Stream m a
408cons x (Stream step state) = Stream step1 Nothing
410 {-# INLINE_LATE step1 #-}
411 step1 _ Nothing = return $ Yield x (Just state)
412 step1 gst (Just st) = do
416 Yield a s -> Yield a (Just s)
417 Skip s -> Skip (Just s)
420-------------------------------------------------------------------------------
422-------------------------------------------------------------------------------
424-- Does not fuse, has the same performance as the StreamK version.
425{-# INLINE_NORMAL uncons #-}
426uncons :: Monad m => Stream m a -> m (Maybe (a, Stream m a))
427uncons (UnStream step state) = go state
430 r <- step defState st
432 Yield x s -> return $ Just (x, Stream step s)
434 Stop -> return Nothing
436------------------------------------------------------------------------------
437-- Generation by unfold
438------------------------------------------------------------------------------
440{-# INLINE_NORMAL unfoldrM #-}
441unfoldrM :: Monad m => (s -> m (Maybe (a, s))) -> s -> Stream m a
442unfoldrM next state = Stream step state
444 {-# INLINE_LATE step #-}
448 Just (x, s) -> Yield x s
451{-# INLINE_LATE unfoldr #-}
452unfoldr :: Monad m => (s -> Maybe (a, s)) -> s -> Stream m a
453unfoldr f = unfoldrM (return . f)
455-- | Convert an 'Unfold' into a 'Stream' by supplying it a seed.
457{-# INLINE_NORMAL unfold #-}
458unfold :: Monad m => Unfold m a b -> a -> Stream m b
459unfold (Unfold ustep inject) seed = Stream step Nothing
461 {-# INLINE_LATE step #-}
462 step _ Nothing = inject seed >>= return . Skip . Just
463 step _ (Just st) = do
466 Yield x s -> Yield x (Just s)
467 Skip s -> Skip (Just s)
470------------------------------------------------------------------------------
471-- Specialized Generation
472------------------------------------------------------------------------------
474{-# INLINE_NORMAL repeatM #-}
475repeatM :: Monad m => m a -> Stream m a
476repeatM x = Stream (\_ _ -> x >>= \r -> return $ Yield r ()) ()
478{-# INLINE_NORMAL repeat #-}
479repeat :: Monad m => a -> Stream m a
480repeat x = Stream (\_ _ -> return $ Yield x ()) ()
482{-# INLINE_NORMAL iterateM #-}
483iterateM :: Monad m => (a -> m a) -> m a -> Stream m a
484iterateM step = Stream (\_ st -> st >>= \x -> return $ Yield x (step x))
486{-# INLINE_NORMAL iterate #-}
487iterate :: Monad m => (a -> a) -> a -> Stream m a
488iterate step st = iterateM (return . step) (return st)
490{-# INLINE_NORMAL replicateM #-}
491replicateM :: forall m a. Monad m => Int -> m a -> Stream m a
492replicateM n p = Stream step n
494 {-# INLINE_LATE step #-}
496 | i <= 0 = return Stop
499 return $ Yield x (i - 1)
501{-# INLINE_NORMAL replicate #-}
502replicate :: Monad m => Int -> a -> Stream m a
503replicate n x = replicateM n (return x)
505-- This would not work properly for floats, therefore we put an Integral
507-- | Can be used to enumerate unbounded integrals. This does not check for
508-- overflow or underflow for bounded integrals.
509{-# INLINE_NORMAL enumerateFromStepIntegral #-}
510enumerateFromStepIntegral :: (Integral a, Monad m) => a -> a -> Stream m a
511enumerateFromStepIntegral from stride =
512 from `seq` stride `seq` Stream step from
514 {-# INLINE_LATE step #-}
515 step _ !x = return $ Yield x $! (x + stride)
517-- We are assuming that "to" is constrained by the type to be within
519{-# INLINE enumerateFromToIntegral #-}
520enumerateFromToIntegral :: (Monad m, Integral a) => a -> a -> Stream m a
521enumerateFromToIntegral from to =
522 takeWhile (<= to) $ enumerateFromStepIntegral from 1
524{-# INLINE enumerateFromIntegral #-}
525enumerateFromIntegral :: (Monad m, Integral a, Bounded a) => a -> Stream m a
526enumerateFromIntegral from = enumerateFromToIntegral from maxBound
528data EnumState a = EnumInit | EnumYield a a a | EnumStop
530{-# INLINE_NORMAL enumerateFromThenToIntegralUp #-}
531enumerateFromThenToIntegralUp
532 :: (Monad m, Integral a)
533 => a -> a -> a -> Stream m a
534enumerateFromThenToIntegralUp from next to = Stream step EnumInit
536 {-# INLINE_LATE step #-}
542 else Yield from EnumStop
543 else -- from <= next <= to
544 let stride = next - from
545 in Skip $ EnumYield from stride (to - stride)
547 step _ (EnumYield x stride toMinus) =
550 then Yield x EnumStop
551 else Yield x $ EnumYield (x + stride) stride toMinus
553 step _ EnumStop = return Stop
555{-# INLINE_NORMAL enumerateFromThenToIntegralDn #-}
556enumerateFromThenToIntegralDn
557 :: (Monad m, Integral a)
558 => a -> a -> a -> Stream m a
559enumerateFromThenToIntegralDn from next to = Stream step EnumInit
561 {-# INLINE_LATE step #-}
563 return $ if to > next
566 else Yield from EnumStop
567 else -- from >= next >= to
568 let stride = next - from
569 in Skip $ EnumYield from stride (to - stride)
571 step _ (EnumYield x stride toMinus) =
574 then Yield x EnumStop
575 else Yield x $ EnumYield (x + stride) stride toMinus
577 step _ EnumStop = return Stop
579{-# INLINE_NORMAL enumerateFromThenToIntegral #-}
580enumerateFromThenToIntegral
581 :: (Monad m, Integral a)
582 => a -> a -> a -> Stream m a
583enumerateFromThenToIntegral from next to
584 | next >= from = enumerateFromThenToIntegralUp from next to
585 | otherwise = enumerateFromThenToIntegralDn from next to
587{-# INLINE_NORMAL enumerateFromThenIntegral #-}
588enumerateFromThenIntegral
589 :: (Monad m, Integral a, Bounded a)
590 => a -> a -> Stream m a
591enumerateFromThenIntegral from next =
593 then enumerateFromThenToIntegralUp from next maxBound
594 else enumerateFromThenToIntegralDn from next minBound
596-- For floating point numbers if the increment is less than the precision then
597-- it just gets lost. Therefore we cannot always increment it correctly by just
599-- 9007199254740992 + 1 + 1 :: Double => 9.007199254740992e15
600-- 9007199254740992 + 2 :: Double => 9.007199254740994e15
602-- Instead we accumulate the increment counter and compute the increment
603-- every time before adding it to the starting number.
605-- This works for Integrals as well as floating point numbers, but
606-- enumerateFromStepIntegral is faster for integrals.
607{-# INLINE_NORMAL enumerateFromStepNum #-}
608enumerateFromStepNum :: (Monad m, Num a) => a -> a -> Stream m a
609enumerateFromStepNum from stride = Stream step 0
611 {-# INLINE_LATE step #-}
612 step _ !i = return $ (Yield $! (from + i * stride)) $! (i + 1)
614{-# INLINE_NORMAL numFrom #-}
615numFrom :: (Monad m, Num a) => a -> Stream m a
616numFrom from = enumerateFromStepNum from 1
618{-# INLINE_NORMAL numFromThen #-}
619numFromThen :: (Monad m, Num a) => a -> a -> Stream m a
620numFromThen from next = enumerateFromStepNum from (next - from)
622-- We cannot write a general function for Num. The only way to write code
623-- portable between the two is to use a 'Real' constraint and convert between
624-- Fractional and Integral using fromRational which is horribly slow.
625{-# INLINE_NORMAL enumerateFromToFractional #-}
626enumerateFromToFractional
627 :: (Monad m, Fractional a, Ord a)
628 => a -> a -> Stream m a
629enumerateFromToFractional from to =
630 takeWhile (<= to + 1 / 2) $ enumerateFromStepNum from 1
632{-# INLINE_NORMAL enumerateFromThenToFractional #-}
633enumerateFromThenToFractional
634 :: (Monad m, Fractional a, Ord a)
635 => a -> a -> a -> Stream m a
636enumerateFromThenToFractional from next to =
637 takeWhile predicate $ numFromThen from next
639 mid = (next - from) / 2
640 predicate | next >= from = (<= to + mid)
641 | otherwise = (>= to + mid)
643-------------------------------------------------------------------------------
644-- Generation by Conversion
645-------------------------------------------------------------------------------
647{-# INLINE_NORMAL fromIndicesM #-}
648fromIndicesM :: Monad m => (Int -> m a) -> Stream m a
649fromIndicesM gen = Stream step 0
651 {-# INLINE_LATE step #-}
654 return $ Yield x (i + 1)
656{-# INLINE fromIndices #-}
657fromIndices :: Monad m => (Int -> a) -> Stream m a
658fromIndices gen = fromIndicesM (return . gen)
660{-# INLINE_NORMAL generateM #-}
661generateM :: Monad m => Int -> (Int -> m a) -> Stream m a
662generateM n gen = n `seq` Stream step 0
664 {-# INLINE_LATE step #-}
665 step _ i | i < n = do
667 return $ Yield x (i + 1)
668 | otherwise = return Stop
670{-# INLINE generate #-}
671generate :: Monad m => Int -> (Int -> a) -> Stream m a
672generate n gen = generateM n (return . gen)
674-- XXX we need the MonadAsync constraint because of a rewrite rule.
675-- | Convert a list of monadic actions to a 'Stream'
676{-# INLINE_LATE fromListM #-}
677fromListM :: MonadAsync m => [m a] -> Stream m a
678fromListM = Stream step
680 {-# INLINE_LATE step #-}
681 step _ (m:ms) = m >>= \x -> return $ Yield x ms
682 step _ [] = return Stop
684{-# INLINE toStreamD #-}
685toStreamD :: (K.IsStream t, Monad m) => t m a -> Stream m a
686toStreamD = fromStreamK . K.toStream
688{-# INLINE_NORMAL fromPrimVar #-}
689fromPrimVar :: (MonadIO m, Prim a) => Var IO a -> Stream m a
690fromPrimVar var = Stream step ()
692 {-# INLINE_LATE step #-}
693 step _ () = liftIO (readVar var) >>= \x -> return $ Yield x ()
695-------------------------------------------------------------------------------
696-- Generation from SVar
697-------------------------------------------------------------------------------
699data FromSVarState t m a =
701 | FromSVarRead (SVar t m a)
702 | FromSVarLoop (SVar t m a) [ChildEvent a]
703 | FromSVarDone (SVar t m a)
705{-# INLINE_NORMAL fromSVar #-}
706fromSVar :: (MonadAsync m) => SVar t m a -> Stream m a
707fromSVar svar = Stream step FromSVarInit
710 {-# INLINE_LATE step #-}
711 step _ FromSVarInit = do
712 ref <- liftIO $ newIORef ()
713 _ <- liftIO $ mkWeakIORef ref hook
714 -- when this copy of svar gets garbage collected "ref" will get
715 -- garbage collected and our GC hook will be called.
716 let sv = svar{svarRef = Just ref}
717 return $ Skip (FromSVarRead sv)
721 {-# NOINLINE hook #-}
723 when (svarInspectMode svar) $ do
724 r <- liftIO $ readIORef (svarStopTime (svarStats svar))
726 printSVar svar "SVar Garbage Collected"
728 -- If there are any SVars referenced by this SVar a GC will prompt
729 -- them to be cleaned up quickly.
730 when (svarInspectMode svar) performMajorGC
732 step _ (FromSVarRead sv) = do
733 list <- readOutputQ sv
734 -- Reversing the output is important to guarantee that we process the
735 -- outputs in the same order as they were generated by the constituent
737 return $ Skip $ FromSVarLoop sv (Prelude.reverse list)
739 step _ (FromSVarLoop sv []) = do
740 done <- postProcess sv
741 return $ Skip $ if done
742 then (FromSVarDone sv)
743 else (FromSVarRead sv)
745 step _ (FromSVarLoop sv (ev : es)) = do
747 ChildYield a -> return $ Yield a (FromSVarLoop sv es)
748 ChildStop tid e -> do
752 stop <- shouldStop tid
755 liftIO (cleanupSVar sv)
756 return $ Skip (FromSVarDone sv)
757 else return $ Skip (FromSVarLoop sv es)
759 case fromException ex of
761 return $ Skip (FromSVarLoop sv es)
762 Nothing -> liftIO (cleanupSVar sv) >> throwM ex
766 case svarStopStyle sv of
767 StopNone -> return False
768 StopAny -> return True
770 sid <- liftIO $ readIORef (svarStopBy sv)
771 return $ if tid == sid then True else False
773 step _ (FromSVarDone sv) = do
774 when (svarInspectMode sv) $ do
775 t <- liftIO $ getTime Monotonic
776 liftIO $ writeIORef (svarStopTime (svarStats sv)) (Just t)
777 liftIO $ printSVar sv "SVar Done"
780-------------------------------------------------------------------------------
781-- Process events received by a fold consumer from a stream producer
782-------------------------------------------------------------------------------
784{-# INLINE_NORMAL fromProducer #-}
785fromProducer :: (MonadAsync m) => SVar t m a -> Stream m a
786fromProducer svar = Stream step (FromSVarRead svar)
789 {-# INLINE_LATE step #-}
790 step _ (FromSVarRead sv) = do
791 list <- readOutputQ sv
792 -- Reversing the output is important to guarantee that we process the
793 -- outputs in the same order as they were generated by the constituent
795 return $ Skip $ FromSVarLoop sv (Prelude.reverse list)
797 step _ (FromSVarLoop sv []) = return $ Skip $ FromSVarRead sv
798 step _ (FromSVarLoop sv (ev : es)) = do
800 ChildYield a -> return $ Yield a (FromSVarLoop sv es)
801 ChildStop tid e -> do
805 sendStopToProducer sv
806 return $ Skip (FromSVarDone sv)
807 Just _ -> error "Bug: fromProducer: received exception"
809 step _ (FromSVarDone sv) = do
810 when (svarInspectMode sv) $ do
811 t <- liftIO $ getTime Monotonic
812 liftIO $ writeIORef (svarStopTime (svarStats sv)) (Just t)
813 liftIO $ printSVar sv "SVar Done"
816 step _ FromSVarInit = undefined
818-------------------------------------------------------------------------------
819-- Hoisting the inner monad
820-------------------------------------------------------------------------------
822{-# INLINE_NORMAL hoist #-}
823hoist :: Monad n => (forall x. m x -> n x) -> Stream m a -> Stream n a
824hoist f (Stream step state) = (Stream step' state)
826 {-# INLINE_LATE step' #-}
828 r <- f $ step (adaptState gst) st
830 Yield x s -> Yield x s
834{-# INLINE generally #-}
835generally :: Monad m => Stream Identity a -> Stream m a
836generally = hoist (return . runIdentity)
838{-# INLINE_NORMAL liftInner #-}
839liftInner :: (Monad m, MonadTrans t, Monad (t m))
840 => Stream m a -> Stream (t m) a
841liftInner (Stream step state) = Stream step' state
843 {-# INLINE_LATE step' #-}
845 r <- lift $ step (adaptState gst) st
847 Yield x s -> Yield x s
851{-# INLINE_NORMAL runReaderT #-}
852runReaderT :: Monad m => s -> Stream (ReaderT s m) a -> Stream m a
853runReaderT sval (Stream step state) = Stream step' state
855 {-# INLINE_LATE step' #-}
857 r <- Reader.runReaderT (step (adaptState gst) st) sval
859 Yield x s -> Yield x s
863{-# INLINE_NORMAL evalStateT #-}
864evalStateT :: Monad m => s -> Stream (StateT s m) a -> Stream m a
865evalStateT sval (Stream step state) = Stream step' (state, sval)
867 {-# INLINE_LATE step' #-}
868 step' gst (st, sv) = do
869 (r, sv') <- State.runStateT (step (adaptState gst) st) sv
871 Yield x s -> Yield x (s, sv')
872 Skip s -> Skip (s, sv')
875{-# INLINE_NORMAL runStateT #-}
876runStateT :: Monad m => s -> Stream (StateT s m) a -> Stream m (s, a)
877runStateT sval (Stream step state) = Stream step' (state, sval)
879 {-# INLINE_LATE step' #-}
880 step' gst (st, sv) = do
881 (r, sv') <- State.runStateT (step (adaptState gst) st) sv
883 Yield x s -> Yield (sv', x) (s, sv')
884 Skip s -> Skip (s, sv')
887------------------------------------------------------------------------------
888-- Elimination by Folds
889------------------------------------------------------------------------------
891------------------------------------------------------------------------------
893------------------------------------------------------------------------------
895{-# INLINE_NORMAL foldr1 #-}
896foldr1 :: Monad m => (a -> a -> a) -> Stream m a -> m (Maybe a)
900 Nothing -> return Nothing
901 Just (h, t) -> fmap Just (foldr f h t)
903------------------------------------------------------------------------------
905------------------------------------------------------------------------------
907{-# INLINE_NORMAL foldlT #-}
908foldlT :: (Monad m, Monad (s m), MonadTrans s)
909 => (s m b -> a -> s m b) -> s m b -> Stream m a -> s m b
910foldlT fstep begin (Stream step state) = go SPEC begin state
913 r <- lift $ step defState st
915 Yield x s -> go SPEC (fstep acc x) s
916 Skip s -> go SPEC acc s
919-- Note, this is going to have horrible performance, because of the nature of
920-- the stream type (i.e. direct stream vs CPS). Its only for reference, it is
921-- likely be practically unusable.
922{-# INLINE_NORMAL foldlS #-}
924 => (Stream m b -> a -> Stream m b) -> Stream m b -> Stream m a -> Stream m b
925foldlS fstep begin (Stream step state) = Stream step' (Left (state, begin))
927 step' gst (Left (st, acc)) = do
928 r <- step (adaptState gst) st
930 Yield x s -> Skip (Left (s, fstep acc x))
931 Skip s -> Skip (Left (s, acc))
932 Stop -> Skip (Right acc)
934 step' gst (Right (Stream stp stt)) = do
935 r <- stp (adaptState gst) stt
937 Yield x s -> Yield x (Right (Stream stp s))
938 Skip s -> Skip (Right (Stream stp s))
941------------------------------------------------------------------------------
943------------------------------------------------------------------------------
945-- | Run a streaming composition, discard the results.
946{-# INLINE_LATE drain #-}
947drain :: Monad m => Stream m a -> m ()
948-- drain = foldrM (\_ xs -> xs) (return ())
949drain (Stream step state) = go SPEC state
952 r <- step defState st
954 Yield _ s -> go SPEC s
958{-# INLINE_NORMAL null #-}
959null :: Monad m => Stream m a -> m Bool
960null m = foldrM (\_ _ -> return False) (return True) m
962{-# INLINE_NORMAL head #-}
963head :: Monad m => Stream m a -> m (Maybe a)
964head m = foldrM (\x _ -> return (Just x)) (return Nothing) m
966{-# INLINE_NORMAL headElse #-}
967headElse :: Monad m => a -> Stream m a -> m a
968headElse a m = foldrM (\x _ -> return x) (return a) m
970-- Does not fuse, has the same performance as the StreamK version.
971{-# INLINE_NORMAL tail #-}
972tail :: Monad m => Stream m a -> m (Maybe (Stream m a))
973tail (UnStream step state) = go state
976 r <- step defState st
978 Yield _ s -> return (Just $ Stream step s)
980 Stop -> return Nothing
982-- XXX will it fuse? need custom impl?
983{-# INLINE_NORMAL last #-}
984last :: Monad m => Stream m a -> m (Maybe a)
985last = foldl' (\_ y -> Just y) Nothing
987{-# INLINE_NORMAL elem #-}
988elem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
989-- elem e m = foldrM (\x xs -> if x == e then return True else xs) (return False) m
990elem e (Stream step state) = go state
993 r <- step defState st
996 | x == e -> return True
1001{-# INLINE_NORMAL notElem #-}
1002notElem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
1003notElem e s = fmap not (elem e s)
1005{-# INLINE_NORMAL all #-}
1006all :: Monad m => (a -> Bool) -> Stream m a -> m Bool
1007-- all p m = foldrM (\x xs -> if p x then xs else return False) (return True) m
1008all p (Stream step state) = go state
1011 r <- step defState st
1015 | otherwise -> return False
1019{-# INLINE_NORMAL any #-}
1020any :: Monad m => (a -> Bool) -> Stream m a -> m Bool
1021-- any p m = foldrM (\x xs -> if p x then return True else xs) (return False) m
1022any p (Stream step state) = go state
1025 r <- step defState st
1028 | p x -> return True
1031 Stop -> return False
1033{-# INLINE_NORMAL maximum #-}
1034maximum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
1035maximum (Stream step state) = go Nothing state
1038 r <- step defState st
1040 Yield x s -> go (Just x) s
1041 Skip s -> go Nothing s
1042 Stop -> return Nothing
1043 go (Just acc) st = do
1044 r <- step defState st
1047 | acc <= x -> go (Just x) s
1048 | otherwise -> go (Just acc) s
1049 Skip s -> go (Just acc) s
1050 Stop -> return (Just acc)
1052{-# INLINE_NORMAL maximumBy #-}
1053maximumBy :: Monad m => (a -> a -> Ordering) -> Stream m a -> m (Maybe a)
1054maximumBy cmp (Stream step state) = go Nothing state
1057 r <- step defState st
1059 Yield x s -> go (Just x) s
1060 Skip s -> go Nothing s
1061 Stop -> return Nothing
1062 go (Just acc) st = do
1063 r <- step defState st
1065 Yield x s -> case cmp acc x of
1066 GT -> go (Just acc) s
1068 Skip s -> go (Just acc) s
1069 Stop -> return (Just acc)
1071{-# INLINE_NORMAL minimum #-}
1072minimum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
1073minimum (Stream step state) = go Nothing state
1076 r <- step defState st
1078 Yield x s -> go (Just x) s
1079 Skip s -> go Nothing s
1080 Stop -> return Nothing
1081 go (Just acc) st = do
1082 r <- step defState st
1085 | acc <= x -> go (Just acc) s
1086 | otherwise -> go (Just x) s
1087 Skip s -> go (Just acc) s
1088 Stop -> return (Just acc)
1090{-# INLINE_NORMAL minimumBy #-}
1091minimumBy :: Monad m => (a -> a -> Ordering) -> Stream m a -> m (Maybe a)
1092minimumBy cmp (Stream step state) = go Nothing state
1095 r <- step defState st
1097 Yield x s -> go (Just x) s
1098 Skip s -> go Nothing s
1099 Stop -> return Nothing
1100 go (Just acc) st = do
1101 r <- step defState st
1103 Yield x s -> case cmp acc x of
1105 _ -> go (Just acc) s
1106 Skip s -> go (Just acc) s
1107 Stop -> return (Just acc)
1109{-# INLINE_NORMAL (!!) #-}
1110(!!) :: (Monad m) => Stream m a -> Int -> m (Maybe a)
1111(Stream step state) !! i = go i state
1114 r <- step defState st
1116 Yield x s | n < 0 -> return Nothing
1117 | n == 0 -> return $ Just x
1118 | otherwise -> go (n - 1) s
1120 Stop -> return Nothing
1122{-# INLINE_NORMAL lookup #-}
1123lookup :: (Monad m, Eq a) => a -> Stream m (a, b) -> m (Maybe b)
1124lookup e m = foldrM (\(a, b) xs -> if e == a then return (Just b) else xs)
1127{-# INLINE_NORMAL findM #-}
1128findM :: Monad m => (a -> m Bool) -> Stream m a -> m (Maybe a)
1129findM p m = foldrM (\x xs -> p x >>= \r -> if r then return (Just x) else xs)
1133find :: Monad m => (a -> Bool) -> Stream m a -> m (Maybe a)
1134find p = findM (return . p)
1136{-# INLINE_NORMAL findIndices #-}
1137findIndices :: Monad m => (a -> Bool) -> Stream m a -> Stream m Int
1138findIndices p (Stream step state) = Stream step' (state, 0)
1140 {-# INLINE_LATE step' #-}
1141 step' gst (st, i) = i `seq` do
1142 r <- step (adaptState gst) st
1144 Yield x s -> if p x then Yield i (s, i+1) else Skip (s, i+1)
1145 Skip s -> Skip (s, i)
1148{-# INLINE toListRev #-}
1149toListRev :: Monad m => Stream m a -> m [a]
1150toListRev = foldl' (flip (:)) []
1152-- We can implement reverse as:
1154-- > reverse = foldlS (flip cons) nil
1156-- However, this implementation is unusable because of the horrible performance
1157-- of cons. So we just convert it to a list first and then stream from the
1160-- XXX Maybe we can use an Array instead of a list here?
1161{-# INLINE_NORMAL reverse #-}
1162reverse :: Monad m => Stream m a -> Stream m a
1163reverse m = Stream step Nothing
1165 {-# INLINE_LATE step #-}
1168 return $ Skip (Just xs)
1169 step _ (Just (x:xs)) = return $ Yield x (Just xs)
1170 step _ (Just []) = return Stop
1172-- Much faster reverse for Storables
1173{-# INLINE_NORMAL reverse' #-}
1174reverse' :: forall m a. (MonadIO m, Storable a) => Stream m a -> Stream m a
1176-- This commented implementation copies the whole stream into one single array
1177-- and then streams from that array, this is 3-4x faster than the chunked code
1178-- that follows. Though this could be problematic due to unbounded large
1179-- allocations. We need to figure out why the chunked code is slower and if we
1180-- can optimize the chunked code to work as fast as this one. It may be a
1182import Foreign.ForeignPtr (touchForeignPtr)
1183import Foreign.ForeignPtr.Unsafe (unsafeForeignPtrToPtr)
1184import Foreign.Ptr (Ptr, plusPtr)
1185reverse' m = Stream step Nothing
1187 {-# INLINE_LATE step #-}
1189 arr <- A.fromStreamD m
1190 let p = aEnd arr `plusPtr` negate (sizeOf (undefined :: a))
1191 return $ Skip $ Just (aStart arr, p)
1193 step _ (Just (start, p)) | p < unsafeForeignPtrToPtr start = return Stop
1195 step _ (Just (start, p)) = do
1196 let !x = A.unsafeInlineIO $ do
1198 touchForeignPtr start
1200 next = p `plusPtr` negate (sizeOf (undefined :: a))
1201 return $ Yield x (Just (start, next))
1208 $ A.fromStreamDArraysOf A.defaultChunkSize m
1211------------------------------------------------------------------------------
1212-- Grouping/Splitting
1213------------------------------------------------------------------------------
1215{-# INLINE_NORMAL splitSuffixBy' #-}
1216splitSuffixBy' :: Monad m
1217 => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1218splitSuffixBy' predicate f (Stream step state) =
1219 Stream (stepOuter f) (Just state)
1223 {-# INLINE_LATE stepOuter #-}
1224 stepOuter (Fold fstep initial done) gst (Just st) = do
1225 res <- step (adaptState gst) st
1231 then done acc' >>= \val -> return $ Yield val (Just s)
1234 Skip s -> return $ Skip $ Just s
1240 res <- step (adaptState gst) stt
1245 then done acc' >>= \val -> return $ Yield val (Just s)
1247 Skip s -> go SPEC s acc
1248 Stop -> done acc >>= \val -> return $ Yield val Nothing
1250 stepOuter _ _ Nothing = return Stop
1252{-# INLINE_NORMAL groupsBy #-}
1258groupsBy cmp f (Stream step state) = Stream (stepOuter f) (Just state, Nothing)
1262 {-# INLINE_LATE stepOuter #-}
1263 stepOuter (Fold fstep initial done) gst (Just st, Nothing) = do
1264 res <- step (adaptState gst) st
1271 Skip s -> return $ Skip $ (Just s, Nothing)
1276 go !_ prev stt !acc = do
1277 res <- step (adaptState gst) stt
1284 else done acc >>= \r -> return $ Yield r (Just s, Just x)
1285 Skip s -> go SPEC prev s acc
1286 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1288 stepOuter (Fold fstep initial done) gst (Just st, Just prev) = do
1290 acc' <- fstep acc prev
1295 -- XXX code duplicated from the previous equation
1297 res <- step (adaptState gst) stt
1304 else done acc >>= \r -> return $ Yield r (Just s, Just x)
1305 Skip s -> go SPEC s acc
1306 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1308 stepOuter _ _ (Nothing,_) = return Stop
1310{-# INLINE_NORMAL groupsRollingBy #-}
1311groupsRollingBy :: Monad m
1316groupsRollingBy cmp f (Stream step state) =
1317 Stream (stepOuter f) (Just state, Nothing)
1320 {-# INLINE_LATE stepOuter #-}
1321 stepOuter (Fold fstep initial done) gst (Just st, Nothing) = do
1322 res <- step (adaptState gst) st
1329 Skip s -> return $ Skip $ (Just s, Nothing)
1333 go !_ prev stt !acc = do
1334 res <- step (adaptState gst) stt
1342 done acc >>= \r -> return $ Yield r (Just s, Just x)
1343 Skip s -> go SPEC prev s acc
1344 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1346 stepOuter (Fold fstep initial done) gst (Just st, Just prev') = do
1348 acc' <- fstep acc prev'
1349 go SPEC prev' st acc'
1352 go !_ prevv stt !acc = do
1353 res <- step (adaptState gst) stt
1360 else done acc >>= \r -> return $ Yield r (Just s, Just x)
1361 Skip s -> go SPEC prevv s acc
1362 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1364 stepOuter _ _ (Nothing, _) = return Stop
1366{-# INLINE_NORMAL splitBy #-}
1367splitBy :: Monad m => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1368splitBy predicate f (Stream step state) = Stream (step' f) (Just state)
1372 {-# INLINE_LATE step' #-}
1373 step' (Fold fstep initial done) gst (Just st) = initial >>= go SPEC st
1378 res <- step (adaptState gst) stt
1382 then done acc >>= \r -> return $ Yield r (Just s)
1386 Skip s -> go SPEC s acc
1387 Stop -> done acc >>= \r -> return $ Yield r Nothing
1389 step' _ _ Nothing = return Stop
1391-- XXX requires -funfolding-use-threshold=150 in lines-unlines benchmark
1392{-# INLINE_NORMAL splitSuffixBy #-}
1393splitSuffixBy :: Monad m
1394 => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1395splitSuffixBy predicate f (Stream step state) = Stream (step' f) (Just state)
1399 {-# INLINE_LATE step' #-}
1400 step' (Fold fstep initial done) gst (Just st) = do
1401 res <- step (adaptState gst) st
1406 then done acc >>= \val -> return $ Yield val (Just s)
1411 Skip s -> return $ Skip $ Just s
1417 res <- step (adaptState gst) stt
1421 then done acc >>= \r -> return $ Yield r (Just s)
1425 Skip s -> go SPEC s acc
1426 Stop -> done acc >>= \r -> return $ Yield r Nothing
1428 step' _ _ Nothing = return Stop
1430{-# INLINE_NORMAL wordsBy #-}
1431wordsBy :: Monad m => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1432wordsBy predicate f (Stream step state) = Stream (stepOuter f) (Just state)
1436 {-# INLINE_LATE stepOuter #-}
1437 stepOuter (Fold fstep initial done) gst (Just st) = do
1438 res <- step (adaptState gst) st
1442 then return $ Skip (Just s)
1448 Skip s -> return $ Skip $ Just s
1454 res <- step (adaptState gst) stt
1458 then done acc >>= \r -> return $ Yield r (Just s)
1462 Skip s -> go SPEC s acc
1463 Stop -> done acc >>= \r -> return $ Yield r Nothing
1465 stepOuter _ _ Nothing = return Stop
1467-- String search algorithms:
1468-- http://www-igm.univ-mlv.fr/~lecroq/string/index.html
1471-- TODO can we unify the splitting operations using a splitting configuration
1472-- like in the split package.
1474data SplitStyle = Infix | Suffix | Prefix deriving (Eq, Show)
1476data SplitOptions = SplitOptions
1477 { style :: SplitStyle
1478 , withSep :: Bool -- ^ keep the separators in output
1479 -- , compact :: Bool -- ^ treat multiple consecutive separators as one
1480 -- , trimHead :: Bool -- ^ drop blank at head
1481 -- , trimTail :: Bool -- ^ drop blank at tail
1485data SplitOnState s a =
1490 | GO_KARP_RABIN s !(RB.Ring a) !(Ptr a)
1493{-# INLINE_NORMAL splitOn #-}
1495 :: forall m a b. (MonadIO m, Storable a, Enum a, Eq a)
1500splitOn patArr@Array{..} (Fold fstep initial done) (Stream step state) =
1501 Stream stepOuter GO_START
1505 patLen = A.length patArr
1506 maxIndex = patLen - 1
1507 elemBits = sizeOf (undefined :: a) * 8
1509 {-# INLINE_LATE stepOuter #-}
1510 stepOuter _ GO_START =
1512 then return $ Skip $ GO_EMPTY_PAT state
1515 r <- liftIO $ (A.unsafeIndexIO patArr 0)
1516 return $ Skip $ GO_SINGLE_PAT state r
1517 else if sizeOf (undefined :: a) * patLen
1518 <= sizeOf (undefined :: Word)
1519 then return $ Skip $ GO_SHORT_PAT state
1521 (rb, rhead) <- liftIO $ RB.new patLen
1522 return $ Skip $ GO_KARP_RABIN state rb rhead
1524 stepOuter gst (GO_SINGLE_PAT stt pat) = initial >>= go SPEC stt
1529 res <- step (adaptState gst) st
1535 return $ Yield r (GO_SINGLE_PAT s pat)
1536 else fstep acc x >>= go SPEC s
1537 Skip s -> go SPEC s acc
1538 Stop -> done acc >>= \r -> return $ Yield r GO_DONE
1540 stepOuter gst (GO_SHORT_PAT stt) = initial >>= go0 SPEC 0 (0 :: Word) stt
1545 mask = (1 `shiftL` (elemBits * patLen)) - 1
1547 addToWord wrd a = (wrd `shiftL` elemBits) .|. fromIntegral (fromEnum a)
1550 patWord = mask .&. A.foldl' addToWord 0 patArr
1552 go0 !_ !idx wrd st !acc = do
1553 res <- step (adaptState gst) st
1556 let wrd' = addToWord wrd x
1559 if wrd' .&. mask == patWord
1562 return $ Yield r (GO_SHORT_PAT s)
1563 else go1 SPEC wrd' s acc
1564 else go0 SPEC (idx + 1) wrd' s acc
1565 Skip s -> go0 SPEC idx wrd s acc
1568 then go2 wrd idx acc
1570 done acc' >>= \r -> return $ Yield r GO_DONE
1573 go1 !_ wrd st !acc = do
1574 res <- step (adaptState gst) st
1577 let wrd' = addToWord wrd x
1578 old = (mask .&. wrd) `shiftR` (elemBits * (patLen - 1))
1579 acc' <- fstep acc (toEnum $ fromIntegral old)
1580 if wrd' .&. mask == patWord
1581 then done acc' >>= \r -> return $ Yield r (GO_SHORT_PAT s)
1582 else go1 SPEC wrd' s acc'
1583 Skip s -> go1 SPEC wrd s acc
1585 acc' <- go2 wrd patLen acc
1586 done acc' >>= \r -> return $ Yield r GO_DONE
1588 go2 !wrd !n !acc | n > 0 = do
1589 let old = (mask .&. wrd) `shiftR` (elemBits * (n - 1))
1590 fstep acc (toEnum $ fromIntegral old) >>= go2 wrd (n - 1)
1591 go2 _ _ acc = return acc
1593 stepOuter gst (GO_KARP_RABIN stt rb rhead) = do
1594 initial >>= go0 SPEC 0 rhead stt
1598 k = 2891336453 :: Word32
1600 addCksum cksum a = cksum * k + fromIntegral (fromEnum a)
1601 deltaCksum cksum old new =
1602 addCksum cksum new - coeff * fromIntegral (fromEnum old)
1604 -- XXX shall we use a random starting hash or 1 instead of 0?
1605 patHash = A.foldl' addCksum 0 patArr
1608 go0 !_ !idx !rh st !acc = do
1609 res <- step (adaptState gst) st
1612 rh' <- liftIO $ RB.unsafeInsert rb rh x
1615 let fold = RB.unsafeFoldRing (RB.ringBound rb)
1616 let !ringHash = fold addCksum 0 rb
1617 if ringHash == patHash
1618 then go2 SPEC ringHash rh' s acc
1619 else go1 SPEC ringHash rh' s acc
1620 else go0 SPEC (idx + 1) rh' s acc
1621 Skip s -> go0 SPEC idx rh s acc
1623 !acc' <- if idx /= 0
1624 then RB.unsafeFoldRingM rh fstep acc rb
1626 done acc' >>= \r -> return $ Yield r GO_DONE
1628 -- XXX Theoretically this code can do 4 times faster if GHC generates
1629 -- optimal code. If we use just "(cksum' == patHash)" condition it goes
1630 -- 4x faster, as soon as we add the "RB.unsafeEqArray rb v" condition
1631 -- the generated code changes drastically and becomes 4x slower. Need
1632 -- to investigate what is going on with GHC.
1634 go1 !_ !cksum !rh st !acc = do
1635 res <- step (adaptState gst) st
1638 old <- liftIO $ peek rh
1639 let cksum' = deltaCksum cksum old x
1640 acc' <- fstep acc old
1642 if (cksum' == patHash)
1644 rh' <- liftIO (RB.unsafeInsert rb rh x)
1645 go2 SPEC cksum' rh' s acc'
1647 rh' <- liftIO (RB.unsafeInsert rb rh x)
1648 go1 SPEC cksum' rh' s acc'
1649 Skip s -> go1 SPEC cksum rh s acc
1651 acc' <- RB.unsafeFoldRingFullM rh fstep acc rb
1652 done acc' >>= \r -> return $ Yield r GO_DONE
1654 go2 !_ !cksum' !rh' s !acc' = do
1655 if RB.unsafeEqArray rb rh' patArr
1658 return $ Yield r (GO_KARP_RABIN s rb rhead)
1659 else go1 SPEC cksum' rh' s acc'
1661 stepOuter gst (GO_EMPTY_PAT st) = do
1662 res <- step (adaptState gst) st
1667 done acc' >>= \r -> return $ Yield r (GO_EMPTY_PAT s)
1668 Skip s -> return $ Skip (GO_EMPTY_PAT s)
1671 stepOuter _ GO_DONE = return Stop
1673{-# INLINE_NORMAL splitSuffixOn #-}
1675 :: forall m a b. (MonadIO m, Storable a, Enum a, Eq a)
1681splitSuffixOn withSep patArr@Array{..} (Fold fstep initial done)
1682 (Stream step state) =
1683 Stream stepOuter GO_START
1687 patLen = A.length patArr
1688 maxIndex = patLen - 1
1689 elemBits = sizeOf (undefined :: a) * 8
1691 {-# INLINE_LATE stepOuter #-}
1692 stepOuter _ GO_START =
1694 then return $ Skip $ GO_EMPTY_PAT state
1697 r <- liftIO $ (A.unsafeIndexIO patArr 0)
1698 return $ Skip $ GO_SINGLE_PAT state r
1699 else if sizeOf (undefined :: a) * patLen
1700 <= sizeOf (undefined :: Word)
1701 then return $ Skip $ GO_SHORT_PAT state
1703 (rb, rhead) <- liftIO $ RB.new patLen
1704 return $ Skip $ GO_KARP_RABIN state rb rhead
1706 stepOuter gst (GO_SINGLE_PAT stt pat) = do
1707 -- This first part is the only difference between splitOn and
1709 -- If the last element is a separator do not issue a blank segment.
1710 res <- step (adaptState gst) stt
1716 acc' <- if withSep then fstep acc x else return acc
1717 done acc' >>= \r -> return $ Yield r (GO_SINGLE_PAT s pat)
1718 else fstep acc x >>= go SPEC s
1719 Skip s -> return $ Skip $ (GO_SINGLE_PAT s pat)
1724 -- This is identical for splitOn and splitSuffixOn
1726 res <- step (adaptState gst) st
1731 acc' <- if withSep then fstep acc x else return acc
1733 return $ Yield r (GO_SINGLE_PAT s pat)
1734 else fstep acc x >>= go SPEC s
1735 Skip s -> go SPEC s acc
1736 Stop -> done acc >>= \r -> return $ Yield r GO_DONE
1738 stepOuter gst (GO_SHORT_PAT stt) = do
1740 -- Call "initial" only if the stream yields an element, otherwise we
1741 -- may call "initial" but never yield anything. initial may produce a
1742 -- side effect, therefore we will end up doing and discard a side
1747 res <- step (adaptState gst) stt
1751 let wrd' = addToWord wrd x
1752 acc' <- if withSep then fstep acc x else return acc
1755 if wrd' .&. mask == patWord
1756 then done acc' >>= \r -> return $ Yield r (GO_SHORT_PAT s)
1757 else go0 SPEC (idx + 1) wrd' s acc'
1758 else go0 SPEC (idx + 1) wrd' s acc'
1759 Skip s -> return $ Skip (GO_SHORT_PAT s)
1765 mask = (1 `shiftL` (elemBits * patLen)) - 1
1767 addToWord wrd a = (wrd `shiftL` elemBits) .|. fromIntegral (fromEnum a)
1770 patWord = mask .&. A.foldl' addToWord 0 patArr
1772 go0 !_ !idx wrd st !acc = do
1773 res <- step (adaptState gst) st
1776 let wrd' = addToWord wrd x
1777 acc' <- if withSep then fstep acc x else return acc
1780 if wrd' .&. mask == patWord
1783 return $ Yield r (GO_SHORT_PAT s)
1784 else go1 SPEC wrd' s acc'
1785 else go0 SPEC (idx + 1) wrd' s acc'
1786 Skip s -> go0 SPEC idx wrd s acc
1788 if (idx == maxIndex) && (wrd .&. mask == patWord)
1791 acc' <- if idx /= 0 && not withSep
1792 then go2 wrd idx acc
1794 done acc' >>= \r -> return $ Yield r GO_DONE
1797 go1 !_ wrd st !acc = do
1798 res <- step (adaptState gst) st
1801 let wrd' = addToWord wrd x
1802 old = (mask .&. wrd) `shiftR` (elemBits * (patLen - 1))
1805 else fstep acc (toEnum $ fromIntegral old)
1806 if wrd' .&. mask == patWord
1807 then done acc' >>= \r -> return $ Yield r (GO_SHORT_PAT s)
1808 else go1 SPEC wrd' s acc'
1809 Skip s -> go1 SPEC wrd s acc
1811 -- If the last sequence is a separator do not issue a blank
1813 if wrd .&. mask == patWord
1818 else go2 wrd patLen acc
1819 done acc' >>= \r -> return $ Yield r GO_DONE
1821 go2 !wrd !n !acc | n > 0 = do
1822 let old = (mask .&. wrd) `shiftR` (elemBits * (n - 1))
1823 fstep acc (toEnum $ fromIntegral old) >>= go2 wrd (n - 1)
1824 go2 _ _ acc = return acc
1826 stepOuter gst (GO_KARP_RABIN stt rb rhead) = do
1828 res <- step (adaptState gst) stt
1832 acc' <- if withSep then fstep acc x else return acc
1833 rh' <- liftIO (RB.unsafeInsert rb rhead x)
1836 let fold = RB.unsafeFoldRing (RB.ringBound rb)
1837 let !ringHash = fold addCksum 0 rb
1838 if ringHash == patHash
1839 then go2 SPEC ringHash rh' s acc'
1840 else go0 SPEC (idx + 1) rh' s acc'
1841 else go0 SPEC (idx + 1) rh' s acc'
1842 Skip s -> return $ Skip (GO_KARP_RABIN s rb rhead)
1847 k = 2891336453 :: Word32
1849 addCksum cksum a = cksum * k + fromIntegral (fromEnum a)
1850 deltaCksum cksum old new =
1851 addCksum cksum new - coeff * fromIntegral (fromEnum old)
1853 -- XXX shall we use a random starting hash or 1 instead of 0?
1854 patHash = A.foldl' addCksum 0 patArr
1857 go0 !_ !idx !rh st !acc = do
1858 res <- step (adaptState gst) st
1861 acc' <- if withSep then fstep acc x else return acc
1862 rh' <- liftIO (RB.unsafeInsert rb rh x)
1865 let fold = RB.unsafeFoldRing (RB.ringBound rb)
1866 let !ringHash = fold addCksum 0 rb
1867 if ringHash == patHash
1868 then go2 SPEC ringHash rh' s acc'
1869 else go1 SPEC ringHash rh' s acc'
1870 else go0 SPEC (idx + 1) rh' s acc'
1871 Skip s -> go0 SPEC idx rh s acc
1873 -- do not issue a blank segment when we end at pattern
1874 if (idx == maxIndex) && RB.unsafeEqArray rb rh patArr
1877 !acc' <- if idx /= 0 && not withSep
1878 then RB.unsafeFoldRingM rh fstep acc rb
1880 done acc' >>= \r -> return $ Yield r GO_DONE
1882 -- XXX Theoretically this code can do 4 times faster if GHC generates
1883 -- optimal code. If we use just "(cksum' == patHash)" condition it goes
1884 -- 4x faster, as soon as we add the "RB.unsafeEqArray rb v" condition
1885 -- the generated code changes drastically and becomes 4x slower. Need
1886 -- to investigate what is going on with GHC.
1888 go1 !_ !cksum !rh st !acc = do
1889 res <- step (adaptState gst) st
1892 old <- liftIO $ peek rh
1893 let cksum' = deltaCksum cksum old x
1898 if (cksum' == patHash)
1900 rh' <- liftIO (RB.unsafeInsert rb rh x)
1901 go2 SPEC cksum' rh' s acc'
1903 rh' <- liftIO (RB.unsafeInsert rb rh x)
1904 go1 SPEC cksum' rh' s acc'
1905 Skip s -> go1 SPEC cksum rh s acc
1907 if RB.unsafeEqArray rb rh patArr
1912 else RB.unsafeFoldRingFullM rh fstep acc rb
1913 done acc' >>= \r -> return $ Yield r GO_DONE
1915 go2 !_ !cksum' !rh' s !acc' = do
1916 if RB.unsafeEqArray rb rh' patArr
1919 return $ Yield r (GO_KARP_RABIN s rb rhead)
1920 else go1 SPEC cksum' rh' s acc'
1922 stepOuter gst (GO_EMPTY_PAT st) = do
1923 res <- step (adaptState gst) st
1928 done acc' >>= \r -> return $ Yield r (GO_EMPTY_PAT s)
1929 Skip s -> return $ Skip (GO_EMPTY_PAT s)
1932 stepOuter _ GO_DONE = return Stop
1934data SplitState s arr
1936 | SplitBuffering s arr
1937 | SplitSplitting s arr
1938 | SplitYielding arr (SplitState s arr)
1941-- XXX An alternative approach would be to use a partial fold (Fold m a b) to
1942-- split using a splitBy like combinator. The Fold would consume upto the
1943-- separator and return any leftover which can then be fed to the next fold.
1945-- We can revisit this once we have partial folds/parsers.
1947-- | Performs infix separator style splitting.
1948{-# INLINE_NORMAL splitInnerBy #-}
1951 => (f a -> m (f a, Maybe (f a))) -- splitter
1952 -> (f a -> f a -> m (f a)) -- joiner
1955splitInnerBy splitter joiner (Stream step1 state1) =
1956 (Stream step (SplitInitial state1))
1960 {-# INLINE_LATE step #-}
1961 step gst (SplitInitial st) = do
1965 (x1, mx2) <- splitter x
1966 return $ case mx2 of
1967 Nothing -> Skip (SplitBuffering s x1)
1968 Just x2 -> Skip (SplitYielding x1 (SplitSplitting s x2))
1969 Skip s -> return $ Skip (SplitInitial s)
1970 Stop -> return $ Stop
1972 step gst (SplitBuffering st buf) = do
1976 (x1, mx2) <- splitter x
1977 buf' <- joiner buf x1
1978 return $ case mx2 of
1979 Nothing -> Skip (SplitBuffering s buf')
1980 Just x2 -> Skip (SplitYielding buf' (SplitSplitting s x2))
1981 Skip s -> return $ Skip (SplitBuffering s buf)
1982 Stop -> return $ Skip (SplitYielding buf SplitFinishing)
1984 step _ (SplitSplitting st buf) = do
1985 (x1, mx2) <- splitter buf
1986 return $ case mx2 of
1987 Nothing -> Skip $ SplitBuffering st x1
1988 Just x2 -> Skip $ SplitYielding x1 (SplitSplitting st x2)
1990 step _ (SplitYielding x next) = return $ Yield x next
1991 step _ SplitFinishing = return $ Stop
1993-- | Performs infix separator style splitting.
1994{-# INLINE_NORMAL splitInnerBySuffix #-}
1996 :: (Monad m, Eq (f a), Monoid (f a))
1997 => (f a -> m (f a, Maybe (f a))) -- splitter
1998 -> (f a -> f a -> m (f a)) -- joiner
2001splitInnerBySuffix splitter joiner (Stream step1 state1) =
2002 (Stream step (SplitInitial state1))
2006 {-# INLINE_LATE step #-}
2007 step gst (SplitInitial st) = do
2011 (x1, mx2) <- splitter x
2012 return $ case mx2 of
2013 Nothing -> Skip (SplitBuffering s x1)
2014 Just x2 -> Skip (SplitYielding x1 (SplitSplitting s x2))
2015 Skip s -> return $ Skip (SplitInitial s)
2016 Stop -> return $ Stop
2018 step gst (SplitBuffering st buf) = do
2022 (x1, mx2) <- splitter x
2023 buf' <- joiner buf x1
2024 return $ case mx2 of
2025 Nothing -> Skip (SplitBuffering s buf')
2026 Just x2 -> Skip (SplitYielding buf' (SplitSplitting s x2))
2027 Skip s -> return $ Skip (SplitBuffering s buf)
2031 else Skip (SplitYielding buf SplitFinishing)
2033 step _ (SplitSplitting st buf) = do
2034 (x1, mx2) <- splitter buf
2035 return $ case mx2 of
2036 Nothing -> Skip $ SplitBuffering st x1
2037 Just x2 -> Skip $ SplitYielding x1 (SplitSplitting st x2)
2039 step _ (SplitYielding x next) = return $ Yield x next
2040 step _ SplitFinishing = return $ Stop
2042------------------------------------------------------------------------------
2044------------------------------------------------------------------------------
2046{-# INLINE_NORMAL isPrefixOf #-}
2047isPrefixOf :: (Eq a, Monad m) => Stream m a -> Stream m a -> m Bool
2048isPrefixOf (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
2050 go (sa, sb, Nothing) = do
2051 r <- stepa defState sa
2053 Yield x sa' -> go (sa', sb, Just x)
2054 Skip sa' -> go (sa', sb, Nothing)
2057 go (sa, sb, Just x) = do
2058 r <- stepb defState sb
2062 then go (sa, sb', Nothing)
2064 Skip sb' -> go (sa, sb', Just x)
2065 Stop -> return False
2067{-# INLINE_NORMAL isSubsequenceOf #-}
2068isSubsequenceOf :: (Eq a, Monad m) => Stream m a -> Stream m a -> m Bool
2069isSubsequenceOf (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
2071 go (sa, sb, Nothing) = do
2072 r <- stepa defState sa
2074 Yield x sa' -> go (sa', sb, Just x)
2075 Skip sa' -> go (sa', sb, Nothing)
2078 go (sa, sb, Just x) = do
2079 r <- stepb defState sb
2083 then go (sa, sb', Nothing)
2084 else go (sa, sb', Just x)
2085 Skip sb' -> go (sa, sb', Just x)
2086 Stop -> return False
2088{-# INLINE_NORMAL stripPrefix #-}
2091 => Stream m a -> Stream m a -> m (Maybe (Stream m a))
2092stripPrefix (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
2094 go (sa, sb, Nothing) = do
2095 r <- stepa defState sa
2097 Yield x sa' -> go (sa', sb, Just x)
2098 Skip sa' -> go (sa', sb, Nothing)
2099 Stop -> return $ Just (Stream stepb sb)
2101 go (sa, sb, Just x) = do
2102 r <- stepb defState sb
2106 then go (sa, sb', Nothing)
2108 Skip sb' -> go (sa, sb', Just x)
2109 Stop -> return Nothing
2111------------------------------------------------------------------------------
2113------------------------------------------------------------------------------
2115-- | Execute a monadic action for each element of the 'Stream'
2116{-# INLINE_NORMAL mapM_ #-}
2117mapM_ :: Monad m => (a -> m b) -> Stream m a -> m ()
2118mapM_ m = drain . mapM m
2120-------------------------------------------------------------------------------
2121-- Stream transformations using Unfolds
2122-------------------------------------------------------------------------------
2124-- Define a unique structure to use in inspection testing
2125data ConcatMapUState o i =
2127 | ConcatMapUInner o i
2129-- | @concatMapU unfold stream@ uses @unfold@ to map the input stream elements
2130-- to streams and then flattens the generated streams into a single output
2133-- This is like 'concatMap' but uses an unfold with an explicit state to
2134-- generate the stream instead of a 'Stream' type generator. This allows better
2135-- optimization via fusion. This can be many times more efficient than
2138{-# INLINE_NORMAL concatMapU #-}
2139concatMapU :: Monad m => Unfold m a b -> Stream m a -> Stream m b
2140concatMapU (Unfold istep inject) (Stream ostep ost) =
2141 Stream step (ConcatMapUOuter ost)
2143 {-# INLINE_LATE step #-}
2144 step gst (ConcatMapUOuter o) = do
2145 r <- ostep (adaptState gst) o
2149 i `seq` return (Skip (ConcatMapUInner o' i))
2150 Skip o' -> return $ Skip (ConcatMapUOuter o')
2151 Stop -> return $ Stop
2153 step _ (ConcatMapUInner o i) = do
2156 Yield x i' -> Yield x (ConcatMapUInner o i')
2157 Skip i' -> Skip (ConcatMapUInner o i')
2158 Stop -> Skip (ConcatMapUOuter o)
2160data ConcatUnfoldInterleaveState o i =
2161 ConcatUnfoldInterleaveOuter o [i]
2162 | ConcatUnfoldInterleaveInner o [i]
2163 | ConcatUnfoldInterleaveInnerL [i] [i]
2164 | ConcatUnfoldInterleaveInnerR [i] [i]
2166-- XXX use arrays to store state instead of lists.
2167-- XXX In general we can use different scheduling strategies e.g. how to
2168-- schedule the outer vs inner loop or assigning weights to different streams
2169-- or outer and inner loops.
2171-- After a yield, switch to the next stream. Do not switch streams on Skip.
2172-- Yield from outer stream switches to the inner stream.
2174-- There are two choices here, (1) exhaust the outer stream first and then
2175-- start yielding from the inner streams, this is much simpler to implement,
2176-- (2) yield at least one element from an inner stream before going back to
2177-- outer stream and opening the next stream from it.
2179-- Ideally, we need some scheduling bias to inner streams vs outer stream.
2180-- Maybe we can configure the behavior.
2182{-# INLINE_NORMAL concatUnfoldInterleave #-}
2183concatUnfoldInterleave :: Monad m => Unfold m a b -> Stream m a -> Stream m b
2184concatUnfoldInterleave (Unfold istep inject) (Stream ostep ost) =
2185 Stream step (ConcatUnfoldInterleaveOuter ost [])
2187 {-# INLINE_LATE step #-}
2188 step gst (ConcatUnfoldInterleaveOuter o ls) = do
2189 r <- ostep (adaptState gst) o
2193 i `seq` return (Skip (ConcatUnfoldInterleaveInner o' (i : ls)))
2194 Skip o' -> return $ Skip (ConcatUnfoldInterleaveOuter o' ls)
2195 Stop -> return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2197 step _ (ConcatUnfoldInterleaveInner _ []) = undefined
2198 step _ (ConcatUnfoldInterleaveInner o (st:ls)) = do
2201 Yield x s -> Yield x (ConcatUnfoldInterleaveOuter o (s:ls))
2202 Skip s -> Skip (ConcatUnfoldInterleaveInner o (s:ls))
2203 Stop -> Skip (ConcatUnfoldInterleaveOuter o ls)
2205 step _ (ConcatUnfoldInterleaveInnerL [] []) = return Stop
2206 step _ (ConcatUnfoldInterleaveInnerL [] rs) =
2207 return $ Skip (ConcatUnfoldInterleaveInnerR [] rs)
2209 step _ (ConcatUnfoldInterleaveInnerL (st:ls) rs) = do
2212 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerL ls (s:rs))
2213 Skip s -> Skip (ConcatUnfoldInterleaveInnerL (s:ls) rs)
2214 Stop -> Skip (ConcatUnfoldInterleaveInnerL ls rs)
2216 step _ (ConcatUnfoldInterleaveInnerR [] []) = return Stop
2217 step _ (ConcatUnfoldInterleaveInnerR ls []) =
2218 return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2220 step _ (ConcatUnfoldInterleaveInnerR ls (st:rs)) = do
2223 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerR (s:ls) rs)
2224 Skip s -> Skip (ConcatUnfoldInterleaveInnerR ls (s:rs))
2225 Stop -> Skip (ConcatUnfoldInterleaveInnerR ls rs)
2227-- XXX In general we can use different scheduling strategies e.g. how to
2228-- schedule the outer vs inner loop or assigning weights to different streams
2229-- or outer and inner loops.
2231-- This could be inefficient if the tasks are too small.
2233-- Compared to concatUnfoldInterleave this one switches streams on Skips.
2235{-# INLINE_NORMAL concatUnfoldRoundrobin #-}
2236concatUnfoldRoundrobin :: Monad m => Unfold m a b -> Stream m a -> Stream m b
2237concatUnfoldRoundrobin (Unfold istep inject) (Stream ostep ost) =
2238 Stream step (ConcatUnfoldInterleaveOuter ost [])
2240 {-# INLINE_LATE step #-}
2241 step gst (ConcatUnfoldInterleaveOuter o ls) = do
2242 r <- ostep (adaptState gst) o
2246 i `seq` return (Skip (ConcatUnfoldInterleaveInner o' (i : ls)))
2247 Skip o' -> return $ Skip (ConcatUnfoldInterleaveInner o' ls)
2248 Stop -> return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2250 step _ (ConcatUnfoldInterleaveInner o []) =
2251 return $ Skip (ConcatUnfoldInterleaveOuter o [])
2253 step _ (ConcatUnfoldInterleaveInner o (st:ls)) = do
2256 Yield x s -> Yield x (ConcatUnfoldInterleaveOuter o (s:ls))
2257 Skip s -> Skip (ConcatUnfoldInterleaveOuter o (s:ls))
2258 Stop -> Skip (ConcatUnfoldInterleaveOuter o ls)
2260 step _ (ConcatUnfoldInterleaveInnerL [] []) = return Stop
2261 step _ (ConcatUnfoldInterleaveInnerL [] rs) =
2262 return $ Skip (ConcatUnfoldInterleaveInnerR [] rs)
2264 step _ (ConcatUnfoldInterleaveInnerL (st:ls) rs) = do
2267 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerL ls (s:rs))
2268 Skip s -> Skip (ConcatUnfoldInterleaveInnerL ls (s:rs))
2269 Stop -> Skip (ConcatUnfoldInterleaveInnerL ls rs)
2271 step _ (ConcatUnfoldInterleaveInnerR [] []) = return Stop
2272 step _ (ConcatUnfoldInterleaveInnerR ls []) =
2273 return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2275 step _ (ConcatUnfoldInterleaveInnerR ls (st:rs)) = do
2278 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerR (s:ls) rs)
2279 Skip s -> Skip (ConcatUnfoldInterleaveInnerR (s:ls) rs)
2280 Stop -> Skip (ConcatUnfoldInterleaveInnerR ls rs)
2282data AppendState s1 s2 = AppendFirst s1 | AppendSecond s2
2284-- Note that this could be much faster compared to the CPS stream. However, as
2285-- the number of streams being composed increases this may become expensive.
2286-- Need to see where the breaking point is between the two.
2288{-# INLINE_NORMAL append #-}
2289append :: Monad m => Stream m a -> Stream m a -> Stream m a
2290append (Stream step1 state1) (Stream step2 state2) =
2291 Stream step (AppendFirst state1)
2295 {-# INLINE_LATE step #-}
2296 step gst (AppendFirst st) = do
2299 Yield a s -> Yield a (AppendFirst s)
2300 Skip s -> Skip (AppendFirst s)
2301 Stop -> Skip (AppendSecond state2)
2303 step gst (AppendSecond st) = do
2306 Yield a s -> Yield a (AppendSecond s)
2307 Skip s -> Skip (AppendSecond s)
2310data InterleaveState s1 s2 = InterleaveFirst s1 s2 | InterleaveSecond s1 s2
2311 | InterleaveSecondOnly s2 | InterleaveFirstOnly s1
2313{-# INLINE_NORMAL interleave #-}
2314interleave :: Monad m => Stream m a -> Stream m a -> Stream m a
2315interleave (Stream step1 state1) (Stream step2 state2) =
2316 Stream step (InterleaveFirst state1 state2)
2320 {-# INLINE_LATE step #-}
2321 step gst (InterleaveFirst st1 st2) = do
2324 Yield a s -> Yield a (InterleaveSecond s st2)
2325 Skip s -> Skip (InterleaveFirst s st2)
2326 Stop -> Skip (InterleaveSecondOnly st2)
2328 step gst (InterleaveSecond st1 st2) = do
2331 Yield a s -> Yield a (InterleaveFirst st1 s)
2332 Skip s -> Skip (InterleaveSecond st1 s)
2333 Stop -> Skip (InterleaveFirstOnly st1)
2335 step gst (InterleaveFirstOnly st1) = do
2338 Yield a s -> Yield a (InterleaveFirstOnly s)
2339 Skip s -> Skip (InterleaveFirstOnly s)
2342 step gst (InterleaveSecondOnly st2) = do
2345 Yield a s -> Yield a (InterleaveSecondOnly s)
2346 Skip s -> Skip (InterleaveSecondOnly s)
2349{-# INLINE_NORMAL interleaveMin #-}
2350interleaveMin :: Monad m => Stream m a -> Stream m a -> Stream m a
2351interleaveMin (Stream step1 state1) (Stream step2 state2) =
2352 Stream step (InterleaveFirst state1 state2)
2356 {-# INLINE_LATE step #-}
2357 step gst (InterleaveFirst st1 st2) = do
2360 Yield a s -> Yield a (InterleaveSecond s st2)
2361 Skip s -> Skip (InterleaveFirst s st2)
2364 step gst (InterleaveSecond st1 st2) = do
2367 Yield a s -> Yield a (InterleaveFirst st1 s)
2368 Skip s -> Skip (InterleaveSecond st1 s)
2371 step _ (InterleaveFirstOnly _) = undefined
2372 step _ (InterleaveSecondOnly _) = undefined
2374{-# INLINE_NORMAL interleaveSuffix #-}
2375interleaveSuffix :: Monad m => Stream m a -> Stream m a -> Stream m a
2376interleaveSuffix (Stream step1 state1) (Stream step2 state2) =
2377 Stream step (InterleaveFirst state1 state2)
2381 {-# INLINE_LATE step #-}
2382 step gst (InterleaveFirst st1 st2) = do
2385 Yield a s -> Yield a (InterleaveSecond s st2)
2386 Skip s -> Skip (InterleaveFirst s st2)
2389 step gst (InterleaveSecond st1 st2) = do
2392 Yield a s -> Yield a (InterleaveFirst st1 s)
2393 Skip s -> Skip (InterleaveSecond st1 s)
2394 Stop -> Skip (InterleaveFirstOnly st1)
2396 step gst (InterleaveFirstOnly st1) = do
2399 Yield a s -> Yield a (InterleaveFirstOnly s)
2400 Skip s -> Skip (InterleaveFirstOnly s)
2403 step _ (InterleaveSecondOnly _) = undefined
2405data InterleaveInfixState s1 s2 a
2406 = InterleaveInfixFirst s1 s2
2407 | InterleaveInfixSecondBuf s1 s2
2408 | InterleaveInfixSecondYield s1 s2 a
2409 | InterleaveInfixFirstYield s1 s2 a
2410 | InterleaveInfixFirstOnly s1
2412{-# INLINE_NORMAL interleaveInfix #-}
2413interleaveInfix :: Monad m => Stream m a -> Stream m a -> Stream m a
2414interleaveInfix (Stream step1 state1) (Stream step2 state2) =
2415 Stream step (InterleaveInfixFirst state1 state2)
2419 {-# INLINE_LATE step #-}
2420 step gst (InterleaveInfixFirst st1 st2) = do
2423 Yield a s -> Yield a (InterleaveInfixSecondBuf s st2)
2424 Skip s -> Skip (InterleaveInfixFirst s st2)
2427 step gst (InterleaveInfixSecondBuf st1 st2) = do
2430 Yield a s -> Skip (InterleaveInfixSecondYield st1 s a)
2431 Skip s -> Skip (InterleaveInfixSecondBuf st1 s)
2432 Stop -> Skip (InterleaveInfixFirstOnly st1)
2434 step gst (InterleaveInfixSecondYield st1 st2 x) = do
2437 Yield a s -> Yield x (InterleaveInfixFirstYield s st2 a)
2438 Skip s -> Skip (InterleaveInfixSecondYield s st2 x)
2441 step _ (InterleaveInfixFirstYield st1 st2 x) = do
2442 return $ Yield x (InterleaveInfixSecondBuf st1 st2)
2444 step gst (InterleaveInfixFirstOnly st1) = do
2447 Yield a s -> Yield a (InterleaveInfixFirstOnly s)
2448 Skip s -> Skip (InterleaveInfixFirstOnly s)
2451{-# INLINE_NORMAL roundRobin #-}
2452roundRobin :: Monad m => Stream m a -> Stream m a -> Stream m a
2453roundRobin (Stream step1 state1) (Stream step2 state2) =
2454 Stream step (InterleaveFirst state1 state2)
2458 {-# INLINE_LATE step #-}
2459 step gst (InterleaveFirst st1 st2) = do
2462 Yield a s -> Yield a (InterleaveSecond s st2)
2463 Skip s -> Skip (InterleaveSecond s st2)
2464 Stop -> Skip (InterleaveSecondOnly st2)
2466 step gst (InterleaveSecond st1 st2) = do
2469 Yield a s -> Yield a (InterleaveFirst st1 s)
2470 Skip s -> Skip (InterleaveFirst st1 s)
2471 Stop -> Skip (InterleaveFirstOnly st1)
2473 step gst (InterleaveSecondOnly st2) = do
2476 Yield a s -> Yield a (InterleaveSecondOnly s)
2477 Skip s -> Skip (InterleaveSecondOnly s)
2480 step gst (InterleaveFirstOnly st1) = do
2483 Yield a s -> Yield a (InterleaveFirstOnly s)
2484 Skip s -> Skip (InterleaveFirstOnly s)
2487data ICUState s1 s2 i1 i2 =
2492 | ICUFirstInner s1 s2 i1
2493 | ICUSecondInner s1 s2 i2
2494 | ICUFirstOnlyInner s1 i1
2495 | ICUSecondOnlyInner s2 i2
2497-- | Interleave streams (full streams, not the elements) unfolded from two
2498-- input streams and concat. Stop when the first stream stops. If the second
2499-- stream ends before the first one then first stream still keeps running alone
2500-- without any interleaving with the second stream.
2502-- [a1, a2, ... an] [b1, b2 ...]
2503-- => [streamA1, streamA2, ... streamAn] [streamB1, streamB2, ...]
2504-- => [streamA1, streamB1, streamA2...StreamAn, streamBn]
2505-- => [a11, a12, ...a1j, b11, b12, ...b1k, a21, a22, ...]
2507{-# INLINE_NORMAL gintercalateSuffix #-}
2510 => Unfold m a c -> Stream m a -> Unfold m b c -> Stream m b -> Stream m c
2512 (Unfold istep1 inject1) (Stream step1 state1)
2513 (Unfold istep2 inject2) (Stream step2 state2) =
2514 Stream step (ICUFirst state1 state2)
2518 {-# INLINE_LATE step #-}
2519 step gst (ICUFirst s1 s2) = do
2520 r <- step1 (adaptState gst) s1
2524 i `seq` return (Skip (ICUFirstInner s s2 i))
2525 Skip s -> return $ Skip (ICUFirst s s2)
2528 step gst (ICUFirstOnly s1) = do
2529 r <- step1 (adaptState gst) s1
2533 i `seq` return (Skip (ICUFirstOnlyInner s i))
2534 Skip s -> return $ Skip (ICUFirstOnly s)
2537 step _ (ICUFirstInner s1 s2 i1) = do
2540 Yield x i' -> Yield x (ICUFirstInner s1 s2 i')
2541 Skip i' -> Skip (ICUFirstInner s1 s2 i')
2542 Stop -> Skip (ICUSecond s1 s2)
2544 step _ (ICUFirstOnlyInner s1 i1) = do
2547 Yield x i' -> Yield x (ICUFirstOnlyInner s1 i')
2548 Skip i' -> Skip (ICUFirstOnlyInner s1 i')
2549 Stop -> Skip (ICUFirstOnly s1)
2551 step gst (ICUSecond s1 s2) = do
2552 r <- step2 (adaptState gst) s2
2556 i `seq` return (Skip (ICUSecondInner s1 s i))
2557 Skip s -> return $ Skip (ICUSecond s1 s)
2558 Stop -> return $ Skip (ICUFirstOnly s1)
2560 step _ (ICUSecondInner s1 s2 i2) = do
2563 Yield x i' -> Yield x (ICUSecondInner s1 s2 i')
2564 Skip i' -> Skip (ICUSecondInner s1 s2 i')
2565 Stop -> Skip (ICUFirst s1 s2)
2567 step _ (ICUSecondOnly _s2) = undefined
2568 step _ (ICUSecondOnlyInner _s2 _i2) = undefined
2570data InterposeSuffixState s1 i1 =
2571 InterposeSuffixFirst s1
2572 -- | InterposeSuffixFirstYield s1 i1
2573 | InterposeSuffixFirstInner s1 i1
2574 | InterposeSuffixSecond s1
2576-- Note that if an unfolded layer turns out to be nil we still emit the
2577-- separator effect. An alternate behavior could be to emit the separator
2578-- effect only if at least one element has been yielded by the unfolding.
2579-- However, that becomes a bit complicated, so we have chosen the former
2581{-# INLINE_NORMAL interposeSuffix #-}
2584 => m c -> Unfold m b c -> Stream m b -> Stream m c
2587 (Unfold istep1 inject1) (Stream step1 state1) =
2588 Stream step (InterposeSuffixFirst state1)
2592 {-# INLINE_LATE step #-}
2593 step gst (InterposeSuffixFirst s1) = do
2594 r <- step1 (adaptState gst) s1
2598 i `seq` return (Skip (InterposeSuffixFirstInner s i))
2599 -- i `seq` return (Skip (InterposeSuffixFirstYield s i))
2600 Skip s -> return $ Skip (InterposeSuffixFirst s)
2604 step _ (InterposeSuffixFirstYield s1 i1) = do
2607 Yield x i' -> Yield x (InterposeSuffixFirstInner s1 i')
2608 Skip i' -> Skip (InterposeSuffixFirstYield s1 i')
2609 Stop -> Skip (InterposeSuffixFirst s1)
2612 step _ (InterposeSuffixFirstInner s1 i1) = do
2615 Yield x i' -> Yield x (InterposeSuffixFirstInner s1 i')
2616 Skip i' -> Skip (InterposeSuffixFirstInner s1 i')
2617 Stop -> Skip (InterposeSuffixSecond s1)
2619 step _ (InterposeSuffixSecond s1) = do
2621 return $ Yield r (InterposeSuffixFirst s1)
2623data ICALState s1 s2 i1 i2 a =
2625 -- | ICALFirstYield s1 s2 i1
2626 | ICALFirstInner s1 s2 i1
2628 | ICALFirstOnlyInner s1 i1
2629 | ICALSecondInject s1 s2
2630 | ICALFirstInject s1 s2 i2
2631 -- | ICALFirstBuf s1 s2 i1 i2
2632 | ICALSecondInner s1 s2 i1 i2
2633 -- -- | ICALSecondInner s1 s2 i1 i2 a
2634 -- -- | ICALFirstResume s1 s2 i1 i2 a
2636-- | Interleave streams (full streams, not the elements) unfolded from two
2637-- input streams and concat. Stop when the first stream stops. If the second
2638-- stream ends before the first one then first stream still keeps running alone
2639-- without any interleaving with the second stream.
2641-- [a1, a2, ... an] [b1, b2 ...]
2642-- => [streamA1, streamA2, ... streamAn] [streamB1, streamB2, ...]
2643-- => [streamA1, streamB1, streamA2...StreamAn, streamBn]
2644-- => [a11, a12, ...a1j, b11, b12, ...b1k, a21, a22, ...]
2646{-# INLINE_NORMAL gintercalate #-}
2649 => Unfold m a c -> Stream m a -> Unfold m b c -> Stream m b -> Stream m c
2651 (Unfold istep1 inject1) (Stream step1 state1)
2652 (Unfold istep2 inject2) (Stream step2 state2) =
2653 Stream step (ICALFirst state1 state2)
2657 {-# INLINE_LATE step #-}
2658 step gst (ICALFirst s1 s2) = do
2659 r <- step1 (adaptState gst) s1
2663 i `seq` return (Skip (ICALFirstInner s s2 i))
2664 -- i `seq` return (Skip (ICALFirstYield s s2 i))
2665 Skip s -> return $ Skip (ICALFirst s s2)
2669 step _ (ICALFirstYield s1 s2 i1) = do
2672 Yield x i' -> Yield x (ICALFirstInner s1 s2 i')
2673 Skip i' -> Skip (ICALFirstYield s1 s2 i')
2674 Stop -> Skip (ICALFirst s1 s2)
2677 step _ (ICALFirstInner s1 s2 i1) = do
2680 Yield x i' -> Yield x (ICALFirstInner s1 s2 i')
2681 Skip i' -> Skip (ICALFirstInner s1 s2 i')
2682 Stop -> Skip (ICALSecondInject s1 s2)
2684 step gst (ICALFirstOnly s1) = do
2685 r <- step1 (adaptState gst) s1
2689 i `seq` return (Skip (ICALFirstOnlyInner s i))
2690 Skip s -> return $ Skip (ICALFirstOnly s)
2693 step _ (ICALFirstOnlyInner s1 i1) = do
2696 Yield x i' -> Yield x (ICALFirstOnlyInner s1 i')
2697 Skip i' -> Skip (ICALFirstOnlyInner s1 i')
2698 Stop -> Skip (ICALFirstOnly s1)
2700 -- We inject the second stream even before checking if the first stream
2701 -- would yield any more elements. There is no clear choice whether we
2702 -- should do this before or after that. Doing it after may make the state
2703 -- machine a bit simpler though.
2704 step gst (ICALSecondInject s1 s2) = do
2705 r <- step2 (adaptState gst) s2
2709 i `seq` return (Skip (ICALFirstInject s1 s i))
2710 Skip s -> return $ Skip (ICALSecondInject s1 s)
2711 Stop -> return $ Skip (ICALFirstOnly s1)
2713 step gst (ICALFirstInject s1 s2 i2) = do
2714 r <- step1 (adaptState gst) s1
2718 i `seq` return (Skip (ICALSecondInner s s2 i i2))
2719 -- i `seq` return (Skip (ICALFirstBuf s s2 i i2))
2720 Skip s -> return $ Skip (ICALFirstInject s s2 i2)
2724 step _ (ICALFirstBuf s1 s2 i1 i2) = do
2727 Yield x i' -> Skip (ICALSecondInner s1 s2 i' i2 x)
2728 Skip i' -> Skip (ICALFirstBuf s1 s2 i' i2)
2731 step _ (ICALSecondInner s1 s2 i1 i2 v) = do
2734 Yield x i' -> Yield x (ICALSecondInner s1 s2 i1 i' v)
2735 Skip i' -> Skip (ICALSecondInner s1 s2 i1 i' v)
2736 Stop -> Skip (ICALFirstResume s1 s2 i1 i2 v)
2739 step _ (ICALSecondInner s1 s2 i1 i2) = do
2742 Yield x i' -> Yield x (ICALSecondInner s1 s2 i1 i')
2743 Skip i' -> Skip (ICALSecondInner s1 s2 i1 i')
2744 Stop -> Skip (ICALFirstInner s1 s2 i1)
2745 -- Stop -> Skip (ICALFirstResume s1 s2 i1 i2)
2748 step _ (ICALFirstResume s1 s2 i1 i2 x) = do
2749 return $ Yield x (ICALFirstInner s1 s2 i1 i2)
2752data InterposeState s1 i1 a =
2754 -- | InterposeFirstYield s1 i1
2755 | InterposeFirstInner s1 i1
2756 | InterposeFirstInject s1
2757 -- | InterposeFirstBuf s1 i1
2758 | InterposeSecondYield s1 i1
2759 -- -- | InterposeSecondYield s1 i1 a
2760 -- -- | InterposeFirstResume s1 i1 a
2762-- Note that this only interposes the pure values, we may run many effects to
2763-- generate those values as some effects may not generate anything (Skip).
2764{-# INLINE_NORMAL interpose #-}
2765interpose :: Monad m => m c -> Unfold m b c -> Stream m b -> Stream m c
2768 (Unfold istep1 inject1) (Stream step1 state1) =
2769 Stream step (InterposeFirst state1)
2773 {-# INLINE_LATE step #-}
2774 step gst (InterposeFirst s1) = do
2775 r <- step1 (adaptState gst) s1
2779 i `seq` return (Skip (InterposeFirstInner s i))
2780 -- i `seq` return (Skip (InterposeFirstYield s i))
2781 Skip s -> return $ Skip (InterposeFirst s)
2785 step _ (InterposeFirstYield s1 i1) = do
2788 Yield x i' -> Yield x (InterposeFirstInner s1 i')
2789 Skip i' -> Skip (InterposeFirstYield s1 i')
2790 Stop -> Skip (InterposeFirst s1)
2793 step _ (InterposeFirstInner s1 i1) = do
2796 Yield x i' -> Yield x (InterposeFirstInner s1 i')
2797 Skip i' -> Skip (InterposeFirstInner s1 i')
2798 Stop -> Skip (InterposeFirstInject s1)
2800 step gst (InterposeFirstInject s1) = do
2801 r <- step1 (adaptState gst) s1
2805 -- i `seq` return (Skip (InterposeFirstBuf s i))
2806 i `seq` return (Skip (InterposeSecondYield s i))
2807 Skip s -> return $ Skip (InterposeFirstInject s)
2811 step _ (InterposeFirstBuf s1 i1) = do
2814 Yield x i' -> Skip (InterposeSecondYield s1 i' x)
2815 Skip i' -> Skip (InterposeFirstBuf s1 i')
2820 step _ (InterposeSecondYield s1 i1 v) = do
2822 return $ Yield r (InterposeFirstResume s1 i1 v)
2824 step _ (InterposeSecondYield s1 i1) = do
2826 return $ Yield r (InterposeFirstInner s1 i1)
2829 step _ (InterposeFirstResume s1 i1 v) = do
2830 return $ Yield v (InterposeFirstInner s1 i1)
2833------------------------------------------------------------------------------
2835------------------------------------------------------------------------------
2837data GbracketState s1 s2 v
2839 | GBracketNormal s1 v
2840 | GBracketException s2
2842-- | The most general bracketing and exception combinator. All other
2843-- combinators can be expressed in terms of this combinator. This can also be
2844-- used for cases which are not covered by the standard combinators.
2848{-# INLINE_NORMAL gbracket #-}
2852 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
2853 -> (c -> m d) -- ^ after, on normal stop
2854 -> (c -> e -> Stream m b) -- ^ on exception
2855 -> (c -> Stream m b) -- ^ stream generator
2857gbracket bef exc aft fexc fnormal =
2858 Stream step GBracketInit
2862 {-# INLINE_LATE step #-}
2863 step _ GBracketInit = do
2865 return $ Skip $ GBracketNormal (fnormal r) r
2867 step gst (GBracketNormal (UnStream step1 st) v) = do
2868 res <- exc $ step1 gst st
2870 Right r -> case r of
2872 return $ Yield x (GBracketNormal (Stream step1 s) v)
2873 Skip s -> return $ Skip (GBracketNormal (Stream step1 s) v)
2874 Stop -> aft v >> return Stop
2875 Left e -> return $ Skip (GBracketException (fexc v e))
2876 step gst (GBracketException (UnStream step1 st)) = do
2879 Yield x s -> return $ Yield x (GBracketException (Stream step1 s))
2880 Skip s -> return $ Skip (GBracketException (Stream step1 s))
2883-- | Create an IORef holding a finalizer that is called automatically when the
2884-- IORef is garbage collected. The IORef can be written to with a 'Nothing'
2885-- value to deactivate the finalizer.
2886newFinalizedIORef :: (MonadIO m, MonadBaseControl IO m)
2887 => m a -> m (IORef (Maybe (IO ())))
2888newFinalizedIORef finalizer = do
2889 mrun <- captureMonadState
2890 ref <- liftIO $ newIORef $ Just $ liftIO $ void $ do
2891 _ <- runInIO mrun finalizer
2894 res <- readIORef ref
2896 Nothing -> return ()
2898 _ <- liftIO $ mkWeakIORef ref finalizer1
2901data GbracketIOState s1 s2 v wref
2903 | GBracketIONormal s1 v wref
2904 | GBracketIOException s2
2906-- | Like gbracket but also uses a finalizer to make sure when the stream is
2907-- garbage collected we run the finalizing action. This requires a MonadIO and
2908-- MonadBaseControl IO constraint.
2910-- | The most general bracketing and exception combinator. All other
2911-- combinators can be expressed in terms of this combinator. This can also be
2912-- used for cases which are not covered by the standard combinators.
2916{-# INLINE_NORMAL gbracketIO #-}
2918 :: (MonadIO m, MonadBaseControl IO m)
2920 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
2921 -> (c -> m d) -- ^ after, on normal stop or GC
2922 -> (c -> e -> Stream m b) -- ^ on exception
2923 -> (c -> Stream m b) -- ^ stream generator
2925gbracketIO bef exc aft fexc fnormal =
2926 Stream step GBracketIOInit
2930 -- If the stream is never evaluated the "aft" action will never be
2931 -- called. For that to occur we will need the user of this API to pass a
2932 -- weak pointer to us.
2933 {-# INLINE_LATE step #-}
2934 step _ GBracketIOInit = do
2936 ref <- newFinalizedIORef (aft r)
2937 return $ Skip $ GBracketIONormal (fnormal r) r ref
2939 step gst (GBracketIONormal (UnStream step1 st) v ref) = do
2940 res <- exc $ step1 gst st
2942 Right r -> case r of
2944 return $ Yield x (GBracketIONormal (Stream step1 s) v ref)
2946 return $ Skip (GBracketIONormal (Stream step1 s) v ref)
2948 liftIO (writeIORef ref Nothing)
2949 aft v >> return Stop
2951 liftIO (writeIORef ref Nothing)
2952 return $ Skip (GBracketIOException (fexc v e))
2953 step gst (GBracketIOException (UnStream step1 st)) = do
2957 return $ Yield x (GBracketIOException (Stream step1 s))
2958 Skip s -> return $ Skip (GBracketIOException (Stream step1 s))
2961-- | Run a side effect before the stream yields its first element.
2962{-# INLINE_NORMAL before #-}
2963before :: Monad m => m b -> Stream m a -> Stream m a
2964before action (Stream step state) = Stream step' Nothing
2968 {-# INLINE_LATE step' #-}
2969 step' _ Nothing = action >> return (Skip (Just state))
2971 step' gst (Just st) = do
2974 Yield x s -> return $ Yield x (Just s)
2975 Skip s -> return $ Skip (Just s)
2978-- | Run a side effect whenever the stream stops normally.
2979{-# INLINE_NORMAL after #-}
2980after :: Monad m => m b -> Stream m a -> Stream m a
2981after action (Stream step state) = Stream step' state
2985 {-# INLINE_LATE step' #-}
2989 Yield x s -> return $ Yield x s
2990 Skip s -> return $ Skip s
2991 Stop -> action >> return Stop
2993{-# INLINE_NORMAL afterIO #-}
2994afterIO :: (MonadIO m, MonadBaseControl IO m)
2995 => m b -> Stream m a -> Stream m a
2996afterIO action (Stream step state) = Stream step' Nothing
3000 {-# INLINE_LATE step' #-}
3001 step' _ Nothing = do
3002 ref <- newFinalizedIORef action
3003 return $ Skip $ Just (state, ref)
3004 step' gst (Just (st, ref)) = do
3007 Yield x s -> return $ Yield x (Just (s, ref))
3008 Skip s -> return $ Skip (Just (s, ref))
3010 Just f <- readIORef ref
3011 writeIORef ref Nothing
3014-- XXX These combinators are expensive due to the call to
3015-- onException/handle/try on each step. Therefore, when possible, they should
3016-- be called in an outer loop where we perform less iterations. For example, we
3017-- cannot call them on each iteration in a char stream, instead we can call
3018-- them when doing an IO on an array.
3020-- XXX For high performance error checks in busy streams we may need another
3021-- Error constructor in step.
3023-- | Run a side effect whenever the stream aborts due to an exception. The
3024-- exception is not caught, simply rethrown.
3025{-# INLINE_NORMAL onException #-}
3026onException :: MonadCatch m => m b -> Stream m a -> Stream m a
3027onException action str =
3028 gbracket (return ()) MC.try return
3029 (\_ (e :: MC.SomeException) -> nilM (action >> MC.throwM e))
3032{-# INLINE_NORMAL _onException #-}
3033_onException :: MonadCatch m => m b -> Stream m a -> Stream m a
3034_onException action (Stream step state) = Stream step' state
3038 {-# INLINE_LATE step' #-}
3040 res <- step gst st `MC.onException` action
3042 Yield x s -> return $ Yield x s
3043 Skip s -> return $ Skip s
3046-- XXX bracket is like concatMap, it generates a stream and then flattens it.
3047-- Like concatMap it has 10x worse performance compared to linear fused
3050-- | Run the first action before the stream starts and remember its output,
3051-- generate a stream using the output, run the second action providing the
3052-- remembered value as an argument whenever the stream ends normally or due to
3054{-# INLINE_NORMAL bracket #-}
3055bracket :: MonadCatch m => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
3056bracket bef aft bet =
3057 gbracket bef MC.try aft
3058 (\a (e :: SomeException) -> nilM (aft a >> MC.throwM e)) bet
3060{-# INLINE_NORMAL bracketIO #-}
3061bracketIO :: (MonadAsync m, MonadCatch m)
3062 => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
3063bracketIO bef aft bet =
3064 gbracketIO bef MC.try aft
3065 (\a (e :: SomeException) -> nilM (aft a >> MC.throwM e)) bet
3067data BracketState s v = BracketInit | BracketRun s v
3069{-# INLINE_NORMAL _bracket #-}
3070_bracket :: MonadCatch m => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
3071_bracket bef aft bet = Stream step' BracketInit
3075 {-# INLINE_LATE step' #-}
3076 step' _ BracketInit = bef >>= \x -> return (Skip (BracketRun (bet x) x))
3078 -- NOTE: It is important to use UnStream instead of the Stream pattern
3079 -- here, otherwise we get huge perf degradation, see note in concatMap.
3080 step' gst (BracketRun (UnStream step state) v) = do
3081 -- res <- step gst state `MC.onException` aft v
3082 res <- MC.try $ step gst state
3084 Left (e :: SomeException) -> aft v >> MC.throwM e >> return Stop
3085 Right r -> case r of
3086 Yield x s -> return $ Yield x (BracketRun (Stream step s) v)
3087 Skip s -> return $ Skip (BracketRun (Stream step s) v)
3088 Stop -> aft v >> return Stop
3090-- | Run a side effect whenever the stream stops normally or aborts due to an
3092{-# INLINE finally #-}
3093finally :: MonadCatch m => m b -> Stream m a -> Stream m a
3094-- finally action xs = after action $ onException action xs
3095finally action xs = bracket (return ()) (\_ -> action) (const xs)
3097{-# INLINE finallyIO #-}
3098finallyIO :: (MonadAsync m, MonadCatch m) => m b -> Stream m a -> Stream m a
3099finallyIO action xs = bracketIO (return ()) (\_ -> action) (const xs)
3101-- | When evaluating a stream if an exception occurs, stream evaluation aborts
3102-- and the specified exception handler is run with the exception as argument.
3103{-# INLINE_NORMAL handle #-}
3104handle :: (MonadCatch m, Exception e)
3105 => (e -> Stream m a) -> Stream m a -> Stream m a
3107 gbracket (return ()) MC.try return (\_ e -> f e) (\_ -> str)
3109{-# INLINE_NORMAL _handle #-}
3110_handle :: (MonadCatch m, Exception e)
3111 => (e -> Stream m a) -> Stream m a -> Stream m a
3112_handle f (Stream step state) = Stream step' (Left state)
3116 {-# INLINE_LATE step' #-}
3117 step' gst (Left st) = do
3118 res <- MC.try $ step gst st
3120 Left e -> return $ Skip $ Right (f e)
3121 Right r -> case r of
3122 Yield x s -> return $ Yield x (Left s)
3123 Skip s -> return $ Skip (Left s)
3126 step' gst (Right (UnStream step1 st)) = do
3129 Yield x s -> return $ Yield x (Right (Stream step1 s))
3130 Skip s -> return $ Skip (Right (Stream step1 s))
3133-------------------------------------------------------------------------------
3134-- General transformation
3135-------------------------------------------------------------------------------
3137{-# INLINE_NORMAL transform #-}
3138transform :: Monad m => Pipe m a b -> Stream m a -> Stream m b
3139transform (Pipe pstep1 pstep2 pstate) (Stream step state) =
3140 Stream step' (Consume pstate, state)
3144 {-# INLINE_LATE step' #-}
3146 step' gst (Consume pst, st) = pst `seq` do
3147 r <- step (adaptState gst) st
3152 Pipe.Yield b pst' -> return $ Yield b (pst', s)
3153 Pipe.Continue pst' -> return $ Skip (pst', s)
3154 Skip s -> return $ Skip (Consume pst, s)
3157 step' _ (Produce pst, st) = pst `seq` do
3160 Pipe.Yield b pst' -> return $ Yield b (pst', st)
3161 Pipe.Continue pst' -> return $ Skip (pst', st)
3163------------------------------------------------------------------------------
3164-- Transformation by Folding (Scans)
3165------------------------------------------------------------------------------
3167------------------------------------------------------------------------------
3169------------------------------------------------------------------------------
3171-- XXX Is a prescan useful, discarding the last step does not sound useful? I
3172-- am not sure about the utility of this function, so this is implemented but
3173-- not exposed. We can expose it if someone provides good reasons why this is
3176-- XXX We have to execute the stream one step ahead to know that we are at the
3177-- last step. The vector implementation of prescan executes the last fold step
3178-- but does not yield the result. This means we have executed the effect but
3179-- discarded value. This does not sound right. In this implementation we are
3180-- not executing the last fold step.
3181{-# INLINE_NORMAL prescanlM' #-}
3182prescanlM' :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
3183prescanlM' f mz (Stream step state) = Stream step' (state, mz)
3185 {-# INLINE_LATE step' #-}
3186 step' gst (st, prev) = do
3187 r <- step (adaptState gst) st
3191 return $ Yield acc (s, f acc x)
3192 Skip s -> return $ Skip (s, prev)
3195{-# INLINE prescanl' #-}
3196prescanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
3197prescanl' f z = prescanlM' (\a b -> return (f a b)) (return z)
3199------------------------------------------------------------------------------
3200-- Monolithic postscans (postscan followed by a map)
3201------------------------------------------------------------------------------
3203-- The performance of a modular postscan followed by a map seems to be
3204-- equivalent to this monolithic scan followed by map therefore we may not need
3205-- this implementation. We just have it for performance comparison and in case
3206-- modular version does not perform well in some situation.
3208{-# INLINE_NORMAL postscanlMx' #-}
3209postscanlMx' :: Monad m
3210 => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> Stream m b
3211postscanlMx' fstep begin done (Stream step state) = do
3212 Stream step' (state, begin)
3214 {-# INLINE_LATE step' #-}
3215 step' gst (st, acc) = do
3216 r <- step (adaptState gst) st
3222 v `seq` y `seq` return (Yield v (s, return y))
3223 Skip s -> return $ Skip (s, acc)
3226{-# INLINE_NORMAL postscanlx' #-}
3227postscanlx' :: Monad m
3228 => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> Stream m b
3229postscanlx' fstep begin done s =
3230 postscanlMx' (\b a -> return (fstep b a)) (return begin) (return . done) s
3232-- XXX do we need consM strict to evaluate the begin value?
3233{-# INLINE scanlMx' #-}
3235 => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> Stream m b
3236scanlMx' fstep begin done s =
3237 (begin >>= \x -> x `seq` done x) `consM` postscanlMx' fstep begin done s
3239{-# INLINE scanlx' #-}
3241 => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> Stream m b
3242scanlx' fstep begin done s =
3243 scanlMx' (\b a -> return (fstep b a)) (return begin) (return . done) s
3245------------------------------------------------------------------------------
3247------------------------------------------------------------------------------
3249{-# INLINE_NORMAL postscanlM' #-}
3250postscanlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
3251postscanlM' fstep begin (Stream step state) =
3252 begin `seq` Stream step' (state, begin)
3254 {-# INLINE_LATE step' #-}
3255 step' gst (st, acc) = acc `seq` do
3256 r <- step (adaptState gst) st
3260 y `seq` return (Yield y (s, y))
3261 Skip s -> return $ Skip (s, acc)
3264{-# INLINE_NORMAL postscanl' #-}
3265postscanl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
3266postscanl' f = postscanlM' (\a b -> return (f a b))
3268{-# INLINE_NORMAL postscanlM #-}
3269postscanlM :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
3270postscanlM fstep begin (Stream step state) = Stream step' (state, begin)
3272 {-# INLINE_LATE step' #-}
3273 step' gst (st, acc) = do
3274 r <- step (adaptState gst) st
3278 return (Yield y (s, y))
3279 Skip s -> return $ Skip (s, acc)
3282{-# INLINE_NORMAL postscanl #-}
3283postscanl :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
3284postscanl f = postscanlM (\a b -> return (f a b))
3286{-# INLINE_NORMAL scanlM' #-}
3287scanlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
3288scanlM' fstep begin s = begin `seq` (begin `cons` postscanlM' fstep begin s)
3290{-# INLINE scanl' #-}
3291scanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
3292scanl' f = scanlM' (\a b -> return (f a b))
3294{-# INLINE_NORMAL scanlM #-}
3295scanlM :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
3296scanlM fstep begin s = begin `cons` postscanlM fstep begin s
3299scanl :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
3300scanl f = scanlM (\a b -> return (f a b))
3302{-# INLINE_NORMAL scanl1M #-}
3303scanl1M :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
3304scanl1M fstep (Stream step state) = Stream step' (state, Nothing)
3306 {-# INLINE_LATE step' #-}
3307 step' gst (st, Nothing) = do
3310 Yield x s -> return $ Yield x (s, Just x)
3311 Skip s -> return $ Skip (s, Nothing)
3314 step' gst (st, Just acc) = do
3319 return $ Yield z (s, Just z)
3320 Skip s -> return $ Skip (s, Just acc)
3323{-# INLINE scanl1 #-}
3324scanl1 :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
3325scanl1 f = scanl1M (\x y -> return (f x y))
3327{-# INLINE_NORMAL scanl1M' #-}
3328scanl1M' :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
3329scanl1M' fstep (Stream step state) = Stream step' (state, Nothing)
3331 {-# INLINE_LATE step' #-}
3332 step' gst (st, Nothing) = do
3335 Yield x s -> x `seq` return $ Yield x (s, Just x)
3336 Skip s -> return $ Skip (s, Nothing)
3339 step' gst (st, Just acc) = acc `seq` do
3344 z `seq` return $ Yield z (s, Just z)
3345 Skip s -> return $ Skip (s, Just acc)
3348{-# INLINE scanl1' #-}
3349scanl1' :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
3350scanl1' f = scanl1M' (\x y -> return (f x y))
3352------------------------------------------------------------------------------
3354------------------------------------------------------------------------------
3356data RollingMapState s a = RollingMapInit s | RollingMapGo s a
3358{-# INLINE rollingMapM #-}
3359rollingMapM :: Monad m => (a -> a -> m b) -> Stream m a -> Stream m b
3360rollingMapM f (Stream step1 state1) = Stream step (RollingMapInit state1)
3362 step gst (RollingMapInit st) = do
3363 r <- step1 (adaptState gst) st
3365 Yield x s -> Skip $ RollingMapGo s x
3366 Skip s -> Skip $ RollingMapInit s
3369 step gst (RollingMapGo s1 x1) = do
3370 r <- step1 (adaptState gst) s1
3374 return $ Yield res $ RollingMapGo s x
3375 Skip s -> return $ Skip $ RollingMapGo s x1
3376 Stop -> return $ Stop
3378{-# INLINE rollingMap #-}
3379rollingMap :: Monad m => (a -> a -> b) -> Stream m a -> Stream m b
3380rollingMap f = rollingMapM (\x y -> return $ f x y)
3382------------------------------------------------------------------------------
3383-- Tapping/Distributing
3384------------------------------------------------------------------------------
3387tap :: Monad m => Fold m a b -> Stream m a -> Stream m a
3388tap (Fold fstep initial extract) (Stream step state) = Stream step' Nothing
3392 step' _ Nothing = do
3394 return $ Skip (Just (r, state))
3396 step' gst (Just (acc, st)) = acc `seq` do
3401 return $ Yield x (Just (acc', s))
3402 Skip s -> return $ Skip (Just (acc, s))
3407{-# INLINE_NORMAL tapOffsetEvery #-}
3408tapOffsetEvery :: Monad m
3409 => Int -> Int -> Fold m a b -> Stream m a -> Stream m a
3410tapOffsetEvery offset n (Fold fstep initial extract) (Stream step state) =
3411 Stream step' Nothing
3415 {-# INLINE_LATE step' #-}
3416 step' _ Nothing = do
3418 return $ Skip (Just (r, state, offset `mod` n))
3420 step' gst (Just (acc, st, count)) | count <= 0 = do
3424 !acc' <- fstep acc x
3425 return $ Yield x (Just (acc', s, n - 1))
3426 Skip s -> return $ Skip (Just (acc, s, count))
3431 step' gst (Just (acc, st, count)) = do
3434 Yield x s -> return $ Yield x (Just (acc, s, count - 1))
3435 Skip s -> return $ Skip (Just (acc, s, count))
3440{-# INLINE_NORMAL pollCounts #-}
3443 => (Stream m Int -> Stream m Int)
3447pollCounts transf fld (Stream step state) = Stream step' Nothing
3450 {-# INLINE_LATE step' #-}
3451 step' _ Nothing = do
3452 countVar <- liftIO $ newVar (0 :: Int)
3454 $ void $ runFold fld
3455 $ transf $ fromPrimVar countVar
3456 return $ Skip (Just (countVar, tid, state))
3458 step' gst (Just (countVar, tid, st)) = do
3462 liftIO $ modifyVar' countVar (+ 1)
3463 return $ Yield x (Just (countVar, tid, s))
3464 Skip s -> return $ Skip (Just (countVar, tid, s))
3466 liftIO $ killThread tid
3469{-# INLINE_NORMAL tapRate #-}
3471 (MonadAsync m, MonadCatch m)
3476tapRate samplingRate action (Stream step state) = Stream step' Nothing
3478 {-# NOINLINE loop #-}
3479 loop countVar prev = do
3482 (do liftIO $ threadDelay (round $ samplingRate * 1000000)
3483 i <- liftIO $ readVar countVar
3484 let !diff = i - prev
3487 (\(e :: AsyncException) -> do
3488 i <- liftIO $ readVar countVar
3489 let !diff = i - prev
3491 throwM (MC.toException e))
3494 {-# INLINE_LATE step' #-}
3495 step' _ Nothing = do
3496 countVar <- liftIO $ newVar 0
3497 tid <- fork $ loop countVar 0
3498 ref <- liftIO $ newIORef ()
3499 _ <- liftIO $ mkWeakIORef ref (killThread tid)
3500 return $ Skip (Just (countVar, tid, state, ref))
3502 step' gst (Just (countVar, tid, st, ref)) = do
3506 liftIO $ modifyVar' countVar (+ 1)
3507 return $ Yield x (Just (countVar, tid, s, ref))
3508 Skip s -> return $ Skip (Just (countVar, tid, s, ref))
3510 liftIO $ killThread tid
3514-------------------------------------------------------------------------------
3516-------------------------------------------------------------------------------
3518{-# INLINE_NORMAL takeWhileM #-}
3519takeWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
3520takeWhileM f (Stream step state) = Stream step' state
3522 {-# INLINE_LATE step' #-}
3528 return $ if b then Yield x s else Stop
3529 Skip s -> return $ Skip s
3532{-# INLINE takeWhile #-}
3533takeWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
3534takeWhile f = takeWhileM (return . f)
3536{-# INLINE_NORMAL drop #-}
3537drop :: Monad m => Int -> Stream m a -> Stream m a
3538drop n (Stream step state) = Stream step' (state, Just n)
3540 {-# INLINE_LATE step' #-}
3541 step' gst (st, Just i)
3546 Yield _ s -> Skip (s, Just (i - 1))
3547 Skip s -> Skip (s, Just i)
3549 | otherwise = return $ Skip (st, Nothing)
3551 step' gst (st, Nothing) = do
3555 Yield x s -> Yield x (s, Nothing)
3556 Skip s -> Skip (s, Nothing)
3559data DropWhileState s a
3561 | DropWhileYield a s
3564{-# INLINE_NORMAL dropWhileM #-}
3565dropWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
3566dropWhileM f (Stream step state) = Stream step' (DropWhileDrop state)
3568 {-# INLINE_LATE step' #-}
3569 step' gst (DropWhileDrop st) = do
3575 then return $ Skip (DropWhileDrop s)
3576 else return $ Skip (DropWhileYield x s)
3577 Skip s -> return $ Skip (DropWhileDrop s)
3580 step' gst (DropWhileNext st) = do
3583 Yield x s -> return $ Skip (DropWhileYield x s)
3584 Skip s -> return $ Skip (DropWhileNext s)
3587 step' _ (DropWhileYield x st) = return $ Yield x (DropWhileNext st)
3589{-# INLINE dropWhile #-}
3590dropWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
3591dropWhile f = dropWhileM (return . f)
3593{-# INLINE_NORMAL filterM #-}
3594filterM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
3595filterM f (Stream step state) = Stream step' state
3597 {-# INLINE_LATE step' #-}
3606 Skip s -> return $ Skip s
3609{-# INLINE filter #-}
3610filter :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
3611filter f = filterM (return . f)
3613{-# INLINE_NORMAL uniq #-}
3614uniq :: (Eq a, Monad m) => Stream m a -> Stream m a
3615uniq (Stream step state) = Stream step' (Nothing, state)
3617 {-# INLINE_LATE step' #-}
3618 step' gst (Nothing, st) = do
3621 Yield x s -> return $ Yield x (Just x, s)
3622 Skip s -> return $ Skip (Nothing, s)
3624 step' gst (Just x, st) = do
3627 Yield y s | x == y -> return $ Skip (Just x, s)
3628 | otherwise -> return $ Yield y (Just y, s)
3629 Skip s -> return $ Skip (Just x, s)
3632------------------------------------------------------------------------------
3633-- Transformation by Mapping
3634------------------------------------------------------------------------------
3636{-# INLINE_NORMAL sequence #-}
3637sequence :: Monad m => Stream m (m a) -> Stream m a
3638sequence (Stream step state) = Stream step' state
3640 {-# INLINE_LATE step' #-}
3642 r <- step (adaptState gst) st
3644 Yield x s -> x >>= \a -> return (Yield a s)
3645 Skip s -> return $ Skip s
3648------------------------------------------------------------------------------
3650------------------------------------------------------------------------------
3652data LoopState x s = FirstYield s
3653 | InterspersingYield s
3656{-# INLINE_NORMAL intersperseM #-}
3657intersperseM :: Monad m => m a -> Stream m a -> Stream m a
3658intersperseM m (Stream step state) = Stream step' (FirstYield state)
3660 {-# INLINE_LATE step' #-}
3661 step' gst (FirstYield st) = do
3665 Yield x s -> Skip (YieldAndCarry x s)
3666 Skip s -> Skip (FirstYield s)
3669 step' gst (InterspersingYield st) = do
3674 return $ Yield a (YieldAndCarry x s)
3675 Skip s -> return $ Skip $ InterspersingYield s
3678 step' _ (YieldAndCarry x st) = return $ Yield x (InterspersingYield st)
3683 | SuffixYield a (SuffixState s a)
3685{-# INLINE_NORMAL intersperseSuffix #-}
3686intersperseSuffix :: forall m a. Monad m => m a -> Stream m a -> Stream m a
3687intersperseSuffix action (Stream step state) = Stream step' (SuffixElem state)
3689 {-# INLINE_LATE step' #-}
3690 step' gst (SuffixElem st) = do
3693 Yield x s -> Skip (SuffixYield x (SuffixSuffix s))
3694 Skip s -> Skip (SuffixElem s)
3697 step' _ (SuffixSuffix st) = do
3698 action >>= \r -> return $ Skip (SuffixYield r (SuffixElem st))
3700 step' _ (SuffixYield x next) = return $ Yield x next
3702data SuffixSpanState s a
3703 = SuffixSpanElem s Int
3704 | SuffixSpanSuffix s
3705 | SuffixSpanYield a (SuffixSpanState s a)
3709-- | intersperse after every n items
3710{-# INLINE_NORMAL intersperseSuffixBySpan #-}
3711intersperseSuffixBySpan :: forall m a. Monad m
3712 => Int -> m a -> Stream m a -> Stream m a
3713intersperseSuffixBySpan n action (Stream step state) =
3714 Stream step' (SuffixSpanElem state n)
3716 {-# INLINE_LATE step' #-}
3717 step' gst (SuffixSpanElem st i) | i > 0 = do
3720 Yield x s -> Skip (SuffixSpanYield x (SuffixSpanElem s (i - 1)))
3721 Skip s -> Skip (SuffixSpanElem s i)
3722 Stop -> if i == n then Stop else Skip SuffixSpanLast
3723 step' _ (SuffixSpanElem st _) = return $ Skip (SuffixSpanSuffix st)
3725 step' _ (SuffixSpanSuffix st) = do
3726 action >>= \r -> return $ Skip (SuffixSpanYield r (SuffixSpanElem st n))
3728 step' _ (SuffixSpanLast) = do
3729 action >>= \r -> return $ Skip (SuffixSpanYield r SuffixSpanStop)
3731 step' _ (SuffixSpanYield x next) = return $ Yield x next
3733 step' _ (SuffixSpanStop) = return Stop
3735{-# INLINE intersperse #-}
3736intersperse :: Monad m => a -> Stream m a -> Stream m a
3737intersperse a = intersperseM (return a)
3739{-# INLINE_NORMAL insertBy #-}
3740insertBy :: Monad m => (a -> a -> Ordering) -> a -> Stream m a -> Stream m a
3741insertBy cmp a (Stream step state) = Stream step' (state, False, Nothing)
3743 {-# INLINE_LATE step' #-}
3744 step' gst (st, False, _) = do
3747 Yield x s -> case cmp a x of
3748 GT -> return $ Yield x (s, False, Nothing)
3749 _ -> return $ Yield a (s, True, Just x)
3750 Skip s -> return $ Skip (s, False, Nothing)
3751 Stop -> return $ Yield a (st, True, Nothing)
3753 step' _ (_, True, Nothing) = return Stop
3755 step' gst (st, True, Just prev) = do
3758 Yield x s -> return $ Yield prev (s, True, Just x)
3759 Skip s -> return $ Skip (s, True, Just prev)
3760 Stop -> return $ Yield prev (st, True, Nothing)
3762------------------------------------------------------------------------------
3764------------------------------------------------------------------------------
3766{-# INLINE_NORMAL deleteBy #-}
3767deleteBy :: Monad m => (a -> a -> Bool) -> a -> Stream m a -> Stream m a
3768deleteBy eq x (Stream step state) = Stream step' (state, False)
3770 {-# INLINE_LATE step' #-}
3771 step' gst (st, False) = do
3774 Yield y s -> return $
3775 if eq x y then Skip (s, True) else Yield y (s, False)
3776 Skip s -> return $ Skip (s, False)
3779 step' gst (st, True) = do
3782 Yield y s -> return $ Yield y (s, True)
3783 Skip s -> return $ Skip (s, True)
3786------------------------------------------------------------------------------
3787-- Transformation by Map and Filter
3788------------------------------------------------------------------------------
3790-- XXX Will this always fuse properly?
3791{-# INLINE_NORMAL mapMaybe #-}
3792mapMaybe :: Monad m => (a -> Maybe b) -> Stream m a -> Stream m b
3793mapMaybe f = fmap fromJust . filter isJust . map f
3795{-# INLINE_NORMAL mapMaybeM #-}
3796mapMaybeM :: Monad m => (a -> m (Maybe b)) -> Stream m a -> Stream m b
3797mapMaybeM f = fmap fromJust . filter isJust . mapM f
3799------------------------------------------------------------------------------
3801------------------------------------------------------------------------------
3803{-# INLINE_NORMAL indexed #-}
3804indexed :: Monad m => Stream m a -> Stream m (Int, a)
3805indexed (Stream step state) = Stream step' (state, 0)
3807 {-# INLINE_LATE step' #-}
3808 step' gst (st, i) = i `seq` do
3809 r <- step (adaptState gst) st
3811 Yield x s -> return $ Yield (i, x) (s, i+1)
3812 Skip s -> return $ Skip (s, i)
3815{-# INLINE_NORMAL indexedR #-}
3816indexedR :: Monad m => Int -> Stream m a -> Stream m (Int, a)
3817indexedR m (Stream step state) = Stream step' (state, m)
3819 {-# INLINE_LATE step' #-}
3820 step' gst (st, i) = i `seq` do
3821 r <- step (adaptState gst) st
3823 Yield x s -> let i' = i - 1
3824 in return $ Yield (i, x) (s, i')
3825 Skip s -> return $ Skip (s, i)
3828{-# INLINE_NORMAL zipWithM #-}
3830 => (a -> b -> m c) -> Stream m a -> Stream m b -> Stream m c
3831zipWithM f (Stream stepa ta) (Stream stepb tb) = Stream step (ta, tb, Nothing)
3833 {-# INLINE_LATE step #-}
3834 step gst (sa, sb, Nothing) = do
3835 r <- stepa (adaptState gst) sa
3838 Yield x sa' -> Skip (sa', sb, Just x)
3839 Skip sa' -> Skip (sa', sb, Nothing)
3842 step gst (sa, sb, Just x) = do
3843 r <- stepb (adaptState gst) sb
3847 return $ Yield z (sa, sb', Nothing)
3848 Skip sb' -> return $ Skip (sa, sb', Just x)
3851#if __GLASGOW_HASKELL__ >= 801
3852{-# RULES "zipWithM xs xs"
3853 forall f xs. zipWithM @Identity f xs xs = mapM (\x -> f x x) xs #-}
3856{-# INLINE zipWith #-}
3857zipWith :: Monad m => (a -> b -> c) -> Stream m a -> Stream m b -> Stream m c
3858zipWith f = zipWithM (\a b -> return (f a b))
3860------------------------------------------------------------------------------
3862------------------------------------------------------------------------------
3864{-# INLINE_NORMAL mergeByM #-}
3867 => (a -> a -> m Ordering) -> Stream m a -> Stream m a -> Stream m a
3868mergeByM cmp (Stream stepa ta) (Stream stepb tb) =
3869 Stream step (Just ta, Just tb, Nothing, Nothing)
3871 {-# INLINE_LATE step #-}
3873 -- one of the values is missing, and the corresponding stream is running
3874 step gst (Just sa, sb, Nothing, b) = do
3877 Yield a sa' -> Skip (Just sa', sb, Just a, b)
3878 Skip sa' -> Skip (Just sa', sb, Nothing, b)
3879 Stop -> Skip (Nothing, sb, Nothing, b)
3881 step gst (sa, Just sb, a, Nothing) = do
3884 Yield b sb' -> Skip (sa, Just sb', a, Just b)
3885 Skip sb' -> Skip (sa, Just sb', a, Nothing)
3886 Stop -> Skip (sa, Nothing, a, Nothing)
3888 -- both the values are available
3889 step _ (sa, sb, Just a, Just b) = do
3891 return $ case res of
3892 GT -> Yield b (sa, sb, Just a, Nothing)
3893 _ -> Yield a (sa, sb, Nothing, Just b)
3895 -- one of the values is missing, corresponding stream is done
3896 step _ (Nothing, sb, Nothing, Just b) =
3897 return $ Yield b (Nothing, sb, Nothing, Nothing)
3899 step _ (sa, Nothing, Just a, Nothing) =
3900 return $ Yield a (sa, Nothing, Nothing, Nothing)
3902 step _ (Nothing, Nothing, Nothing, Nothing) = return Stop
3904{-# INLINE mergeBy #-}
3907 => (a -> a -> Ordering) -> Stream m a -> Stream m a -> Stream m a
3908mergeBy cmp = mergeByM (\a b -> return $ cmp a b)
3910------------------------------------------------------------------------------
3911-- Transformation comprehensions
3912------------------------------------------------------------------------------
3914{-# INLINE_NORMAL the #-}
3915the :: (Eq a, Monad m) => Stream m a -> m (Maybe a)
3916the (Stream step state) = go state
3919 r <- step defState st
3921 Yield x s -> go' x s
3923 Stop -> return Nothing
3925 r <- step defState st
3927 Yield x s | x == n -> go' n s
3928 | otherwise -> return Nothing
3930 Stop -> return (Just n)
3932{-# INLINE runFold #-}
3933runFold :: (Monad m) => Fold m a b -> Stream m a -> m b
3934runFold (Fold step begin done) = foldlMx' step begin done
3936-------------------------------------------------------------------------------
3937-- Concurrent application and fold
3938-------------------------------------------------------------------------------
3940-- XXX These functions should be moved to Stream/Parallel.hs
3942-- Using StreamD the worker stream producing code can fuse with the code to
3943-- queue output to the SVar giving some perf boost.
3945-- Note that StreamD can only be used in limited situations, specifically, we
3946-- cannot implement joinStreamVarPar using this.
3948-- XXX make sure that the SVar passed is a Parallel style SVar.
3950-- | Fold the supplied stream to the SVar asynchronously using Parallel
3951-- concurrency style.
3952-- {-# INLINE_NORMAL toSVarParallel #-}
3953{-# INLINE toSVarParallel #-}
3954toSVarParallel :: MonadAsync m
3955 => State t m a -> SVar t m a -> Stream m a -> m ()
3956toSVarParallel st sv xs =
3957 if svarInspectMode sv
3961 case getYieldLimit st of
3962 Nothing -> doFork (work Nothing)
3964 (handleChildException sv)
3965 Just _ -> doFork (workLim Nothing)
3967 (handleChildException sv)
3972 {-# NOINLINE work #-}
3973 work info = (runFold (FL.toParallelSVar sv info) xs)
3975 {-# NOINLINE workLim #-}
3976 workLim info = runFold (FL.toParallelSVarLimited sv info) xs
3978 {-# NOINLINE forkWithDiag #-}
3980 -- We do not use workerCount in case of ParallelVar but still there is
3981 -- no harm in maintaining it correctly.
3982 liftIO $ atomicModifyIORefCAS_ (workerCount sv) $ \n -> n + 1
3984 -- This allocation matters when significant number of workers are being
3985 -- sent. We allocate it only when needed. The overhead increases by 4x.
3987 case yieldRateInfo sv of
3988 Nothing -> return Nothing
3989 Just _ -> liftIO $ do
3990 cntRef <- newIORef 0
3991 t <- getTime Monotonic
3992 lat <- newIORef (0, t)
3993 return $ Just WorkerInfo
3994 { workerYieldMax = 0
3995 , workerYieldCount = cntRef
3996 , workerLatencyStart = lat
3999 case getYieldLimit st of
4000 Nothing -> doFork (work winfo)
4002 (handleChildException sv)
4003 Just _ -> doFork (workLim winfo)
4005 (handleChildException sv)
4008{-# INLINE_NORMAL mkParallelD #-}
4009mkParallelD :: MonadAsync m => Stream m a -> Stream m a
4010mkParallelD m = Stream step Nothing
4013 step gst Nothing = do
4014 sv <- newParallelVar StopNone gst
4015 toSVarParallel gst sv m
4016 -- XXX use unfold instead?
4017 return $ Skip $ Just $ fromSVar sv
4019 step gst (Just (UnStream step1 st)) = do
4022 Yield a s -> Yield a (Just $ Stream step1 s)
4023 Skip s -> Skip (Just $ Stream step1 s)
4026-- Compare with mkAsync. mkAsync uses an Async style SVar whereas this uses a
4027-- parallel style SVar for evaluation. Currently, parallel style cannot use
4028-- rate control whereas Async style can use rate control. In async style SVar
4029-- the worker thread terminates when the buffer is full whereas in Parallel
4032-- | Make the stream producer and consumer run concurrently by introducing a
4033-- buffer between them. The producer thread evaluates the input stream until
4034-- the buffer fills, it blocks if the buffer is full until there is space in
4035-- the buffer. The consumer consumes the stream lazily from the buffer.
4039{-# INLINE_NORMAL mkParallel #-}
4040mkParallel :: (K.IsStream t, MonadAsync m) => t m a -> t m a
4041mkParallel = fromStreamD . mkParallelD . toStreamD
4043-------------------------------------------------------------------------------
4045-------------------------------------------------------------------------------
4047-- | Create an SVar with a fold consumer that will fold any elements sent to it
4048-- using the supplied fold function.
4049{-# INLINE newFoldSVar #-}
4050newFoldSVar :: MonadAsync m => State t m a -> Fold m a b -> m (SVar t m a)
4051newFoldSVar stt f = do
4052 -- Buffer size for the SVar is derived from the current state
4053 sv <- newParallelVar StopAny (adaptState stt)
4054 -- Add the producer thread-id to the SVar.
4055 liftIO myThreadId >>= modifyThread sv
4056 void $ doFork (work sv) (svarMrun sv) (handleFoldException sv)
4061 {-# NOINLINE work #-}
4062 work sv = void $ runFold f $ fromProducer sv
4064data TapState sv st = TapInit | Tapping sv st | TapDone st
4066{-# INLINE_NORMAL tapAsync #-}
4067tapAsync :: MonadAsync m => Fold m a b -> Stream m a -> Stream m a
4068tapAsync f (Stream step1 state1) = Stream step TapInit
4072 -- In general, a Stop event would come equipped with the result
4073 -- of the fold. It is not used here but it would be useful in
4074 -- applicative and distribute.
4075 done <- fromConsumer svr
4076 when (not done) $ do
4077 liftIO $ withDiagMVar svr "teeToSVar: waiting to drain"
4078 $ takeMVar (outputDoorBellFromConsumer svr)
4082 liftIO $ sendStop svr Nothing
4083 -- drain/wait until a stop event arrives from the fold.
4086 {-# INLINE_LATE step #-}
4087 step gst TapInit = do
4088 sv <- newFoldSVar gst f
4089 return $ Skip (Tapping sv state1)
4091 step gst (Tapping sv st) = do
4095 done <- pushToFold sv a
4098 -- XXX we do not need to wait synchronously here
4100 return $ Yield a (TapDone s)
4101 else return $ Yield a (Tapping sv s)
4102 Skip s -> return $ Skip (Tapping sv s)
4107 step gst (TapDone st) = do
4110 Yield a s -> Yield a (TapDone s)
4111 Skip s -> Skip (TapDone s)
4114-- XXX Exported from Array again as this fold is specific to Array
4115-- | Take last 'n' elements from the stream and discard the rest.
4117lastN :: (Storable a, MonadIO m) => Int -> Fold m a (Array a)
4118lastN n = Fold step initial done
4120 step (Tuple3' rb rh i) a = do
4121 rh1 <- liftIO $ RB.unsafeInsert rb rh a
4122 return $ Tuple3' rb rh1 (i + 1)
4123 initial = fmap (\(a, b) -> Tuple3' a b (0 :: Int)) $ liftIO $ RB.new n
4124 done (Tuple3' rb rh i) = do
4125 arr <- liftIO $ A.newArray n
4126 foldFunc i rh snoc' arr rb
4127 snoc' b a = liftIO $ A.unsafeSnoc b a
4129 | i < n = RB.unsafeFoldRingM
4130 | otherwise = RB.unsafeFoldRingFullM
4132------------------------------------------------------------------------------
4134------------------------------------------------------------------------------
4136-- XXX using getTime in the loop can be pretty expensive especially for
4137-- computations where iterations are lightweight. We have the following
4140-- 1) Run a timeout thread updating a flag asynchronously and check that
4141-- flag here, that way we can have a cheap termination check.
4143-- 2) Use COARSE clock to get time with lower resolution but more efficiently.
4145-- 3) Use rdtscp/rdtsc to get time directly from the processor, compute the
4146-- termination value of rdtsc in the beginning and then in each iteration just
4147-- get rdtsc and check if we should terminate.
4151 | TakeByTimeCheck st s
4152 | TakeByTimeYield st s
4154{-# INLINE_NORMAL takeByTime #-}
4155takeByTime :: (MonadIO m, TimeUnit64 t) => t -> Stream m a -> Stream m a
4156takeByTime duration (Stream step1 state1) = Stream step (TakeByTimeInit state1)
4159 lim = toRelTime64 duration
4161 {-# INLINE_LATE step #-}
4162 step _ (TakeByTimeInit _) | lim == 0 = return Stop
4163 step _ (TakeByTimeInit st) = do
4164 t0 <- liftIO $ getTime Monotonic
4165 return $ Skip (TakeByTimeYield st t0)
4166 step _ (TakeByTimeCheck st t0) = do
4167 t <- liftIO $ getTime Monotonic
4169 if diffAbsTime64 t t0 > lim
4171 else Skip (TakeByTimeYield st t0)
4172 step gst (TakeByTimeYield st t0) = do
4175 Yield x s -> Yield x (TakeByTimeCheck s t0)
4176 Skip s -> Skip (TakeByTimeCheck s t0)
4179data DropByTime st s x
4181 | DropByTimeGen st s
4182 | DropByTimeCheck st s x
4183 | DropByTimeYield st
4185{-# INLINE_NORMAL dropByTime #-}
4186dropByTime :: (MonadIO m, TimeUnit64 t) => t -> Stream m a -> Stream m a
4187dropByTime duration (Stream step1 state1) = Stream step (DropByTimeInit state1)
4190 lim = toRelTime64 duration
4192 {-# INLINE_LATE step #-}
4193 step _ (DropByTimeInit st) = do
4194 t0 <- liftIO $ getTime Monotonic
4195 return $ Skip (DropByTimeGen st t0)
4196 step gst (DropByTimeGen st t0) = do
4199 Yield x s -> Skip (DropByTimeCheck s t0 x)
4200 Skip s -> Skip (DropByTimeGen s t0)
4202 step _ (DropByTimeCheck st t0 x) = do
4203 t <- liftIO $ getTime Monotonic
4204 if diffAbsTime64 t t0 <= lim
4205 then return $ Skip $ DropByTimeGen st t0
4206 else return $ Yield x $ DropByTimeYield st
4207 step gst (DropByTimeYield st) = do
4210 Yield x s -> Yield x (DropByTimeYield s)
4211 Skip s -> Skip (DropByTimeYield s)
4214{-# INLINE_NORMAL currentTime #-}
4215currentTime :: MonadAsync m => Double -> Stream m AbsTime
4216currentTime g = Stream step Nothing
4219 threadDelay $ delayTime
4220 MicroSecond64 t <- fromAbsTime <$> getTime Monotonic
4221 modifyVar' timeVar (const t)
4223 g' = g * 10 ^ (6 :: Int)
4224 {-# INLINE delayTime #-}
4225 delayTime = if g' >= fromIntegral (maxBound :: Int)
4228 {-# INLINE_LATE step #-}
4230 timeVar <- liftIO $ newVar (0 :: Int64)
4231 tid <- forkManaged $ liftIO $ void $ next timeVar
4232 return $ Skip $ Just (timeVar, tid)
4233 step _ s@(Just (timeVar, _)) = do
4234 a <- liftIO $ readVar timeVar
4235 return $ Yield (toAbsTime (MicroSecond64 a)) s
keysesc to close
- j k
- move down, up
- enter
- open
- esc
- back, or up a level
- 1 2 3 4
- readme, files, commits, refs
- /
- filter
- y
- copy the clone command
- c
- compare the chosen commit with another (or shift-click both)
- ?
- show or hide this list