1{-# LANGUAGE BangPatterns #-}
3{-# LANGUAGE ConstraintKinds #-}
4{-# LANGUAGE ExistentialQuantification #-}
5{-# LANGUAGE FlexibleContexts #-}
6{-# LANGUAGE FlexibleInstances #-}
7{-# LANGUAGE MultiParamTypeClasses #-}
8{-# LANGUAGE PatternSynonyms #-}
9{-# LANGUAGE RecordWildCards #-}
10{-# LANGUAGE ScopedTypeVariables #-}
11{-# LANGUAGE ViewPatterns #-}
12{-# LANGUAGE RankNTypes #-}
13{-# LANGUAGE MagicHash #-}
15#if __GLASGOW_HASKELL__ >= 801
16{-# LANGUAGE TypeApplications #-}
22-- Module : Streamly.Internal.Data.Stream.StreamD
23-- Copyright : (c) 2018 Harendra Kumar
24-- (c) Roman Leshchinskiy 2008-2010
25-- (c) The University of Glasgow, 2009
28-- Maintainer : streamly@composewell.com
29-- Stability : experimental
32-- Direct style re-implementation of CPS style stream in StreamK module. The
33-- symbol or suffix 'D' in this module denotes the "Direct" style. GHC is able
34-- to INLINE and fuse direct style better, providing better performance than
38-- import qualified Streamly.Internal.Data.Stream.StreamD as D
41-- Some of the functions in this file have been adapted from the vector
42-- library, https://hackage.haskell.org/package/vector.
44module Streamly.Internal.Data.Stream.StreamD
49#if __GLASGOW_HASKELL__ >= 800
50 , Stream (Stream, UnStream)
70 -- ** Specialized Generation
71 -- | Generate a monadic stream from a seed.
84 , enumerateFromStepIntegral
85 , enumerateFromIntegral
86 , enumerateFromThenIntegral
87 , enumerateFromToIntegral
88 , enumerateFromThenToIntegral
90 , enumerateFromStepNum
93 , enumerateFromToFractional
94 , enumerateFromThenToFractional
100 -- | Transform an input structure into a stream.
101 -- | Direct style stream does not support @fromFoldable@.
131 -- ** Specialized Folds
158 -- ** Flattening nested streams
161 , ConcatMapUState (..)
163 , ConcatUnfoldInterleaveState (..)
164 , concatUnfoldInterleave
165 , concatUnfoldRoundrobin
168 , InterleaveState(..)
173 , roundRobin -- interleaveFair?/ParallelFair
206 -- | Transform a stream into another type.
223 -- ** By folding (scans)
270 , intersperseSuffixBySpan
294 -- * Transformation comprehensions
300 , clearIORefFinalizer
312 -- * Concurrent Application
320import Control.Concurrent (killThread, myThreadId, takeMVar, threadDelay)
321import Control.Exception
322 (Exception, SomeException, AsyncException, fromException)
323import Control.Monad (void, when, forever)
324import Control.Monad.Catch (MonadCatch, throwM)
325import Control.Monad.IO.Class (MonadIO(..))
326import Control.Monad.Reader (ReaderT)
327import Control.Monad.State.Strict (StateT)
328import Control.Monad.Trans (MonadTrans(lift))
329import Control.Monad.Trans.Control (MonadBaseControl)
330import Data.Bits (shiftR, shiftL, (.|.), (.&.))
331import Data.Functor.Identity (Identity(..))
332import Data.Int (Int64)
333import Data.IORef (newIORef, readIORef, mkWeakIORef, writeIORef, IORef)
334import Data.Maybe (fromJust, isJust, isNothing)
335import Data.Word (Word32)
336import Foreign.Ptr (Ptr)
337import Foreign.Storable (Storable(..))
338import GHC.Types (SPEC(..))
339import System.Mem (performMajorGC)
341 hiding (map, mapM, mapM_, repeat, foldr, last, take, filter,
342 takeWhile, drop, dropWhile, all, any, maximum, minimum, elem,
343 notElem, null, head, tail, zipWith, lookup, foldr1, sequence,
344 (!!), scanl, scanl1, concatMap, replicate, enumFromTo, concat,
347import qualified Control.Monad.Catch as MC
348import qualified Control.Monad.Reader as Reader
349import qualified Control.Monad.State.Strict as State
350import qualified Prelude
352import Streamly.Internal.Mutable.Prim.Var
353 (Prim, Var, readVar, newVar, modifyVar')
354import Streamly.Internal.Data.Time.Units
355 (TimeUnit64, toRelTime64, diffAbsTime64)
357import Streamly.Internal.Data.Atomics (atomicModifyIORefCAS_)
358import Streamly.Internal.Memory.Array.Types (Array(..))
359import Streamly.Internal.Data.Fold.Types (Fold(..))
360import Streamly.Internal.Data.Pipe.Types (Pipe(..), PipeState(..))
361import Streamly.Internal.Data.Time.Clock (Clock(Monotonic), getTime)
362import Streamly.Internal.Data.Time.Units
363 (MicroSecond64(..), fromAbsTime, toAbsTime, AbsTime)
364import Streamly.Internal.Data.Unfold.Types (Unfold(..))
365import Streamly.Internal.Data.Strict (Tuple3'(..))
367import Streamly.Internal.Data.Stream.StreamD.Type
368import Streamly.Internal.Data.SVar
369import Streamly.Internal.Data.Stream.SVar (fromConsumer, pushToFold)
371import qualified Streamly.Internal.Data.Pipe.Types as Pipe
372import qualified Streamly.Internal.Memory.Array.Types as A
373import qualified Streamly.Internal.Data.Fold as FL
374import qualified Streamly.Memory.Ring as RB
375import qualified Streamly.Internal.Data.Stream.StreamK as K
377------------------------------------------------------------------------------
379------------------------------------------------------------------------------
381-- | An empty 'Stream'.
382{-# INLINE_NORMAL nil #-}
383nil :: Monad m => Stream m a
384nil = Stream (\_ _ -> return Stop) ()
386-- | An empty 'Stream' with a side effect.
387{-# INLINE_NORMAL nilM #-}
388nilM :: Monad m => m b -> Stream m a
389nilM m = Stream (\_ _ -> m >> return Stop) ()
391{-# INLINE_NORMAL consM #-}
392consM :: Monad m => m a -> Stream m a -> Stream m a
393consM m (Stream step state) = Stream step1 Nothing
395 {-# INLINE_LATE step1 #-}
396 step1 _ Nothing = m >>= \x -> return $ Yield x (Just state)
397 step1 gst (Just st) = do
401 Yield a s -> Yield a (Just s)
402 Skip s -> Skip (Just s)
405-- XXX implement in terms of consM?
406-- cons x = consM (return x)
408-- | Can fuse but has O(n^2) complexity.
409{-# INLINE_NORMAL cons #-}
410cons :: Monad m => a -> Stream m a -> Stream m a
411cons x (Stream step state) = Stream step1 Nothing
413 {-# INLINE_LATE step1 #-}
414 step1 _ Nothing = return $ Yield x (Just state)
415 step1 gst (Just st) = do
419 Yield a s -> Yield a (Just s)
420 Skip s -> Skip (Just s)
423-------------------------------------------------------------------------------
425-------------------------------------------------------------------------------
427-- Does not fuse, has the same performance as the StreamK version.
428{-# INLINE_NORMAL uncons #-}
429uncons :: Monad m => Stream m a -> m (Maybe (a, Stream m a))
430uncons (UnStream step state) = go state
433 r <- step defState st
435 Yield x s -> return $ Just (x, Stream step s)
437 Stop -> return Nothing
439------------------------------------------------------------------------------
440-- Generation by unfold
441------------------------------------------------------------------------------
443{-# INLINE_NORMAL unfoldrM #-}
444unfoldrM :: Monad m => (s -> m (Maybe (a, s))) -> s -> Stream m a
445unfoldrM next state = Stream step state
447 {-# INLINE_LATE step #-}
451 Just (x, s) -> Yield x s
454{-# INLINE_LATE unfoldr #-}
455unfoldr :: Monad m => (s -> Maybe (a, s)) -> s -> Stream m a
456unfoldr f = unfoldrM (return . f)
458-- | Convert an 'Unfold' into a 'Stream' by supplying it a seed.
460{-# INLINE_NORMAL unfold #-}
461unfold :: Monad m => Unfold m a b -> a -> Stream m b
462unfold (Unfold ustep inject) seed = Stream step Nothing
464 {-# INLINE_LATE step #-}
465 step _ Nothing = inject seed >>= return . Skip . Just
466 step _ (Just st) = do
469 Yield x s -> Yield x (Just s)
470 Skip s -> Skip (Just s)
473------------------------------------------------------------------------------
474-- Specialized Generation
475------------------------------------------------------------------------------
477{-# INLINE_NORMAL repeatM #-}
478repeatM :: Monad m => m a -> Stream m a
479repeatM x = Stream (\_ _ -> x >>= \r -> return $ Yield r ()) ()
481{-# INLINE_NORMAL repeat #-}
482repeat :: Monad m => a -> Stream m a
483repeat x = Stream (\_ _ -> return $ Yield x ()) ()
485{-# INLINE_NORMAL iterateM #-}
486iterateM :: Monad m => (a -> m a) -> m a -> Stream m a
487iterateM step = Stream (\_ st -> st >>= \x -> return $ Yield x (step x))
489{-# INLINE_NORMAL iterate #-}
490iterate :: Monad m => (a -> a) -> a -> Stream m a
491iterate step st = iterateM (return . step) (return st)
493{-# INLINE_NORMAL replicateM #-}
494replicateM :: forall m a. Monad m => Int -> m a -> Stream m a
495replicateM n p = Stream step n
497 {-# INLINE_LATE step #-}
499 | i <= 0 = return Stop
502 return $ Yield x (i - 1)
504{-# INLINE_NORMAL replicate #-}
505replicate :: Monad m => Int -> a -> Stream m a
506replicate n x = replicateM n (return x)
508-- This would not work properly for floats, therefore we put an Integral
510-- | Can be used to enumerate unbounded integrals. This does not check for
511-- overflow or underflow for bounded integrals.
512{-# INLINE_NORMAL enumerateFromStepIntegral #-}
513enumerateFromStepIntegral :: (Integral a, Monad m) => a -> a -> Stream m a
514enumerateFromStepIntegral from stride =
515 from `seq` stride `seq` Stream step from
517 {-# INLINE_LATE step #-}
518 step _ !x = return $ Yield x $! (x + stride)
520-- We are assuming that "to" is constrained by the type to be within
522{-# INLINE enumerateFromToIntegral #-}
523enumerateFromToIntegral :: (Monad m, Integral a) => a -> a -> Stream m a
524enumerateFromToIntegral from to =
525 takeWhile (<= to) $ enumerateFromStepIntegral from 1
527{-# INLINE enumerateFromIntegral #-}
528enumerateFromIntegral :: (Monad m, Integral a, Bounded a) => a -> Stream m a
529enumerateFromIntegral from = enumerateFromToIntegral from maxBound
531data EnumState a = EnumInit | EnumYield a a a | EnumStop
533{-# INLINE_NORMAL enumerateFromThenToIntegralUp #-}
534enumerateFromThenToIntegralUp
535 :: (Monad m, Integral a)
536 => a -> a -> a -> Stream m a
537enumerateFromThenToIntegralUp from next to = Stream step EnumInit
539 {-# INLINE_LATE step #-}
545 else Yield from EnumStop
546 else -- from <= next <= to
547 let stride = next - from
548 in Skip $ EnumYield from stride (to - stride)
550 step _ (EnumYield x stride toMinus) =
553 then Yield x EnumStop
554 else Yield x $ EnumYield (x + stride) stride toMinus
556 step _ EnumStop = return Stop
558{-# INLINE_NORMAL enumerateFromThenToIntegralDn #-}
559enumerateFromThenToIntegralDn
560 :: (Monad m, Integral a)
561 => a -> a -> a -> Stream m a
562enumerateFromThenToIntegralDn from next to = Stream step EnumInit
564 {-# INLINE_LATE step #-}
566 return $ if to > next
569 else Yield from EnumStop
570 else -- from >= next >= to
571 let stride = next - from
572 in Skip $ EnumYield from stride (to - stride)
574 step _ (EnumYield x stride toMinus) =
577 then Yield x EnumStop
578 else Yield x $ EnumYield (x + stride) stride toMinus
580 step _ EnumStop = return Stop
582{-# INLINE_NORMAL enumerateFromThenToIntegral #-}
583enumerateFromThenToIntegral
584 :: (Monad m, Integral a)
585 => a -> a -> a -> Stream m a
586enumerateFromThenToIntegral from next to
587 | next >= from = enumerateFromThenToIntegralUp from next to
588 | otherwise = enumerateFromThenToIntegralDn from next to
590{-# INLINE_NORMAL enumerateFromThenIntegral #-}
591enumerateFromThenIntegral
592 :: (Monad m, Integral a, Bounded a)
593 => a -> a -> Stream m a
594enumerateFromThenIntegral from next =
596 then enumerateFromThenToIntegralUp from next maxBound
597 else enumerateFromThenToIntegralDn from next minBound
599-- For floating point numbers if the increment is less than the precision then
600-- it just gets lost. Therefore we cannot always increment it correctly by just
602-- 9007199254740992 + 1 + 1 :: Double => 9.007199254740992e15
603-- 9007199254740992 + 2 :: Double => 9.007199254740994e15
605-- Instead we accumulate the increment counter and compute the increment
606-- every time before adding it to the starting number.
608-- This works for Integrals as well as floating point numbers, but
609-- enumerateFromStepIntegral is faster for integrals.
610{-# INLINE_NORMAL enumerateFromStepNum #-}
611enumerateFromStepNum :: (Monad m, Num a) => a -> a -> Stream m a
612enumerateFromStepNum from stride = Stream step 0
614 {-# INLINE_LATE step #-}
615 step _ !i = return $ (Yield $! (from + i * stride)) $! (i + 1)
617{-# INLINE_NORMAL numFrom #-}
618numFrom :: (Monad m, Num a) => a -> Stream m a
619numFrom from = enumerateFromStepNum from 1
621{-# INLINE_NORMAL numFromThen #-}
622numFromThen :: (Monad m, Num a) => a -> a -> Stream m a
623numFromThen from next = enumerateFromStepNum from (next - from)
625-- We cannot write a general function for Num. The only way to write code
626-- portable between the two is to use a 'Real' constraint and convert between
627-- Fractional and Integral using fromRational which is horribly slow.
628{-# INLINE_NORMAL enumerateFromToFractional #-}
629enumerateFromToFractional
630 :: (Monad m, Fractional a, Ord a)
631 => a -> a -> Stream m a
632enumerateFromToFractional from to =
633 takeWhile (<= to + 1 / 2) $ enumerateFromStepNum from 1
635{-# INLINE_NORMAL enumerateFromThenToFractional #-}
636enumerateFromThenToFractional
637 :: (Monad m, Fractional a, Ord a)
638 => a -> a -> a -> Stream m a
639enumerateFromThenToFractional from next to =
640 takeWhile predicate $ numFromThen from next
642 mid = (next - from) / 2
643 predicate | next >= from = (<= to + mid)
644 | otherwise = (>= to + mid)
646-------------------------------------------------------------------------------
647-- Generation by Conversion
648-------------------------------------------------------------------------------
650{-# INLINE_NORMAL fromIndicesM #-}
651fromIndicesM :: Monad m => (Int -> m a) -> Stream m a
652fromIndicesM gen = Stream step 0
654 {-# INLINE_LATE step #-}
657 return $ Yield x (i + 1)
659{-# INLINE fromIndices #-}
660fromIndices :: Monad m => (Int -> a) -> Stream m a
661fromIndices gen = fromIndicesM (return . gen)
663{-# INLINE_NORMAL generateM #-}
664generateM :: Monad m => Int -> (Int -> m a) -> Stream m a
665generateM n gen = n `seq` Stream step 0
667 {-# INLINE_LATE step #-}
668 step _ i | i < n = do
670 return $ Yield x (i + 1)
671 | otherwise = return Stop
673{-# INLINE generate #-}
674generate :: Monad m => Int -> (Int -> a) -> Stream m a
675generate n gen = generateM n (return . gen)
677-- XXX we need the MonadAsync constraint because of a rewrite rule.
678-- | Convert a list of monadic actions to a 'Stream'
679{-# INLINE_LATE fromListM #-}
680fromListM :: MonadAsync m => [m a] -> Stream m a
681fromListM = Stream step
683 {-# INLINE_LATE step #-}
684 step _ (m:ms) = m >>= \x -> return $ Yield x ms
685 step _ [] = return Stop
687{-# INLINE toStreamD #-}
688toStreamD :: (K.IsStream t, Monad m) => t m a -> Stream m a
689toStreamD = fromStreamK . K.toStream
691{-# INLINE_NORMAL fromPrimVar #-}
692fromPrimVar :: (MonadIO m, Prim a) => Var IO a -> Stream m a
693fromPrimVar var = Stream step ()
695 {-# INLINE_LATE step #-}
696 step _ () = liftIO (readVar var) >>= \x -> return $ Yield x ()
698-------------------------------------------------------------------------------
699-- Generation from SVar
700-------------------------------------------------------------------------------
702data FromSVarState t m a =
704 | FromSVarRead (SVar t m a)
705 | FromSVarLoop (SVar t m a) [ChildEvent a]
706 | FromSVarDone (SVar t m a)
708{-# INLINE_NORMAL fromSVar #-}
709fromSVar :: (MonadAsync m) => SVar t m a -> Stream m a
710fromSVar svar = Stream step FromSVarInit
713 {-# INLINE_LATE step #-}
714 step _ FromSVarInit = do
715 ref <- liftIO $ newIORef ()
716 _ <- liftIO $ mkWeakIORef ref hook
717 -- when this copy of svar gets garbage collected "ref" will get
718 -- garbage collected and our GC hook will be called.
719 let sv = svar{svarRef = Just ref}
720 return $ Skip (FromSVarRead sv)
724 {-# NOINLINE hook #-}
726 when (svarInspectMode svar) $ do
727 r <- liftIO $ readIORef (svarStopTime (svarStats svar))
729 printSVar svar "SVar Garbage Collected"
731 -- If there are any SVars referenced by this SVar a GC will prompt
732 -- them to be cleaned up quickly.
733 when (svarInspectMode svar) performMajorGC
735 step _ (FromSVarRead sv) = do
736 list <- readOutputQ sv
737 -- Reversing the output is important to guarantee that we process the
738 -- outputs in the same order as they were generated by the constituent
740 return $ Skip $ FromSVarLoop sv (Prelude.reverse list)
742 step _ (FromSVarLoop sv []) = do
743 done <- postProcess sv
744 return $ Skip $ if done
745 then (FromSVarDone sv)
746 else (FromSVarRead sv)
748 step _ (FromSVarLoop sv (ev : es)) = do
750 ChildYield a -> return $ Yield a (FromSVarLoop sv es)
751 ChildStop tid e -> do
755 stop <- shouldStop tid
758 liftIO (cleanupSVar sv)
759 return $ Skip (FromSVarDone sv)
760 else return $ Skip (FromSVarLoop sv es)
762 case fromException ex of
764 return $ Skip (FromSVarLoop sv es)
765 Nothing -> liftIO (cleanupSVar sv) >> throwM ex
769 case svarStopStyle sv of
770 StopNone -> return False
771 StopAny -> return True
773 sid <- liftIO $ readIORef (svarStopBy sv)
774 return $ if tid == sid then True else False
776 step _ (FromSVarDone sv) = do
777 when (svarInspectMode sv) $ do
778 t <- liftIO $ getTime Monotonic
779 liftIO $ writeIORef (svarStopTime (svarStats sv)) (Just t)
780 liftIO $ printSVar sv "SVar Done"
783-------------------------------------------------------------------------------
784-- Process events received by a fold consumer from a stream producer
785-------------------------------------------------------------------------------
787{-# INLINE_NORMAL fromProducer #-}
788fromProducer :: (MonadAsync m) => SVar t m a -> Stream m a
789fromProducer svar = Stream step (FromSVarRead svar)
792 {-# INLINE_LATE step #-}
793 step _ (FromSVarRead sv) = do
794 list <- readOutputQ sv
795 -- Reversing the output is important to guarantee that we process the
796 -- outputs in the same order as they were generated by the constituent
798 return $ Skip $ FromSVarLoop sv (Prelude.reverse list)
800 step _ (FromSVarLoop sv []) = return $ Skip $ FromSVarRead sv
801 step _ (FromSVarLoop sv (ev : es)) = do
803 ChildYield a -> return $ Yield a (FromSVarLoop sv es)
804 ChildStop tid e -> do
808 sendStopToProducer sv
809 return $ Skip (FromSVarDone sv)
810 Just _ -> error "Bug: fromProducer: received exception"
812 step _ (FromSVarDone sv) = do
813 when (svarInspectMode sv) $ do
814 t <- liftIO $ getTime Monotonic
815 liftIO $ writeIORef (svarStopTime (svarStats sv)) (Just t)
816 liftIO $ printSVar sv "SVar Done"
819 step _ FromSVarInit = undefined
821-------------------------------------------------------------------------------
822-- Hoisting the inner monad
823-------------------------------------------------------------------------------
825{-# INLINE_NORMAL hoist #-}
826hoist :: Monad n => (forall x. m x -> n x) -> Stream m a -> Stream n a
827hoist f (Stream step state) = (Stream step' state)
829 {-# INLINE_LATE step' #-}
831 r <- f $ step (adaptState gst) st
833 Yield x s -> Yield x s
837{-# INLINE generally #-}
838generally :: Monad m => Stream Identity a -> Stream m a
839generally = hoist (return . runIdentity)
841{-# INLINE_NORMAL liftInner #-}
842liftInner :: (Monad m, MonadTrans t, Monad (t m))
843 => Stream m a -> Stream (t m) a
844liftInner (Stream step state) = Stream step' state
846 {-# INLINE_LATE step' #-}
848 r <- lift $ step (adaptState gst) st
850 Yield x s -> Yield x s
854{-# INLINE_NORMAL runReaderT #-}
855runReaderT :: Monad m => s -> Stream (ReaderT s m) a -> Stream m a
856runReaderT sval (Stream step state) = Stream step' state
858 {-# INLINE_LATE step' #-}
860 r <- Reader.runReaderT (step (adaptState gst) st) sval
862 Yield x s -> Yield x s
866{-# INLINE_NORMAL evalStateT #-}
867evalStateT :: Monad m => s -> Stream (StateT s m) a -> Stream m a
868evalStateT sval (Stream step state) = Stream step' (state, sval)
870 {-# INLINE_LATE step' #-}
871 step' gst (st, sv) = do
872 (r, sv') <- State.runStateT (step (adaptState gst) st) sv
874 Yield x s -> Yield x (s, sv')
875 Skip s -> Skip (s, sv')
878{-# INLINE_NORMAL runStateT #-}
879runStateT :: Monad m => s -> Stream (StateT s m) a -> Stream m (s, a)
880runStateT sval (Stream step state) = Stream step' (state, sval)
882 {-# INLINE_LATE step' #-}
883 step' gst (st, sv) = do
884 (r, sv') <- State.runStateT (step (adaptState gst) st) sv
886 Yield x s -> Yield (sv', x) (s, sv')
887 Skip s -> Skip (s, sv')
890------------------------------------------------------------------------------
891-- Elimination by Folds
892------------------------------------------------------------------------------
894------------------------------------------------------------------------------
896------------------------------------------------------------------------------
898{-# INLINE_NORMAL foldr1 #-}
899foldr1 :: Monad m => (a -> a -> a) -> Stream m a -> m (Maybe a)
903 Nothing -> return Nothing
904 Just (h, t) -> fmap Just (foldr f h t)
906------------------------------------------------------------------------------
908------------------------------------------------------------------------------
910{-# INLINE_NORMAL foldlT #-}
911foldlT :: (Monad m, Monad (s m), MonadTrans s)
912 => (s m b -> a -> s m b) -> s m b -> Stream m a -> s m b
913foldlT fstep begin (Stream step state) = go SPEC begin state
916 r <- lift $ step defState st
918 Yield x s -> go SPEC (fstep acc x) s
919 Skip s -> go SPEC acc s
922-- Note, this is going to have horrible performance, because of the nature of
923-- the stream type (i.e. direct stream vs CPS). Its only for reference, it is
924-- likely be practically unusable.
925{-# INLINE_NORMAL foldlS #-}
927 => (Stream m b -> a -> Stream m b) -> Stream m b -> Stream m a -> Stream m b
928foldlS fstep begin (Stream step state) = Stream step' (Left (state, begin))
930 step' gst (Left (st, acc)) = do
931 r <- step (adaptState gst) st
933 Yield x s -> Skip (Left (s, fstep acc x))
934 Skip s -> Skip (Left (s, acc))
935 Stop -> Skip (Right acc)
937 step' gst (Right (Stream stp stt)) = do
938 r <- stp (adaptState gst) stt
940 Yield x s -> Yield x (Right (Stream stp s))
941 Skip s -> Skip (Right (Stream stp s))
944------------------------------------------------------------------------------
946------------------------------------------------------------------------------
948-- | Run a streaming composition, discard the results.
949{-# INLINE_LATE drain #-}
950drain :: Monad m => Stream m a -> m ()
951-- drain = foldrM (\_ xs -> xs) (return ())
952drain (Stream step state) = go SPEC state
955 r <- step defState st
957 Yield _ s -> go SPEC s
961{-# INLINE_NORMAL null #-}
962null :: Monad m => Stream m a -> m Bool
963null m = foldrM (\_ _ -> return False) (return True) m
965{-# INLINE_NORMAL head #-}
966head :: Monad m => Stream m a -> m (Maybe a)
967head m = foldrM (\x _ -> return (Just x)) (return Nothing) m
969{-# INLINE_NORMAL headElse #-}
970headElse :: Monad m => a -> Stream m a -> m a
971headElse a m = foldrM (\x _ -> return x) (return a) m
973-- Does not fuse, has the same performance as the StreamK version.
974{-# INLINE_NORMAL tail #-}
975tail :: Monad m => Stream m a -> m (Maybe (Stream m a))
976tail (UnStream step state) = go state
979 r <- step defState st
981 Yield _ s -> return (Just $ Stream step s)
983 Stop -> return Nothing
985-- XXX will it fuse? need custom impl?
986{-# INLINE_NORMAL last #-}
987last :: Monad m => Stream m a -> m (Maybe a)
988last = foldl' (\_ y -> Just y) Nothing
990{-# INLINE_NORMAL elem #-}
991elem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
992-- elem e m = foldrM (\x xs -> if x == e then return True else xs) (return False) m
993elem e (Stream step state) = go state
996 r <- step defState st
999 | x == e -> return True
1002 Stop -> return False
1004{-# INLINE_NORMAL notElem #-}
1005notElem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
1006notElem e s = fmap not (elem e s)
1008{-# INLINE_NORMAL all #-}
1009all :: Monad m => (a -> Bool) -> Stream m a -> m Bool
1010-- all p m = foldrM (\x xs -> if p x then xs else return False) (return True) m
1011all p (Stream step state) = go state
1014 r <- step defState st
1018 | otherwise -> return False
1022{-# INLINE_NORMAL any #-}
1023any :: Monad m => (a -> Bool) -> Stream m a -> m Bool
1024-- any p m = foldrM (\x xs -> if p x then return True else xs) (return False) m
1025any p (Stream step state) = go state
1028 r <- step defState st
1031 | p x -> return True
1034 Stop -> return False
1036{-# INLINE_NORMAL maximum #-}
1037maximum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
1038maximum (Stream step state) = go Nothing state
1041 r <- step defState st
1043 Yield x s -> go (Just x) s
1044 Skip s -> go Nothing s
1045 Stop -> return Nothing
1046 go (Just acc) st = do
1047 r <- step defState st
1050 | acc <= x -> go (Just x) s
1051 | otherwise -> go (Just acc) s
1052 Skip s -> go (Just acc) s
1053 Stop -> return (Just acc)
1055{-# INLINE_NORMAL maximumBy #-}
1056maximumBy :: Monad m => (a -> a -> Ordering) -> Stream m a -> m (Maybe a)
1057maximumBy cmp (Stream step state) = go Nothing state
1060 r <- step defState st
1062 Yield x s -> go (Just x) s
1063 Skip s -> go Nothing s
1064 Stop -> return Nothing
1065 go (Just acc) st = do
1066 r <- step defState st
1068 Yield x s -> case cmp acc x of
1069 GT -> go (Just acc) s
1071 Skip s -> go (Just acc) s
1072 Stop -> return (Just acc)
1074{-# INLINE_NORMAL minimum #-}
1075minimum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
1076minimum (Stream step state) = go Nothing state
1079 r <- step defState st
1081 Yield x s -> go (Just x) s
1082 Skip s -> go Nothing s
1083 Stop -> return Nothing
1084 go (Just acc) st = do
1085 r <- step defState st
1088 | acc <= x -> go (Just acc) s
1089 | otherwise -> go (Just x) s
1090 Skip s -> go (Just acc) s
1091 Stop -> return (Just acc)
1093{-# INLINE_NORMAL minimumBy #-}
1094minimumBy :: Monad m => (a -> a -> Ordering) -> Stream m a -> m (Maybe a)
1095minimumBy cmp (Stream step state) = go Nothing state
1098 r <- step defState st
1100 Yield x s -> go (Just x) s
1101 Skip s -> go Nothing s
1102 Stop -> return Nothing
1103 go (Just acc) st = do
1104 r <- step defState st
1106 Yield x s -> case cmp acc x of
1108 _ -> go (Just acc) s
1109 Skip s -> go (Just acc) s
1110 Stop -> return (Just acc)
1112{-# INLINE_NORMAL (!!) #-}
1113(!!) :: (Monad m) => Stream m a -> Int -> m (Maybe a)
1114(Stream step state) !! i = go i state
1117 r <- step defState st
1119 Yield x s | n < 0 -> return Nothing
1120 | n == 0 -> return $ Just x
1121 | otherwise -> go (n - 1) s
1123 Stop -> return Nothing
1125{-# INLINE_NORMAL lookup #-}
1126lookup :: (Monad m, Eq a) => a -> Stream m (a, b) -> m (Maybe b)
1127lookup e m = foldrM (\(a, b) xs -> if e == a then return (Just b) else xs)
1130{-# INLINE_NORMAL findM #-}
1131findM :: Monad m => (a -> m Bool) -> Stream m a -> m (Maybe a)
1132findM p m = foldrM (\x xs -> p x >>= \r -> if r then return (Just x) else xs)
1136find :: Monad m => (a -> Bool) -> Stream m a -> m (Maybe a)
1137find p = findM (return . p)
1139{-# INLINE_NORMAL findIndices #-}
1140findIndices :: Monad m => (a -> Bool) -> Stream m a -> Stream m Int
1141findIndices p (Stream step state) = Stream step' (state, 0)
1143 {-# INLINE_LATE step' #-}
1144 step' gst (st, i) = i `seq` do
1145 r <- step (adaptState gst) st
1147 Yield x s -> if p x then Yield i (s, i+1) else Skip (s, i+1)
1148 Skip s -> Skip (s, i)
1151{-# INLINE toListRev #-}
1152toListRev :: Monad m => Stream m a -> m [a]
1153toListRev = foldl' (flip (:)) []
1155-- We can implement reverse as:
1157-- > reverse = foldlS (flip cons) nil
1159-- However, this implementation is unusable because of the horrible performance
1160-- of cons. So we just convert it to a list first and then stream from the
1163-- XXX Maybe we can use an Array instead of a list here?
1164{-# INLINE_NORMAL reverse #-}
1165reverse :: Monad m => Stream m a -> Stream m a
1166reverse m = Stream step Nothing
1168 {-# INLINE_LATE step #-}
1171 return $ Skip (Just xs)
1172 step _ (Just (x:xs)) = return $ Yield x (Just xs)
1173 step _ (Just []) = return Stop
1175-- Much faster reverse for Storables
1176{-# INLINE_NORMAL reverse' #-}
1177reverse' :: forall m a. (MonadIO m, Storable a) => Stream m a -> Stream m a
1179-- This commented implementation copies the whole stream into one single array
1180-- and then streams from that array, this is 3-4x faster than the chunked code
1181-- that follows. Though this could be problematic due to unbounded large
1182-- allocations. We need to figure out why the chunked code is slower and if we
1183-- can optimize the chunked code to work as fast as this one. It may be a
1185import Foreign.ForeignPtr (touchForeignPtr)
1186import Foreign.ForeignPtr.Unsafe (unsafeForeignPtrToPtr)
1187import Foreign.Ptr (Ptr, plusPtr)
1188reverse' m = Stream step Nothing
1190 {-# INLINE_LATE step #-}
1192 arr <- A.fromStreamD m
1193 let p = aEnd arr `plusPtr` negate (sizeOf (undefined :: a))
1194 return $ Skip $ Just (aStart arr, p)
1196 step _ (Just (start, p)) | p < unsafeForeignPtrToPtr start = return Stop
1198 step _ (Just (start, p)) = do
1199 let !x = A.unsafeInlineIO $ do
1201 touchForeignPtr start
1203 next = p `plusPtr` negate (sizeOf (undefined :: a))
1204 return $ Yield x (Just (start, next))
1211 $ A.fromStreamDArraysOf A.defaultChunkSize m
1214------------------------------------------------------------------------------
1215-- Grouping/Splitting
1216------------------------------------------------------------------------------
1218{-# INLINE_NORMAL splitSuffixBy' #-}
1219splitSuffixBy' :: Monad m
1220 => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1221splitSuffixBy' predicate f (Stream step state) =
1222 Stream (stepOuter f) (Just state)
1226 {-# INLINE_LATE stepOuter #-}
1227 stepOuter (Fold fstep initial done) gst (Just st) = do
1228 res <- step (adaptState gst) st
1234 then done acc' >>= \val -> return $ Yield val (Just s)
1237 Skip s -> return $ Skip $ Just s
1243 res <- step (adaptState gst) stt
1248 then done acc' >>= \val -> return $ Yield val (Just s)
1250 Skip s -> go SPEC s acc
1251 Stop -> done acc >>= \val -> return $ Yield val Nothing
1253 stepOuter _ _ Nothing = return Stop
1255{-# INLINE_NORMAL groupsBy #-}
1261groupsBy cmp f (Stream step state) = Stream (stepOuter f) (Just state, Nothing)
1265 {-# INLINE_LATE stepOuter #-}
1266 stepOuter (Fold fstep initial done) gst (Just st, Nothing) = do
1267 res <- step (adaptState gst) st
1274 Skip s -> return $ Skip $ (Just s, Nothing)
1279 go !_ prev stt !acc = do
1280 res <- step (adaptState gst) stt
1287 else done acc >>= \r -> return $ Yield r (Just s, Just x)
1288 Skip s -> go SPEC prev s acc
1289 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1291 stepOuter (Fold fstep initial done) gst (Just st, Just prev) = do
1293 acc' <- fstep acc prev
1298 -- XXX code duplicated from the previous equation
1300 res <- step (adaptState gst) stt
1307 else done acc >>= \r -> return $ Yield r (Just s, Just x)
1308 Skip s -> go SPEC s acc
1309 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1311 stepOuter _ _ (Nothing,_) = return Stop
1313{-# INLINE_NORMAL groupsRollingBy #-}
1314groupsRollingBy :: Monad m
1319groupsRollingBy cmp f (Stream step state) =
1320 Stream (stepOuter f) (Just state, Nothing)
1323 {-# INLINE_LATE stepOuter #-}
1324 stepOuter (Fold fstep initial done) gst (Just st, Nothing) = do
1325 res <- step (adaptState gst) st
1332 Skip s -> return $ Skip $ (Just s, Nothing)
1336 go !_ prev stt !acc = do
1337 res <- step (adaptState gst) stt
1345 done acc >>= \r -> return $ Yield r (Just s, Just x)
1346 Skip s -> go SPEC prev s acc
1347 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1349 stepOuter (Fold fstep initial done) gst (Just st, Just prev') = do
1351 acc' <- fstep acc prev'
1352 go SPEC prev' st acc'
1355 go !_ prevv stt !acc = do
1356 res <- step (adaptState gst) stt
1363 else done acc >>= \r -> return $ Yield r (Just s, Just x)
1364 Skip s -> go SPEC prevv s acc
1365 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1367 stepOuter _ _ (Nothing, _) = return Stop
1369{-# INLINE_NORMAL splitBy #-}
1370splitBy :: Monad m => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1371splitBy predicate f (Stream step state) = Stream (step' f) (Just state)
1375 {-# INLINE_LATE step' #-}
1376 step' (Fold fstep initial done) gst (Just st) = initial >>= go SPEC st
1381 res <- step (adaptState gst) stt
1385 then done acc >>= \r -> return $ Yield r (Just s)
1389 Skip s -> go SPEC s acc
1390 Stop -> done acc >>= \r -> return $ Yield r Nothing
1392 step' _ _ Nothing = return Stop
1394-- XXX requires -funfolding-use-threshold=150 in lines-unlines benchmark
1395{-# INLINE_NORMAL splitSuffixBy #-}
1396splitSuffixBy :: Monad m
1397 => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1398splitSuffixBy predicate f (Stream step state) = Stream (step' f) (Just state)
1402 {-# INLINE_LATE step' #-}
1403 step' (Fold fstep initial done) gst (Just st) = do
1404 res <- step (adaptState gst) st
1409 then done acc >>= \val -> return $ Yield val (Just s)
1414 Skip s -> return $ Skip $ Just s
1420 res <- step (adaptState gst) stt
1424 then done acc >>= \r -> return $ Yield r (Just s)
1428 Skip s -> go SPEC s acc
1429 Stop -> done acc >>= \r -> return $ Yield r Nothing
1431 step' _ _ Nothing = return Stop
1433{-# INLINE_NORMAL wordsBy #-}
1434wordsBy :: Monad m => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1435wordsBy predicate f (Stream step state) = Stream (stepOuter f) (Just state)
1439 {-# INLINE_LATE stepOuter #-}
1440 stepOuter (Fold fstep initial done) gst (Just st) = do
1441 res <- step (adaptState gst) st
1445 then return $ Skip (Just s)
1451 Skip s -> return $ Skip $ Just s
1457 res <- step (adaptState gst) stt
1461 then done acc >>= \r -> return $ Yield r (Just s)
1465 Skip s -> go SPEC s acc
1466 Stop -> done acc >>= \r -> return $ Yield r Nothing
1468 stepOuter _ _ Nothing = return Stop
1470-- String search algorithms:
1471-- http://www-igm.univ-mlv.fr/~lecroq/string/index.html
1474-- TODO can we unify the splitting operations using a splitting configuration
1475-- like in the split package.
1477data SplitStyle = Infix | Suffix | Prefix deriving (Eq, Show)
1479data SplitOptions = SplitOptions
1480 { style :: SplitStyle
1481 , withSep :: Bool -- ^ keep the separators in output
1482 -- , compact :: Bool -- ^ treat multiple consecutive separators as one
1483 -- , trimHead :: Bool -- ^ drop blank at head
1484 -- , trimTail :: Bool -- ^ drop blank at tail
1488data SplitOnState s a =
1493 | GO_KARP_RABIN s !(RB.Ring a) !(Ptr a)
1496{-# INLINE_NORMAL splitOn #-}
1498 :: forall m a b. (MonadIO m, Storable a, Enum a, Eq a)
1503splitOn patArr@Array{..} (Fold fstep initial done) (Stream step state) =
1504 Stream stepOuter GO_START
1508 patLen = A.length patArr
1509 maxIndex = patLen - 1
1510 elemBits = sizeOf (undefined :: a) * 8
1512 {-# INLINE_LATE stepOuter #-}
1513 stepOuter _ GO_START =
1515 then return $ Skip $ GO_EMPTY_PAT state
1518 r <- liftIO $ (A.unsafeIndexIO patArr 0)
1519 return $ Skip $ GO_SINGLE_PAT state r
1520 else if sizeOf (undefined :: a) * patLen
1521 <= sizeOf (undefined :: Word)
1522 then return $ Skip $ GO_SHORT_PAT state
1524 (rb, rhead) <- liftIO $ RB.new patLen
1525 return $ Skip $ GO_KARP_RABIN state rb rhead
1527 stepOuter gst (GO_SINGLE_PAT stt pat) = initial >>= go SPEC stt
1532 res <- step (adaptState gst) st
1538 return $ Yield r (GO_SINGLE_PAT s pat)
1539 else fstep acc x >>= go SPEC s
1540 Skip s -> go SPEC s acc
1541 Stop -> done acc >>= \r -> return $ Yield r GO_DONE
1543 stepOuter gst (GO_SHORT_PAT stt) = initial >>= go0 SPEC 0 (0 :: Word) stt
1548 mask = (1 `shiftL` (elemBits * patLen)) - 1
1550 addToWord wrd a = (wrd `shiftL` elemBits) .|. fromIntegral (fromEnum a)
1553 patWord = mask .&. A.foldl' addToWord 0 patArr
1555 go0 !_ !idx wrd st !acc = do
1556 res <- step (adaptState gst) st
1559 let wrd' = addToWord wrd x
1562 if wrd' .&. mask == patWord
1565 return $ Yield r (GO_SHORT_PAT s)
1566 else go1 SPEC wrd' s acc
1567 else go0 SPEC (idx + 1) wrd' s acc
1568 Skip s -> go0 SPEC idx wrd s acc
1571 then go2 wrd idx acc
1573 done acc' >>= \r -> return $ Yield r GO_DONE
1576 go1 !_ wrd st !acc = do
1577 res <- step (adaptState gst) st
1580 let wrd' = addToWord wrd x
1581 old = (mask .&. wrd) `shiftR` (elemBits * (patLen - 1))
1582 acc' <- fstep acc (toEnum $ fromIntegral old)
1583 if wrd' .&. mask == patWord
1584 then done acc' >>= \r -> return $ Yield r (GO_SHORT_PAT s)
1585 else go1 SPEC wrd' s acc'
1586 Skip s -> go1 SPEC wrd s acc
1588 acc' <- go2 wrd patLen acc
1589 done acc' >>= \r -> return $ Yield r GO_DONE
1591 go2 !wrd !n !acc | n > 0 = do
1592 let old = (mask .&. wrd) `shiftR` (elemBits * (n - 1))
1593 fstep acc (toEnum $ fromIntegral old) >>= go2 wrd (n - 1)
1594 go2 _ _ acc = return acc
1596 stepOuter gst (GO_KARP_RABIN stt rb rhead) = do
1597 initial >>= go0 SPEC 0 rhead stt
1601 k = 2891336453 :: Word32
1603 addCksum cksum a = cksum * k + fromIntegral (fromEnum a)
1604 deltaCksum cksum old new =
1605 addCksum cksum new - coeff * fromIntegral (fromEnum old)
1607 -- XXX shall we use a random starting hash or 1 instead of 0?
1608 patHash = A.foldl' addCksum 0 patArr
1611 go0 !_ !idx !rh st !acc = do
1612 res <- step (adaptState gst) st
1615 rh' <- liftIO $ RB.unsafeInsert rb rh x
1618 let fold = RB.unsafeFoldRing (RB.ringBound rb)
1619 let !ringHash = fold addCksum 0 rb
1620 if ringHash == patHash
1621 then go2 SPEC ringHash rh' s acc
1622 else go1 SPEC ringHash rh' s acc
1623 else go0 SPEC (idx + 1) rh' s acc
1624 Skip s -> go0 SPEC idx rh s acc
1626 !acc' <- if idx /= 0
1627 then RB.unsafeFoldRingM rh fstep acc rb
1629 done acc' >>= \r -> return $ Yield r GO_DONE
1631 -- XXX Theoretically this code can do 4 times faster if GHC generates
1632 -- optimal code. If we use just "(cksum' == patHash)" condition it goes
1633 -- 4x faster, as soon as we add the "RB.unsafeEqArray rb v" condition
1634 -- the generated code changes drastically and becomes 4x slower. Need
1635 -- to investigate what is going on with GHC.
1637 go1 !_ !cksum !rh st !acc = do
1638 res <- step (adaptState gst) st
1641 old <- liftIO $ peek rh
1642 let cksum' = deltaCksum cksum old x
1643 acc' <- fstep acc old
1645 if (cksum' == patHash)
1647 rh' <- liftIO (RB.unsafeInsert rb rh x)
1648 go2 SPEC cksum' rh' s acc'
1650 rh' <- liftIO (RB.unsafeInsert rb rh x)
1651 go1 SPEC cksum' rh' s acc'
1652 Skip s -> go1 SPEC cksum rh s acc
1654 acc' <- RB.unsafeFoldRingFullM rh fstep acc rb
1655 done acc' >>= \r -> return $ Yield r GO_DONE
1657 go2 !_ !cksum' !rh' s !acc' = do
1658 if RB.unsafeEqArray rb rh' patArr
1661 return $ Yield r (GO_KARP_RABIN s rb rhead)
1662 else go1 SPEC cksum' rh' s acc'
1664 stepOuter gst (GO_EMPTY_PAT st) = do
1665 res <- step (adaptState gst) st
1670 done acc' >>= \r -> return $ Yield r (GO_EMPTY_PAT s)
1671 Skip s -> return $ Skip (GO_EMPTY_PAT s)
1674 stepOuter _ GO_DONE = return Stop
1676{-# INLINE_NORMAL splitSuffixOn #-}
1678 :: forall m a b. (MonadIO m, Storable a, Enum a, Eq a)
1684splitSuffixOn withSep patArr@Array{..} (Fold fstep initial done)
1685 (Stream step state) =
1686 Stream stepOuter GO_START
1690 patLen = A.length patArr
1691 maxIndex = patLen - 1
1692 elemBits = sizeOf (undefined :: a) * 8
1694 {-# INLINE_LATE stepOuter #-}
1695 stepOuter _ GO_START =
1697 then return $ Skip $ GO_EMPTY_PAT state
1700 r <- liftIO $ (A.unsafeIndexIO patArr 0)
1701 return $ Skip $ GO_SINGLE_PAT state r
1702 else if sizeOf (undefined :: a) * patLen
1703 <= sizeOf (undefined :: Word)
1704 then return $ Skip $ GO_SHORT_PAT state
1706 (rb, rhead) <- liftIO $ RB.new patLen
1707 return $ Skip $ GO_KARP_RABIN state rb rhead
1709 stepOuter gst (GO_SINGLE_PAT stt pat) = do
1710 -- This first part is the only difference between splitOn and
1712 -- If the last element is a separator do not issue a blank segment.
1713 res <- step (adaptState gst) stt
1719 acc' <- if withSep then fstep acc x else return acc
1720 done acc' >>= \r -> return $ Yield r (GO_SINGLE_PAT s pat)
1721 else fstep acc x >>= go SPEC s
1722 Skip s -> return $ Skip $ (GO_SINGLE_PAT s pat)
1727 -- This is identical for splitOn and splitSuffixOn
1729 res <- step (adaptState gst) st
1734 acc' <- if withSep then fstep acc x else return acc
1736 return $ Yield r (GO_SINGLE_PAT s pat)
1737 else fstep acc x >>= go SPEC s
1738 Skip s -> go SPEC s acc
1739 Stop -> done acc >>= \r -> return $ Yield r GO_DONE
1741 stepOuter gst (GO_SHORT_PAT stt) = do
1743 -- Call "initial" only if the stream yields an element, otherwise we
1744 -- may call "initial" but never yield anything. initial may produce a
1745 -- side effect, therefore we will end up doing and discard a side
1750 res <- step (adaptState gst) stt
1754 let wrd' = addToWord wrd x
1755 acc' <- if withSep then fstep acc x else return acc
1758 if wrd' .&. mask == patWord
1759 then done acc' >>= \r -> return $ Yield r (GO_SHORT_PAT s)
1760 else go0 SPEC (idx + 1) wrd' s acc'
1761 else go0 SPEC (idx + 1) wrd' s acc'
1762 Skip s -> return $ Skip (GO_SHORT_PAT s)
1768 mask = (1 `shiftL` (elemBits * patLen)) - 1
1770 addToWord wrd a = (wrd `shiftL` elemBits) .|. fromIntegral (fromEnum a)
1773 patWord = mask .&. A.foldl' addToWord 0 patArr
1775 go0 !_ !idx wrd st !acc = do
1776 res <- step (adaptState gst) st
1779 let wrd' = addToWord wrd x
1780 acc' <- if withSep then fstep acc x else return acc
1783 if wrd' .&. mask == patWord
1786 return $ Yield r (GO_SHORT_PAT s)
1787 else go1 SPEC wrd' s acc'
1788 else go0 SPEC (idx + 1) wrd' s acc'
1789 Skip s -> go0 SPEC idx wrd s acc
1791 if (idx == maxIndex) && (wrd .&. mask == patWord)
1794 acc' <- if idx /= 0 && not withSep
1795 then go2 wrd idx acc
1797 done acc' >>= \r -> return $ Yield r GO_DONE
1800 go1 !_ wrd st !acc = do
1801 res <- step (adaptState gst) st
1804 let wrd' = addToWord wrd x
1805 old = (mask .&. wrd) `shiftR` (elemBits * (patLen - 1))
1808 else fstep acc (toEnum $ fromIntegral old)
1809 if wrd' .&. mask == patWord
1810 then done acc' >>= \r -> return $ Yield r (GO_SHORT_PAT s)
1811 else go1 SPEC wrd' s acc'
1812 Skip s -> go1 SPEC wrd s acc
1814 -- If the last sequence is a separator do not issue a blank
1816 if wrd .&. mask == patWord
1821 else go2 wrd patLen acc
1822 done acc' >>= \r -> return $ Yield r GO_DONE
1824 go2 !wrd !n !acc | n > 0 = do
1825 let old = (mask .&. wrd) `shiftR` (elemBits * (n - 1))
1826 fstep acc (toEnum $ fromIntegral old) >>= go2 wrd (n - 1)
1827 go2 _ _ acc = return acc
1829 stepOuter gst (GO_KARP_RABIN stt rb rhead) = do
1831 res <- step (adaptState gst) stt
1835 acc' <- if withSep then fstep acc x else return acc
1836 rh' <- liftIO (RB.unsafeInsert rb rhead x)
1839 let fold = RB.unsafeFoldRing (RB.ringBound rb)
1840 let !ringHash = fold addCksum 0 rb
1841 if ringHash == patHash
1842 then go2 SPEC ringHash rh' s acc'
1843 else go0 SPEC (idx + 1) rh' s acc'
1844 else go0 SPEC (idx + 1) rh' s acc'
1845 Skip s -> return $ Skip (GO_KARP_RABIN s rb rhead)
1850 k = 2891336453 :: Word32
1852 addCksum cksum a = cksum * k + fromIntegral (fromEnum a)
1853 deltaCksum cksum old new =
1854 addCksum cksum new - coeff * fromIntegral (fromEnum old)
1856 -- XXX shall we use a random starting hash or 1 instead of 0?
1857 patHash = A.foldl' addCksum 0 patArr
1860 go0 !_ !idx !rh st !acc = do
1861 res <- step (adaptState gst) st
1864 acc' <- if withSep then fstep acc x else return acc
1865 rh' <- liftIO (RB.unsafeInsert rb rh x)
1868 let fold = RB.unsafeFoldRing (RB.ringBound rb)
1869 let !ringHash = fold addCksum 0 rb
1870 if ringHash == patHash
1871 then go2 SPEC ringHash rh' s acc'
1872 else go1 SPEC ringHash rh' s acc'
1873 else go0 SPEC (idx + 1) rh' s acc'
1874 Skip s -> go0 SPEC idx rh s acc
1876 -- do not issue a blank segment when we end at pattern
1877 if (idx == maxIndex) && RB.unsafeEqArray rb rh patArr
1880 !acc' <- if idx /= 0 && not withSep
1881 then RB.unsafeFoldRingM rh fstep acc rb
1883 done acc' >>= \r -> return $ Yield r GO_DONE
1885 -- XXX Theoretically this code can do 4 times faster if GHC generates
1886 -- optimal code. If we use just "(cksum' == patHash)" condition it goes
1887 -- 4x faster, as soon as we add the "RB.unsafeEqArray rb v" condition
1888 -- the generated code changes drastically and becomes 4x slower. Need
1889 -- to investigate what is going on with GHC.
1891 go1 !_ !cksum !rh st !acc = do
1892 res <- step (adaptState gst) st
1895 old <- liftIO $ peek rh
1896 let cksum' = deltaCksum cksum old x
1901 if (cksum' == patHash)
1903 rh' <- liftIO (RB.unsafeInsert rb rh x)
1904 go2 SPEC cksum' rh' s acc'
1906 rh' <- liftIO (RB.unsafeInsert rb rh x)
1907 go1 SPEC cksum' rh' s acc'
1908 Skip s -> go1 SPEC cksum rh s acc
1910 if RB.unsafeEqArray rb rh patArr
1915 else RB.unsafeFoldRingFullM rh fstep acc rb
1916 done acc' >>= \r -> return $ Yield r GO_DONE
1918 go2 !_ !cksum' !rh' s !acc' = do
1919 if RB.unsafeEqArray rb rh' patArr
1922 return $ Yield r (GO_KARP_RABIN s rb rhead)
1923 else go1 SPEC cksum' rh' s acc'
1925 stepOuter gst (GO_EMPTY_PAT st) = do
1926 res <- step (adaptState gst) st
1931 done acc' >>= \r -> return $ Yield r (GO_EMPTY_PAT s)
1932 Skip s -> return $ Skip (GO_EMPTY_PAT s)
1935 stepOuter _ GO_DONE = return Stop
1937data SplitState s arr
1939 | SplitBuffering s arr
1940 | SplitSplitting s arr
1941 | SplitYielding arr (SplitState s arr)
1944-- XXX An alternative approach would be to use a partial fold (Fold m a b) to
1945-- split using a splitBy like combinator. The Fold would consume upto the
1946-- separator and return any leftover which can then be fed to the next fold.
1948-- We can revisit this once we have partial folds/parsers.
1950-- | Performs infix separator style splitting.
1951{-# INLINE_NORMAL splitInnerBy #-}
1954 => (f a -> m (f a, Maybe (f a))) -- splitter
1955 -> (f a -> f a -> m (f a)) -- joiner
1958splitInnerBy splitter joiner (Stream step1 state1) =
1959 (Stream step (SplitInitial state1))
1963 {-# INLINE_LATE step #-}
1964 step gst (SplitInitial st) = do
1968 (x1, mx2) <- splitter x
1969 return $ case mx2 of
1970 Nothing -> Skip (SplitBuffering s x1)
1971 Just x2 -> Skip (SplitYielding x1 (SplitSplitting s x2))
1972 Skip s -> return $ Skip (SplitInitial s)
1973 Stop -> return $ Stop
1975 step gst (SplitBuffering st buf) = do
1979 (x1, mx2) <- splitter x
1980 buf' <- joiner buf x1
1981 return $ case mx2 of
1982 Nothing -> Skip (SplitBuffering s buf')
1983 Just x2 -> Skip (SplitYielding buf' (SplitSplitting s x2))
1984 Skip s -> return $ Skip (SplitBuffering s buf)
1985 Stop -> return $ Skip (SplitYielding buf SplitFinishing)
1987 step _ (SplitSplitting st buf) = do
1988 (x1, mx2) <- splitter buf
1989 return $ case mx2 of
1990 Nothing -> Skip $ SplitBuffering st x1
1991 Just x2 -> Skip $ SplitYielding x1 (SplitSplitting st x2)
1993 step _ (SplitYielding x next) = return $ Yield x next
1994 step _ SplitFinishing = return $ Stop
1996-- | Performs infix separator style splitting.
1997{-# INLINE_NORMAL splitInnerBySuffix #-}
1999 :: (Monad m, Eq (f a), Monoid (f a))
2000 => (f a -> m (f a, Maybe (f a))) -- splitter
2001 -> (f a -> f a -> m (f a)) -- joiner
2004splitInnerBySuffix splitter joiner (Stream step1 state1) =
2005 (Stream step (SplitInitial state1))
2009 {-# INLINE_LATE step #-}
2010 step gst (SplitInitial st) = do
2014 (x1, mx2) <- splitter x
2015 return $ case mx2 of
2016 Nothing -> Skip (SplitBuffering s x1)
2017 Just x2 -> Skip (SplitYielding x1 (SplitSplitting s x2))
2018 Skip s -> return $ Skip (SplitInitial s)
2019 Stop -> return $ Stop
2021 step gst (SplitBuffering st buf) = do
2025 (x1, mx2) <- splitter x
2026 buf' <- joiner buf x1
2027 return $ case mx2 of
2028 Nothing -> Skip (SplitBuffering s buf')
2029 Just x2 -> Skip (SplitYielding buf' (SplitSplitting s x2))
2030 Skip s -> return $ Skip (SplitBuffering s buf)
2034 else Skip (SplitYielding buf SplitFinishing)
2036 step _ (SplitSplitting st buf) = do
2037 (x1, mx2) <- splitter buf
2038 return $ case mx2 of
2039 Nothing -> Skip $ SplitBuffering st x1
2040 Just x2 -> Skip $ SplitYielding x1 (SplitSplitting st x2)
2042 step _ (SplitYielding x next) = return $ Yield x next
2043 step _ SplitFinishing = return $ Stop
2045------------------------------------------------------------------------------
2047------------------------------------------------------------------------------
2049{-# INLINE_NORMAL isPrefixOf #-}
2050isPrefixOf :: (Eq a, Monad m) => Stream m a -> Stream m a -> m Bool
2051isPrefixOf (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
2053 go (sa, sb, Nothing) = do
2054 r <- stepa defState sa
2056 Yield x sa' -> go (sa', sb, Just x)
2057 Skip sa' -> go (sa', sb, Nothing)
2060 go (sa, sb, Just x) = do
2061 r <- stepb defState sb
2065 then go (sa, sb', Nothing)
2067 Skip sb' -> go (sa, sb', Just x)
2068 Stop -> return False
2070{-# INLINE_NORMAL isSubsequenceOf #-}
2071isSubsequenceOf :: (Eq a, Monad m) => Stream m a -> Stream m a -> m Bool
2072isSubsequenceOf (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
2074 go (sa, sb, Nothing) = do
2075 r <- stepa defState sa
2077 Yield x sa' -> go (sa', sb, Just x)
2078 Skip sa' -> go (sa', sb, Nothing)
2081 go (sa, sb, Just x) = do
2082 r <- stepb defState sb
2086 then go (sa, sb', Nothing)
2087 else go (sa, sb', Just x)
2088 Skip sb' -> go (sa, sb', Just x)
2089 Stop -> return False
2091{-# INLINE_NORMAL stripPrefix #-}
2094 => Stream m a -> Stream m a -> m (Maybe (Stream m a))
2095stripPrefix (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
2097 go (sa, sb, Nothing) = do
2098 r <- stepa defState sa
2100 Yield x sa' -> go (sa', sb, Just x)
2101 Skip sa' -> go (sa', sb, Nothing)
2102 Stop -> return $ Just (Stream stepb sb)
2104 go (sa, sb, Just x) = do
2105 r <- stepb defState sb
2109 then go (sa, sb', Nothing)
2111 Skip sb' -> go (sa, sb', Just x)
2112 Stop -> return Nothing
2114------------------------------------------------------------------------------
2116------------------------------------------------------------------------------
2118-- | Execute a monadic action for each element of the 'Stream'
2119{-# INLINE_NORMAL mapM_ #-}
2120mapM_ :: Monad m => (a -> m b) -> Stream m a -> m ()
2121mapM_ m = drain . mapM m
2123-------------------------------------------------------------------------------
2124-- Stream transformations using Unfolds
2125-------------------------------------------------------------------------------
2127-- Define a unique structure to use in inspection testing
2128data ConcatMapUState o i =
2130 | ConcatMapUInner o i
2132-- | @concatMapU unfold stream@ uses @unfold@ to map the input stream elements
2133-- to streams and then flattens the generated streams into a single output
2136-- This is like 'concatMap' but uses an unfold with an explicit state to
2137-- generate the stream instead of a 'Stream' type generator. This allows better
2138-- optimization via fusion. This can be many times more efficient than
2141{-# INLINE_NORMAL concatMapU #-}
2142concatMapU :: Monad m => Unfold m a b -> Stream m a -> Stream m b
2143concatMapU (Unfold istep inject) (Stream ostep ost) =
2144 Stream step (ConcatMapUOuter ost)
2146 {-# INLINE_LATE step #-}
2147 step gst (ConcatMapUOuter o) = do
2148 r <- ostep (adaptState gst) o
2152 i `seq` return (Skip (ConcatMapUInner o' i))
2153 Skip o' -> return $ Skip (ConcatMapUOuter o')
2154 Stop -> return $ Stop
2156 step _ (ConcatMapUInner o i) = do
2159 Yield x i' -> Yield x (ConcatMapUInner o i')
2160 Skip i' -> Skip (ConcatMapUInner o i')
2161 Stop -> Skip (ConcatMapUOuter o)
2163data ConcatUnfoldInterleaveState o i =
2164 ConcatUnfoldInterleaveOuter o [i]
2165 | ConcatUnfoldInterleaveInner o [i]
2166 | ConcatUnfoldInterleaveInnerL [i] [i]
2167 | ConcatUnfoldInterleaveInnerR [i] [i]
2169-- XXX use arrays to store state instead of lists.
2170-- XXX In general we can use different scheduling strategies e.g. how to
2171-- schedule the outer vs inner loop or assigning weights to different streams
2172-- or outer and inner loops.
2174-- After a yield, switch to the next stream. Do not switch streams on Skip.
2175-- Yield from outer stream switches to the inner stream.
2177-- There are two choices here, (1) exhaust the outer stream first and then
2178-- start yielding from the inner streams, this is much simpler to implement,
2179-- (2) yield at least one element from an inner stream before going back to
2180-- outer stream and opening the next stream from it.
2182-- Ideally, we need some scheduling bias to inner streams vs outer stream.
2183-- Maybe we can configure the behavior.
2185{-# INLINE_NORMAL concatUnfoldInterleave #-}
2186concatUnfoldInterleave :: Monad m => Unfold m a b -> Stream m a -> Stream m b
2187concatUnfoldInterleave (Unfold istep inject) (Stream ostep ost) =
2188 Stream step (ConcatUnfoldInterleaveOuter ost [])
2190 {-# INLINE_LATE step #-}
2191 step gst (ConcatUnfoldInterleaveOuter o ls) = do
2192 r <- ostep (adaptState gst) o
2196 i `seq` return (Skip (ConcatUnfoldInterleaveInner o' (i : ls)))
2197 Skip o' -> return $ Skip (ConcatUnfoldInterleaveOuter o' ls)
2198 Stop -> return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2200 step _ (ConcatUnfoldInterleaveInner _ []) = undefined
2201 step _ (ConcatUnfoldInterleaveInner o (st:ls)) = do
2204 Yield x s -> Yield x (ConcatUnfoldInterleaveOuter o (s:ls))
2205 Skip s -> Skip (ConcatUnfoldInterleaveInner o (s:ls))
2206 Stop -> Skip (ConcatUnfoldInterleaveOuter o ls)
2208 step _ (ConcatUnfoldInterleaveInnerL [] []) = return Stop
2209 step _ (ConcatUnfoldInterleaveInnerL [] rs) =
2210 return $ Skip (ConcatUnfoldInterleaveInnerR [] rs)
2212 step _ (ConcatUnfoldInterleaveInnerL (st:ls) rs) = do
2215 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerL ls (s:rs))
2216 Skip s -> Skip (ConcatUnfoldInterleaveInnerL (s:ls) rs)
2217 Stop -> Skip (ConcatUnfoldInterleaveInnerL ls rs)
2219 step _ (ConcatUnfoldInterleaveInnerR [] []) = return Stop
2220 step _ (ConcatUnfoldInterleaveInnerR ls []) =
2221 return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2223 step _ (ConcatUnfoldInterleaveInnerR ls (st:rs)) = do
2226 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerR (s:ls) rs)
2227 Skip s -> Skip (ConcatUnfoldInterleaveInnerR ls (s:rs))
2228 Stop -> Skip (ConcatUnfoldInterleaveInnerR ls rs)
2230-- XXX In general we can use different scheduling strategies e.g. how to
2231-- schedule the outer vs inner loop or assigning weights to different streams
2232-- or outer and inner loops.
2234-- This could be inefficient if the tasks are too small.
2236-- Compared to concatUnfoldInterleave this one switches streams on Skips.
2238{-# INLINE_NORMAL concatUnfoldRoundrobin #-}
2239concatUnfoldRoundrobin :: Monad m => Unfold m a b -> Stream m a -> Stream m b
2240concatUnfoldRoundrobin (Unfold istep inject) (Stream ostep ost) =
2241 Stream step (ConcatUnfoldInterleaveOuter ost [])
2243 {-# INLINE_LATE step #-}
2244 step gst (ConcatUnfoldInterleaveOuter o ls) = do
2245 r <- ostep (adaptState gst) o
2249 i `seq` return (Skip (ConcatUnfoldInterleaveInner o' (i : ls)))
2250 Skip o' -> return $ Skip (ConcatUnfoldInterleaveInner o' ls)
2251 Stop -> return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2253 step _ (ConcatUnfoldInterleaveInner o []) =
2254 return $ Skip (ConcatUnfoldInterleaveOuter o [])
2256 step _ (ConcatUnfoldInterleaveInner o (st:ls)) = do
2259 Yield x s -> Yield x (ConcatUnfoldInterleaveOuter o (s:ls))
2260 Skip s -> Skip (ConcatUnfoldInterleaveOuter o (s:ls))
2261 Stop -> Skip (ConcatUnfoldInterleaveOuter o ls)
2263 step _ (ConcatUnfoldInterleaveInnerL [] []) = return Stop
2264 step _ (ConcatUnfoldInterleaveInnerL [] rs) =
2265 return $ Skip (ConcatUnfoldInterleaveInnerR [] rs)
2267 step _ (ConcatUnfoldInterleaveInnerL (st:ls) rs) = do
2270 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerL ls (s:rs))
2271 Skip s -> Skip (ConcatUnfoldInterleaveInnerL ls (s:rs))
2272 Stop -> Skip (ConcatUnfoldInterleaveInnerL ls rs)
2274 step _ (ConcatUnfoldInterleaveInnerR [] []) = return Stop
2275 step _ (ConcatUnfoldInterleaveInnerR ls []) =
2276 return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2278 step _ (ConcatUnfoldInterleaveInnerR ls (st:rs)) = do
2281 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerR (s:ls) rs)
2282 Skip s -> Skip (ConcatUnfoldInterleaveInnerR (s:ls) rs)
2283 Stop -> Skip (ConcatUnfoldInterleaveInnerR ls rs)
2285data AppendState s1 s2 = AppendFirst s1 | AppendSecond s2
2287-- Note that this could be much faster compared to the CPS stream. However, as
2288-- the number of streams being composed increases this may become expensive.
2289-- Need to see where the breaking point is between the two.
2291{-# INLINE_NORMAL append #-}
2292append :: Monad m => Stream m a -> Stream m a -> Stream m a
2293append (Stream step1 state1) (Stream step2 state2) =
2294 Stream step (AppendFirst state1)
2298 {-# INLINE_LATE step #-}
2299 step gst (AppendFirst st) = do
2302 Yield a s -> Yield a (AppendFirst s)
2303 Skip s -> Skip (AppendFirst s)
2304 Stop -> Skip (AppendSecond state2)
2306 step gst (AppendSecond st) = do
2309 Yield a s -> Yield a (AppendSecond s)
2310 Skip s -> Skip (AppendSecond s)
2313data InterleaveState s1 s2 = InterleaveFirst s1 s2 | InterleaveSecond s1 s2
2314 | InterleaveSecondOnly s2 | InterleaveFirstOnly s1
2316{-# INLINE_NORMAL interleave #-}
2317interleave :: Monad m => Stream m a -> Stream m a -> Stream m a
2318interleave (Stream step1 state1) (Stream step2 state2) =
2319 Stream step (InterleaveFirst state1 state2)
2323 {-# INLINE_LATE step #-}
2324 step gst (InterleaveFirst st1 st2) = do
2327 Yield a s -> Yield a (InterleaveSecond s st2)
2328 Skip s -> Skip (InterleaveFirst s st2)
2329 Stop -> Skip (InterleaveSecondOnly st2)
2331 step gst (InterleaveSecond st1 st2) = do
2334 Yield a s -> Yield a (InterleaveFirst st1 s)
2335 Skip s -> Skip (InterleaveSecond st1 s)
2336 Stop -> Skip (InterleaveFirstOnly st1)
2338 step gst (InterleaveFirstOnly st1) = do
2341 Yield a s -> Yield a (InterleaveFirstOnly s)
2342 Skip s -> Skip (InterleaveFirstOnly s)
2345 step gst (InterleaveSecondOnly st2) = do
2348 Yield a s -> Yield a (InterleaveSecondOnly s)
2349 Skip s -> Skip (InterleaveSecondOnly s)
2352{-# INLINE_NORMAL interleaveMin #-}
2353interleaveMin :: Monad m => Stream m a -> Stream m a -> Stream m a
2354interleaveMin (Stream step1 state1) (Stream step2 state2) =
2355 Stream step (InterleaveFirst state1 state2)
2359 {-# INLINE_LATE step #-}
2360 step gst (InterleaveFirst st1 st2) = do
2363 Yield a s -> Yield a (InterleaveSecond s st2)
2364 Skip s -> Skip (InterleaveFirst s st2)
2367 step gst (InterleaveSecond st1 st2) = do
2370 Yield a s -> Yield a (InterleaveFirst st1 s)
2371 Skip s -> Skip (InterleaveSecond st1 s)
2374 step _ (InterleaveFirstOnly _) = undefined
2375 step _ (InterleaveSecondOnly _) = undefined
2377{-# INLINE_NORMAL interleaveSuffix #-}
2378interleaveSuffix :: Monad m => Stream m a -> Stream m a -> Stream m a
2379interleaveSuffix (Stream step1 state1) (Stream step2 state2) =
2380 Stream step (InterleaveFirst state1 state2)
2384 {-# INLINE_LATE step #-}
2385 step gst (InterleaveFirst st1 st2) = do
2388 Yield a s -> Yield a (InterleaveSecond s st2)
2389 Skip s -> Skip (InterleaveFirst s st2)
2392 step gst (InterleaveSecond st1 st2) = do
2395 Yield a s -> Yield a (InterleaveFirst st1 s)
2396 Skip s -> Skip (InterleaveSecond st1 s)
2397 Stop -> Skip (InterleaveFirstOnly st1)
2399 step gst (InterleaveFirstOnly st1) = do
2402 Yield a s -> Yield a (InterleaveFirstOnly s)
2403 Skip s -> Skip (InterleaveFirstOnly s)
2406 step _ (InterleaveSecondOnly _) = undefined
2408data InterleaveInfixState s1 s2 a
2409 = InterleaveInfixFirst s1 s2
2410 | InterleaveInfixSecondBuf s1 s2
2411 | InterleaveInfixSecondYield s1 s2 a
2412 | InterleaveInfixFirstYield s1 s2 a
2413 | InterleaveInfixFirstOnly s1
2415{-# INLINE_NORMAL interleaveInfix #-}
2416interleaveInfix :: Monad m => Stream m a -> Stream m a -> Stream m a
2417interleaveInfix (Stream step1 state1) (Stream step2 state2) =
2418 Stream step (InterleaveInfixFirst state1 state2)
2422 {-# INLINE_LATE step #-}
2423 step gst (InterleaveInfixFirst st1 st2) = do
2426 Yield a s -> Yield a (InterleaveInfixSecondBuf s st2)
2427 Skip s -> Skip (InterleaveInfixFirst s st2)
2430 step gst (InterleaveInfixSecondBuf st1 st2) = do
2433 Yield a s -> Skip (InterleaveInfixSecondYield st1 s a)
2434 Skip s -> Skip (InterleaveInfixSecondBuf st1 s)
2435 Stop -> Skip (InterleaveInfixFirstOnly st1)
2437 step gst (InterleaveInfixSecondYield st1 st2 x) = do
2440 Yield a s -> Yield x (InterleaveInfixFirstYield s st2 a)
2441 Skip s -> Skip (InterleaveInfixSecondYield s st2 x)
2444 step _ (InterleaveInfixFirstYield st1 st2 x) = do
2445 return $ Yield x (InterleaveInfixSecondBuf st1 st2)
2447 step gst (InterleaveInfixFirstOnly st1) = do
2450 Yield a s -> Yield a (InterleaveInfixFirstOnly s)
2451 Skip s -> Skip (InterleaveInfixFirstOnly s)
2454{-# INLINE_NORMAL roundRobin #-}
2455roundRobin :: Monad m => Stream m a -> Stream m a -> Stream m a
2456roundRobin (Stream step1 state1) (Stream step2 state2) =
2457 Stream step (InterleaveFirst state1 state2)
2461 {-# INLINE_LATE step #-}
2462 step gst (InterleaveFirst st1 st2) = do
2465 Yield a s -> Yield a (InterleaveSecond s st2)
2466 Skip s -> Skip (InterleaveSecond s st2)
2467 Stop -> Skip (InterleaveSecondOnly st2)
2469 step gst (InterleaveSecond st1 st2) = do
2472 Yield a s -> Yield a (InterleaveFirst st1 s)
2473 Skip s -> Skip (InterleaveFirst st1 s)
2474 Stop -> Skip (InterleaveFirstOnly st1)
2476 step gst (InterleaveSecondOnly st2) = do
2479 Yield a s -> Yield a (InterleaveSecondOnly s)
2480 Skip s -> Skip (InterleaveSecondOnly s)
2483 step gst (InterleaveFirstOnly st1) = do
2486 Yield a s -> Yield a (InterleaveFirstOnly s)
2487 Skip s -> Skip (InterleaveFirstOnly s)
2490data ICUState s1 s2 i1 i2 =
2495 | ICUFirstInner s1 s2 i1
2496 | ICUSecondInner s1 s2 i2
2497 | ICUFirstOnlyInner s1 i1
2498 | ICUSecondOnlyInner s2 i2
2500-- | Interleave streams (full streams, not the elements) unfolded from two
2501-- input streams and concat. Stop when the first stream stops. If the second
2502-- stream ends before the first one then first stream still keeps running alone
2503-- without any interleaving with the second stream.
2505-- [a1, a2, ... an] [b1, b2 ...]
2506-- => [streamA1, streamA2, ... streamAn] [streamB1, streamB2, ...]
2507-- => [streamA1, streamB1, streamA2...StreamAn, streamBn]
2508-- => [a11, a12, ...a1j, b11, b12, ...b1k, a21, a22, ...]
2510{-# INLINE_NORMAL gintercalateSuffix #-}
2513 => Unfold m a c -> Stream m a -> Unfold m b c -> Stream m b -> Stream m c
2515 (Unfold istep1 inject1) (Stream step1 state1)
2516 (Unfold istep2 inject2) (Stream step2 state2) =
2517 Stream step (ICUFirst state1 state2)
2521 {-# INLINE_LATE step #-}
2522 step gst (ICUFirst s1 s2) = do
2523 r <- step1 (adaptState gst) s1
2527 i `seq` return (Skip (ICUFirstInner s s2 i))
2528 Skip s -> return $ Skip (ICUFirst s s2)
2531 step gst (ICUFirstOnly s1) = do
2532 r <- step1 (adaptState gst) s1
2536 i `seq` return (Skip (ICUFirstOnlyInner s i))
2537 Skip s -> return $ Skip (ICUFirstOnly s)
2540 step _ (ICUFirstInner s1 s2 i1) = do
2543 Yield x i' -> Yield x (ICUFirstInner s1 s2 i')
2544 Skip i' -> Skip (ICUFirstInner s1 s2 i')
2545 Stop -> Skip (ICUSecond s1 s2)
2547 step _ (ICUFirstOnlyInner s1 i1) = do
2550 Yield x i' -> Yield x (ICUFirstOnlyInner s1 i')
2551 Skip i' -> Skip (ICUFirstOnlyInner s1 i')
2552 Stop -> Skip (ICUFirstOnly s1)
2554 step gst (ICUSecond s1 s2) = do
2555 r <- step2 (adaptState gst) s2
2559 i `seq` return (Skip (ICUSecondInner s1 s i))
2560 Skip s -> return $ Skip (ICUSecond s1 s)
2561 Stop -> return $ Skip (ICUFirstOnly s1)
2563 step _ (ICUSecondInner s1 s2 i2) = do
2566 Yield x i' -> Yield x (ICUSecondInner s1 s2 i')
2567 Skip i' -> Skip (ICUSecondInner s1 s2 i')
2568 Stop -> Skip (ICUFirst s1 s2)
2570 step _ (ICUSecondOnly _s2) = undefined
2571 step _ (ICUSecondOnlyInner _s2 _i2) = undefined
2573data InterposeSuffixState s1 i1 =
2574 InterposeSuffixFirst s1
2575 -- | InterposeSuffixFirstYield s1 i1
2576 | InterposeSuffixFirstInner s1 i1
2577 | InterposeSuffixSecond s1
2579-- Note that if an unfolded layer turns out to be nil we still emit the
2580-- separator effect. An alternate behavior could be to emit the separator
2581-- effect only if at least one element has been yielded by the unfolding.
2582-- However, that becomes a bit complicated, so we have chosen the former
2584{-# INLINE_NORMAL interposeSuffix #-}
2587 => m c -> Unfold m b c -> Stream m b -> Stream m c
2590 (Unfold istep1 inject1) (Stream step1 state1) =
2591 Stream step (InterposeSuffixFirst state1)
2595 {-# INLINE_LATE step #-}
2596 step gst (InterposeSuffixFirst s1) = do
2597 r <- step1 (adaptState gst) s1
2601 i `seq` return (Skip (InterposeSuffixFirstInner s i))
2602 -- i `seq` return (Skip (InterposeSuffixFirstYield s i))
2603 Skip s -> return $ Skip (InterposeSuffixFirst s)
2607 step _ (InterposeSuffixFirstYield s1 i1) = do
2610 Yield x i' -> Yield x (InterposeSuffixFirstInner s1 i')
2611 Skip i' -> Skip (InterposeSuffixFirstYield s1 i')
2612 Stop -> Skip (InterposeSuffixFirst s1)
2615 step _ (InterposeSuffixFirstInner s1 i1) = do
2618 Yield x i' -> Yield x (InterposeSuffixFirstInner s1 i')
2619 Skip i' -> Skip (InterposeSuffixFirstInner s1 i')
2620 Stop -> Skip (InterposeSuffixSecond s1)
2622 step _ (InterposeSuffixSecond s1) = do
2624 return $ Yield r (InterposeSuffixFirst s1)
2626data ICALState s1 s2 i1 i2 a =
2628 -- | ICALFirstYield s1 s2 i1
2629 | ICALFirstInner s1 s2 i1
2631 | ICALFirstOnlyInner s1 i1
2632 | ICALSecondInject s1 s2
2633 | ICALFirstInject s1 s2 i2
2634 -- | ICALFirstBuf s1 s2 i1 i2
2635 | ICALSecondInner s1 s2 i1 i2
2636 -- -- | ICALSecondInner s1 s2 i1 i2 a
2637 -- -- | ICALFirstResume s1 s2 i1 i2 a
2639-- | Interleave streams (full streams, not the elements) unfolded from two
2640-- input streams and concat. Stop when the first stream stops. If the second
2641-- stream ends before the first one then first stream still keeps running alone
2642-- without any interleaving with the second stream.
2644-- [a1, a2, ... an] [b1, b2 ...]
2645-- => [streamA1, streamA2, ... streamAn] [streamB1, streamB2, ...]
2646-- => [streamA1, streamB1, streamA2...StreamAn, streamBn]
2647-- => [a11, a12, ...a1j, b11, b12, ...b1k, a21, a22, ...]
2649{-# INLINE_NORMAL gintercalate #-}
2652 => Unfold m a c -> Stream m a -> Unfold m b c -> Stream m b -> Stream m c
2654 (Unfold istep1 inject1) (Stream step1 state1)
2655 (Unfold istep2 inject2) (Stream step2 state2) =
2656 Stream step (ICALFirst state1 state2)
2660 {-# INLINE_LATE step #-}
2661 step gst (ICALFirst s1 s2) = do
2662 r <- step1 (adaptState gst) s1
2666 i `seq` return (Skip (ICALFirstInner s s2 i))
2667 -- i `seq` return (Skip (ICALFirstYield s s2 i))
2668 Skip s -> return $ Skip (ICALFirst s s2)
2672 step _ (ICALFirstYield s1 s2 i1) = do
2675 Yield x i' -> Yield x (ICALFirstInner s1 s2 i')
2676 Skip i' -> Skip (ICALFirstYield s1 s2 i')
2677 Stop -> Skip (ICALFirst s1 s2)
2680 step _ (ICALFirstInner s1 s2 i1) = do
2683 Yield x i' -> Yield x (ICALFirstInner s1 s2 i')
2684 Skip i' -> Skip (ICALFirstInner s1 s2 i')
2685 Stop -> Skip (ICALSecondInject s1 s2)
2687 step gst (ICALFirstOnly s1) = do
2688 r <- step1 (adaptState gst) s1
2692 i `seq` return (Skip (ICALFirstOnlyInner s i))
2693 Skip s -> return $ Skip (ICALFirstOnly s)
2696 step _ (ICALFirstOnlyInner s1 i1) = do
2699 Yield x i' -> Yield x (ICALFirstOnlyInner s1 i')
2700 Skip i' -> Skip (ICALFirstOnlyInner s1 i')
2701 Stop -> Skip (ICALFirstOnly s1)
2703 -- We inject the second stream even before checking if the first stream
2704 -- would yield any more elements. There is no clear choice whether we
2705 -- should do this before or after that. Doing it after may make the state
2706 -- machine a bit simpler though.
2707 step gst (ICALSecondInject s1 s2) = do
2708 r <- step2 (adaptState gst) s2
2712 i `seq` return (Skip (ICALFirstInject s1 s i))
2713 Skip s -> return $ Skip (ICALSecondInject s1 s)
2714 Stop -> return $ Skip (ICALFirstOnly s1)
2716 step gst (ICALFirstInject s1 s2 i2) = do
2717 r <- step1 (adaptState gst) s1
2721 i `seq` return (Skip (ICALSecondInner s s2 i i2))
2722 -- i `seq` return (Skip (ICALFirstBuf s s2 i i2))
2723 Skip s -> return $ Skip (ICALFirstInject s s2 i2)
2727 step _ (ICALFirstBuf s1 s2 i1 i2) = do
2730 Yield x i' -> Skip (ICALSecondInner s1 s2 i' i2 x)
2731 Skip i' -> Skip (ICALFirstBuf s1 s2 i' i2)
2734 step _ (ICALSecondInner s1 s2 i1 i2 v) = do
2737 Yield x i' -> Yield x (ICALSecondInner s1 s2 i1 i' v)
2738 Skip i' -> Skip (ICALSecondInner s1 s2 i1 i' v)
2739 Stop -> Skip (ICALFirstResume s1 s2 i1 i2 v)
2742 step _ (ICALSecondInner s1 s2 i1 i2) = do
2745 Yield x i' -> Yield x (ICALSecondInner s1 s2 i1 i')
2746 Skip i' -> Skip (ICALSecondInner s1 s2 i1 i')
2747 Stop -> Skip (ICALFirstInner s1 s2 i1)
2748 -- Stop -> Skip (ICALFirstResume s1 s2 i1 i2)
2751 step _ (ICALFirstResume s1 s2 i1 i2 x) = do
2752 return $ Yield x (ICALFirstInner s1 s2 i1 i2)
2755data InterposeState s1 i1 a =
2757 -- | InterposeFirstYield s1 i1
2758 | InterposeFirstInner s1 i1
2759 | InterposeFirstInject s1
2760 -- | InterposeFirstBuf s1 i1
2761 | InterposeSecondYield s1 i1
2762 -- -- | InterposeSecondYield s1 i1 a
2763 -- -- | InterposeFirstResume s1 i1 a
2765-- Note that this only interposes the pure values, we may run many effects to
2766-- generate those values as some effects may not generate anything (Skip).
2767{-# INLINE_NORMAL interpose #-}
2768interpose :: Monad m => m c -> Unfold m b c -> Stream m b -> Stream m c
2771 (Unfold istep1 inject1) (Stream step1 state1) =
2772 Stream step (InterposeFirst state1)
2776 {-# INLINE_LATE step #-}
2777 step gst (InterposeFirst s1) = do
2778 r <- step1 (adaptState gst) s1
2782 i `seq` return (Skip (InterposeFirstInner s i))
2783 -- i `seq` return (Skip (InterposeFirstYield s i))
2784 Skip s -> return $ Skip (InterposeFirst s)
2788 step _ (InterposeFirstYield s1 i1) = do
2791 Yield x i' -> Yield x (InterposeFirstInner s1 i')
2792 Skip i' -> Skip (InterposeFirstYield s1 i')
2793 Stop -> Skip (InterposeFirst s1)
2796 step _ (InterposeFirstInner s1 i1) = do
2799 Yield x i' -> Yield x (InterposeFirstInner s1 i')
2800 Skip i' -> Skip (InterposeFirstInner s1 i')
2801 Stop -> Skip (InterposeFirstInject s1)
2803 step gst (InterposeFirstInject s1) = do
2804 r <- step1 (adaptState gst) s1
2808 -- i `seq` return (Skip (InterposeFirstBuf s i))
2809 i `seq` return (Skip (InterposeSecondYield s i))
2810 Skip s -> return $ Skip (InterposeFirstInject s)
2814 step _ (InterposeFirstBuf s1 i1) = do
2817 Yield x i' -> Skip (InterposeSecondYield s1 i' x)
2818 Skip i' -> Skip (InterposeFirstBuf s1 i')
2823 step _ (InterposeSecondYield s1 i1 v) = do
2825 return $ Yield r (InterposeFirstResume s1 i1 v)
2827 step _ (InterposeSecondYield s1 i1) = do
2829 return $ Yield r (InterposeFirstInner s1 i1)
2832 step _ (InterposeFirstResume s1 i1 v) = do
2833 return $ Yield v (InterposeFirstInner s1 i1)
2836------------------------------------------------------------------------------
2838------------------------------------------------------------------------------
2840data GbracketState s1 s2 v
2842 | GBracketNormal s1 v
2843 | GBracketException s2
2845-- | The most general bracketing and exception combinator. All other
2846-- combinators can be expressed in terms of this combinator. This can also be
2847-- used for cases which are not covered by the standard combinators.
2851{-# INLINE_NORMAL gbracket #-}
2855 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
2856 -> (c -> m d) -- ^ after, on normal stop
2857 -> (c -> e -> Stream m b) -- ^ on exception
2858 -> (c -> Stream m b) -- ^ stream generator
2860gbracket bef exc aft fexc fnormal =
2861 Stream step GBracketInit
2865 {-# INLINE_LATE step #-}
2866 step _ GBracketInit = do
2868 return $ Skip $ GBracketNormal (fnormal r) r
2870 step gst (GBracketNormal (UnStream step1 st) v) = do
2871 res <- exc $ step1 gst st
2873 Right r -> case r of
2875 return $ Yield x (GBracketNormal (Stream step1 s) v)
2876 Skip s -> return $ Skip (GBracketNormal (Stream step1 s) v)
2877 Stop -> aft v >> return Stop
2878 Left e -> return $ Skip (GBracketException (fexc v e))
2879 step gst (GBracketException (UnStream step1 st)) = do
2882 Yield x s -> return $ Yield x (GBracketException (Stream step1 s))
2883 Skip s -> return $ Skip (GBracketException (Stream step1 s))
2886-- | Create an IORef holding a finalizer that is called automatically when the
2887-- IORef is garbage collected. The IORef can be written to with a 'Nothing'
2888-- value to deactivate the finalizer.
2889newFinalizedIORef :: (MonadIO m, MonadBaseControl IO m)
2890 => m a -> m (IORef (Maybe (IO ())))
2891newFinalizedIORef finalizer = do
2892 mrun <- captureMonadState
2893 ref <- liftIO $ newIORef $ Just $ liftIO $ void $ do
2894 _ <- runInIO mrun finalizer
2897 res <- readIORef ref
2899 Nothing -> return ()
2901 _ <- liftIO $ mkWeakIORef ref finalizer1
2904-- | Run the finalizer stored in an IORef and deactivate it so that it is run
2907runIORefFinalizer :: MonadIO m => IORef (Maybe (IO ())) -> m ()
2908runIORefFinalizer ref = liftIO $ do
2909 res <- readIORef ref
2911 Nothing -> return ()
2912 Just f -> writeIORef ref Nothing >> f
2914-- | Deactivate the finalizer stored in an IORef without running it.
2916clearIORefFinalizer :: MonadIO m => IORef (Maybe (IO ())) -> m ()
2917clearIORefFinalizer ref = liftIO $ writeIORef ref Nothing
2919data GbracketIOState s1 s2 v wref
2921 | GBracketIONormal s1 v wref
2922 | GBracketIOException s2
2924-- | Like gbracket but also uses a finalizer to make sure when the stream is
2925-- garbage collected we run the finalizing action. This requires a MonadIO and
2926-- MonadBaseControl IO constraint.
2928-- | The most general bracketing and exception combinator. All other
2929-- combinators can be expressed in terms of this combinator. This can also be
2930-- used for cases which are not covered by the standard combinators.
2934{-# INLINE_NORMAL gbracketIO #-}
2936 :: (MonadIO m, MonadBaseControl IO m)
2938 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
2939 -> (c -> m d) -- ^ after, on normal stop or GC
2940 -> (c -> e -> Stream m b) -- ^ on exception
2941 -> (c -> Stream m b) -- ^ stream generator
2943gbracketIO bef exc aft fexc fnormal =
2944 Stream step GBracketIOInit
2948 -- If the stream is never evaluated the "aft" action will never be
2949 -- called. For that to occur we will need the user of this API to pass a
2950 -- weak pointer to us.
2951 {-# INLINE_LATE step #-}
2952 step _ GBracketIOInit = do
2954 ref <- newFinalizedIORef (aft r)
2955 return $ Skip $ GBracketIONormal (fnormal r) r ref
2957 step gst (GBracketIONormal (UnStream step1 st) v ref) = do
2958 res <- exc $ step1 gst st
2960 Right r -> case r of
2962 return $ Yield x (GBracketIONormal (Stream step1 s) v ref)
2964 return $ Skip (GBracketIONormal (Stream step1 s) v ref)
2966 runIORefFinalizer ref
2969 clearIORefFinalizer ref
2970 return $ Skip (GBracketIOException (fexc v e))
2971 step gst (GBracketIOException (UnStream step1 st)) = do
2975 return $ Yield x (GBracketIOException (Stream step1 s))
2976 Skip s -> return $ Skip (GBracketIOException (Stream step1 s))
2979-- | Run a side effect before the stream yields its first element.
2980{-# INLINE_NORMAL before #-}
2981before :: Monad m => m b -> Stream m a -> Stream m a
2982before action (Stream step state) = Stream step' Nothing
2986 {-# INLINE_LATE step' #-}
2987 step' _ Nothing = action >> return (Skip (Just state))
2989 step' gst (Just st) = do
2992 Yield x s -> return $ Yield x (Just s)
2993 Skip s -> return $ Skip (Just s)
2996-- | Run a side effect whenever the stream stops normally.
2997{-# INLINE_NORMAL after #-}
2998after :: Monad m => m b -> Stream m a -> Stream m a
2999after action (Stream step state) = Stream step' state
3003 {-# INLINE_LATE step' #-}
3007 Yield x s -> return $ Yield x s
3008 Skip s -> return $ Skip s
3009 Stop -> action >> return Stop
3011{-# INLINE_NORMAL afterIO #-}
3012afterIO :: (MonadIO m, MonadBaseControl IO m)
3013 => m b -> Stream m a -> Stream m a
3014afterIO action (Stream step state) = Stream step' Nothing
3018 {-# INLINE_LATE step' #-}
3019 step' _ Nothing = do
3020 ref <- newFinalizedIORef action
3021 return $ Skip $ Just (state, ref)
3022 step' gst (Just (st, ref)) = do
3025 Yield x s -> return $ Yield x (Just (s, ref))
3026 Skip s -> return $ Skip (Just (s, ref))
3028 runIORefFinalizer ref
3031-- XXX These combinators are expensive due to the call to
3032-- onException/handle/try on each step. Therefore, when possible, they should
3033-- be called in an outer loop where we perform less iterations. For example, we
3034-- cannot call them on each iteration in a char stream, instead we can call
3035-- them when doing an IO on an array.
3037-- XXX For high performance error checks in busy streams we may need another
3038-- Error constructor in step.
3040-- | Run a side effect whenever the stream aborts due to an exception. The
3041-- exception is not caught, simply rethrown.
3042{-# INLINE_NORMAL onException #-}
3043onException :: MonadCatch m => m b -> Stream m a -> Stream m a
3044onException action str =
3045 gbracket (return ()) MC.try return
3046 (\_ (e :: MC.SomeException) -> nilM (action >> MC.throwM e))
3049{-# INLINE_NORMAL _onException #-}
3050_onException :: MonadCatch m => m b -> Stream m a -> Stream m a
3051_onException action (Stream step state) = Stream step' state
3055 {-# INLINE_LATE step' #-}
3057 res <- step gst st `MC.onException` action
3059 Yield x s -> return $ Yield x s
3060 Skip s -> return $ Skip s
3063-- XXX bracket is like concatMap, it generates a stream and then flattens it.
3064-- Like concatMap it has 10x worse performance compared to linear fused
3067-- | Run the first action before the stream starts and remember its output,
3068-- generate a stream using the output, run the second action providing the
3069-- remembered value as an argument whenever the stream ends normally or due to
3071{-# INLINE_NORMAL bracket #-}
3072bracket :: MonadCatch m => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
3073bracket bef aft bet =
3074 gbracket bef MC.try aft
3075 (\a (e :: SomeException) -> nilM (aft a >> MC.throwM e)) bet
3077{-# INLINE_NORMAL bracketIO #-}
3078bracketIO :: (MonadAsync m, MonadCatch m)
3079 => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
3080bracketIO bef aft bet =
3081 gbracketIO bef MC.try aft
3082 (\a (e :: SomeException) -> nilM (aft a >> MC.throwM e)) bet
3084data BracketState s v = BracketInit | BracketRun s v
3086{-# INLINE_NORMAL _bracket #-}
3087_bracket :: MonadCatch m => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
3088_bracket bef aft bet = Stream step' BracketInit
3092 {-# INLINE_LATE step' #-}
3093 step' _ BracketInit = bef >>= \x -> return (Skip (BracketRun (bet x) x))
3095 -- NOTE: It is important to use UnStream instead of the Stream pattern
3096 -- here, otherwise we get huge perf degradation, see note in concatMap.
3097 step' gst (BracketRun (UnStream step state) v) = do
3098 -- res <- step gst state `MC.onException` aft v
3099 res <- MC.try $ step gst state
3101 Left (e :: SomeException) -> aft v >> MC.throwM e >> return Stop
3102 Right r -> case r of
3103 Yield x s -> return $ Yield x (BracketRun (Stream step s) v)
3104 Skip s -> return $ Skip (BracketRun (Stream step s) v)
3105 Stop -> aft v >> return Stop
3107-- | Run a side effect whenever the stream stops normally or aborts due to an
3109{-# INLINE finally #-}
3110finally :: MonadCatch m => m b -> Stream m a -> Stream m a
3111-- finally action xs = after action $ onException action xs
3112finally action xs = bracket (return ()) (\_ -> action) (const xs)
3114{-# INLINE finallyIO #-}
3115finallyIO :: (MonadAsync m, MonadCatch m) => m b -> Stream m a -> Stream m a
3116finallyIO action xs = bracketIO (return ()) (\_ -> action) (const xs)
3118-- | When evaluating a stream if an exception occurs, stream evaluation aborts
3119-- and the specified exception handler is run with the exception as argument.
3120{-# INLINE_NORMAL handle #-}
3121handle :: (MonadCatch m, Exception e)
3122 => (e -> Stream m a) -> Stream m a -> Stream m a
3124 gbracket (return ()) MC.try return (\_ e -> f e) (\_ -> str)
3126{-# INLINE_NORMAL _handle #-}
3127_handle :: (MonadCatch m, Exception e)
3128 => (e -> Stream m a) -> Stream m a -> Stream m a
3129_handle f (Stream step state) = Stream step' (Left state)
3133 {-# INLINE_LATE step' #-}
3134 step' gst (Left st) = do
3135 res <- MC.try $ step gst st
3137 Left e -> return $ Skip $ Right (f e)
3138 Right r -> case r of
3139 Yield x s -> return $ Yield x (Left s)
3140 Skip s -> return $ Skip (Left s)
3143 step' gst (Right (UnStream step1 st)) = do
3146 Yield x s -> return $ Yield x (Right (Stream step1 s))
3147 Skip s -> return $ Skip (Right (Stream step1 s))
3150-------------------------------------------------------------------------------
3151-- General transformation
3152-------------------------------------------------------------------------------
3154{-# INLINE_NORMAL transform #-}
3155transform :: Monad m => Pipe m a b -> Stream m a -> Stream m b
3156transform (Pipe pstep1 pstep2 pstate) (Stream step state) =
3157 Stream step' (Consume pstate, state)
3161 {-# INLINE_LATE step' #-}
3163 step' gst (Consume pst, st) = pst `seq` do
3164 r <- step (adaptState gst) st
3169 Pipe.Yield b pst' -> return $ Yield b (pst', s)
3170 Pipe.Continue pst' -> return $ Skip (pst', s)
3171 Skip s -> return $ Skip (Consume pst, s)
3174 step' _ (Produce pst, st) = pst `seq` do
3177 Pipe.Yield b pst' -> return $ Yield b (pst', st)
3178 Pipe.Continue pst' -> return $ Skip (pst', st)
3180------------------------------------------------------------------------------
3181-- Transformation by Folding (Scans)
3182------------------------------------------------------------------------------
3184------------------------------------------------------------------------------
3186------------------------------------------------------------------------------
3188-- XXX Is a prescan useful, discarding the last step does not sound useful? I
3189-- am not sure about the utility of this function, so this is implemented but
3190-- not exposed. We can expose it if someone provides good reasons why this is
3193-- XXX We have to execute the stream one step ahead to know that we are at the
3194-- last step. The vector implementation of prescan executes the last fold step
3195-- but does not yield the result. This means we have executed the effect but
3196-- discarded value. This does not sound right. In this implementation we are
3197-- not executing the last fold step.
3198{-# INLINE_NORMAL prescanlM' #-}
3199prescanlM' :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
3200prescanlM' f mz (Stream step state) = Stream step' (state, mz)
3202 {-# INLINE_LATE step' #-}
3203 step' gst (st, prev) = do
3204 r <- step (adaptState gst) st
3208 return $ Yield acc (s, f acc x)
3209 Skip s -> return $ Skip (s, prev)
3212{-# INLINE prescanl' #-}
3213prescanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
3214prescanl' f z = prescanlM' (\a b -> return (f a b)) (return z)
3216------------------------------------------------------------------------------
3217-- Monolithic postscans (postscan followed by a map)
3218------------------------------------------------------------------------------
3220-- The performance of a modular postscan followed by a map seems to be
3221-- equivalent to this monolithic scan followed by map therefore we may not need
3222-- this implementation. We just have it for performance comparison and in case
3223-- modular version does not perform well in some situation.
3225{-# INLINE_NORMAL postscanlMx' #-}
3226postscanlMx' :: Monad m
3227 => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> Stream m b
3228postscanlMx' fstep begin done (Stream step state) = do
3229 Stream step' (state, begin)
3231 {-# INLINE_LATE step' #-}
3232 step' gst (st, acc) = do
3233 r <- step (adaptState gst) st
3239 v `seq` y `seq` return (Yield v (s, return y))
3240 Skip s -> return $ Skip (s, acc)
3243{-# INLINE_NORMAL postscanlx' #-}
3244postscanlx' :: Monad m
3245 => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> Stream m b
3246postscanlx' fstep begin done s =
3247 postscanlMx' (\b a -> return (fstep b a)) (return begin) (return . done) s
3249-- XXX do we need consM strict to evaluate the begin value?
3250{-# INLINE scanlMx' #-}
3252 => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> Stream m b
3253scanlMx' fstep begin done s =
3254 (begin >>= \x -> x `seq` done x) `consM` postscanlMx' fstep begin done s
3256{-# INLINE scanlx' #-}
3258 => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> Stream m b
3259scanlx' fstep begin done s =
3260 scanlMx' (\b a -> return (fstep b a)) (return begin) (return . done) s
3262------------------------------------------------------------------------------
3264------------------------------------------------------------------------------
3266{-# INLINE_NORMAL postscanlM' #-}
3267postscanlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
3268postscanlM' fstep begin (Stream step state) =
3269 begin `seq` Stream step' (state, begin)
3271 {-# INLINE_LATE step' #-}
3272 step' gst (st, acc) = acc `seq` do
3273 r <- step (adaptState gst) st
3277 y `seq` return (Yield y (s, y))
3278 Skip s -> return $ Skip (s, acc)
3281{-# INLINE_NORMAL postscanl' #-}
3282postscanl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
3283postscanl' f = postscanlM' (\a b -> return (f a b))
3285{-# INLINE_NORMAL postscanlM #-}
3286postscanlM :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
3287postscanlM fstep begin (Stream step state) = Stream step' (state, begin)
3289 {-# INLINE_LATE step' #-}
3290 step' gst (st, acc) = do
3291 r <- step (adaptState gst) st
3295 return (Yield y (s, y))
3296 Skip s -> return $ Skip (s, acc)
3299{-# INLINE_NORMAL postscanl #-}
3300postscanl :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
3301postscanl f = postscanlM (\a b -> return (f a b))
3303{-# INLINE_NORMAL scanlM' #-}
3304scanlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
3305scanlM' fstep begin s = begin `seq` (begin `cons` postscanlM' fstep begin s)
3307{-# INLINE scanl' #-}
3308scanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
3309scanl' f = scanlM' (\a b -> return (f a b))
3311{-# INLINE_NORMAL scanlM #-}
3312scanlM :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
3313scanlM fstep begin s = begin `cons` postscanlM fstep begin s
3316scanl :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
3317scanl f = scanlM (\a b -> return (f a b))
3319{-# INLINE_NORMAL scanl1M #-}
3320scanl1M :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
3321scanl1M fstep (Stream step state) = Stream step' (state, Nothing)
3323 {-# INLINE_LATE step' #-}
3324 step' gst (st, Nothing) = do
3327 Yield x s -> return $ Yield x (s, Just x)
3328 Skip s -> return $ Skip (s, Nothing)
3331 step' gst (st, Just acc) = do
3336 return $ Yield z (s, Just z)
3337 Skip s -> return $ Skip (s, Just acc)
3340{-# INLINE scanl1 #-}
3341scanl1 :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
3342scanl1 f = scanl1M (\x y -> return (f x y))
3344{-# INLINE_NORMAL scanl1M' #-}
3345scanl1M' :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
3346scanl1M' fstep (Stream step state) = Stream step' (state, Nothing)
3348 {-# INLINE_LATE step' #-}
3349 step' gst (st, Nothing) = do
3352 Yield x s -> x `seq` return $ Yield x (s, Just x)
3353 Skip s -> return $ Skip (s, Nothing)
3356 step' gst (st, Just acc) = acc `seq` do
3361 z `seq` return $ Yield z (s, Just z)
3362 Skip s -> return $ Skip (s, Just acc)
3365{-# INLINE scanl1' #-}
3366scanl1' :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
3367scanl1' f = scanl1M' (\x y -> return (f x y))
3369------------------------------------------------------------------------------
3371------------------------------------------------------------------------------
3373data RollingMapState s a = RollingMapInit s | RollingMapGo s a
3375{-# INLINE rollingMapM #-}
3376rollingMapM :: Monad m => (a -> a -> m b) -> Stream m a -> Stream m b
3377rollingMapM f (Stream step1 state1) = Stream step (RollingMapInit state1)
3379 step gst (RollingMapInit st) = do
3380 r <- step1 (adaptState gst) st
3382 Yield x s -> Skip $ RollingMapGo s x
3383 Skip s -> Skip $ RollingMapInit s
3386 step gst (RollingMapGo s1 x1) = do
3387 r <- step1 (adaptState gst) s1
3391 return $ Yield res $ RollingMapGo s x
3392 Skip s -> return $ Skip $ RollingMapGo s x1
3393 Stop -> return $ Stop
3395{-# INLINE rollingMap #-}
3396rollingMap :: Monad m => (a -> a -> b) -> Stream m a -> Stream m b
3397rollingMap f = rollingMapM (\x y -> return $ f x y)
3399------------------------------------------------------------------------------
3400-- Tapping/Distributing
3401------------------------------------------------------------------------------
3404tap :: Monad m => Fold m a b -> Stream m a -> Stream m a
3405tap (Fold fstep initial extract) (Stream step state) = Stream step' Nothing
3409 step' _ Nothing = do
3411 return $ Skip (Just (r, state))
3413 step' gst (Just (acc, st)) = acc `seq` do
3418 return $ Yield x (Just (acc', s))
3419 Skip s -> return $ Skip (Just (acc, s))
3424{-# INLINE_NORMAL tapOffsetEvery #-}
3425tapOffsetEvery :: Monad m
3426 => Int -> Int -> Fold m a b -> Stream m a -> Stream m a
3427tapOffsetEvery offset n (Fold fstep initial extract) (Stream step state) =
3428 Stream step' Nothing
3432 {-# INLINE_LATE step' #-}
3433 step' _ Nothing = do
3435 return $ Skip (Just (r, state, offset `mod` n))
3437 step' gst (Just (acc, st, count)) | count <= 0 = do
3441 !acc' <- fstep acc x
3442 return $ Yield x (Just (acc', s, n - 1))
3443 Skip s -> return $ Skip (Just (acc, s, count))
3448 step' gst (Just (acc, st, count)) = do
3451 Yield x s -> return $ Yield x (Just (acc, s, count - 1))
3452 Skip s -> return $ Skip (Just (acc, s, count))
3457{-# INLINE_NORMAL pollCounts #-}
3460 => (Stream m Int -> Stream m Int)
3464pollCounts transf fld (Stream step state) = Stream step' Nothing
3467 {-# INLINE_LATE step' #-}
3468 step' _ Nothing = do
3469 -- As long as we are using an "Int" for counts lockfree reads from
3470 -- Var should work correctly on both 32-bit and 64-bit machines.
3471 -- However, an Int on a 32-bit machine may overflow quickly.
3472 countVar <- liftIO $ newVar (0 :: Int)
3474 $ void $ runFold fld
3475 $ transf $ fromPrimVar countVar
3476 return $ Skip (Just (countVar, tid, state))
3478 step' gst (Just (countVar, tid, st)) = do
3482 liftIO $ modifyVar' countVar (+ 1)
3483 return $ Yield x (Just (countVar, tid, s))
3484 Skip s -> return $ Skip (Just (countVar, tid, s))
3486 liftIO $ killThread tid
3489{-# INLINE_NORMAL tapRate #-}
3491 (MonadAsync m, MonadCatch m)
3496tapRate samplingRate action (Stream step state) = Stream step' Nothing
3498 {-# NOINLINE loop #-}
3499 loop countVar prev = do
3502 (do liftIO $ threadDelay (round $ samplingRate * 1000000)
3503 i <- liftIO $ readVar countVar
3504 let !diff = i - prev
3507 (\(e :: AsyncException) -> do
3508 i <- liftIO $ readVar countVar
3509 let !diff = i - prev
3511 throwM (MC.toException e))
3514 {-# INLINE_LATE step' #-}
3515 step' _ Nothing = do
3516 countVar <- liftIO $ newVar 0
3517 tid <- fork $ loop countVar 0
3518 ref <- liftIO $ newIORef ()
3519 _ <- liftIO $ mkWeakIORef ref (killThread tid)
3520 return $ Skip (Just (countVar, tid, state, ref))
3522 step' gst (Just (countVar, tid, st, ref)) = do
3526 liftIO $ modifyVar' countVar (+ 1)
3527 return $ Yield x (Just (countVar, tid, s, ref))
3528 Skip s -> return $ Skip (Just (countVar, tid, s, ref))
3530 liftIO $ killThread tid
3534-------------------------------------------------------------------------------
3536-------------------------------------------------------------------------------
3538{-# INLINE_NORMAL takeWhileM #-}
3539takeWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
3540takeWhileM f (Stream step state) = Stream step' state
3542 {-# INLINE_LATE step' #-}
3548 return $ if b then Yield x s else Stop
3549 Skip s -> return $ Skip s
3552{-# INLINE takeWhile #-}
3553takeWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
3554takeWhile f = takeWhileM (return . f)
3556{-# INLINE_NORMAL drop #-}
3557drop :: Monad m => Int -> Stream m a -> Stream m a
3558drop n (Stream step state) = Stream step' (state, Just n)
3560 {-# INLINE_LATE step' #-}
3561 step' gst (st, Just i)
3566 Yield _ s -> Skip (s, Just (i - 1))
3567 Skip s -> Skip (s, Just i)
3569 | otherwise = return $ Skip (st, Nothing)
3571 step' gst (st, Nothing) = do
3575 Yield x s -> Yield x (s, Nothing)
3576 Skip s -> Skip (s, Nothing)
3579data DropWhileState s a
3581 | DropWhileYield a s
3584{-# INLINE_NORMAL dropWhileM #-}
3585dropWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
3586dropWhileM f (Stream step state) = Stream step' (DropWhileDrop state)
3588 {-# INLINE_LATE step' #-}
3589 step' gst (DropWhileDrop st) = do
3595 then return $ Skip (DropWhileDrop s)
3596 else return $ Skip (DropWhileYield x s)
3597 Skip s -> return $ Skip (DropWhileDrop s)
3600 step' gst (DropWhileNext st) = do
3603 Yield x s -> return $ Skip (DropWhileYield x s)
3604 Skip s -> return $ Skip (DropWhileNext s)
3607 step' _ (DropWhileYield x st) = return $ Yield x (DropWhileNext st)
3609{-# INLINE dropWhile #-}
3610dropWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
3611dropWhile f = dropWhileM (return . f)
3613{-# INLINE_NORMAL filterM #-}
3614filterM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
3615filterM f (Stream step state) = Stream step' state
3617 {-# INLINE_LATE step' #-}
3626 Skip s -> return $ Skip s
3629{-# INLINE filter #-}
3630filter :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
3631filter f = filterM (return . f)
3633{-# INLINE_NORMAL uniq #-}
3634uniq :: (Eq a, Monad m) => Stream m a -> Stream m a
3635uniq (Stream step state) = Stream step' (Nothing, state)
3637 {-# INLINE_LATE step' #-}
3638 step' gst (Nothing, st) = do
3641 Yield x s -> return $ Yield x (Just x, s)
3642 Skip s -> return $ Skip (Nothing, s)
3644 step' gst (Just x, st) = do
3647 Yield y s | x == y -> return $ Skip (Just x, s)
3648 | otherwise -> return $ Yield y (Just y, s)
3649 Skip s -> return $ Skip (Just x, s)
3652------------------------------------------------------------------------------
3653-- Transformation by Mapping
3654------------------------------------------------------------------------------
3656{-# INLINE_NORMAL sequence #-}
3657sequence :: Monad m => Stream m (m a) -> Stream m a
3658sequence (Stream step state) = Stream step' state
3660 {-# INLINE_LATE step' #-}
3662 r <- step (adaptState gst) st
3664 Yield x s -> x >>= \a -> return (Yield a s)
3665 Skip s -> return $ Skip s
3668------------------------------------------------------------------------------
3670------------------------------------------------------------------------------
3672data LoopState x s = FirstYield s
3673 | InterspersingYield s
3676{-# INLINE_NORMAL intersperseM #-}
3677intersperseM :: Monad m => m a -> Stream m a -> Stream m a
3678intersperseM m (Stream step state) = Stream step' (FirstYield state)
3680 {-# INLINE_LATE step' #-}
3681 step' gst (FirstYield st) = do
3685 Yield x s -> Skip (YieldAndCarry x s)
3686 Skip s -> Skip (FirstYield s)
3689 step' gst (InterspersingYield st) = do
3694 return $ Yield a (YieldAndCarry x s)
3695 Skip s -> return $ Skip $ InterspersingYield s
3698 step' _ (YieldAndCarry x st) = return $ Yield x (InterspersingYield st)
3703 | SuffixYield a (SuffixState s a)
3705{-# INLINE_NORMAL intersperseSuffix #-}
3706intersperseSuffix :: forall m a. Monad m => m a -> Stream m a -> Stream m a
3707intersperseSuffix action (Stream step state) = Stream step' (SuffixElem state)
3709 {-# INLINE_LATE step' #-}
3710 step' gst (SuffixElem st) = do
3713 Yield x s -> Skip (SuffixYield x (SuffixSuffix s))
3714 Skip s -> Skip (SuffixElem s)
3717 step' _ (SuffixSuffix st) = do
3718 action >>= \r -> return $ Skip (SuffixYield r (SuffixElem st))
3720 step' _ (SuffixYield x next) = return $ Yield x next
3722data SuffixSpanState s a
3723 = SuffixSpanElem s Int
3724 | SuffixSpanSuffix s
3725 | SuffixSpanYield a (SuffixSpanState s a)
3729-- | intersperse after every n items
3730{-# INLINE_NORMAL intersperseSuffixBySpan #-}
3731intersperseSuffixBySpan :: forall m a. Monad m
3732 => Int -> m a -> Stream m a -> Stream m a
3733intersperseSuffixBySpan n action (Stream step state) =
3734 Stream step' (SuffixSpanElem state n)
3736 {-# INLINE_LATE step' #-}
3737 step' gst (SuffixSpanElem st i) | i > 0 = do
3740 Yield x s -> Skip (SuffixSpanYield x (SuffixSpanElem s (i - 1)))
3741 Skip s -> Skip (SuffixSpanElem s i)
3742 Stop -> if i == n then Stop else Skip SuffixSpanLast
3743 step' _ (SuffixSpanElem st _) = return $ Skip (SuffixSpanSuffix st)
3745 step' _ (SuffixSpanSuffix st) = do
3746 action >>= \r -> return $ Skip (SuffixSpanYield r (SuffixSpanElem st n))
3748 step' _ (SuffixSpanLast) = do
3749 action >>= \r -> return $ Skip (SuffixSpanYield r SuffixSpanStop)
3751 step' _ (SuffixSpanYield x next) = return $ Yield x next
3753 step' _ (SuffixSpanStop) = return Stop
3755{-# INLINE intersperse #-}
3756intersperse :: Monad m => a -> Stream m a -> Stream m a
3757intersperse a = intersperseM (return a)
3759{-# INLINE_NORMAL insertBy #-}
3760insertBy :: Monad m => (a -> a -> Ordering) -> a -> Stream m a -> Stream m a
3761insertBy cmp a (Stream step state) = Stream step' (state, False, Nothing)
3763 {-# INLINE_LATE step' #-}
3764 step' gst (st, False, _) = do
3767 Yield x s -> case cmp a x of
3768 GT -> return $ Yield x (s, False, Nothing)
3769 _ -> return $ Yield a (s, True, Just x)
3770 Skip s -> return $ Skip (s, False, Nothing)
3771 Stop -> return $ Yield a (st, True, Nothing)
3773 step' _ (_, True, Nothing) = return Stop
3775 step' gst (st, True, Just prev) = do
3778 Yield x s -> return $ Yield prev (s, True, Just x)
3779 Skip s -> return $ Skip (s, True, Just prev)
3780 Stop -> return $ Yield prev (st, True, Nothing)
3782------------------------------------------------------------------------------
3784------------------------------------------------------------------------------
3786{-# INLINE_NORMAL deleteBy #-}
3787deleteBy :: Monad m => (a -> a -> Bool) -> a -> Stream m a -> Stream m a
3788deleteBy eq x (Stream step state) = Stream step' (state, False)
3790 {-# INLINE_LATE step' #-}
3791 step' gst (st, False) = do
3794 Yield y s -> return $
3795 if eq x y then Skip (s, True) else Yield y (s, False)
3796 Skip s -> return $ Skip (s, False)
3799 step' gst (st, True) = do
3802 Yield y s -> return $ Yield y (s, True)
3803 Skip s -> return $ Skip (s, True)
3806------------------------------------------------------------------------------
3807-- Transformation by Map and Filter
3808------------------------------------------------------------------------------
3810-- XXX Will this always fuse properly?
3811{-# INLINE_NORMAL mapMaybe #-}
3812mapMaybe :: Monad m => (a -> Maybe b) -> Stream m a -> Stream m b
3813mapMaybe f = fmap fromJust . filter isJust . map f
3815{-# INLINE_NORMAL mapMaybeM #-}
3816mapMaybeM :: Monad m => (a -> m (Maybe b)) -> Stream m a -> Stream m b
3817mapMaybeM f = fmap fromJust . filter isJust . mapM f
3819------------------------------------------------------------------------------
3821------------------------------------------------------------------------------
3823{-# INLINE_NORMAL indexed #-}
3824indexed :: Monad m => Stream m a -> Stream m (Int, a)
3825indexed (Stream step state) = Stream step' (state, 0)
3827 {-# INLINE_LATE step' #-}
3828 step' gst (st, i) = i `seq` do
3829 r <- step (adaptState gst) st
3831 Yield x s -> return $ Yield (i, x) (s, i+1)
3832 Skip s -> return $ Skip (s, i)
3835{-# INLINE_NORMAL indexedR #-}
3836indexedR :: Monad m => Int -> Stream m a -> Stream m (Int, a)
3837indexedR m (Stream step state) = Stream step' (state, m)
3839 {-# INLINE_LATE step' #-}
3840 step' gst (st, i) = i `seq` do
3841 r <- step (adaptState gst) st
3843 Yield x s -> let i' = i - 1
3844 in return $ Yield (i, x) (s, i')
3845 Skip s -> return $ Skip (s, i)
3848{-# INLINE_NORMAL zipWithM #-}
3850 => (a -> b -> m c) -> Stream m a -> Stream m b -> Stream m c
3851zipWithM f (Stream stepa ta) (Stream stepb tb) = Stream step (ta, tb, Nothing)
3853 {-# INLINE_LATE step #-}
3854 step gst (sa, sb, Nothing) = do
3855 r <- stepa (adaptState gst) sa
3858 Yield x sa' -> Skip (sa', sb, Just x)
3859 Skip sa' -> Skip (sa', sb, Nothing)
3862 step gst (sa, sb, Just x) = do
3863 r <- stepb (adaptState gst) sb
3867 return $ Yield z (sa, sb', Nothing)
3868 Skip sb' -> return $ Skip (sa, sb', Just x)
3871#if __GLASGOW_HASKELL__ >= 801
3872{-# RULES "zipWithM xs xs"
3873 forall f xs. zipWithM @Identity f xs xs = mapM (\x -> f x x) xs #-}
3876{-# INLINE zipWith #-}
3877zipWith :: Monad m => (a -> b -> c) -> Stream m a -> Stream m b -> Stream m c
3878zipWith f = zipWithM (\a b -> return (f a b))
3880------------------------------------------------------------------------------
3882------------------------------------------------------------------------------
3884{-# INLINE_NORMAL mergeByM #-}
3887 => (a -> a -> m Ordering) -> Stream m a -> Stream m a -> Stream m a
3888mergeByM cmp (Stream stepa ta) (Stream stepb tb) =
3889 Stream step (Just ta, Just tb, Nothing, Nothing)
3891 {-# INLINE_LATE step #-}
3893 -- one of the values is missing, and the corresponding stream is running
3894 step gst (Just sa, sb, Nothing, b) = do
3897 Yield a sa' -> Skip (Just sa', sb, Just a, b)
3898 Skip sa' -> Skip (Just sa', sb, Nothing, b)
3899 Stop -> Skip (Nothing, sb, Nothing, b)
3901 step gst (sa, Just sb, a, Nothing) = do
3904 Yield b sb' -> Skip (sa, Just sb', a, Just b)
3905 Skip sb' -> Skip (sa, Just sb', a, Nothing)
3906 Stop -> Skip (sa, Nothing, a, Nothing)
3908 -- both the values are available
3909 step _ (sa, sb, Just a, Just b) = do
3911 return $ case res of
3912 GT -> Yield b (sa, sb, Just a, Nothing)
3913 _ -> Yield a (sa, sb, Nothing, Just b)
3915 -- one of the values is missing, corresponding stream is done
3916 step _ (Nothing, sb, Nothing, Just b) =
3917 return $ Yield b (Nothing, sb, Nothing, Nothing)
3919 step _ (sa, Nothing, Just a, Nothing) =
3920 return $ Yield a (sa, Nothing, Nothing, Nothing)
3922 step _ (Nothing, Nothing, Nothing, Nothing) = return Stop
3924{-# INLINE mergeBy #-}
3927 => (a -> a -> Ordering) -> Stream m a -> Stream m a -> Stream m a
3928mergeBy cmp = mergeByM (\a b -> return $ cmp a b)
3930------------------------------------------------------------------------------
3931-- Transformation comprehensions
3932------------------------------------------------------------------------------
3934{-# INLINE_NORMAL the #-}
3935the :: (Eq a, Monad m) => Stream m a -> m (Maybe a)
3936the (Stream step state) = go state
3939 r <- step defState st
3941 Yield x s -> go' x s
3943 Stop -> return Nothing
3945 r <- step defState st
3947 Yield x s | x == n -> go' n s
3948 | otherwise -> return Nothing
3950 Stop -> return (Just n)
3952{-# INLINE runFold #-}
3953runFold :: (Monad m) => Fold m a b -> Stream m a -> m b
3954runFold (Fold step begin done) = foldlMx' step begin done
3956-------------------------------------------------------------------------------
3957-- Concurrent application and fold
3958-------------------------------------------------------------------------------
3960-- XXX These functions should be moved to Stream/Parallel.hs
3962-- Using StreamD the worker stream producing code can fuse with the code to
3963-- queue output to the SVar giving some perf boost.
3965-- Note that StreamD can only be used in limited situations, specifically, we
3966-- cannot implement joinStreamVarPar using this.
3968-- XXX make sure that the SVar passed is a Parallel style SVar.
3970-- | Fold the supplied stream to the SVar asynchronously using Parallel
3971-- concurrency style.
3972-- {-# INLINE_NORMAL toSVarParallel #-}
3973{-# INLINE toSVarParallel #-}
3974toSVarParallel :: MonadAsync m
3975 => State t m a -> SVar t m a -> Stream m a -> m ()
3976toSVarParallel st sv xs =
3977 if svarInspectMode sv
3981 case getYieldLimit st of
3982 Nothing -> doFork (work Nothing)
3984 (handleChildException sv)
3985 Just _ -> doFork (workLim Nothing)
3987 (handleChildException sv)
3992 {-# NOINLINE work #-}
3993 work info = (runFold (FL.toParallelSVar sv info) xs)
3995 {-# NOINLINE workLim #-}
3996 workLim info = runFold (FL.toParallelSVarLimited sv info) xs
3998 {-# NOINLINE forkWithDiag #-}
4000 -- We do not use workerCount in case of ParallelVar but still there is
4001 -- no harm in maintaining it correctly.
4002 liftIO $ atomicModifyIORefCAS_ (workerCount sv) $ \n -> n + 1
4004 -- This allocation matters when significant number of workers are being
4005 -- sent. We allocate it only when needed. The overhead increases by 4x.
4007 case yieldRateInfo sv of
4008 Nothing -> return Nothing
4009 Just _ -> liftIO $ do
4010 cntRef <- newIORef 0
4011 t <- getTime Monotonic
4012 lat <- newIORef (0, t)
4013 return $ Just WorkerInfo
4014 { workerYieldMax = 0
4015 , workerYieldCount = cntRef
4016 , workerLatencyStart = lat
4019 case getYieldLimit st of
4020 Nothing -> doFork (work winfo)
4022 (handleChildException sv)
4023 Just _ -> doFork (workLim winfo)
4025 (handleChildException sv)
4028{-# INLINE_NORMAL mkParallelD #-}
4029mkParallelD :: MonadAsync m => Stream m a -> Stream m a
4030mkParallelD m = Stream step Nothing
4033 step gst Nothing = do
4034 sv <- newParallelVar StopNone gst
4035 toSVarParallel gst sv m
4036 -- XXX use unfold instead?
4037 return $ Skip $ Just $ fromSVar sv
4039 step gst (Just (UnStream step1 st)) = do
4042 Yield a s -> Yield a (Just $ Stream step1 s)
4043 Skip s -> Skip (Just $ Stream step1 s)
4046-- Compare with mkAsync. mkAsync uses an Async style SVar whereas this uses a
4047-- parallel style SVar for evaluation. Currently, parallel style cannot use
4048-- rate control whereas Async style can use rate control. In async style SVar
4049-- the worker thread terminates when the buffer is full whereas in Parallel
4052-- | Make the stream producer and consumer run concurrently by introducing a
4053-- buffer between them. The producer thread evaluates the input stream until
4054-- the buffer fills, it blocks if the buffer is full until there is space in
4055-- the buffer. The consumer consumes the stream lazily from the buffer.
4059{-# INLINE_NORMAL mkParallel #-}
4060mkParallel :: (K.IsStream t, MonadAsync m) => t m a -> t m a
4061mkParallel = fromStreamD . mkParallelD . toStreamD
4063-------------------------------------------------------------------------------
4065-------------------------------------------------------------------------------
4067-- | Create an SVar with a fold consumer that will fold any elements sent to it
4068-- using the supplied fold function.
4069{-# INLINE newFoldSVar #-}
4070newFoldSVar :: MonadAsync m => State t m a -> Fold m a b -> m (SVar t m a)
4071newFoldSVar stt f = do
4072 -- Buffer size for the SVar is derived from the current state
4073 sv <- newParallelVar StopAny (adaptState stt)
4074 -- Add the producer thread-id to the SVar.
4075 liftIO myThreadId >>= modifyThread sv
4076 void $ doFork (work sv) (svarMrun sv) (handleFoldException sv)
4081 {-# NOINLINE work #-}
4082 work sv = void $ runFold f $ fromProducer sv
4084data TapState sv st = TapInit | Tapping sv st | TapDone st
4086{-# INLINE_NORMAL tapAsync #-}
4087tapAsync :: MonadAsync m => Fold m a b -> Stream m a -> Stream m a
4088tapAsync f (Stream step1 state1) = Stream step TapInit
4092 -- In general, a Stop event would come equipped with the result
4093 -- of the fold. It is not used here but it would be useful in
4094 -- applicative and distribute.
4095 done <- fromConsumer svr
4096 when (not done) $ do
4097 liftIO $ withDiagMVar svr "teeToSVar: waiting to drain"
4098 $ takeMVar (outputDoorBellFromConsumer svr)
4102 liftIO $ sendStop svr Nothing
4103 -- drain/wait until a stop event arrives from the fold.
4106 {-# INLINE_LATE step #-}
4107 step gst TapInit = do
4108 sv <- newFoldSVar gst f
4109 return $ Skip (Tapping sv state1)
4111 step gst (Tapping sv st) = do
4115 done <- pushToFold sv a
4118 -- XXX we do not need to wait synchronously here
4120 return $ Yield a (TapDone s)
4121 else return $ Yield a (Tapping sv s)
4122 Skip s -> return $ Skip (Tapping sv s)
4127 step gst (TapDone st) = do
4130 Yield a s -> Yield a (TapDone s)
4131 Skip s -> Skip (TapDone s)
4134-- XXX Exported from Array again as this fold is specific to Array
4135-- | Take last 'n' elements from the stream and discard the rest.
4137lastN :: (Storable a, MonadIO m) => Int -> Fold m a (Array a)
4138lastN n = Fold step initial done
4140 step (Tuple3' rb rh i) a = do
4141 rh1 <- liftIO $ RB.unsafeInsert rb rh a
4142 return $ Tuple3' rb rh1 (i + 1)
4143 initial = fmap (\(a, b) -> Tuple3' a b (0 :: Int)) $ liftIO $ RB.new n
4144 done (Tuple3' rb rh i) = do
4145 arr <- liftIO $ A.newArray n
4146 foldFunc i rh snoc' arr rb
4147 snoc' b a = liftIO $ A.unsafeSnoc b a
4149 | i < n = RB.unsafeFoldRingM
4150 | otherwise = RB.unsafeFoldRingFullM
4152------------------------------------------------------------------------------
4154------------------------------------------------------------------------------
4156-- XXX using getTime in the loop can be pretty expensive especially for
4157-- computations where iterations are lightweight. We have the following
4160-- 1) Run a timeout thread updating a flag asynchronously and check that
4161-- flag here, that way we can have a cheap termination check.
4163-- 2) Use COARSE clock to get time with lower resolution but more efficiently.
4165-- 3) Use rdtscp/rdtsc to get time directly from the processor, compute the
4166-- termination value of rdtsc in the beginning and then in each iteration just
4167-- get rdtsc and check if we should terminate.
4171 | TakeByTimeCheck st s
4172 | TakeByTimeYield st s
4174{-# INLINE_NORMAL takeByTime #-}
4175takeByTime :: (MonadIO m, TimeUnit64 t) => t -> Stream m a -> Stream m a
4176takeByTime duration (Stream step1 state1) = Stream step (TakeByTimeInit state1)
4179 lim = toRelTime64 duration
4181 {-# INLINE_LATE step #-}
4182 step _ (TakeByTimeInit _) | lim == 0 = return Stop
4183 step _ (TakeByTimeInit st) = do
4184 t0 <- liftIO $ getTime Monotonic
4185 return $ Skip (TakeByTimeYield st t0)
4186 step _ (TakeByTimeCheck st t0) = do
4187 t <- liftIO $ getTime Monotonic
4189 if diffAbsTime64 t t0 > lim
4191 else Skip (TakeByTimeYield st t0)
4192 step gst (TakeByTimeYield st t0) = do
4195 Yield x s -> Yield x (TakeByTimeCheck s t0)
4196 Skip s -> Skip (TakeByTimeCheck s t0)
4199data DropByTime st s x
4201 | DropByTimeGen st s
4202 | DropByTimeCheck st s x
4203 | DropByTimeYield st
4205{-# INLINE_NORMAL dropByTime #-}
4206dropByTime :: (MonadIO m, TimeUnit64 t) => t -> Stream m a -> Stream m a
4207dropByTime duration (Stream step1 state1) = Stream step (DropByTimeInit state1)
4210 lim = toRelTime64 duration
4212 {-# INLINE_LATE step #-}
4213 step _ (DropByTimeInit st) = do
4214 t0 <- liftIO $ getTime Monotonic
4215 return $ Skip (DropByTimeGen st t0)
4216 step gst (DropByTimeGen st t0) = do
4219 Yield x s -> Skip (DropByTimeCheck s t0 x)
4220 Skip s -> Skip (DropByTimeGen s t0)
4222 step _ (DropByTimeCheck st t0 x) = do
4223 t <- liftIO $ getTime Monotonic
4224 if diffAbsTime64 t t0 <= lim
4225 then return $ Skip $ DropByTimeGen st t0
4226 else return $ Yield x $ DropByTimeYield st
4227 step gst (DropByTimeYield st) = do
4230 Yield x s -> Yield x (DropByTimeYield s)
4231 Skip s -> Skip (DropByTimeYield s)
4234-- XXX we should move this to stream generation section of this file. Also, the
4235-- take/drop combinators above should be moved to filtering section.
4236{-# INLINE_NORMAL currentTime #-}
4237currentTime :: MonadAsync m => Double -> Stream m AbsTime
4238currentTime g = Stream step Nothing
4242 g' = g * 10 ^ (6 :: Int)
4244 -- XXX should have a minimum granularity to avoid high CPU usage?
4245 {-# INLINE delayTime #-}
4247 if g' >= fromIntegral (maxBound :: Int)
4251 updateTimeVar timeVar = do
4252 threadDelay $ delayTime
4253 MicroSecond64 t <- fromAbsTime <$> getTime Monotonic
4254 modifyVar' timeVar (const t)
4256 {-# INLINE_LATE step #-}
4258 -- XXX note that this is safe only on a 64-bit machine. On a 32-bit
4259 -- machine a 64-bit 'Var' cannot be read consistently without a lock
4260 -- while another thread is writing to it.
4261 timeVar <- liftIO $ newVar (0 :: Int64)
4262 tid <- forkManaged $ liftIO $ forever (updateTimeVar timeVar)
4263 return $ Skip $ Just (timeVar, tid)
4265 step _ s@(Just (timeVar, _)) = do
4266 a <- liftIO $ readVar timeVar
4267 -- XXX we can perhaps use an AbsTime64 using a 64 bit Int for
4268 -- efficiency. or maybe we can use a representation using Double for
4269 -- floating precision time
4270 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