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@.
134 -- ** Specialized Folds
161 -- ** Flattening nested streams
164 , ConcatMapUState (..)
166 , ConcatUnfoldInterleaveState (..)
167 , concatUnfoldInterleave
168 , concatUnfoldRoundrobin
171 , InterleaveState(..)
176 , roundRobin -- interleaveFair?/ParallelFair
209 -- | Transform a stream into another type.
226 -- ** By folding (scans)
273 , intersperseSuffixBySpan
297 -- * Transformation comprehensions
303 , clearIORefFinalizer
315 -- * Concurrent Application
324import Control.Concurrent (killThread, myThreadId, takeMVar, threadDelay)
325import Control.Exception
326 (assert, Exception, SomeException, AsyncException, fromException, mask_)
327import Control.Monad (void, when, forever)
328import Control.Monad.Catch (MonadCatch, throwM)
329import Control.Monad.IO.Class (MonadIO(..))
330import Control.Monad.Reader (ReaderT)
331import Control.Monad.State.Strict (StateT)
332import Control.Monad.Trans (MonadTrans(lift))
333import Control.Monad.Trans.Control (MonadBaseControl, liftBaseOp_)
334import Data.Bits (shiftR, shiftL, (.|.), (.&.))
335import Data.Functor.Identity (Identity(..))
336import Data.Int (Int64)
337import Data.IORef (newIORef, readIORef, mkWeakIORef, writeIORef, IORef)
338import Data.Maybe (fromJust, isJust, isNothing)
339import Data.Word (Word32)
340import Foreign.Ptr (Ptr)
341import Foreign.Storable (Storable(..))
342import GHC.Types (SPEC(..))
343import System.Mem (performMajorGC)
345 hiding (map, mapM, mapM_, repeat, foldr, last, take, filter,
346 takeWhile, drop, dropWhile, all, any, maximum, minimum, elem,
347 notElem, null, head, tail, zipWith, lookup, foldr1, sequence,
348 (!!), scanl, scanl1, concatMap, replicate, enumFromTo, concat,
349 reverse, iterate, splitAt)
351import qualified Control.Monad.Catch as MC
352import qualified Control.Monad.Reader as Reader
353import qualified Control.Monad.State.Strict as State
354import qualified Prelude
356import Fusion.Plugin.Types (Fuse(..))
357import Streamly.Internal.Mutable.Prim.Var
358 (Prim, Var, readVar, newVar, modifyVar')
359import Streamly.Internal.Data.Time.Units
360 (TimeUnit64, toRelTime64, diffAbsTime64)
362import Streamly.Internal.Data.Atomics (atomicModifyIORefCAS_)
363import Streamly.Internal.Memory.Array.Types (Array(..))
364import Streamly.Internal.Data.Fold.Types (Fold(..))
365import Streamly.Internal.Data.Parser.Types (Parser(..))
366import Streamly.Internal.Data.Pipe.Types (Pipe(..), PipeState(..))
367import Streamly.Internal.Data.Time.Clock (Clock(Monotonic), getTime)
368import Streamly.Internal.Data.Time.Units
369 (MicroSecond64(..), fromAbsTime, toAbsTime, AbsTime)
370import Streamly.Internal.Data.Unfold.Types (Unfold(..))
371import Streamly.Internal.Data.Strict (Tuple3'(..))
373import Streamly.Internal.Data.Stream.StreamD.Type
374import Streamly.Internal.Data.SVar
375import Streamly.Internal.Data.Stream.SVar (fromConsumer, pushToFold)
377import qualified Streamly.Internal.Data.Pipe.Types as Pipe
378import qualified Streamly.Internal.Memory.Array.Types as A
379import qualified Streamly.Internal.Data.Fold as FL
380import qualified Streamly.Memory.Ring as RB
381import qualified Streamly.Internal.Data.Stream.StreamK as K
382import qualified Streamly.Internal.Data.Parser.Types as PR
384------------------------------------------------------------------------------
386------------------------------------------------------------------------------
388-- | An empty 'Stream'.
389{-# INLINE_NORMAL nil #-}
390nil :: Monad m => Stream m a
391nil = Stream (\_ _ -> return Stop) ()
393-- | An empty 'Stream' with a side effect.
394{-# INLINE_NORMAL nilM #-}
395nilM :: Monad m => m b -> Stream m a
396nilM m = Stream (\_ _ -> m >> return Stop) ()
398{-# INLINE_NORMAL consM #-}
399consM :: Monad m => m a -> Stream m a -> Stream m a
400consM m (Stream step state) = Stream step1 Nothing
402 {-# INLINE_LATE step1 #-}
403 step1 _ Nothing = m >>= \x -> return $ Yield x (Just state)
404 step1 gst (Just st) = do
408 Yield a s -> Yield a (Just s)
409 Skip s -> Skip (Just s)
412-- XXX implement in terms of consM?
413-- cons x = consM (return x)
415-- | Can fuse but has O(n^2) complexity.
416{-# INLINE_NORMAL cons #-}
417cons :: Monad m => a -> Stream m a -> Stream m a
418cons x (Stream step state) = Stream step1 Nothing
420 {-# INLINE_LATE step1 #-}
421 step1 _ Nothing = return $ Yield x (Just state)
422 step1 gst (Just st) = do
426 Yield a s -> Yield a (Just s)
427 Skip s -> Skip (Just s)
430-------------------------------------------------------------------------------
432-------------------------------------------------------------------------------
434-- Does not fuse, has the same performance as the StreamK version.
435{-# INLINE_NORMAL uncons #-}
436uncons :: Monad m => Stream m a -> m (Maybe (a, Stream m a))
437uncons (UnStream step state) = go state
440 r <- step defState st
442 Yield x s -> return $ Just (x, Stream step s)
444 Stop -> return Nothing
446------------------------------------------------------------------------------
447-- Generation by unfold
448------------------------------------------------------------------------------
450{-# INLINE_NORMAL unfoldrM #-}
451unfoldrM :: Monad m => (s -> m (Maybe (a, s))) -> s -> Stream m a
452unfoldrM next state = Stream step state
454 {-# INLINE_LATE step #-}
458 Just (x, s) -> Yield x s
461{-# INLINE_LATE unfoldr #-}
462unfoldr :: Monad m => (s -> Maybe (a, s)) -> s -> Stream m a
463unfoldr f = unfoldrM (return . f)
465-- | Convert an 'Unfold' into a 'Stream' by supplying it a seed.
467{-# INLINE_NORMAL unfold #-}
468unfold :: Monad m => Unfold m a b -> a -> Stream m b
469unfold (Unfold ustep inject) seed = Stream step Nothing
471 {-# INLINE_LATE step #-}
472 step _ Nothing = inject seed >>= return . Skip . Just
473 step _ (Just st) = do
476 Yield x s -> Yield x (Just s)
477 Skip s -> Skip (Just s)
480------------------------------------------------------------------------------
481-- Specialized Generation
482------------------------------------------------------------------------------
484{-# INLINE_NORMAL repeatM #-}
485repeatM :: Monad m => m a -> Stream m a
486repeatM x = Stream (\_ _ -> x >>= \r -> return $ Yield r ()) ()
488{-# INLINE_NORMAL repeat #-}
489repeat :: Monad m => a -> Stream m a
490repeat x = Stream (\_ _ -> return $ Yield x ()) ()
492{-# INLINE_NORMAL iterateM #-}
493iterateM :: Monad m => (a -> m a) -> m a -> Stream m a
494iterateM step = Stream (\_ st -> st >>= \x -> return $ Yield x (step x))
496{-# INLINE_NORMAL iterate #-}
497iterate :: Monad m => (a -> a) -> a -> Stream m a
498iterate step st = iterateM (return . step) (return st)
500{-# INLINE_NORMAL replicateM #-}
501replicateM :: forall m a. Monad m => Int -> m a -> Stream m a
502replicateM n p = Stream step n
504 {-# INLINE_LATE step #-}
506 | i <= 0 = return Stop
509 return $ Yield x (i - 1)
511{-# INLINE_NORMAL replicate #-}
512replicate :: Monad m => Int -> a -> Stream m a
513replicate n x = replicateM n (return x)
515-- This would not work properly for floats, therefore we put an Integral
517-- | Can be used to enumerate unbounded integrals. This does not check for
518-- overflow or underflow for bounded integrals.
519{-# INLINE_NORMAL enumerateFromStepIntegral #-}
520enumerateFromStepIntegral :: (Integral a, Monad m) => a -> a -> Stream m a
521enumerateFromStepIntegral from stride =
522 from `seq` stride `seq` Stream step from
524 {-# INLINE_LATE step #-}
525 step _ !x = return $ Yield x $! (x + stride)
527-- We are assuming that "to" is constrained by the type to be within
529{-# INLINE enumerateFromToIntegral #-}
530enumerateFromToIntegral :: (Monad m, Integral a) => a -> a -> Stream m a
531enumerateFromToIntegral from to =
532 takeWhile (<= to) $ enumerateFromStepIntegral from 1
534{-# INLINE enumerateFromIntegral #-}
535enumerateFromIntegral :: (Monad m, Integral a, Bounded a) => a -> Stream m a
536enumerateFromIntegral from = enumerateFromToIntegral from maxBound
538data EnumState a = EnumInit | EnumYield a a a | EnumStop
540{-# INLINE_NORMAL enumerateFromThenToIntegralUp #-}
541enumerateFromThenToIntegralUp
542 :: (Monad m, Integral a)
543 => a -> a -> a -> Stream m a
544enumerateFromThenToIntegralUp from next to = Stream step EnumInit
546 {-# INLINE_LATE step #-}
552 else Yield from EnumStop
553 else -- from <= next <= to
554 let stride = next - from
555 in Skip $ EnumYield from stride (to - stride)
557 step _ (EnumYield x stride toMinus) =
560 then Yield x EnumStop
561 else Yield x $ EnumYield (x + stride) stride toMinus
563 step _ EnumStop = return Stop
565{-# INLINE_NORMAL enumerateFromThenToIntegralDn #-}
566enumerateFromThenToIntegralDn
567 :: (Monad m, Integral a)
568 => a -> a -> a -> Stream m a
569enumerateFromThenToIntegralDn from next to = Stream step EnumInit
571 {-# INLINE_LATE step #-}
573 return $ if to > next
576 else Yield from EnumStop
577 else -- from >= next >= to
578 let stride = next - from
579 in Skip $ EnumYield from stride (to - stride)
581 step _ (EnumYield x stride toMinus) =
584 then Yield x EnumStop
585 else Yield x $ EnumYield (x + stride) stride toMinus
587 step _ EnumStop = return Stop
589{-# INLINE_NORMAL enumerateFromThenToIntegral #-}
590enumerateFromThenToIntegral
591 :: (Monad m, Integral a)
592 => a -> a -> a -> Stream m a
593enumerateFromThenToIntegral from next to
594 | next >= from = enumerateFromThenToIntegralUp from next to
595 | otherwise = enumerateFromThenToIntegralDn from next to
597{-# INLINE_NORMAL enumerateFromThenIntegral #-}
598enumerateFromThenIntegral
599 :: (Monad m, Integral a, Bounded a)
600 => a -> a -> Stream m a
601enumerateFromThenIntegral from next =
603 then enumerateFromThenToIntegralUp from next maxBound
604 else enumerateFromThenToIntegralDn from next minBound
606-- For floating point numbers if the increment is less than the precision then
607-- it just gets lost. Therefore we cannot always increment it correctly by just
609-- 9007199254740992 + 1 + 1 :: Double => 9.007199254740992e15
610-- 9007199254740992 + 2 :: Double => 9.007199254740994e15
612-- Instead we accumulate the increment counter and compute the increment
613-- every time before adding it to the starting number.
615-- This works for Integrals as well as floating point numbers, but
616-- enumerateFromStepIntegral is faster for integrals.
617{-# INLINE_NORMAL enumerateFromStepNum #-}
618enumerateFromStepNum :: (Monad m, Num a) => a -> a -> Stream m a
619enumerateFromStepNum from stride = Stream step 0
621 {-# INLINE_LATE step #-}
622 step _ !i = return $ (Yield $! (from + i * stride)) $! (i + 1)
624{-# INLINE_NORMAL numFrom #-}
625numFrom :: (Monad m, Num a) => a -> Stream m a
626numFrom from = enumerateFromStepNum from 1
628{-# INLINE_NORMAL numFromThen #-}
629numFromThen :: (Monad m, Num a) => a -> a -> Stream m a
630numFromThen from next = enumerateFromStepNum from (next - from)
632-- We cannot write a general function for Num. The only way to write code
633-- portable between the two is to use a 'Real' constraint and convert between
634-- Fractional and Integral using fromRational which is horribly slow.
635{-# INLINE_NORMAL enumerateFromToFractional #-}
636enumerateFromToFractional
637 :: (Monad m, Fractional a, Ord a)
638 => a -> a -> Stream m a
639enumerateFromToFractional from to =
640 takeWhile (<= to + 1 / 2) $ enumerateFromStepNum from 1
642{-# INLINE_NORMAL enumerateFromThenToFractional #-}
643enumerateFromThenToFractional
644 :: (Monad m, Fractional a, Ord a)
645 => a -> a -> a -> Stream m a
646enumerateFromThenToFractional from next to =
647 takeWhile predicate $ numFromThen from next
649 mid = (next - from) / 2
650 predicate | next >= from = (<= to + mid)
651 | otherwise = (>= to + mid)
653-------------------------------------------------------------------------------
654-- Generation by Conversion
655-------------------------------------------------------------------------------
657{-# INLINE_NORMAL fromIndicesM #-}
658fromIndicesM :: Monad m => (Int -> m a) -> Stream m a
659fromIndicesM gen = Stream step 0
661 {-# INLINE_LATE step #-}
664 return $ Yield x (i + 1)
666{-# INLINE fromIndices #-}
667fromIndices :: Monad m => (Int -> a) -> Stream m a
668fromIndices gen = fromIndicesM (return . gen)
670{-# INLINE_NORMAL generateM #-}
671generateM :: Monad m => Int -> (Int -> m a) -> Stream m a
672generateM n gen = n `seq` Stream step 0
674 {-# INLINE_LATE step #-}
675 step _ i | i < n = do
677 return $ Yield x (i + 1)
678 | otherwise = return Stop
680{-# INLINE generate #-}
681generate :: Monad m => Int -> (Int -> a) -> Stream m a
682generate n gen = generateM n (return . gen)
684-- XXX we need the MonadAsync constraint because of a rewrite rule.
685-- | Convert a list of monadic actions to a 'Stream'
686{-# INLINE_LATE fromListM #-}
687fromListM :: MonadAsync m => [m a] -> Stream m a
688fromListM = Stream step
690 {-# INLINE_LATE step #-}
691 step _ (m:ms) = m >>= \x -> return $ Yield x ms
692 step _ [] = return Stop
694{-# INLINE toStreamD #-}
695toStreamD :: (K.IsStream t, Monad m) => t m a -> Stream m a
696toStreamD = fromStreamK . K.toStream
698{-# INLINE_NORMAL fromPrimVar #-}
699fromPrimVar :: (MonadIO m, Prim a) => Var IO a -> Stream m a
700fromPrimVar var = Stream step ()
702 {-# INLINE_LATE step #-}
703 step _ () = liftIO (readVar var) >>= \x -> return $ Yield x ()
705-------------------------------------------------------------------------------
706-- Generation from SVar
707-------------------------------------------------------------------------------
709data FromSVarState t m a =
711 | FromSVarRead (SVar t m a)
712 | FromSVarLoop (SVar t m a) [ChildEvent a]
713 | FromSVarDone (SVar t m a)
715{-# INLINE_NORMAL fromSVar #-}
716fromSVar :: (MonadAsync m) => SVar t m a -> Stream m a
717fromSVar svar = Stream step FromSVarInit
720 {-# INLINE_LATE step #-}
721 step _ FromSVarInit = do
722 ref <- liftIO $ newIORef ()
723 _ <- liftIO $ mkWeakIORef ref hook
724 -- when this copy of svar gets garbage collected "ref" will get
725 -- garbage collected and our GC hook will be called.
726 let sv = svar{svarRef = Just ref}
727 return $ Skip (FromSVarRead sv)
731 {-# NOINLINE hook #-}
733 when (svarInspectMode svar) $ do
734 r <- liftIO $ readIORef (svarStopTime (svarStats svar))
736 printSVar svar "SVar Garbage Collected"
738 -- If there are any SVars referenced by this SVar a GC will prompt
739 -- them to be cleaned up quickly.
740 when (svarInspectMode svar) performMajorGC
742 step _ (FromSVarRead sv) = do
743 list <- readOutputQ sv
744 -- Reversing the output is important to guarantee that we process the
745 -- outputs in the same order as they were generated by the constituent
747 return $ Skip $ FromSVarLoop sv (Prelude.reverse list)
749 step _ (FromSVarLoop sv []) = do
750 done <- postProcess sv
751 return $ Skip $ if done
752 then (FromSVarDone sv)
753 else (FromSVarRead sv)
755 step _ (FromSVarLoop sv (ev : es)) = do
757 ChildYield a -> return $ Yield a (FromSVarLoop sv es)
758 ChildStop tid e -> do
762 stop <- shouldStop tid
765 liftIO (cleanupSVar sv)
766 return $ Skip (FromSVarDone sv)
767 else return $ Skip (FromSVarLoop sv es)
769 case fromException ex of
771 return $ Skip (FromSVarLoop sv es)
772 Nothing -> liftIO (cleanupSVar sv) >> throwM ex
776 case svarStopStyle sv of
777 StopNone -> return False
778 StopAny -> return True
780 sid <- liftIO $ readIORef (svarStopBy sv)
781 return $ if tid == sid then True else False
783 step _ (FromSVarDone sv) = do
784 when (svarInspectMode sv) $ do
785 t <- liftIO $ getTime Monotonic
786 liftIO $ writeIORef (svarStopTime (svarStats sv)) (Just t)
787 liftIO $ printSVar sv "SVar Done"
790-------------------------------------------------------------------------------
791-- Process events received by a fold consumer from a stream producer
792-------------------------------------------------------------------------------
794{-# INLINE_NORMAL fromProducer #-}
795fromProducer :: (MonadAsync m) => SVar t m a -> Stream m a
796fromProducer svar = Stream step (FromSVarRead svar)
799 {-# INLINE_LATE step #-}
800 step _ (FromSVarRead sv) = do
801 list <- readOutputQ sv
802 -- Reversing the output is important to guarantee that we process the
803 -- outputs in the same order as they were generated by the constituent
805 return $ Skip $ FromSVarLoop sv (Prelude.reverse list)
807 step _ (FromSVarLoop sv []) = return $ Skip $ FromSVarRead sv
808 step _ (FromSVarLoop sv (ev : es)) = do
810 ChildYield a -> return $ Yield a (FromSVarLoop sv es)
811 ChildStop tid e -> do
815 sendStopToProducer sv
816 return $ Skip (FromSVarDone sv)
817 Just _ -> error "Bug: fromProducer: received exception"
819 step _ (FromSVarDone sv) = do
820 when (svarInspectMode sv) $ do
821 t <- liftIO $ getTime Monotonic
822 liftIO $ writeIORef (svarStopTime (svarStats sv)) (Just t)
823 liftIO $ printSVar sv "SVar Done"
826 step _ FromSVarInit = undefined
828-------------------------------------------------------------------------------
829-- Hoisting the inner monad
830-------------------------------------------------------------------------------
832{-# INLINE_NORMAL hoist #-}
833hoist :: Monad n => (forall x. m x -> n x) -> Stream m a -> Stream n a
834hoist f (Stream step state) = (Stream step' state)
836 {-# INLINE_LATE step' #-}
838 r <- f $ step (adaptState gst) st
840 Yield x s -> Yield x s
844{-# INLINE generally #-}
845generally :: Monad m => Stream Identity a -> Stream m a
846generally = hoist (return . runIdentity)
848{-# INLINE_NORMAL liftInner #-}
849liftInner :: (Monad m, MonadTrans t, Monad (t m))
850 => Stream m a -> Stream (t m) a
851liftInner (Stream step state) = Stream step' state
853 {-# INLINE_LATE step' #-}
855 r <- lift $ step (adaptState gst) st
857 Yield x s -> Yield x s
861{-# INLINE_NORMAL runReaderT #-}
862runReaderT :: Monad m => s -> Stream (ReaderT s m) a -> Stream m a
863runReaderT sval (Stream step state) = Stream step' state
865 {-# INLINE_LATE step' #-}
867 r <- Reader.runReaderT (step (adaptState gst) st) sval
869 Yield x s -> Yield x s
873{-# INLINE_NORMAL evalStateT #-}
874evalStateT :: Monad m => s -> Stream (StateT s m) a -> Stream m a
875evalStateT sval (Stream step state) = Stream step' (state, sval)
877 {-# INLINE_LATE step' #-}
878 step' gst (st, sv) = do
879 (r, sv') <- State.runStateT (step (adaptState gst) st) sv
881 Yield x s -> Yield x (s, sv')
882 Skip s -> Skip (s, sv')
885{-# INLINE_NORMAL runStateT #-}
886runStateT :: Monad m => s -> Stream (StateT s m) a -> Stream m (s, a)
887runStateT sval (Stream step state) = Stream step' (state, sval)
889 {-# INLINE_LATE step' #-}
890 step' gst (st, sv) = do
891 (r, sv') <- State.runStateT (step (adaptState gst) st) sv
893 Yield x s -> Yield (sv', x) (s, sv')
894 Skip s -> Skip (s, sv')
897------------------------------------------------------------------------------
898-- Elimination by Folds
899------------------------------------------------------------------------------
901------------------------------------------------------------------------------
903------------------------------------------------------------------------------
905{-# INLINE_NORMAL foldr1 #-}
906foldr1 :: Monad m => (a -> a -> a) -> Stream m a -> m (Maybe a)
910 Nothing -> return Nothing
911 Just (h, t) -> fmap Just (foldr f h t)
913------------------------------------------------------------------------------
915------------------------------------------------------------------------------
917{-# INLINE_NORMAL foldlT #-}
918foldlT :: (Monad m, Monad (s m), MonadTrans s)
919 => (s m b -> a -> s m b) -> s m b -> Stream m a -> s m b
920foldlT fstep begin (Stream step state) = go SPEC begin state
923 r <- lift $ step defState st
925 Yield x s -> go SPEC (fstep acc x) s
926 Skip s -> go SPEC acc s
929-- Note, this is going to have horrible performance, because of the nature of
930-- the stream type (i.e. direct stream vs CPS). Its only for reference, it is
931-- likely be practically unusable.
932{-# INLINE_NORMAL foldlS #-}
934 => (Stream m b -> a -> Stream m b) -> Stream m b -> Stream m a -> Stream m b
935foldlS fstep begin (Stream step state) = Stream step' (Left (state, begin))
937 step' gst (Left (st, acc)) = do
938 r <- step (adaptState gst) st
940 Yield x s -> Skip (Left (s, fstep acc x))
941 Skip s -> Skip (Left (s, acc))
942 Stop -> Skip (Right acc)
944 step' gst (Right (Stream stp stt)) = do
945 r <- stp (adaptState gst) stt
947 Yield x s -> Yield x (Right (Stream stp s))
948 Skip s -> Skip (Right (Stream stp s))
951------------------------------------------------------------------------------
953------------------------------------------------------------------------------
955-- Inlined definition. Without the inline "serially/parser/take" benchmark
956-- degrades and parseChunks does not fuse. Even using "inline" at the callsite
958{-# INLINE splitAt #-}
959splitAt :: Int -> [a] -> ([a],[a])
962 | otherwise = splitAt' n ls
964 splitAt' :: Int -> [a] -> ([a], [a])
965 splitAt' _ [] = ([], [])
966 splitAt' 1 (x:xs) = ([x], xs)
967 splitAt' m (x:xs) = (x:xs', xs'')
969 (xs', xs'') = splitAt' (m - 1) xs
971-- | Run a 'Parse' over a stream.
972{-# INLINE_NORMAL parselMx' #-}
975 => (s -> a -> m (PR.Step s))
977 -> (s -> m (Either String b))
979 -> m (Either String b)
980parselMx' pstep initial extract (Stream step state) = do
981 initial >>= go SPEC state []
985 -- XXX currently we are using a dumb list based approach for backtracking
986 -- buffer. This can be replaced by a sliding/ring buffer using Data.Array.
987 -- That will allow us more efficient random back and forth movement.
989 go !_ st buf !pst = do
990 r <- step defState st
995 -- PR.Yield 0 pst1 -> go SPEC s [] pst1
996 PR.Yield n pst1 -> do
997 assert (n <= length (x:buf)) (return ())
998 go SPEC s (Prelude.take n (x:buf)) pst1
999 -- PR.Skip 0 pst1 -> go SPEC s (x:buf) pst1
1000 PR.Skip n pst1 -> do
1001 assert (n <= length (x:buf)) (return ())
1002 let (src0, buf1) = splitAt n (x:buf)
1003 src = Prelude.reverse src0
1004 gobuf SPEC s buf1 src pst1
1005 PR.Stop _ pst1 -> extract pst1
1006 Skip s -> go SPEC s buf pst
1009 gobuf !_ s buf [] !pst = go SPEC s buf pst
1010 gobuf !_ s buf (x:xs) !pst = do
1013 -- PR.Yield 0 pst1 -> go SPEC s [] pst1
1014 PR.Yield n pst1 -> do
1015 assert (n <= length (x:buf)) (return ())
1016 gobuf SPEC s (Prelude.take n (x:buf)) xs pst1
1017 -- PR.Skip 0 pst1 -> gobuf SPEC s (x:buf) xs pst1
1018 PR.Skip n pst1 -> do
1019 assert (n <= length (x:buf)) (return ())
1020 let (src0, buf1) = splitAt n (x:buf)
1021 src = Prelude.reverse src0 ++ xs
1022 gobuf SPEC s buf1 src pst1
1023 PR.Stop _ pst1 -> extract pst1
1025------------------------------------------------------------------------------
1027------------------------------------------------------------------------------
1029{-# ANN type ParseChunksState Fuse #-}
1030data ParseChunksState x inpBuf st pst =
1031 ParseChunksInit inpBuf st
1032 | ParseChunksInitLeftOver inpBuf
1033 | ParseChunksStream st inpBuf pst
1034 | ParseChunksBuf inpBuf st inpBuf pst
1035 | ParseChunksYield x (ParseChunksState x inpBuf st pst)
1037{-# INLINE_NORMAL parseChunks #-}
1043parseChunks (Parser pstep initial extract) (Stream step state) =
1044 Stream stepOuter (ParseChunksInit [] state)
1048 {-# INLINE_LATE stepOuter #-}
1049 -- Buffer is empty, go to stream processing loop
1050 stepOuter _ (ParseChunksInit [] st) = do
1051 initial >>= return . Skip . ParseChunksStream st []
1053 -- Buffer is not empty, go to buffered processing loop
1054 stepOuter _ (ParseChunksInit src st) = do
1055 initial >>= return . Skip . ParseChunksBuf src st []
1057 -- XXX we just discard any leftover input at the end
1058 stepOuter _ (ParseChunksInitLeftOver _) = return Stop
1060 -- Buffer is empty process elements from the stream
1061 stepOuter gst (ParseChunksStream st buf pst) = do
1062 r <- step (adaptState gst) st
1067 -- PR.Yield 0 pst1 -> go SPEC s [] pst1
1068 PR.Yield n pst1 -> do
1069 assert (n <= length (x:buf)) (return ())
1070 let buf1 = Prelude.take n (x:buf)
1071 return $ Skip $ ParseChunksStream s buf1 pst1
1072 -- PR.Skip 0 pst1 ->
1073 -- return $ Skip $ ParseChunksStream s (x:buf) pst1
1074 PR.Skip n pst1 -> do
1075 assert (n <= length (x:buf)) (return ())
1076 let (src0, buf1) = splitAt n (x:buf)
1077 src = Prelude.reverse src0
1078 return $ Skip $ ParseChunksBuf src s buf1 pst1
1079 -- XXX Halt 0 common case?
1080 PR.Stop n pst1 -> do
1083 Left err -> error err
1085 assert (n <= length (x:buf)) (return ())
1086 let src = Prelude.reverse (Prelude.take n (x:buf))
1087 return $ Skip $ ParseChunksYield a (ParseChunksInit src s)
1088 Skip s -> return $ Skip $ ParseChunksStream s buf pst
1092 Left err -> error err
1094 let src = Prelude.reverse buf
1095 return $ Skip $ ParseChunksYield a (ParseChunksInitLeftOver src)
1097 -- go back to stream processing mode
1098 stepOuter _ (ParseChunksBuf [] s buf pst) =
1099 return $ Skip $ ParseChunksStream s buf pst
1101 -- buffered processing loop
1102 stepOuter _ (ParseChunksBuf (x:xs) s buf pst) = do
1105 -- PR.Yield 0 pst1 ->
1106 PR.Yield n pst1 -> do
1107 assert (n <= length (x:buf)) (return ())
1108 let buf1 = Prelude.take n (x:buf)
1109 return $ Skip $ ParseChunksBuf xs s buf1 pst1
1110 -- PR.Skip 0 pst1 -> return $ Skip $ ParseChunksBuf xs s (x:buf) pst1
1111 PR.Skip n pst1 -> do
1112 assert (n <= length (x:buf)) (return ())
1113 let (src0, buf1) = splitAt n (x:buf)
1114 src = Prelude.reverse src0 ++ xs
1115 return $ Skip $ ParseChunksBuf src s buf1 pst1
1117 PR.Stop n pst1 -> do
1120 Left err -> error err
1122 assert (n <= length (x:buf)) (return ())
1123 let src = Prelude.reverse (Prelude.take n (x:buf)) ++ xs
1124 return $ Skip $ ParseChunksYield a (ParseChunksInit src s)
1126 stepOuter _ (ParseChunksYield a next) = return $ Yield a next
1128------------------------------------------------------------------------------
1130------------------------------------------------------------------------------
1132-- | Run a streaming composition, discard the results.
1133{-# INLINE_LATE drain #-}
1134drain :: Monad m => Stream m a -> m ()
1135-- drain = foldrM (\_ xs -> xs) (return ())
1136drain (Stream step state) = go SPEC state
1139 r <- step defState st
1141 Yield _ s -> go SPEC s
1145{-# INLINE_NORMAL null #-}
1146null :: Monad m => Stream m a -> m Bool
1147null m = foldrM (\_ _ -> return False) (return True) m
1149{-# INLINE_NORMAL head #-}
1150head :: Monad m => Stream m a -> m (Maybe a)
1151head m = foldrM (\x _ -> return (Just x)) (return Nothing) m
1153{-# INLINE_NORMAL headElse #-}
1154headElse :: Monad m => a -> Stream m a -> m a
1155headElse a m = foldrM (\x _ -> return x) (return a) m
1157-- Does not fuse, has the same performance as the StreamK version.
1158{-# INLINE_NORMAL tail #-}
1159tail :: Monad m => Stream m a -> m (Maybe (Stream m a))
1160tail (UnStream step state) = go state
1163 r <- step defState st
1165 Yield _ s -> return (Just $ Stream step s)
1167 Stop -> return Nothing
1169-- XXX will it fuse? need custom impl?
1170{-# INLINE_NORMAL last #-}
1171last :: Monad m => Stream m a -> m (Maybe a)
1172last = foldl' (\_ y -> Just y) Nothing
1174{-# INLINE_NORMAL elem #-}
1175elem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
1176-- elem e m = foldrM (\x xs -> if x == e then return True else xs) (return False) m
1177elem e (Stream step state) = go state
1180 r <- step defState st
1183 | x == e -> return True
1186 Stop -> return False
1188{-# INLINE_NORMAL notElem #-}
1189notElem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
1190notElem e s = fmap not (elem e s)
1192{-# INLINE_NORMAL all #-}
1193all :: Monad m => (a -> Bool) -> Stream m a -> m Bool
1194-- all p m = foldrM (\x xs -> if p x then xs else return False) (return True) m
1195all p (Stream step state) = go state
1198 r <- step defState st
1202 | otherwise -> return False
1206{-# INLINE_NORMAL any #-}
1207any :: Monad m => (a -> Bool) -> Stream m a -> m Bool
1208-- any p m = foldrM (\x xs -> if p x then return True else xs) (return False) m
1209any p (Stream step state) = go state
1212 r <- step defState st
1215 | p x -> return True
1218 Stop -> return False
1220{-# INLINE_NORMAL maximum #-}
1221maximum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
1222maximum (Stream step state) = go Nothing state
1225 r <- step defState st
1227 Yield x s -> go (Just x) s
1228 Skip s -> go Nothing s
1229 Stop -> return Nothing
1230 go (Just acc) st = do
1231 r <- step defState st
1234 | acc <= x -> go (Just x) s
1235 | otherwise -> go (Just acc) s
1236 Skip s -> go (Just acc) s
1237 Stop -> return (Just acc)
1239{-# INLINE_NORMAL maximumBy #-}
1240maximumBy :: Monad m => (a -> a -> Ordering) -> Stream m a -> m (Maybe a)
1241maximumBy cmp (Stream step state) = go Nothing state
1244 r <- step defState st
1246 Yield x s -> go (Just x) s
1247 Skip s -> go Nothing s
1248 Stop -> return Nothing
1249 go (Just acc) st = do
1250 r <- step defState st
1252 Yield x s -> case cmp acc x of
1253 GT -> go (Just acc) s
1255 Skip s -> go (Just acc) s
1256 Stop -> return (Just acc)
1258{-# INLINE_NORMAL minimum #-}
1259minimum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
1260minimum (Stream step state) = go Nothing state
1263 r <- step defState st
1265 Yield x s -> go (Just x) s
1266 Skip s -> go Nothing s
1267 Stop -> return Nothing
1268 go (Just acc) st = do
1269 r <- step defState st
1272 | acc <= x -> go (Just acc) s
1273 | otherwise -> go (Just x) s
1274 Skip s -> go (Just acc) s
1275 Stop -> return (Just acc)
1277{-# INLINE_NORMAL minimumBy #-}
1278minimumBy :: Monad m => (a -> a -> Ordering) -> Stream m a -> m (Maybe a)
1279minimumBy cmp (Stream step state) = go Nothing state
1282 r <- step defState st
1284 Yield x s -> go (Just x) s
1285 Skip s -> go Nothing s
1286 Stop -> return Nothing
1287 go (Just acc) st = do
1288 r <- step defState st
1290 Yield x s -> case cmp acc x of
1292 _ -> go (Just acc) s
1293 Skip s -> go (Just acc) s
1294 Stop -> return (Just acc)
1296{-# INLINE_NORMAL (!!) #-}
1297(!!) :: (Monad m) => Stream m a -> Int -> m (Maybe a)
1298(Stream step state) !! i = go i state
1301 r <- step defState st
1303 Yield x s | n < 0 -> return Nothing
1304 | n == 0 -> return $ Just x
1305 | otherwise -> go (n - 1) s
1307 Stop -> return Nothing
1309{-# INLINE_NORMAL lookup #-}
1310lookup :: (Monad m, Eq a) => a -> Stream m (a, b) -> m (Maybe b)
1311lookup e m = foldrM (\(a, b) xs -> if e == a then return (Just b) else xs)
1314{-# INLINE_NORMAL findM #-}
1315findM :: Monad m => (a -> m Bool) -> Stream m a -> m (Maybe a)
1316findM p m = foldrM (\x xs -> p x >>= \r -> if r then return (Just x) else xs)
1320find :: Monad m => (a -> Bool) -> Stream m a -> m (Maybe a)
1321find p = findM (return . p)
1323{-# INLINE_NORMAL findIndices #-}
1324findIndices :: Monad m => (a -> Bool) -> Stream m a -> Stream m Int
1325findIndices p (Stream step state) = Stream step' (state, 0)
1327 {-# INLINE_LATE step' #-}
1328 step' gst (st, i) = i `seq` do
1329 r <- step (adaptState gst) st
1331 Yield x s -> if p x then Yield i (s, i+1) else Skip (s, i+1)
1332 Skip s -> Skip (s, i)
1335{-# INLINE toListRev #-}
1336toListRev :: Monad m => Stream m a -> m [a]
1337toListRev = foldl' (flip (:)) []
1339-- We can implement reverse as:
1341-- > reverse = foldlS (flip cons) nil
1343-- However, this implementation is unusable because of the horrible performance
1344-- of cons. So we just convert it to a list first and then stream from the
1347-- XXX Maybe we can use an Array instead of a list here?
1348{-# INLINE_NORMAL reverse #-}
1349reverse :: Monad m => Stream m a -> Stream m a
1350reverse m = Stream step Nothing
1352 {-# INLINE_LATE step #-}
1355 return $ Skip (Just xs)
1356 step _ (Just (x:xs)) = return $ Yield x (Just xs)
1357 step _ (Just []) = return Stop
1359-- Much faster reverse for Storables
1360{-# INLINE_NORMAL reverse' #-}
1361reverse' :: forall m a. (MonadIO m, Storable a) => Stream m a -> Stream m a
1363-- This commented implementation copies the whole stream into one single array
1364-- and then streams from that array, this is 3-4x faster than the chunked code
1365-- that follows. Though this could be problematic due to unbounded large
1366-- allocations. We need to figure out why the chunked code is slower and if we
1367-- can optimize the chunked code to work as fast as this one. It may be a
1369import Foreign.ForeignPtr (touchForeignPtr)
1370import Foreign.ForeignPtr.Unsafe (unsafeForeignPtrToPtr)
1371import Foreign.Ptr (Ptr, plusPtr)
1372reverse' m = Stream step Nothing
1374 {-# INLINE_LATE step #-}
1376 arr <- A.fromStreamD m
1377 let p = aEnd arr `plusPtr` negate (sizeOf (undefined :: a))
1378 return $ Skip $ Just (aStart arr, p)
1380 step _ (Just (start, p)) | p < unsafeForeignPtrToPtr start = return Stop
1382 step _ (Just (start, p)) = do
1383 let !x = A.unsafeInlineIO $ do
1385 touchForeignPtr start
1387 next = p `plusPtr` negate (sizeOf (undefined :: a))
1388 return $ Yield x (Just (start, next))
1395 $ A.fromStreamDArraysOf A.defaultChunkSize m
1398------------------------------------------------------------------------------
1399-- Grouping/Splitting
1400------------------------------------------------------------------------------
1402{-# INLINE_NORMAL splitSuffixBy' #-}
1403splitSuffixBy' :: Monad m
1404 => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1405splitSuffixBy' predicate f (Stream step state) =
1406 Stream (stepOuter f) (Just state)
1410 {-# INLINE_LATE stepOuter #-}
1411 stepOuter (Fold fstep initial done) gst (Just st) = do
1412 res <- step (adaptState gst) st
1418 then done acc' >>= \val -> return $ Yield val (Just s)
1421 Skip s -> return $ Skip $ Just s
1427 res <- step (adaptState gst) stt
1432 then done acc' >>= \val -> return $ Yield val (Just s)
1434 Skip s -> go SPEC s acc
1435 Stop -> done acc >>= \val -> return $ Yield val Nothing
1437 stepOuter _ _ Nothing = return Stop
1439{-# INLINE_NORMAL groupsBy #-}
1445groupsBy cmp f (Stream step state) = Stream (stepOuter f) (Just state, Nothing)
1449 {-# INLINE_LATE stepOuter #-}
1450 stepOuter (Fold fstep initial done) gst (Just st, Nothing) = do
1451 res <- step (adaptState gst) st
1458 Skip s -> return $ Skip $ (Just s, Nothing)
1463 go !_ prev stt !acc = do
1464 res <- step (adaptState gst) stt
1471 else done acc >>= \r -> return $ Yield r (Just s, Just x)
1472 Skip s -> go SPEC prev s acc
1473 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1475 stepOuter (Fold fstep initial done) gst (Just st, Just prev) = do
1477 acc' <- fstep acc prev
1482 -- XXX code duplicated from the previous equation
1484 res <- step (adaptState gst) stt
1491 else done acc >>= \r -> return $ Yield r (Just s, Just x)
1492 Skip s -> go SPEC s acc
1493 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1495 stepOuter _ _ (Nothing,_) = return Stop
1497{-# INLINE_NORMAL groupsRollingBy #-}
1498groupsRollingBy :: Monad m
1503groupsRollingBy cmp f (Stream step state) =
1504 Stream (stepOuter f) (Just state, Nothing)
1507 {-# INLINE_LATE stepOuter #-}
1508 stepOuter (Fold fstep initial done) gst (Just st, Nothing) = do
1509 res <- step (adaptState gst) st
1516 Skip s -> return $ Skip $ (Just s, Nothing)
1520 go !_ prev stt !acc = do
1521 res <- step (adaptState gst) stt
1529 done acc >>= \r -> return $ Yield r (Just s, Just x)
1530 Skip s -> go SPEC prev s acc
1531 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1533 stepOuter (Fold fstep initial done) gst (Just st, Just prev') = do
1535 acc' <- fstep acc prev'
1536 go SPEC prev' st acc'
1539 go !_ prevv stt !acc = do
1540 res <- step (adaptState gst) stt
1547 else done acc >>= \r -> return $ Yield r (Just s, Just x)
1548 Skip s -> go SPEC prevv s acc
1549 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1551 stepOuter _ _ (Nothing, _) = return Stop
1553{-# INLINE_NORMAL splitBy #-}
1554splitBy :: Monad m => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1555splitBy predicate f (Stream step state) = Stream (step' f) (Just state)
1559 {-# INLINE_LATE step' #-}
1560 step' (Fold fstep initial done) gst (Just st) = initial >>= go SPEC st
1565 res <- step (adaptState gst) stt
1569 then done acc >>= \r -> return $ Yield r (Just s)
1573 Skip s -> go SPEC s acc
1574 Stop -> done acc >>= \r -> return $ Yield r Nothing
1576 step' _ _ Nothing = return Stop
1578-- XXX requires -funfolding-use-threshold=150 in lines-unlines benchmark
1579{-# INLINE_NORMAL splitSuffixBy #-}
1580splitSuffixBy :: Monad m
1581 => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1582splitSuffixBy predicate f (Stream step state) = Stream (step' f) (Just state)
1586 {-# INLINE_LATE step' #-}
1587 step' (Fold fstep initial done) gst (Just st) = do
1588 res <- step (adaptState gst) st
1593 then done acc >>= \val -> return $ Yield val (Just s)
1598 Skip s -> return $ Skip $ Just s
1604 res <- step (adaptState gst) stt
1608 then done acc >>= \r -> return $ Yield r (Just s)
1612 Skip s -> go SPEC s acc
1613 Stop -> done acc >>= \r -> return $ Yield r Nothing
1615 step' _ _ Nothing = return Stop
1617{-# INLINE_NORMAL wordsBy #-}
1618wordsBy :: Monad m => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1619wordsBy predicate f (Stream step state) = Stream (stepOuter f) (Just state)
1623 {-# INLINE_LATE stepOuter #-}
1624 stepOuter (Fold fstep initial done) gst (Just st) = do
1625 res <- step (adaptState gst) st
1629 then return $ Skip (Just s)
1635 Skip s -> return $ Skip $ Just s
1641 res <- step (adaptState gst) stt
1645 then done acc >>= \r -> return $ Yield r (Just s)
1649 Skip s -> go SPEC s acc
1650 Stop -> done acc >>= \r -> return $ Yield r Nothing
1652 stepOuter _ _ Nothing = return Stop
1654-- String search algorithms:
1655-- http://www-igm.univ-mlv.fr/~lecroq/string/index.html
1658-- TODO can we unify the splitting operations using a splitting configuration
1659-- like in the split package.
1661data SplitStyle = Infix | Suffix | Prefix deriving (Eq, Show)
1663data SplitOptions = SplitOptions
1664 { style :: SplitStyle
1665 , withSep :: Bool -- ^ keep the separators in output
1666 -- , compact :: Bool -- ^ treat multiple consecutive separators as one
1667 -- , trimHead :: Bool -- ^ drop blank at head
1668 -- , trimTail :: Bool -- ^ drop blank at tail
1672data SplitOnState s a =
1677 | GO_KARP_RABIN s !(RB.Ring a) !(Ptr a)
1680{-# INLINE_NORMAL splitOn #-}
1682 :: forall m a b. (MonadIO m, Storable a, Enum a, Eq a)
1687splitOn patArr@Array{..} (Fold fstep initial done) (Stream step state) =
1688 Stream stepOuter GO_START
1692 patLen = A.length patArr
1693 maxIndex = patLen - 1
1694 elemBits = sizeOf (undefined :: a) * 8
1696 {-# INLINE_LATE stepOuter #-}
1697 stepOuter _ GO_START =
1699 then return $ Skip $ GO_EMPTY_PAT state
1702 r <- liftIO $ (A.unsafeIndexIO patArr 0)
1703 return $ Skip $ GO_SINGLE_PAT state r
1704 else if sizeOf (undefined :: a) * patLen
1705 <= sizeOf (undefined :: Word)
1706 then return $ Skip $ GO_SHORT_PAT state
1708 (rb, rhead) <- liftIO $ RB.new patLen
1709 return $ Skip $ GO_KARP_RABIN state rb rhead
1711 stepOuter gst (GO_SINGLE_PAT stt pat) = initial >>= go SPEC stt
1716 res <- step (adaptState gst) st
1722 return $ Yield r (GO_SINGLE_PAT s pat)
1723 else fstep acc x >>= go SPEC s
1724 Skip s -> go SPEC s acc
1725 Stop -> done acc >>= \r -> return $ Yield r GO_DONE
1727 stepOuter gst (GO_SHORT_PAT stt) = initial >>= go0 SPEC 0 (0 :: Word) stt
1732 mask = (1 `shiftL` (elemBits * patLen)) - 1
1734 addToWord wrd a = (wrd `shiftL` elemBits) .|. fromIntegral (fromEnum a)
1737 patWord = mask .&. A.foldl' addToWord 0 patArr
1739 go0 !_ !idx wrd st !acc = do
1740 res <- step (adaptState gst) st
1743 let wrd' = addToWord wrd x
1746 if wrd' .&. mask == patWord
1749 return $ Yield r (GO_SHORT_PAT s)
1750 else go1 SPEC wrd' s acc
1751 else go0 SPEC (idx + 1) wrd' s acc
1752 Skip s -> go0 SPEC idx wrd s acc
1755 then go2 wrd idx acc
1757 done acc' >>= \r -> return $ Yield r GO_DONE
1760 go1 !_ wrd st !acc = do
1761 res <- step (adaptState gst) st
1764 let wrd' = addToWord wrd x
1765 old = (mask .&. wrd) `shiftR` (elemBits * (patLen - 1))
1766 acc' <- fstep acc (toEnum $ fromIntegral old)
1767 if wrd' .&. mask == patWord
1768 then done acc' >>= \r -> return $ Yield r (GO_SHORT_PAT s)
1769 else go1 SPEC wrd' s acc'
1770 Skip s -> go1 SPEC wrd s acc
1772 acc' <- go2 wrd patLen acc
1773 done acc' >>= \r -> return $ Yield r GO_DONE
1775 go2 !wrd !n !acc | n > 0 = do
1776 let old = (mask .&. wrd) `shiftR` (elemBits * (n - 1))
1777 fstep acc (toEnum $ fromIntegral old) >>= go2 wrd (n - 1)
1778 go2 _ _ acc = return acc
1780 stepOuter gst (GO_KARP_RABIN stt rb rhead) = do
1781 initial >>= go0 SPEC 0 rhead stt
1785 k = 2891336453 :: Word32
1787 addCksum cksum a = cksum * k + fromIntegral (fromEnum a)
1788 deltaCksum cksum old new =
1789 addCksum cksum new - coeff * fromIntegral (fromEnum old)
1791 -- XXX shall we use a random starting hash or 1 instead of 0?
1792 patHash = A.foldl' addCksum 0 patArr
1795 go0 !_ !idx !rh st !acc = do
1796 res <- step (adaptState gst) st
1799 rh' <- liftIO $ RB.unsafeInsert rb rh x
1802 let fold = RB.unsafeFoldRing (RB.ringBound rb)
1803 let !ringHash = fold addCksum 0 rb
1804 if ringHash == patHash
1805 then go2 SPEC ringHash rh' s acc
1806 else go1 SPEC ringHash rh' s acc
1807 else go0 SPEC (idx + 1) rh' s acc
1808 Skip s -> go0 SPEC idx rh s acc
1810 !acc' <- if idx /= 0
1811 then RB.unsafeFoldRingM rh fstep acc rb
1813 done acc' >>= \r -> return $ Yield r GO_DONE
1815 -- XXX Theoretically this code can do 4 times faster if GHC generates
1816 -- optimal code. If we use just "(cksum' == patHash)" condition it goes
1817 -- 4x faster, as soon as we add the "RB.unsafeEqArray rb v" condition
1818 -- the generated code changes drastically and becomes 4x slower. Need
1819 -- to investigate what is going on with GHC.
1821 go1 !_ !cksum !rh st !acc = do
1822 res <- step (adaptState gst) st
1825 old <- liftIO $ peek rh
1826 let cksum' = deltaCksum cksum old x
1827 acc' <- fstep acc old
1829 if (cksum' == patHash)
1831 rh' <- liftIO (RB.unsafeInsert rb rh x)
1832 go2 SPEC cksum' rh' s acc'
1834 rh' <- liftIO (RB.unsafeInsert rb rh x)
1835 go1 SPEC cksum' rh' s acc'
1836 Skip s -> go1 SPEC cksum rh s acc
1838 acc' <- RB.unsafeFoldRingFullM rh fstep acc rb
1839 done acc' >>= \r -> return $ Yield r GO_DONE
1841 go2 !_ !cksum' !rh' s !acc' = do
1842 if RB.unsafeEqArray rb rh' patArr
1845 return $ Yield r (GO_KARP_RABIN s rb rhead)
1846 else go1 SPEC cksum' rh' s acc'
1848 stepOuter gst (GO_EMPTY_PAT st) = do
1849 res <- step (adaptState gst) st
1854 done acc' >>= \r -> return $ Yield r (GO_EMPTY_PAT s)
1855 Skip s -> return $ Skip (GO_EMPTY_PAT s)
1858 stepOuter _ GO_DONE = return Stop
1860{-# INLINE_NORMAL splitSuffixOn #-}
1862 :: forall m a b. (MonadIO m, Storable a, Enum a, Eq a)
1868splitSuffixOn withSep patArr@Array{..} (Fold fstep initial done)
1869 (Stream step state) =
1870 Stream stepOuter GO_START
1874 patLen = A.length patArr
1875 maxIndex = patLen - 1
1876 elemBits = sizeOf (undefined :: a) * 8
1878 {-# INLINE_LATE stepOuter #-}
1879 stepOuter _ GO_START =
1881 then return $ Skip $ GO_EMPTY_PAT state
1884 r <- liftIO $ (A.unsafeIndexIO patArr 0)
1885 return $ Skip $ GO_SINGLE_PAT state r
1886 else if sizeOf (undefined :: a) * patLen
1887 <= sizeOf (undefined :: Word)
1888 then return $ Skip $ GO_SHORT_PAT state
1890 (rb, rhead) <- liftIO $ RB.new patLen
1891 return $ Skip $ GO_KARP_RABIN state rb rhead
1893 stepOuter gst (GO_SINGLE_PAT stt pat) = do
1894 -- This first part is the only difference between splitOn and
1896 -- If the last element is a separator do not issue a blank segment.
1897 res <- step (adaptState gst) stt
1903 acc' <- if withSep then fstep acc x else return acc
1904 done acc' >>= \r -> return $ Yield r (GO_SINGLE_PAT s pat)
1905 else fstep acc x >>= go SPEC s
1906 Skip s -> return $ Skip $ (GO_SINGLE_PAT s pat)
1911 -- This is identical for splitOn and splitSuffixOn
1913 res <- step (adaptState gst) st
1918 acc' <- if withSep then fstep acc x else return acc
1920 return $ Yield r (GO_SINGLE_PAT s pat)
1921 else fstep acc x >>= go SPEC s
1922 Skip s -> go SPEC s acc
1923 Stop -> done acc >>= \r -> return $ Yield r GO_DONE
1925 stepOuter gst (GO_SHORT_PAT stt) = do
1927 -- Call "initial" only if the stream yields an element, otherwise we
1928 -- may call "initial" but never yield anything. initial may produce a
1929 -- side effect, therefore we will end up doing and discard a side
1934 res <- step (adaptState gst) stt
1938 let wrd' = addToWord wrd x
1939 acc' <- if withSep then fstep acc x else return acc
1942 if wrd' .&. mask == patWord
1943 then done acc' >>= \r -> return $ Yield r (GO_SHORT_PAT s)
1944 else go0 SPEC (idx + 1) wrd' s acc'
1945 else go0 SPEC (idx + 1) wrd' s acc'
1946 Skip s -> return $ Skip (GO_SHORT_PAT s)
1952 mask = (1 `shiftL` (elemBits * patLen)) - 1
1954 addToWord wrd a = (wrd `shiftL` elemBits) .|. fromIntegral (fromEnum a)
1957 patWord = mask .&. A.foldl' addToWord 0 patArr
1959 go0 !_ !idx wrd st !acc = do
1960 res <- step (adaptState gst) st
1963 let wrd' = addToWord wrd x
1964 acc' <- if withSep then fstep acc x else return acc
1967 if wrd' .&. mask == patWord
1970 return $ Yield r (GO_SHORT_PAT s)
1971 else go1 SPEC wrd' s acc'
1972 else go0 SPEC (idx + 1) wrd' s acc'
1973 Skip s -> go0 SPEC idx wrd s acc
1975 if (idx == maxIndex) && (wrd .&. mask == patWord)
1978 acc' <- if idx /= 0 && not withSep
1979 then go2 wrd idx acc
1981 done acc' >>= \r -> return $ Yield r GO_DONE
1984 go1 !_ wrd st !acc = do
1985 res <- step (adaptState gst) st
1988 let wrd' = addToWord wrd x
1989 old = (mask .&. wrd) `shiftR` (elemBits * (patLen - 1))
1992 else fstep acc (toEnum $ fromIntegral old)
1993 if wrd' .&. mask == patWord
1994 then done acc' >>= \r -> return $ Yield r (GO_SHORT_PAT s)
1995 else go1 SPEC wrd' s acc'
1996 Skip s -> go1 SPEC wrd s acc
1998 -- If the last sequence is a separator do not issue a blank
2000 if wrd .&. mask == patWord
2005 else go2 wrd patLen acc
2006 done acc' >>= \r -> return $ Yield r GO_DONE
2008 go2 !wrd !n !acc | n > 0 = do
2009 let old = (mask .&. wrd) `shiftR` (elemBits * (n - 1))
2010 fstep acc (toEnum $ fromIntegral old) >>= go2 wrd (n - 1)
2011 go2 _ _ acc = return acc
2013 stepOuter gst (GO_KARP_RABIN stt rb rhead) = do
2015 res <- step (adaptState gst) stt
2019 acc' <- if withSep then fstep acc x else return acc
2020 rh' <- liftIO (RB.unsafeInsert rb rhead x)
2023 let fold = RB.unsafeFoldRing (RB.ringBound rb)
2024 let !ringHash = fold addCksum 0 rb
2025 if ringHash == patHash
2026 then go2 SPEC ringHash rh' s acc'
2027 else go0 SPEC (idx + 1) rh' s acc'
2028 else go0 SPEC (idx + 1) rh' s acc'
2029 Skip s -> return $ Skip (GO_KARP_RABIN s rb rhead)
2034 k = 2891336453 :: Word32
2036 addCksum cksum a = cksum * k + fromIntegral (fromEnum a)
2037 deltaCksum cksum old new =
2038 addCksum cksum new - coeff * fromIntegral (fromEnum old)
2040 -- XXX shall we use a random starting hash or 1 instead of 0?
2041 patHash = A.foldl' addCksum 0 patArr
2044 go0 !_ !idx !rh st !acc = do
2045 res <- step (adaptState gst) st
2048 acc' <- if withSep then fstep acc x else return acc
2049 rh' <- liftIO (RB.unsafeInsert rb rh x)
2052 let fold = RB.unsafeFoldRing (RB.ringBound rb)
2053 let !ringHash = fold addCksum 0 rb
2054 if ringHash == patHash
2055 then go2 SPEC ringHash rh' s acc'
2056 else go1 SPEC ringHash rh' s acc'
2057 else go0 SPEC (idx + 1) rh' s acc'
2058 Skip s -> go0 SPEC idx rh s acc
2060 -- do not issue a blank segment when we end at pattern
2061 if (idx == maxIndex) && RB.unsafeEqArray rb rh patArr
2064 !acc' <- if idx /= 0 && not withSep
2065 then RB.unsafeFoldRingM rh fstep acc rb
2067 done acc' >>= \r -> return $ Yield r GO_DONE
2069 -- XXX Theoretically this code can do 4 times faster if GHC generates
2070 -- optimal code. If we use just "(cksum' == patHash)" condition it goes
2071 -- 4x faster, as soon as we add the "RB.unsafeEqArray rb v" condition
2072 -- the generated code changes drastically and becomes 4x slower. Need
2073 -- to investigate what is going on with GHC.
2075 go1 !_ !cksum !rh st !acc = do
2076 res <- step (adaptState gst) st
2079 old <- liftIO $ peek rh
2080 let cksum' = deltaCksum cksum old x
2085 if (cksum' == patHash)
2087 rh' <- liftIO (RB.unsafeInsert rb rh x)
2088 go2 SPEC cksum' rh' s acc'
2090 rh' <- liftIO (RB.unsafeInsert rb rh x)
2091 go1 SPEC cksum' rh' s acc'
2092 Skip s -> go1 SPEC cksum rh s acc
2094 if RB.unsafeEqArray rb rh patArr
2099 else RB.unsafeFoldRingFullM rh fstep acc rb
2100 done acc' >>= \r -> return $ Yield r GO_DONE
2102 go2 !_ !cksum' !rh' s !acc' = do
2103 if RB.unsafeEqArray rb rh' patArr
2106 return $ Yield r (GO_KARP_RABIN s rb rhead)
2107 else go1 SPEC cksum' rh' s acc'
2109 stepOuter gst (GO_EMPTY_PAT st) = do
2110 res <- step (adaptState gst) st
2115 done acc' >>= \r -> return $ Yield r (GO_EMPTY_PAT s)
2116 Skip s -> return $ Skip (GO_EMPTY_PAT s)
2119 stepOuter _ GO_DONE = return Stop
2121data SplitState s arr
2123 | SplitBuffering s arr
2124 | SplitSplitting s arr
2125 | SplitYielding arr (SplitState s arr)
2128-- XXX An alternative approach would be to use a partial fold (Fold m a b) to
2129-- split using a splitBy like combinator. The Fold would consume upto the
2130-- separator and return any leftover which can then be fed to the next fold.
2132-- We can revisit this once we have partial folds/parsers.
2134-- | Performs infix separator style splitting.
2135{-# INLINE_NORMAL splitInnerBy #-}
2138 => (f a -> m (f a, Maybe (f a))) -- splitter
2139 -> (f a -> f a -> m (f a)) -- joiner
2142splitInnerBy splitter joiner (Stream step1 state1) =
2143 (Stream step (SplitInitial state1))
2147 {-# INLINE_LATE step #-}
2148 step gst (SplitInitial st) = do
2152 (x1, mx2) <- splitter x
2153 return $ case mx2 of
2154 Nothing -> Skip (SplitBuffering s x1)
2155 Just x2 -> Skip (SplitYielding x1 (SplitSplitting s x2))
2156 Skip s -> return $ Skip (SplitInitial s)
2157 Stop -> return $ Stop
2159 step gst (SplitBuffering st buf) = do
2163 (x1, mx2) <- splitter x
2164 buf' <- joiner buf x1
2165 return $ case mx2 of
2166 Nothing -> Skip (SplitBuffering s buf')
2167 Just x2 -> Skip (SplitYielding buf' (SplitSplitting s x2))
2168 Skip s -> return $ Skip (SplitBuffering s buf)
2169 Stop -> return $ Skip (SplitYielding buf SplitFinishing)
2171 step _ (SplitSplitting st buf) = do
2172 (x1, mx2) <- splitter buf
2173 return $ case mx2 of
2174 Nothing -> Skip $ SplitBuffering st x1
2175 Just x2 -> Skip $ SplitYielding x1 (SplitSplitting st x2)
2177 step _ (SplitYielding x next) = return $ Yield x next
2178 step _ SplitFinishing = return $ Stop
2180-- | Performs infix separator style splitting.
2181{-# INLINE_NORMAL splitInnerBySuffix #-}
2183 :: (Monad m, Eq (f a), Monoid (f a))
2184 => (f a -> m (f a, Maybe (f a))) -- splitter
2185 -> (f a -> f a -> m (f a)) -- joiner
2188splitInnerBySuffix splitter joiner (Stream step1 state1) =
2189 (Stream step (SplitInitial state1))
2193 {-# INLINE_LATE step #-}
2194 step gst (SplitInitial st) = do
2198 (x1, mx2) <- splitter x
2199 return $ case mx2 of
2200 Nothing -> Skip (SplitBuffering s x1)
2201 Just x2 -> Skip (SplitYielding x1 (SplitSplitting s x2))
2202 Skip s -> return $ Skip (SplitInitial s)
2203 Stop -> return $ Stop
2205 step gst (SplitBuffering st buf) = do
2209 (x1, mx2) <- splitter x
2210 buf' <- joiner buf x1
2211 return $ case mx2 of
2212 Nothing -> Skip (SplitBuffering s buf')
2213 Just x2 -> Skip (SplitYielding buf' (SplitSplitting s x2))
2214 Skip s -> return $ Skip (SplitBuffering s buf)
2218 else Skip (SplitYielding buf SplitFinishing)
2220 step _ (SplitSplitting st buf) = do
2221 (x1, mx2) <- splitter buf
2222 return $ case mx2 of
2223 Nothing -> Skip $ SplitBuffering st x1
2224 Just x2 -> Skip $ SplitYielding x1 (SplitSplitting st x2)
2226 step _ (SplitYielding x next) = return $ Yield x next
2227 step _ SplitFinishing = return $ Stop
2229------------------------------------------------------------------------------
2231------------------------------------------------------------------------------
2233{-# INLINE_NORMAL isPrefixOf #-}
2234isPrefixOf :: (Eq a, Monad m) => Stream m a -> Stream m a -> m Bool
2235isPrefixOf (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
2237 go (sa, sb, Nothing) = do
2238 r <- stepa defState sa
2240 Yield x sa' -> go (sa', sb, Just x)
2241 Skip sa' -> go (sa', sb, Nothing)
2244 go (sa, sb, Just x) = do
2245 r <- stepb defState sb
2249 then go (sa, sb', Nothing)
2251 Skip sb' -> go (sa, sb', Just x)
2252 Stop -> return False
2254{-# INLINE_NORMAL isSubsequenceOf #-}
2255isSubsequenceOf :: (Eq a, Monad m) => Stream m a -> Stream m a -> m Bool
2256isSubsequenceOf (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
2258 go (sa, sb, Nothing) = do
2259 r <- stepa defState sa
2261 Yield x sa' -> go (sa', sb, Just x)
2262 Skip sa' -> go (sa', sb, Nothing)
2265 go (sa, sb, Just x) = do
2266 r <- stepb defState sb
2270 then go (sa, sb', Nothing)
2271 else go (sa, sb', Just x)
2272 Skip sb' -> go (sa, sb', Just x)
2273 Stop -> return False
2275{-# INLINE_NORMAL stripPrefix #-}
2278 => Stream m a -> Stream m a -> m (Maybe (Stream m a))
2279stripPrefix (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
2281 go (sa, sb, Nothing) = do
2282 r <- stepa defState sa
2284 Yield x sa' -> go (sa', sb, Just x)
2285 Skip sa' -> go (sa', sb, Nothing)
2286 Stop -> return $ Just (Stream stepb sb)
2288 go (sa, sb, Just x) = do
2289 r <- stepb defState sb
2293 then go (sa, sb', Nothing)
2295 Skip sb' -> go (sa, sb', Just x)
2296 Stop -> return Nothing
2298------------------------------------------------------------------------------
2300------------------------------------------------------------------------------
2302-- | Execute a monadic action for each element of the 'Stream'
2303{-# INLINE_NORMAL mapM_ #-}
2304mapM_ :: Monad m => (a -> m b) -> Stream m a -> m ()
2305mapM_ m = drain . mapM m
2307-------------------------------------------------------------------------------
2308-- Stream transformations using Unfolds
2309-------------------------------------------------------------------------------
2311-- Define a unique structure to use in inspection testing
2312data ConcatMapUState o i =
2314 | ConcatMapUInner o i
2316-- | @concatMapU unfold stream@ uses @unfold@ to map the input stream elements
2317-- to streams and then flattens the generated streams into a single output
2320-- This is like 'concatMap' but uses an unfold with an explicit state to
2321-- generate the stream instead of a 'Stream' type generator. This allows better
2322-- optimization via fusion. This can be many times more efficient than
2325{-# INLINE_NORMAL concatMapU #-}
2326concatMapU :: Monad m => Unfold m a b -> Stream m a -> Stream m b
2327concatMapU (Unfold istep inject) (Stream ostep ost) =
2328 Stream step (ConcatMapUOuter ost)
2330 {-# INLINE_LATE step #-}
2331 step gst (ConcatMapUOuter o) = do
2332 r <- ostep (adaptState gst) o
2336 i `seq` return (Skip (ConcatMapUInner o' i))
2337 Skip o' -> return $ Skip (ConcatMapUOuter o')
2338 Stop -> return $ Stop
2340 step _ (ConcatMapUInner o i) = do
2343 Yield x i' -> Yield x (ConcatMapUInner o i')
2344 Skip i' -> Skip (ConcatMapUInner o i')
2345 Stop -> Skip (ConcatMapUOuter o)
2347data ConcatUnfoldInterleaveState o i =
2348 ConcatUnfoldInterleaveOuter o [i]
2349 | ConcatUnfoldInterleaveInner o [i]
2350 | ConcatUnfoldInterleaveInnerL [i] [i]
2351 | ConcatUnfoldInterleaveInnerR [i] [i]
2353-- XXX use arrays to store state instead of lists.
2354-- XXX In general we can use different scheduling strategies e.g. how to
2355-- schedule the outer vs inner loop or assigning weights to different streams
2356-- or outer and inner loops.
2358-- After a yield, switch to the next stream. Do not switch streams on Skip.
2359-- Yield from outer stream switches to the inner stream.
2361-- There are two choices here, (1) exhaust the outer stream first and then
2362-- start yielding from the inner streams, this is much simpler to implement,
2363-- (2) yield at least one element from an inner stream before going back to
2364-- outer stream and opening the next stream from it.
2366-- Ideally, we need some scheduling bias to inner streams vs outer stream.
2367-- Maybe we can configure the behavior.
2369{-# INLINE_NORMAL concatUnfoldInterleave #-}
2370concatUnfoldInterleave :: Monad m => Unfold m a b -> Stream m a -> Stream m b
2371concatUnfoldInterleave (Unfold istep inject) (Stream ostep ost) =
2372 Stream step (ConcatUnfoldInterleaveOuter ost [])
2374 {-# INLINE_LATE step #-}
2375 step gst (ConcatUnfoldInterleaveOuter o ls) = do
2376 r <- ostep (adaptState gst) o
2380 i `seq` return (Skip (ConcatUnfoldInterleaveInner o' (i : ls)))
2381 Skip o' -> return $ Skip (ConcatUnfoldInterleaveOuter o' ls)
2382 Stop -> return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2384 step _ (ConcatUnfoldInterleaveInner _ []) = undefined
2385 step _ (ConcatUnfoldInterleaveInner o (st:ls)) = do
2388 Yield x s -> Yield x (ConcatUnfoldInterleaveOuter o (s:ls))
2389 Skip s -> Skip (ConcatUnfoldInterleaveInner o (s:ls))
2390 Stop -> Skip (ConcatUnfoldInterleaveOuter o ls)
2392 step _ (ConcatUnfoldInterleaveInnerL [] []) = return Stop
2393 step _ (ConcatUnfoldInterleaveInnerL [] rs) =
2394 return $ Skip (ConcatUnfoldInterleaveInnerR [] rs)
2396 step _ (ConcatUnfoldInterleaveInnerL (st:ls) rs) = do
2399 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerL ls (s:rs))
2400 Skip s -> Skip (ConcatUnfoldInterleaveInnerL (s:ls) rs)
2401 Stop -> Skip (ConcatUnfoldInterleaveInnerL ls rs)
2403 step _ (ConcatUnfoldInterleaveInnerR [] []) = return Stop
2404 step _ (ConcatUnfoldInterleaveInnerR ls []) =
2405 return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2407 step _ (ConcatUnfoldInterleaveInnerR ls (st:rs)) = do
2410 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerR (s:ls) rs)
2411 Skip s -> Skip (ConcatUnfoldInterleaveInnerR ls (s:rs))
2412 Stop -> Skip (ConcatUnfoldInterleaveInnerR ls rs)
2414-- XXX In general we can use different scheduling strategies e.g. how to
2415-- schedule the outer vs inner loop or assigning weights to different streams
2416-- or outer and inner loops.
2418-- This could be inefficient if the tasks are too small.
2420-- Compared to concatUnfoldInterleave this one switches streams on Skips.
2422{-# INLINE_NORMAL concatUnfoldRoundrobin #-}
2423concatUnfoldRoundrobin :: Monad m => Unfold m a b -> Stream m a -> Stream m b
2424concatUnfoldRoundrobin (Unfold istep inject) (Stream ostep ost) =
2425 Stream step (ConcatUnfoldInterleaveOuter ost [])
2427 {-# INLINE_LATE step #-}
2428 step gst (ConcatUnfoldInterleaveOuter o ls) = do
2429 r <- ostep (adaptState gst) o
2433 i `seq` return (Skip (ConcatUnfoldInterleaveInner o' (i : ls)))
2434 Skip o' -> return $ Skip (ConcatUnfoldInterleaveInner o' ls)
2435 Stop -> return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2437 step _ (ConcatUnfoldInterleaveInner o []) =
2438 return $ Skip (ConcatUnfoldInterleaveOuter o [])
2440 step _ (ConcatUnfoldInterleaveInner o (st:ls)) = do
2443 Yield x s -> Yield x (ConcatUnfoldInterleaveOuter o (s:ls))
2444 Skip s -> Skip (ConcatUnfoldInterleaveOuter o (s:ls))
2445 Stop -> Skip (ConcatUnfoldInterleaveOuter o ls)
2447 step _ (ConcatUnfoldInterleaveInnerL [] []) = return Stop
2448 step _ (ConcatUnfoldInterleaveInnerL [] rs) =
2449 return $ Skip (ConcatUnfoldInterleaveInnerR [] rs)
2451 step _ (ConcatUnfoldInterleaveInnerL (st:ls) rs) = do
2454 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerL ls (s:rs))
2455 Skip s -> Skip (ConcatUnfoldInterleaveInnerL ls (s:rs))
2456 Stop -> Skip (ConcatUnfoldInterleaveInnerL ls rs)
2458 step _ (ConcatUnfoldInterleaveInnerR [] []) = return Stop
2459 step _ (ConcatUnfoldInterleaveInnerR ls []) =
2460 return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2462 step _ (ConcatUnfoldInterleaveInnerR ls (st:rs)) = do
2465 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerR (s:ls) rs)
2466 Skip s -> Skip (ConcatUnfoldInterleaveInnerR (s:ls) rs)
2467 Stop -> Skip (ConcatUnfoldInterleaveInnerR ls rs)
2469data AppendState s1 s2 = AppendFirst s1 | AppendSecond s2
2471-- Note that this could be much faster compared to the CPS stream. However, as
2472-- the number of streams being composed increases this may become expensive.
2473-- Need to see where the breaking point is between the two.
2475{-# INLINE_NORMAL append #-}
2476append :: Monad m => Stream m a -> Stream m a -> Stream m a
2477append (Stream step1 state1) (Stream step2 state2) =
2478 Stream step (AppendFirst state1)
2482 {-# INLINE_LATE step #-}
2483 step gst (AppendFirst st) = do
2486 Yield a s -> Yield a (AppendFirst s)
2487 Skip s -> Skip (AppendFirst s)
2488 Stop -> Skip (AppendSecond state2)
2490 step gst (AppendSecond st) = do
2493 Yield a s -> Yield a (AppendSecond s)
2494 Skip s -> Skip (AppendSecond s)
2497data InterleaveState s1 s2 = InterleaveFirst s1 s2 | InterleaveSecond s1 s2
2498 | InterleaveSecondOnly s2 | InterleaveFirstOnly s1
2500{-# INLINE_NORMAL interleave #-}
2501interleave :: Monad m => Stream m a -> Stream m a -> Stream m a
2502interleave (Stream step1 state1) (Stream step2 state2) =
2503 Stream step (InterleaveFirst state1 state2)
2507 {-# INLINE_LATE step #-}
2508 step gst (InterleaveFirst st1 st2) = do
2511 Yield a s -> Yield a (InterleaveSecond s st2)
2512 Skip s -> Skip (InterleaveFirst s st2)
2513 Stop -> Skip (InterleaveSecondOnly st2)
2515 step gst (InterleaveSecond st1 st2) = do
2518 Yield a s -> Yield a (InterleaveFirst st1 s)
2519 Skip s -> Skip (InterleaveSecond st1 s)
2520 Stop -> Skip (InterleaveFirstOnly st1)
2522 step gst (InterleaveFirstOnly st1) = do
2525 Yield a s -> Yield a (InterleaveFirstOnly s)
2526 Skip s -> Skip (InterleaveFirstOnly s)
2529 step gst (InterleaveSecondOnly st2) = do
2532 Yield a s -> Yield a (InterleaveSecondOnly s)
2533 Skip s -> Skip (InterleaveSecondOnly s)
2536{-# INLINE_NORMAL interleaveMin #-}
2537interleaveMin :: Monad m => Stream m a -> Stream m a -> Stream m a
2538interleaveMin (Stream step1 state1) (Stream step2 state2) =
2539 Stream step (InterleaveFirst state1 state2)
2543 {-# INLINE_LATE step #-}
2544 step gst (InterleaveFirst st1 st2) = do
2547 Yield a s -> Yield a (InterleaveSecond s st2)
2548 Skip s -> Skip (InterleaveFirst s st2)
2551 step gst (InterleaveSecond st1 st2) = do
2554 Yield a s -> Yield a (InterleaveFirst st1 s)
2555 Skip s -> Skip (InterleaveSecond st1 s)
2558 step _ (InterleaveFirstOnly _) = undefined
2559 step _ (InterleaveSecondOnly _) = undefined
2561{-# INLINE_NORMAL interleaveSuffix #-}
2562interleaveSuffix :: Monad m => Stream m a -> Stream m a -> Stream m a
2563interleaveSuffix (Stream step1 state1) (Stream step2 state2) =
2564 Stream step (InterleaveFirst state1 state2)
2568 {-# INLINE_LATE step #-}
2569 step gst (InterleaveFirst st1 st2) = do
2572 Yield a s -> Yield a (InterleaveSecond s st2)
2573 Skip s -> Skip (InterleaveFirst s st2)
2576 step gst (InterleaveSecond st1 st2) = do
2579 Yield a s -> Yield a (InterleaveFirst st1 s)
2580 Skip s -> Skip (InterleaveSecond st1 s)
2581 Stop -> Skip (InterleaveFirstOnly st1)
2583 step gst (InterleaveFirstOnly st1) = do
2586 Yield a s -> Yield a (InterleaveFirstOnly s)
2587 Skip s -> Skip (InterleaveFirstOnly s)
2590 step _ (InterleaveSecondOnly _) = undefined
2592data InterleaveInfixState s1 s2 a
2593 = InterleaveInfixFirst s1 s2
2594 | InterleaveInfixSecondBuf s1 s2
2595 | InterleaveInfixSecondYield s1 s2 a
2596 | InterleaveInfixFirstYield s1 s2 a
2597 | InterleaveInfixFirstOnly s1
2599{-# INLINE_NORMAL interleaveInfix #-}
2600interleaveInfix :: Monad m => Stream m a -> Stream m a -> Stream m a
2601interleaveInfix (Stream step1 state1) (Stream step2 state2) =
2602 Stream step (InterleaveInfixFirst state1 state2)
2606 {-# INLINE_LATE step #-}
2607 step gst (InterleaveInfixFirst st1 st2) = do
2610 Yield a s -> Yield a (InterleaveInfixSecondBuf s st2)
2611 Skip s -> Skip (InterleaveInfixFirst s st2)
2614 step gst (InterleaveInfixSecondBuf st1 st2) = do
2617 Yield a s -> Skip (InterleaveInfixSecondYield st1 s a)
2618 Skip s -> Skip (InterleaveInfixSecondBuf st1 s)
2619 Stop -> Skip (InterleaveInfixFirstOnly st1)
2621 step gst (InterleaveInfixSecondYield st1 st2 x) = do
2624 Yield a s -> Yield x (InterleaveInfixFirstYield s st2 a)
2625 Skip s -> Skip (InterleaveInfixSecondYield s st2 x)
2628 step _ (InterleaveInfixFirstYield st1 st2 x) = do
2629 return $ Yield x (InterleaveInfixSecondBuf st1 st2)
2631 step gst (InterleaveInfixFirstOnly st1) = do
2634 Yield a s -> Yield a (InterleaveInfixFirstOnly s)
2635 Skip s -> Skip (InterleaveInfixFirstOnly s)
2638{-# INLINE_NORMAL roundRobin #-}
2639roundRobin :: Monad m => Stream m a -> Stream m a -> Stream m a
2640roundRobin (Stream step1 state1) (Stream step2 state2) =
2641 Stream step (InterleaveFirst state1 state2)
2645 {-# INLINE_LATE step #-}
2646 step gst (InterleaveFirst st1 st2) = do
2649 Yield a s -> Yield a (InterleaveSecond s st2)
2650 Skip s -> Skip (InterleaveSecond s st2)
2651 Stop -> Skip (InterleaveSecondOnly st2)
2653 step gst (InterleaveSecond st1 st2) = do
2656 Yield a s -> Yield a (InterleaveFirst st1 s)
2657 Skip s -> Skip (InterleaveFirst st1 s)
2658 Stop -> Skip (InterleaveFirstOnly st1)
2660 step gst (InterleaveSecondOnly st2) = do
2663 Yield a s -> Yield a (InterleaveSecondOnly s)
2664 Skip s -> Skip (InterleaveSecondOnly s)
2667 step gst (InterleaveFirstOnly st1) = do
2670 Yield a s -> Yield a (InterleaveFirstOnly s)
2671 Skip s -> Skip (InterleaveFirstOnly s)
2674data ICUState s1 s2 i1 i2 =
2679 | ICUFirstInner s1 s2 i1
2680 | ICUSecondInner s1 s2 i2
2681 | ICUFirstOnlyInner s1 i1
2682 | ICUSecondOnlyInner s2 i2
2684-- | Interleave streams (full streams, not the elements) unfolded from two
2685-- input streams and concat. Stop when the first stream stops. If the second
2686-- stream ends before the first one then first stream still keeps running alone
2687-- without any interleaving with the second stream.
2689-- [a1, a2, ... an] [b1, b2 ...]
2690-- => [streamA1, streamA2, ... streamAn] [streamB1, streamB2, ...]
2691-- => [streamA1, streamB1, streamA2...StreamAn, streamBn]
2692-- => [a11, a12, ...a1j, b11, b12, ...b1k, a21, a22, ...]
2694{-# INLINE_NORMAL gintercalateSuffix #-}
2697 => Unfold m a c -> Stream m a -> Unfold m b c -> Stream m b -> Stream m c
2699 (Unfold istep1 inject1) (Stream step1 state1)
2700 (Unfold istep2 inject2) (Stream step2 state2) =
2701 Stream step (ICUFirst state1 state2)
2705 {-# INLINE_LATE step #-}
2706 step gst (ICUFirst s1 s2) = do
2707 r <- step1 (adaptState gst) s1
2711 i `seq` return (Skip (ICUFirstInner s s2 i))
2712 Skip s -> return $ Skip (ICUFirst s s2)
2715 step gst (ICUFirstOnly s1) = do
2716 r <- step1 (adaptState gst) s1
2720 i `seq` return (Skip (ICUFirstOnlyInner s i))
2721 Skip s -> return $ Skip (ICUFirstOnly s)
2724 step _ (ICUFirstInner s1 s2 i1) = do
2727 Yield x i' -> Yield x (ICUFirstInner s1 s2 i')
2728 Skip i' -> Skip (ICUFirstInner s1 s2 i')
2729 Stop -> Skip (ICUSecond s1 s2)
2731 step _ (ICUFirstOnlyInner s1 i1) = do
2734 Yield x i' -> Yield x (ICUFirstOnlyInner s1 i')
2735 Skip i' -> Skip (ICUFirstOnlyInner s1 i')
2736 Stop -> Skip (ICUFirstOnly s1)
2738 step gst (ICUSecond s1 s2) = do
2739 r <- step2 (adaptState gst) s2
2743 i `seq` return (Skip (ICUSecondInner s1 s i))
2744 Skip s -> return $ Skip (ICUSecond s1 s)
2745 Stop -> return $ Skip (ICUFirstOnly s1)
2747 step _ (ICUSecondInner s1 s2 i2) = do
2750 Yield x i' -> Yield x (ICUSecondInner s1 s2 i')
2751 Skip i' -> Skip (ICUSecondInner s1 s2 i')
2752 Stop -> Skip (ICUFirst s1 s2)
2754 step _ (ICUSecondOnly _s2) = undefined
2755 step _ (ICUSecondOnlyInner _s2 _i2) = undefined
2757data InterposeSuffixState s1 i1 =
2758 InterposeSuffixFirst s1
2759 -- | InterposeSuffixFirstYield s1 i1
2760 | InterposeSuffixFirstInner s1 i1
2761 | InterposeSuffixSecond s1
2763-- Note that if an unfolded layer turns out to be nil we still emit the
2764-- separator effect. An alternate behavior could be to emit the separator
2765-- effect only if at least one element has been yielded by the unfolding.
2766-- However, that becomes a bit complicated, so we have chosen the former
2768{-# INLINE_NORMAL interposeSuffix #-}
2771 => m c -> Unfold m b c -> Stream m b -> Stream m c
2774 (Unfold istep1 inject1) (Stream step1 state1) =
2775 Stream step (InterposeSuffixFirst state1)
2779 {-# INLINE_LATE step #-}
2780 step gst (InterposeSuffixFirst s1) = do
2781 r <- step1 (adaptState gst) s1
2785 i `seq` return (Skip (InterposeSuffixFirstInner s i))
2786 -- i `seq` return (Skip (InterposeSuffixFirstYield s i))
2787 Skip s -> return $ Skip (InterposeSuffixFirst s)
2791 step _ (InterposeSuffixFirstYield s1 i1) = do
2794 Yield x i' -> Yield x (InterposeSuffixFirstInner s1 i')
2795 Skip i' -> Skip (InterposeSuffixFirstYield s1 i')
2796 Stop -> Skip (InterposeSuffixFirst s1)
2799 step _ (InterposeSuffixFirstInner s1 i1) = do
2802 Yield x i' -> Yield x (InterposeSuffixFirstInner s1 i')
2803 Skip i' -> Skip (InterposeSuffixFirstInner s1 i')
2804 Stop -> Skip (InterposeSuffixSecond s1)
2806 step _ (InterposeSuffixSecond s1) = do
2808 return $ Yield r (InterposeSuffixFirst s1)
2810data ICALState s1 s2 i1 i2 a =
2812 -- | ICALFirstYield s1 s2 i1
2813 | ICALFirstInner s1 s2 i1
2815 | ICALFirstOnlyInner s1 i1
2816 | ICALSecondInject s1 s2
2817 | ICALFirstInject s1 s2 i2
2818 -- | ICALFirstBuf s1 s2 i1 i2
2819 | ICALSecondInner s1 s2 i1 i2
2820 -- -- | ICALSecondInner s1 s2 i1 i2 a
2821 -- -- | ICALFirstResume s1 s2 i1 i2 a
2823-- | Interleave streams (full streams, not the elements) unfolded from two
2824-- input streams and concat. Stop when the first stream stops. If the second
2825-- stream ends before the first one then first stream still keeps running alone
2826-- without any interleaving with the second stream.
2828-- [a1, a2, ... an] [b1, b2 ...]
2829-- => [streamA1, streamA2, ... streamAn] [streamB1, streamB2, ...]
2830-- => [streamA1, streamB1, streamA2...StreamAn, streamBn]
2831-- => [a11, a12, ...a1j, b11, b12, ...b1k, a21, a22, ...]
2833{-# INLINE_NORMAL gintercalate #-}
2836 => Unfold m a c -> Stream m a -> Unfold m b c -> Stream m b -> Stream m c
2838 (Unfold istep1 inject1) (Stream step1 state1)
2839 (Unfold istep2 inject2) (Stream step2 state2) =
2840 Stream step (ICALFirst state1 state2)
2844 {-# INLINE_LATE step #-}
2845 step gst (ICALFirst s1 s2) = do
2846 r <- step1 (adaptState gst) s1
2850 i `seq` return (Skip (ICALFirstInner s s2 i))
2851 -- i `seq` return (Skip (ICALFirstYield s s2 i))
2852 Skip s -> return $ Skip (ICALFirst s s2)
2856 step _ (ICALFirstYield s1 s2 i1) = do
2859 Yield x i' -> Yield x (ICALFirstInner s1 s2 i')
2860 Skip i' -> Skip (ICALFirstYield s1 s2 i')
2861 Stop -> Skip (ICALFirst s1 s2)
2864 step _ (ICALFirstInner s1 s2 i1) = do
2867 Yield x i' -> Yield x (ICALFirstInner s1 s2 i')
2868 Skip i' -> Skip (ICALFirstInner s1 s2 i')
2869 Stop -> Skip (ICALSecondInject s1 s2)
2871 step gst (ICALFirstOnly s1) = do
2872 r <- step1 (adaptState gst) s1
2876 i `seq` return (Skip (ICALFirstOnlyInner s i))
2877 Skip s -> return $ Skip (ICALFirstOnly s)
2880 step _ (ICALFirstOnlyInner s1 i1) = do
2883 Yield x i' -> Yield x (ICALFirstOnlyInner s1 i')
2884 Skip i' -> Skip (ICALFirstOnlyInner s1 i')
2885 Stop -> Skip (ICALFirstOnly s1)
2887 -- We inject the second stream even before checking if the first stream
2888 -- would yield any more elements. There is no clear choice whether we
2889 -- should do this before or after that. Doing it after may make the state
2890 -- machine a bit simpler though.
2891 step gst (ICALSecondInject s1 s2) = do
2892 r <- step2 (adaptState gst) s2
2896 i `seq` return (Skip (ICALFirstInject s1 s i))
2897 Skip s -> return $ Skip (ICALSecondInject s1 s)
2898 Stop -> return $ Skip (ICALFirstOnly s1)
2900 step gst (ICALFirstInject s1 s2 i2) = do
2901 r <- step1 (adaptState gst) s1
2905 i `seq` return (Skip (ICALSecondInner s s2 i i2))
2906 -- i `seq` return (Skip (ICALFirstBuf s s2 i i2))
2907 Skip s -> return $ Skip (ICALFirstInject s s2 i2)
2911 step _ (ICALFirstBuf s1 s2 i1 i2) = do
2914 Yield x i' -> Skip (ICALSecondInner s1 s2 i' i2 x)
2915 Skip i' -> Skip (ICALFirstBuf s1 s2 i' i2)
2918 step _ (ICALSecondInner s1 s2 i1 i2 v) = do
2921 Yield x i' -> Yield x (ICALSecondInner s1 s2 i1 i' v)
2922 Skip i' -> Skip (ICALSecondInner s1 s2 i1 i' v)
2923 Stop -> Skip (ICALFirstResume s1 s2 i1 i2 v)
2926 step _ (ICALSecondInner s1 s2 i1 i2) = do
2929 Yield x i' -> Yield x (ICALSecondInner s1 s2 i1 i')
2930 Skip i' -> Skip (ICALSecondInner s1 s2 i1 i')
2931 Stop -> Skip (ICALFirstInner s1 s2 i1)
2932 -- Stop -> Skip (ICALFirstResume s1 s2 i1 i2)
2935 step _ (ICALFirstResume s1 s2 i1 i2 x) = do
2936 return $ Yield x (ICALFirstInner s1 s2 i1 i2)
2939data InterposeState s1 i1 a =
2941 -- | InterposeFirstYield s1 i1
2942 | InterposeFirstInner s1 i1
2943 | InterposeFirstInject s1
2944 -- | InterposeFirstBuf s1 i1
2945 | InterposeSecondYield s1 i1
2946 -- -- | InterposeSecondYield s1 i1 a
2947 -- -- | InterposeFirstResume s1 i1 a
2949-- Note that this only interposes the pure values, we may run many effects to
2950-- generate those values as some effects may not generate anything (Skip).
2951{-# INLINE_NORMAL interpose #-}
2952interpose :: Monad m => m c -> Unfold m b c -> Stream m b -> Stream m c
2955 (Unfold istep1 inject1) (Stream step1 state1) =
2956 Stream step (InterposeFirst state1)
2960 {-# INLINE_LATE step #-}
2961 step gst (InterposeFirst s1) = do
2962 r <- step1 (adaptState gst) s1
2966 i `seq` return (Skip (InterposeFirstInner s i))
2967 -- i `seq` return (Skip (InterposeFirstYield s i))
2968 Skip s -> return $ Skip (InterposeFirst s)
2972 step _ (InterposeFirstYield s1 i1) = do
2975 Yield x i' -> Yield x (InterposeFirstInner s1 i')
2976 Skip i' -> Skip (InterposeFirstYield s1 i')
2977 Stop -> Skip (InterposeFirst s1)
2980 step _ (InterposeFirstInner s1 i1) = do
2983 Yield x i' -> Yield x (InterposeFirstInner s1 i')
2984 Skip i' -> Skip (InterposeFirstInner s1 i')
2985 Stop -> Skip (InterposeFirstInject s1)
2987 step gst (InterposeFirstInject s1) = do
2988 r <- step1 (adaptState gst) s1
2992 -- i `seq` return (Skip (InterposeFirstBuf s i))
2993 i `seq` return (Skip (InterposeSecondYield s i))
2994 Skip s -> return $ Skip (InterposeFirstInject s)
2998 step _ (InterposeFirstBuf s1 i1) = do
3001 Yield x i' -> Skip (InterposeSecondYield s1 i' x)
3002 Skip i' -> Skip (InterposeFirstBuf s1 i')
3007 step _ (InterposeSecondYield s1 i1 v) = do
3009 return $ Yield r (InterposeFirstResume s1 i1 v)
3011 step _ (InterposeSecondYield s1 i1) = do
3013 return $ Yield r (InterposeFirstInner s1 i1)
3016 step _ (InterposeFirstResume s1 i1 v) = do
3017 return $ Yield v (InterposeFirstInner s1 i1)
3020------------------------------------------------------------------------------
3022------------------------------------------------------------------------------
3024data GbracketState s1 s2 v
3026 | GBracketNormal s1 v
3027 | GBracketException s2
3029-- | The most general bracketing and exception combinator. All other
3030-- combinators can be expressed in terms of this combinator. This can also be
3031-- used for cases which are not covered by the standard combinators.
3035{-# INLINE_NORMAL gbracket #-}
3039 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
3040 -> (c -> m d) -- ^ after, on normal stop
3041 -> (c -> e -> Stream m b) -- ^ on exception
3042 -> (c -> Stream m b) -- ^ stream generator
3044gbracket bef exc aft fexc fnormal =
3045 Stream step GBracketInit
3049 {-# INLINE_LATE step #-}
3050 step _ GBracketInit = do
3052 return $ Skip $ GBracketNormal (fnormal r) r
3054 step gst (GBracketNormal (UnStream step1 st) v) = do
3055 res <- exc $ step1 gst st
3057 Right r -> case r of
3059 return $ Yield x (GBracketNormal (Stream step1 s) v)
3060 Skip s -> return $ Skip (GBracketNormal (Stream step1 s) v)
3061 Stop -> aft v >> return Stop
3062 Left e -> return $ Skip (GBracketException (fexc v e))
3063 step gst (GBracketException (UnStream step1 st)) = do
3066 Yield x s -> return $ Yield x (GBracketException (Stream step1 s))
3067 Skip s -> return $ Skip (GBracketException (Stream step1 s))
3070-- | Create an IORef holding a finalizer that is called automatically when the
3071-- IORef is garbage collected. The IORef can be written to with a 'Nothing'
3072-- value to deactivate the finalizer.
3073newFinalizedIORef :: (MonadIO m, MonadBaseControl IO m)
3074 => m a -> m (IORef (Maybe (IO ())))
3075newFinalizedIORef finalizer = do
3076 mrun <- captureMonadState
3077 ref <- liftIO $ newIORef $ Just $ liftIO $ void $ do
3078 _ <- runInIO mrun finalizer
3081 res <- readIORef ref
3083 Nothing -> return ()
3085 _ <- liftIO $ mkWeakIORef ref finalizer1
3088-- | Run the finalizer stored in an IORef and deactivate it so that it is run
3091runIORefFinalizer :: MonadIO m => IORef (Maybe (IO ())) -> m ()
3092runIORefFinalizer ref = liftIO $ do
3093 res <- readIORef ref
3095 Nothing -> return ()
3096 Just f -> writeIORef ref Nothing >> f
3098-- | Deactivate the finalizer stored in an IORef without running it.
3100clearIORefFinalizer :: MonadIO m => IORef (Maybe (IO ())) -> m ()
3101clearIORefFinalizer ref = liftIO $ writeIORef ref Nothing
3103data GbracketIOState s1 s2 v wref
3105 | GBracketIONormal s1 v wref
3106 | GBracketIOException s2
3108-- | Like gbracket but also uses a finalizer to make sure when the stream is
3109-- garbage collected we run the finalizing action. This requires a MonadIO and
3110-- MonadBaseControl IO constraint.
3112-- | The most general bracketing and exception combinator. All other
3113-- combinators can be expressed in terms of this combinator. This can also be
3114-- used for cases which are not covered by the standard combinators.
3118{-# INLINE_NORMAL gbracketIO #-}
3120 :: (MonadIO m, MonadBaseControl IO m)
3122 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
3123 -> (c -> m d) -- ^ after, on normal stop or GC
3124 -> (c -> e -> Stream m b) -- ^ on exception
3125 -> (c -> Stream m b) -- ^ stream generator
3127gbracketIO bef exc aft fexc fnormal =
3128 Stream step GBracketIOInit
3132 -- If the stream is never evaluated the "aft" action will never be
3133 -- called. For that to occur we will need the user of this API to pass a
3134 -- weak pointer to us.
3135 {-# INLINE_LATE step #-}
3136 step _ GBracketIOInit = do
3137 -- We mask asynchronous exceptions to make the execution
3138 -- of 'bef' and the registration of 'aft' atomic.
3139 -- A similar thing is done in the resourcet package: https://git.io/JvKV3
3140 -- Tutorial: https://markkarpov.com/tutorial/exceptions.html
3141 (r, ref) <- liftBaseOp_ mask_ $ do
3143 ref <- newFinalizedIORef (aft r)
3145 return $ Skip $ GBracketIONormal (fnormal r) r ref
3147 step gst (GBracketIONormal (UnStream step1 st) v ref) = do
3148 res <- exc $ step1 gst st
3150 Right r -> case r of
3152 return $ Yield x (GBracketIONormal (Stream step1 s) v ref)
3154 return $ Skip (GBracketIONormal (Stream step1 s) v ref)
3156 runIORefFinalizer ref
3159 clearIORefFinalizer ref
3160 return $ Skip (GBracketIOException (fexc v e))
3161 step gst (GBracketIOException (UnStream step1 st)) = do
3165 return $ Yield x (GBracketIOException (Stream step1 s))
3166 Skip s -> return $ Skip (GBracketIOException (Stream step1 s))
3169-- | Run a side effect before the stream yields its first element.
3170{-# INLINE_NORMAL before #-}
3171before :: Monad m => m b -> Stream m a -> Stream m a
3172before action (Stream step state) = Stream step' Nothing
3176 {-# INLINE_LATE step' #-}
3177 step' _ Nothing = action >> return (Skip (Just state))
3179 step' gst (Just st) = do
3182 Yield x s -> return $ Yield x (Just s)
3183 Skip s -> return $ Skip (Just s)
3186-- | Run a side effect whenever the stream stops normally.
3187{-# INLINE_NORMAL after #-}
3188after :: Monad m => m b -> Stream m a -> Stream m a
3189after action (Stream step state) = Stream step' state
3193 {-# INLINE_LATE step' #-}
3197 Yield x s -> return $ Yield x s
3198 Skip s -> return $ Skip s
3199 Stop -> action >> return Stop
3201{-# INLINE_NORMAL afterIO #-}
3202afterIO :: (MonadIO m, MonadBaseControl IO m)
3203 => m b -> Stream m a -> Stream m a
3204afterIO action (Stream step state) = Stream step' Nothing
3208 {-# INLINE_LATE step' #-}
3209 step' _ Nothing = do
3210 ref <- newFinalizedIORef action
3211 return $ Skip $ Just (state, ref)
3212 step' gst (Just (st, ref)) = do
3215 Yield x s -> return $ Yield x (Just (s, ref))
3216 Skip s -> return $ Skip (Just (s, ref))
3218 runIORefFinalizer ref
3221-- XXX These combinators are expensive due to the call to
3222-- onException/handle/try on each step. Therefore, when possible, they should
3223-- be called in an outer loop where we perform less iterations. For example, we
3224-- cannot call them on each iteration in a char stream, instead we can call
3225-- them when doing an IO on an array.
3227-- XXX For high performance error checks in busy streams we may need another
3228-- Error constructor in step.
3230-- | Run a side effect whenever the stream aborts due to an exception. The
3231-- exception is not caught, simply rethrown.
3232{-# INLINE_NORMAL onException #-}
3233onException :: MonadCatch m => m b -> Stream m a -> Stream m a
3234onException action str =
3235 gbracket (return ()) MC.try return
3236 (\_ (e :: MC.SomeException) -> nilM (action >> MC.throwM e))
3239{-# INLINE_NORMAL _onException #-}
3240_onException :: MonadCatch m => m b -> Stream m a -> Stream m a
3241_onException action (Stream step state) = Stream step' state
3245 {-# INLINE_LATE step' #-}
3247 res <- step gst st `MC.onException` action
3249 Yield x s -> return $ Yield x s
3250 Skip s -> return $ Skip s
3253-- XXX bracket is like concatMap, it generates a stream and then flattens it.
3254-- Like concatMap it has 10x worse performance compared to linear fused
3257-- | Run the first action before the stream starts and remember its output,
3258-- generate a stream using the output, run the second action providing the
3259-- remembered value as an argument whenever the stream ends normally or due to
3261{-# INLINE_NORMAL bracket #-}
3262bracket :: MonadCatch m => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
3263bracket bef aft bet =
3264 gbracket bef MC.try aft
3265 (\a (e :: SomeException) -> nilM (aft a >> MC.throwM e)) bet
3267{-# INLINE_NORMAL bracketIO #-}
3268bracketIO :: (MonadAsync m, MonadCatch m)
3269 => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
3270bracketIO bef aft bet =
3271 gbracketIO bef MC.try aft
3272 (\a (e :: SomeException) -> nilM (aft a >> MC.throwM e)) bet
3274data BracketState s v = BracketInit | BracketRun s v
3276{-# INLINE_NORMAL _bracket #-}
3277_bracket :: MonadCatch m => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
3278_bracket bef aft bet = Stream step' BracketInit
3282 {-# INLINE_LATE step' #-}
3283 step' _ BracketInit = bef >>= \x -> return (Skip (BracketRun (bet x) x))
3285 -- NOTE: It is important to use UnStream instead of the Stream pattern
3286 -- here, otherwise we get huge perf degradation, see note in concatMap.
3287 step' gst (BracketRun (UnStream step state) v) = do
3288 -- res <- step gst state `MC.onException` aft v
3289 res <- MC.try $ step gst state
3291 Left (e :: SomeException) -> aft v >> MC.throwM e >> return Stop
3292 Right r -> case r of
3293 Yield x s -> return $ Yield x (BracketRun (Stream step s) v)
3294 Skip s -> return $ Skip (BracketRun (Stream step s) v)
3295 Stop -> aft v >> return Stop
3297-- | Run a side effect whenever the stream stops normally or aborts due to an
3299{-# INLINE finally #-}
3300finally :: MonadCatch m => m b -> Stream m a -> Stream m a
3301-- finally action xs = after action $ onException action xs
3302finally action xs = bracket (return ()) (\_ -> action) (const xs)
3304{-# INLINE finallyIO #-}
3305finallyIO :: (MonadAsync m, MonadCatch m) => m b -> Stream m a -> Stream m a
3306finallyIO action xs = bracketIO (return ()) (\_ -> action) (const xs)
3308-- | When evaluating a stream if an exception occurs, stream evaluation aborts
3309-- and the specified exception handler is run with the exception as argument.
3310{-# INLINE_NORMAL handle #-}
3311handle :: (MonadCatch m, Exception e)
3312 => (e -> Stream m a) -> Stream m a -> Stream m a
3314 gbracket (return ()) MC.try return (\_ e -> f e) (\_ -> str)
3316{-# INLINE_NORMAL _handle #-}
3317_handle :: (MonadCatch m, Exception e)
3318 => (e -> Stream m a) -> Stream m a -> Stream m a
3319_handle f (Stream step state) = Stream step' (Left state)
3323 {-# INLINE_LATE step' #-}
3324 step' gst (Left st) = do
3325 res <- MC.try $ step gst st
3327 Left e -> return $ Skip $ Right (f e)
3328 Right r -> case r of
3329 Yield x s -> return $ Yield x (Left s)
3330 Skip s -> return $ Skip (Left s)
3333 step' gst (Right (UnStream step1 st)) = do
3336 Yield x s -> return $ Yield x (Right (Stream step1 s))
3337 Skip s -> return $ Skip (Right (Stream step1 s))
3340-------------------------------------------------------------------------------
3341-- General transformation
3342-------------------------------------------------------------------------------
3344{-# INLINE_NORMAL transform #-}
3345transform :: Monad m => Pipe m a b -> Stream m a -> Stream m b
3346transform (Pipe pstep1 pstep2 pstate) (Stream step state) =
3347 Stream step' (Consume pstate, state)
3351 {-# INLINE_LATE step' #-}
3353 step' gst (Consume pst, st) = pst `seq` do
3354 r <- step (adaptState gst) st
3359 Pipe.Yield b pst' -> return $ Yield b (pst', s)
3360 Pipe.Continue pst' -> return $ Skip (pst', s)
3361 Skip s -> return $ Skip (Consume pst, s)
3364 step' _ (Produce pst, st) = pst `seq` do
3367 Pipe.Yield b pst' -> return $ Yield b (pst', st)
3368 Pipe.Continue pst' -> return $ Skip (pst', st)
3370------------------------------------------------------------------------------
3371-- Transformation by Folding (Scans)
3372------------------------------------------------------------------------------
3374------------------------------------------------------------------------------
3376------------------------------------------------------------------------------
3378-- XXX Is a prescan useful, discarding the last step does not sound useful? I
3379-- am not sure about the utility of this function, so this is implemented but
3380-- not exposed. We can expose it if someone provides good reasons why this is
3383-- XXX We have to execute the stream one step ahead to know that we are at the
3384-- last step. The vector implementation of prescan executes the last fold step
3385-- but does not yield the result. This means we have executed the effect but
3386-- discarded value. This does not sound right. In this implementation we are
3387-- not executing the last fold step.
3388{-# INLINE_NORMAL prescanlM' #-}
3389prescanlM' :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
3390prescanlM' f mz (Stream step state) = Stream step' (state, mz)
3392 {-# INLINE_LATE step' #-}
3393 step' gst (st, prev) = do
3394 r <- step (adaptState gst) st
3398 return $ Yield acc (s, f acc x)
3399 Skip s -> return $ Skip (s, prev)
3402{-# INLINE prescanl' #-}
3403prescanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
3404prescanl' f z = prescanlM' (\a b -> return (f a b)) (return z)
3406------------------------------------------------------------------------------
3407-- Monolithic postscans (postscan followed by a map)
3408------------------------------------------------------------------------------
3410-- The performance of a modular postscan followed by a map seems to be
3411-- equivalent to this monolithic scan followed by map therefore we may not need
3412-- this implementation. We just have it for performance comparison and in case
3413-- modular version does not perform well in some situation.
3415{-# INLINE_NORMAL postscanlMx' #-}
3416postscanlMx' :: Monad m
3417 => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> Stream m b
3418postscanlMx' fstep begin done (Stream step state) = do
3419 Stream step' (state, begin)
3421 {-# INLINE_LATE step' #-}
3422 step' gst (st, acc) = do
3423 r <- step (adaptState gst) st
3429 v `seq` y `seq` return (Yield v (s, return y))
3430 Skip s -> return $ Skip (s, acc)
3433{-# INLINE_NORMAL postscanlx' #-}
3434postscanlx' :: Monad m
3435 => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> Stream m b
3436postscanlx' fstep begin done s =
3437 postscanlMx' (\b a -> return (fstep b a)) (return begin) (return . done) s
3439-- XXX do we need consM strict to evaluate the begin value?
3440{-# INLINE scanlMx' #-}
3442 => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> Stream m b
3443scanlMx' fstep begin done s =
3444 (begin >>= \x -> x `seq` done x) `consM` postscanlMx' fstep begin done s
3446{-# INLINE scanlx' #-}
3448 => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> Stream m b
3449scanlx' fstep begin done s =
3450 scanlMx' (\b a -> return (fstep b a)) (return begin) (return . done) s
3452------------------------------------------------------------------------------
3454------------------------------------------------------------------------------
3456{-# INLINE_NORMAL postscanlM' #-}
3457postscanlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
3458postscanlM' fstep begin (Stream step state) =
3459 begin `seq` Stream step' (state, begin)
3461 {-# INLINE_LATE step' #-}
3462 step' gst (st, acc) = acc `seq` do
3463 r <- step (adaptState gst) st
3467 y `seq` return (Yield y (s, y))
3468 Skip s -> return $ Skip (s, acc)
3471{-# INLINE_NORMAL postscanl' #-}
3472postscanl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
3473postscanl' f = postscanlM' (\a b -> return (f a b))
3475{-# INLINE_NORMAL postscanlM #-}
3476postscanlM :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
3477postscanlM fstep begin (Stream step state) = Stream step' (state, begin)
3479 {-# INLINE_LATE step' #-}
3480 step' gst (st, acc) = do
3481 r <- step (adaptState gst) st
3485 return (Yield y (s, y))
3486 Skip s -> return $ Skip (s, acc)
3489{-# INLINE_NORMAL postscanl #-}
3490postscanl :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
3491postscanl f = postscanlM (\a b -> return (f a b))
3493{-# INLINE_NORMAL scanlM' #-}
3494scanlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
3495scanlM' fstep begin s = begin `seq` (begin `cons` postscanlM' fstep begin s)
3497{-# INLINE scanl' #-}
3498scanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
3499scanl' f = scanlM' (\a b -> return (f a b))
3501{-# INLINE_NORMAL scanlM #-}
3502scanlM :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
3503scanlM fstep begin s = begin `cons` postscanlM fstep begin s
3506scanl :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
3507scanl f = scanlM (\a b -> return (f a b))
3509{-# INLINE_NORMAL scanl1M #-}
3510scanl1M :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
3511scanl1M fstep (Stream step state) = Stream step' (state, Nothing)
3513 {-# INLINE_LATE step' #-}
3514 step' gst (st, Nothing) = do
3517 Yield x s -> return $ Yield x (s, Just x)
3518 Skip s -> return $ Skip (s, Nothing)
3521 step' gst (st, Just acc) = do
3526 return $ Yield z (s, Just z)
3527 Skip s -> return $ Skip (s, Just acc)
3530{-# INLINE scanl1 #-}
3531scanl1 :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
3532scanl1 f = scanl1M (\x y -> return (f x y))
3534{-# INLINE_NORMAL scanl1M' #-}
3535scanl1M' :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
3536scanl1M' fstep (Stream step state) = Stream step' (state, Nothing)
3538 {-# INLINE_LATE step' #-}
3539 step' gst (st, Nothing) = do
3542 Yield x s -> x `seq` return $ Yield x (s, Just x)
3543 Skip s -> return $ Skip (s, Nothing)
3546 step' gst (st, Just acc) = acc `seq` do
3551 z `seq` return $ Yield z (s, Just z)
3552 Skip s -> return $ Skip (s, Just acc)
3555{-# INLINE scanl1' #-}
3556scanl1' :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
3557scanl1' f = scanl1M' (\x y -> return (f x y))
3559------------------------------------------------------------------------------
3561------------------------------------------------------------------------------
3563data RollingMapState s a = RollingMapInit s | RollingMapGo s a
3565{-# INLINE rollingMapM #-}
3566rollingMapM :: Monad m => (a -> a -> m b) -> Stream m a -> Stream m b
3567rollingMapM f (Stream step1 state1) = Stream step (RollingMapInit state1)
3569 step gst (RollingMapInit st) = do
3570 r <- step1 (adaptState gst) st
3572 Yield x s -> Skip $ RollingMapGo s x
3573 Skip s -> Skip $ RollingMapInit s
3576 step gst (RollingMapGo s1 x1) = do
3577 r <- step1 (adaptState gst) s1
3581 return $ Yield res $ RollingMapGo s x
3582 Skip s -> return $ Skip $ RollingMapGo s x1
3583 Stop -> return $ Stop
3585{-# INLINE rollingMap #-}
3586rollingMap :: Monad m => (a -> a -> b) -> Stream m a -> Stream m b
3587rollingMap f = rollingMapM (\x y -> return $ f x y)
3589------------------------------------------------------------------------------
3590-- Tapping/Distributing
3591------------------------------------------------------------------------------
3594tap :: Monad m => Fold m a b -> Stream m a -> Stream m a
3595tap (Fold fstep initial extract) (Stream step state) = Stream step' Nothing
3599 step' _ Nothing = do
3601 return $ Skip (Just (r, state))
3603 step' gst (Just (acc, st)) = acc `seq` do
3608 return $ Yield x (Just (acc', s))
3609 Skip s -> return $ Skip (Just (acc, s))
3614{-# INLINE_NORMAL tapOffsetEvery #-}
3615tapOffsetEvery :: Monad m
3616 => Int -> Int -> Fold m a b -> Stream m a -> Stream m a
3617tapOffsetEvery offset n (Fold fstep initial extract) (Stream step state) =
3618 Stream step' Nothing
3622 {-# INLINE_LATE step' #-}
3623 step' _ Nothing = do
3625 return $ Skip (Just (r, state, offset `mod` n))
3627 step' gst (Just (acc, st, count)) | count <= 0 = do
3631 !acc' <- fstep acc x
3632 return $ Yield x (Just (acc', s, n - 1))
3633 Skip s -> return $ Skip (Just (acc, s, count))
3638 step' gst (Just (acc, st, count)) = do
3641 Yield x s -> return $ Yield x (Just (acc, s, count - 1))
3642 Skip s -> return $ Skip (Just (acc, s, count))
3647{-# INLINE_NORMAL pollCounts #-}
3651 -> (Stream m Int -> Stream m Int)
3655pollCounts predicate transf fld (Stream step state) = Stream step' Nothing
3658 {-# INLINE_LATE step' #-}
3659 step' _ Nothing = do
3660 -- As long as we are using an "Int" for counts lockfree reads from
3661 -- Var should work correctly on both 32-bit and 64-bit machines.
3662 -- However, an Int on a 32-bit machine may overflow quickly.
3663 countVar <- liftIO $ newVar (0 :: Int)
3665 $ void $ runFold fld
3666 $ transf $ fromPrimVar countVar
3667 return $ Skip (Just (countVar, tid, state))
3669 step' gst (Just (countVar, tid, st)) = do
3673 when (predicate x) $ liftIO $ modifyVar' countVar (+ 1)
3674 return $ Yield x (Just (countVar, tid, s))
3675 Skip s -> return $ Skip (Just (countVar, tid, s))
3677 liftIO $ killThread tid
3680{-# INLINE_NORMAL tapRate #-}
3682 (MonadAsync m, MonadCatch m)
3687tapRate samplingRate action (Stream step state) = Stream step' Nothing
3689 {-# NOINLINE loop #-}
3690 loop countVar prev = do
3693 (do liftIO $ threadDelay (round $ samplingRate * 1000000)
3694 i <- liftIO $ readVar countVar
3695 let !diff = i - prev
3698 (\(e :: AsyncException) -> do
3699 i <- liftIO $ readVar countVar
3700 let !diff = i - prev
3702 throwM (MC.toException e))
3705 {-# INLINE_LATE step' #-}
3706 step' _ Nothing = do
3707 countVar <- liftIO $ newVar 0
3708 tid <- fork $ loop countVar 0
3709 ref <- liftIO $ newIORef ()
3710 _ <- liftIO $ mkWeakIORef ref (killThread tid)
3711 return $ Skip (Just (countVar, tid, state, ref))
3713 step' gst (Just (countVar, tid, st, ref)) = do
3717 liftIO $ modifyVar' countVar (+ 1)
3718 return $ Yield x (Just (countVar, tid, s, ref))
3719 Skip s -> return $ Skip (Just (countVar, tid, s, ref))
3721 liftIO $ killThread tid
3725-------------------------------------------------------------------------------
3727-------------------------------------------------------------------------------
3729{-# INLINE_NORMAL takeWhileM #-}
3730takeWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
3731takeWhileM f (Stream step state) = Stream step' state
3733 {-# INLINE_LATE step' #-}
3739 return $ if b then Yield x s else Stop
3740 Skip s -> return $ Skip s
3743{-# INLINE takeWhile #-}
3744takeWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
3745takeWhile f = takeWhileM (return . f)
3747{-# INLINE_NORMAL drop #-}
3748drop :: Monad m => Int -> Stream m a -> Stream m a
3749drop n (Stream step state) = Stream step' (state, Just n)
3751 {-# INLINE_LATE step' #-}
3752 step' gst (st, Just i)
3757 Yield _ s -> Skip (s, Just (i - 1))
3758 Skip s -> Skip (s, Just i)
3760 | otherwise = return $ Skip (st, Nothing)
3762 step' gst (st, Nothing) = do
3766 Yield x s -> Yield x (s, Nothing)
3767 Skip s -> Skip (s, Nothing)
3770data DropWhileState s a
3772 | DropWhileYield a s
3775{-# INLINE_NORMAL dropWhileM #-}
3776dropWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
3777dropWhileM f (Stream step state) = Stream step' (DropWhileDrop state)
3779 {-# INLINE_LATE step' #-}
3780 step' gst (DropWhileDrop st) = do
3786 then return $ Skip (DropWhileDrop s)
3787 else return $ Skip (DropWhileYield x s)
3788 Skip s -> return $ Skip (DropWhileDrop s)
3791 step' gst (DropWhileNext st) = do
3794 Yield x s -> return $ Skip (DropWhileYield x s)
3795 Skip s -> return $ Skip (DropWhileNext s)
3798 step' _ (DropWhileYield x st) = return $ Yield x (DropWhileNext st)
3800{-# INLINE dropWhile #-}
3801dropWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
3802dropWhile f = dropWhileM (return . f)
3804{-# INLINE_NORMAL filterM #-}
3805filterM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
3806filterM f (Stream step state) = Stream step' state
3808 {-# INLINE_LATE step' #-}
3817 Skip s -> return $ Skip s
3820{-# INLINE filter #-}
3821filter :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
3822filter f = filterM (return . f)
3824{-# INLINE_NORMAL uniq #-}
3825uniq :: (Eq a, Monad m) => Stream m a -> Stream m a
3826uniq (Stream step state) = Stream step' (Nothing, state)
3828 {-# INLINE_LATE step' #-}
3829 step' gst (Nothing, st) = do
3832 Yield x s -> return $ Yield x (Just x, s)
3833 Skip s -> return $ Skip (Nothing, s)
3835 step' gst (Just x, st) = do
3838 Yield y s | x == y -> return $ Skip (Just x, s)
3839 | otherwise -> return $ Yield y (Just y, s)
3840 Skip s -> return $ Skip (Just x, s)
3843------------------------------------------------------------------------------
3844-- Transformation by Mapping
3845------------------------------------------------------------------------------
3847{-# INLINE_NORMAL sequence #-}
3848sequence :: Monad m => Stream m (m a) -> Stream m a
3849sequence (Stream step state) = Stream step' state
3851 {-# INLINE_LATE step' #-}
3853 r <- step (adaptState gst) st
3855 Yield x s -> x >>= \a -> return (Yield a s)
3856 Skip s -> return $ Skip s
3859------------------------------------------------------------------------------
3861------------------------------------------------------------------------------
3863data LoopState x s = FirstYield s
3864 | InterspersingYield s
3867{-# INLINE_NORMAL intersperseM #-}
3868intersperseM :: Monad m => m a -> Stream m a -> Stream m a
3869intersperseM m (Stream step state) = Stream step' (FirstYield state)
3871 {-# INLINE_LATE step' #-}
3872 step' gst (FirstYield st) = do
3876 Yield x s -> Skip (YieldAndCarry x s)
3877 Skip s -> Skip (FirstYield s)
3880 step' gst (InterspersingYield st) = do
3885 return $ Yield a (YieldAndCarry x s)
3886 Skip s -> return $ Skip $ InterspersingYield s
3889 step' _ (YieldAndCarry x st) = return $ Yield x (InterspersingYield st)
3894 | SuffixYield a (SuffixState s a)
3896{-# INLINE_NORMAL intersperseSuffix #-}
3897intersperseSuffix :: forall m a. Monad m => m a -> Stream m a -> Stream m a
3898intersperseSuffix action (Stream step state) = Stream step' (SuffixElem state)
3900 {-# INLINE_LATE step' #-}
3901 step' gst (SuffixElem st) = do
3904 Yield x s -> Skip (SuffixYield x (SuffixSuffix s))
3905 Skip s -> Skip (SuffixElem s)
3908 step' _ (SuffixSuffix st) = do
3909 action >>= \r -> return $ Skip (SuffixYield r (SuffixElem st))
3911 step' _ (SuffixYield x next) = return $ Yield x next
3913data SuffixSpanState s a
3914 = SuffixSpanElem s Int
3915 | SuffixSpanSuffix s
3916 | SuffixSpanYield a (SuffixSpanState s a)
3920-- | intersperse after every n items
3921{-# INLINE_NORMAL intersperseSuffixBySpan #-}
3922intersperseSuffixBySpan :: forall m a. Monad m
3923 => Int -> m a -> Stream m a -> Stream m a
3924intersperseSuffixBySpan n action (Stream step state) =
3925 Stream step' (SuffixSpanElem state n)
3927 {-# INLINE_LATE step' #-}
3928 step' gst (SuffixSpanElem st i) | i > 0 = do
3931 Yield x s -> Skip (SuffixSpanYield x (SuffixSpanElem s (i - 1)))
3932 Skip s -> Skip (SuffixSpanElem s i)
3933 Stop -> if i == n then Stop else Skip SuffixSpanLast
3934 step' _ (SuffixSpanElem st _) = return $ Skip (SuffixSpanSuffix st)
3936 step' _ (SuffixSpanSuffix st) = do
3937 action >>= \r -> return $ Skip (SuffixSpanYield r (SuffixSpanElem st n))
3939 step' _ (SuffixSpanLast) = do
3940 action >>= \r -> return $ Skip (SuffixSpanYield r SuffixSpanStop)
3942 step' _ (SuffixSpanYield x next) = return $ Yield x next
3944 step' _ (SuffixSpanStop) = return Stop
3946{-# INLINE intersperse #-}
3947intersperse :: Monad m => a -> Stream m a -> Stream m a
3948intersperse a = intersperseM (return a)
3950{-# INLINE_NORMAL insertBy #-}
3951insertBy :: Monad m => (a -> a -> Ordering) -> a -> Stream m a -> Stream m a
3952insertBy cmp a (Stream step state) = Stream step' (state, False, Nothing)
3954 {-# INLINE_LATE step' #-}
3955 step' gst (st, False, _) = do
3958 Yield x s -> case cmp a x of
3959 GT -> return $ Yield x (s, False, Nothing)
3960 _ -> return $ Yield a (s, True, Just x)
3961 Skip s -> return $ Skip (s, False, Nothing)
3962 Stop -> return $ Yield a (st, True, Nothing)
3964 step' _ (_, True, Nothing) = return Stop
3966 step' gst (st, True, Just prev) = do
3969 Yield x s -> return $ Yield prev (s, True, Just x)
3970 Skip s -> return $ Skip (s, True, Just prev)
3971 Stop -> return $ Yield prev (st, True, Nothing)
3973------------------------------------------------------------------------------
3975------------------------------------------------------------------------------
3977{-# INLINE_NORMAL deleteBy #-}
3978deleteBy :: Monad m => (a -> a -> Bool) -> a -> Stream m a -> Stream m a
3979deleteBy eq x (Stream step state) = Stream step' (state, False)
3981 {-# INLINE_LATE step' #-}
3982 step' gst (st, False) = do
3985 Yield y s -> return $
3986 if eq x y then Skip (s, True) else Yield y (s, False)
3987 Skip s -> return $ Skip (s, False)
3990 step' gst (st, True) = do
3993 Yield y s -> return $ Yield y (s, True)
3994 Skip s -> return $ Skip (s, True)
3997------------------------------------------------------------------------------
3998-- Transformation by Map and Filter
3999------------------------------------------------------------------------------
4001-- XXX Will this always fuse properly?
4002{-# INLINE_NORMAL mapMaybe #-}
4003mapMaybe :: Monad m => (a -> Maybe b) -> Stream m a -> Stream m b
4004mapMaybe f = fmap fromJust . filter isJust . map f
4006{-# INLINE_NORMAL mapMaybeM #-}
4007mapMaybeM :: Monad m => (a -> m (Maybe b)) -> Stream m a -> Stream m b
4008mapMaybeM f = fmap fromJust . filter isJust . mapM f
4010------------------------------------------------------------------------------
4012------------------------------------------------------------------------------
4014{-# INLINE_NORMAL indexed #-}
4015indexed :: Monad m => Stream m a -> Stream m (Int, a)
4016indexed (Stream step state) = Stream step' (state, 0)
4018 {-# INLINE_LATE step' #-}
4019 step' gst (st, i) = i `seq` do
4020 r <- step (adaptState gst) st
4022 Yield x s -> return $ Yield (i, x) (s, i+1)
4023 Skip s -> return $ Skip (s, i)
4026{-# INLINE_NORMAL indexedR #-}
4027indexedR :: Monad m => Int -> Stream m a -> Stream m (Int, a)
4028indexedR m (Stream step state) = Stream step' (state, m)
4030 {-# INLINE_LATE step' #-}
4031 step' gst (st, i) = i `seq` do
4032 r <- step (adaptState gst) st
4034 Yield x s -> let i' = i - 1
4035 in return $ Yield (i, x) (s, i')
4036 Skip s -> return $ Skip (s, i)
4039{-# INLINE_NORMAL zipWithM #-}
4041 => (a -> b -> m c) -> Stream m a -> Stream m b -> Stream m c
4042zipWithM f (Stream stepa ta) (Stream stepb tb) = Stream step (ta, tb, Nothing)
4044 {-# INLINE_LATE step #-}
4045 step gst (sa, sb, Nothing) = do
4046 r <- stepa (adaptState gst) sa
4049 Yield x sa' -> Skip (sa', sb, Just x)
4050 Skip sa' -> Skip (sa', sb, Nothing)
4053 step gst (sa, sb, Just x) = do
4054 r <- stepb (adaptState gst) sb
4058 return $ Yield z (sa, sb', Nothing)
4059 Skip sb' -> return $ Skip (sa, sb', Just x)
4062#if __GLASGOW_HASKELL__ >= 801
4063{-# RULES "zipWithM xs xs"
4064 forall f xs. zipWithM @Identity f xs xs = mapM (\x -> f x x) xs #-}
4067{-# INLINE zipWith #-}
4068zipWith :: Monad m => (a -> b -> c) -> Stream m a -> Stream m b -> Stream m c
4069zipWith f = zipWithM (\a b -> return (f a b))
4071------------------------------------------------------------------------------
4073------------------------------------------------------------------------------
4075{-# INLINE_NORMAL mergeByM #-}
4078 => (a -> a -> m Ordering) -> Stream m a -> Stream m a -> Stream m a
4079mergeByM cmp (Stream stepa ta) (Stream stepb tb) =
4080 Stream step (Just ta, Just tb, Nothing, Nothing)
4082 {-# INLINE_LATE step #-}
4084 -- one of the values is missing, and the corresponding stream is running
4085 step gst (Just sa, sb, Nothing, b) = do
4088 Yield a sa' -> Skip (Just sa', sb, Just a, b)
4089 Skip sa' -> Skip (Just sa', sb, Nothing, b)
4090 Stop -> Skip (Nothing, sb, Nothing, b)
4092 step gst (sa, Just sb, a, Nothing) = do
4095 Yield b sb' -> Skip (sa, Just sb', a, Just b)
4096 Skip sb' -> Skip (sa, Just sb', a, Nothing)
4097 Stop -> Skip (sa, Nothing, a, Nothing)
4099 -- both the values are available
4100 step _ (sa, sb, Just a, Just b) = do
4102 return $ case res of
4103 GT -> Yield b (sa, sb, Just a, Nothing)
4104 _ -> Yield a (sa, sb, Nothing, Just b)
4106 -- one of the values is missing, corresponding stream is done
4107 step _ (Nothing, sb, Nothing, Just b) =
4108 return $ Yield b (Nothing, sb, Nothing, Nothing)
4110 step _ (sa, Nothing, Just a, Nothing) =
4111 return $ Yield a (sa, Nothing, Nothing, Nothing)
4113 step _ (Nothing, Nothing, Nothing, Nothing) = return Stop
4115{-# INLINE mergeBy #-}
4118 => (a -> a -> Ordering) -> Stream m a -> Stream m a -> Stream m a
4119mergeBy cmp = mergeByM (\a b -> return $ cmp a b)
4121------------------------------------------------------------------------------
4122-- Transformation comprehensions
4123------------------------------------------------------------------------------
4125{-# INLINE_NORMAL the #-}
4126the :: (Eq a, Monad m) => Stream m a -> m (Maybe a)
4127the (Stream step state) = go state
4130 r <- step defState st
4132 Yield x s -> go' x s
4134 Stop -> return Nothing
4136 r <- step defState st
4138 Yield x s | x == n -> go' n s
4139 | otherwise -> return Nothing
4141 Stop -> return (Just n)
4143{-# INLINE runFold #-}
4144runFold :: (Monad m) => Fold m a b -> Stream m a -> m b
4145runFold (Fold step begin done) = foldlMx' step begin done
4147-------------------------------------------------------------------------------
4148-- Concurrent application and fold
4149-------------------------------------------------------------------------------
4151-- XXX These functions should be moved to Stream/Parallel.hs
4153-- Using StreamD the worker stream producing code can fuse with the code to
4154-- queue output to the SVar giving some perf boost.
4156-- Note that StreamD can only be used in limited situations, specifically, we
4157-- cannot implement joinStreamVarPar using this.
4159-- XXX make sure that the SVar passed is a Parallel style SVar.
4161-- | Fold the supplied stream to the SVar asynchronously using Parallel
4162-- concurrency style.
4163-- {-# INLINE_NORMAL toSVarParallel #-}
4164{-# INLINE toSVarParallel #-}
4165toSVarParallel :: MonadAsync m
4166 => State t m a -> SVar t m a -> Stream m a -> m ()
4167toSVarParallel st sv xs =
4168 if svarInspectMode sv
4172 case getYieldLimit st of
4173 Nothing -> doFork (work Nothing)
4175 (handleChildException sv)
4176 Just _ -> doFork (workLim Nothing)
4178 (handleChildException sv)
4183 {-# NOINLINE work #-}
4184 work info = (runFold (FL.toParallelSVar sv info) xs)
4186 {-# NOINLINE workLim #-}
4187 workLim info = runFold (FL.toParallelSVarLimited sv info) xs
4189 {-# NOINLINE forkWithDiag #-}
4191 -- We do not use workerCount in case of ParallelVar but still there is
4192 -- no harm in maintaining it correctly.
4193 liftIO $ atomicModifyIORefCAS_ (workerCount sv) $ \n -> n + 1
4195 -- This allocation matters when significant number of workers are being
4196 -- sent. We allocate it only when needed. The overhead increases by 4x.
4198 case yieldRateInfo sv of
4199 Nothing -> return Nothing
4200 Just _ -> liftIO $ do
4201 cntRef <- newIORef 0
4202 t <- getTime Monotonic
4203 lat <- newIORef (0, t)
4204 return $ Just WorkerInfo
4205 { workerYieldMax = 0
4206 , workerYieldCount = cntRef
4207 , workerLatencyStart = lat
4210 case getYieldLimit st of
4211 Nothing -> doFork (work winfo)
4213 (handleChildException sv)
4214 Just _ -> doFork (workLim winfo)
4216 (handleChildException sv)
4219{-# INLINE_NORMAL mkParallelD #-}
4220mkParallelD :: MonadAsync m => Stream m a -> Stream m a
4221mkParallelD m = Stream step Nothing
4224 step gst Nothing = do
4225 sv <- newParallelVar StopNone gst
4226 toSVarParallel gst sv m
4227 -- XXX use unfold instead?
4228 return $ Skip $ Just $ fromSVar sv
4230 step gst (Just (UnStream step1 st)) = do
4233 Yield a s -> Yield a (Just $ Stream step1 s)
4234 Skip s -> Skip (Just $ Stream step1 s)
4237-- Compare with mkAsync. mkAsync uses an Async style SVar whereas this uses a
4238-- parallel style SVar for evaluation. Currently, parallel style cannot use
4239-- rate control whereas Async style can use rate control. In async style SVar
4240-- the worker thread terminates when the buffer is full whereas in Parallel
4243-- | Make the stream producer and consumer run concurrently by introducing a
4244-- buffer between them. The producer thread evaluates the input stream until
4245-- the buffer fills, it blocks if the buffer is full until there is space in
4246-- the buffer. The consumer consumes the stream lazily from the buffer.
4250{-# INLINE_NORMAL mkParallel #-}
4251mkParallel :: (K.IsStream t, MonadAsync m) => t m a -> t m a
4252mkParallel = fromStreamD . mkParallelD . toStreamD
4254-- Note: we can use another API with two callbacks stop and yield if we want
4255-- the callback to be able to indicate end of stream.
4257-- | Generates a callback and a stream pair. The callback returned is used to
4258-- queue values to the stream. The stream is infinite, there is no way for the
4259-- callback to indicate that it is done now.
4263{-# INLINE_NORMAL newCallbackStream #-}
4264newCallbackStream :: (K.IsStream t, MonadAsync m) => m ((a -> m ()), t m a)
4265newCallbackStream = do
4266 sv <- newParallelVar StopNone defState
4268 -- XXX Add our own thread-id to the SVar as we can not know the callback's
4269 -- thread-id and the callback is not run in a managed worker. We need to
4270 -- handle this better.
4271 liftIO myThreadId >>= modifyThread sv
4273 let callback a = liftIO $ void $ send sv (ChildYield a)
4274 -- XXX we can return an SVar and then the consumer can unfold from the
4276 return (callback, fromStreamD (fromSVar sv))
4278-------------------------------------------------------------------------------
4280-------------------------------------------------------------------------------
4282-- | Create an SVar with a fold consumer that will fold any elements sent to it
4283-- using the supplied fold function.
4284{-# INLINE newFoldSVar #-}
4285newFoldSVar :: MonadAsync m => State t m a -> Fold m a b -> m (SVar t m a)
4286newFoldSVar stt f = do
4287 -- Buffer size for the SVar is derived from the current state
4288 sv <- newParallelVar StopAny (adaptState stt)
4289 -- Add the producer thread-id to the SVar.
4290 liftIO myThreadId >>= modifyThread sv
4291 void $ doFork (work sv) (svarMrun sv) (handleFoldException sv)
4296 {-# NOINLINE work #-}
4297 work sv = void $ runFold f $ fromProducer sv
4299data TapState sv st = TapInit | Tapping sv st | TapDone st
4301{-# INLINE_NORMAL tapAsync #-}
4302tapAsync :: MonadAsync m => Fold m a b -> Stream m a -> Stream m a
4303tapAsync f (Stream step1 state1) = Stream step TapInit
4307 -- In general, a Stop event would come equipped with the result
4308 -- of the fold. It is not used here but it would be useful in
4309 -- applicative and distribute.
4310 done <- fromConsumer svr
4311 when (not done) $ do
4312 liftIO $ withDiagMVar svr "teeToSVar: waiting to drain"
4313 $ takeMVar (outputDoorBellFromConsumer svr)
4317 liftIO $ sendStop svr Nothing
4318 -- drain/wait until a stop event arrives from the fold.
4321 {-# INLINE_LATE step #-}
4322 step gst TapInit = do
4323 sv <- newFoldSVar gst f
4324 return $ Skip (Tapping sv state1)
4326 step gst (Tapping sv st) = do
4330 done <- pushToFold sv a
4333 -- XXX we do not need to wait synchronously here
4335 return $ Yield a (TapDone s)
4336 else return $ Yield a (Tapping sv s)
4337 Skip s -> return $ Skip (Tapping sv s)
4342 step gst (TapDone st) = do
4345 Yield a s -> Yield a (TapDone s)
4346 Skip s -> Skip (TapDone s)
4349-- XXX Exported from Array again as this fold is specific to Array
4350-- | Take last 'n' elements from the stream and discard the rest.
4352lastN :: (Storable a, MonadIO m) => Int -> Fold m a (Array a)
4354 | n <= 0 = fmap (const mempty) FL.drain
4355 | otherwise = Fold step initial done
4357 step (Tuple3' rb rh i) a = do
4358 rh1 <- liftIO $ RB.unsafeInsert rb rh a
4359 return $ Tuple3' rb rh1 (i + 1)
4360 initial = fmap (\(a, b) -> Tuple3' a b (0 :: Int)) $ liftIO $ RB.new n
4361 done (Tuple3' rb rh i) = do
4362 arr <- liftIO $ A.newArray n
4363 foldFunc i rh snoc' arr rb
4364 snoc' b a = liftIO $ A.unsafeSnoc b a
4366 | i < n = RB.unsafeFoldRingM
4367 | otherwise = RB.unsafeFoldRingFullM
4369------------------------------------------------------------------------------
4371------------------------------------------------------------------------------
4373-- XXX using getTime in the loop can be pretty expensive especially for
4374-- computations where iterations are lightweight. We have the following
4377-- 1) Run a timeout thread updating a flag asynchronously and check that
4378-- flag here, that way we can have a cheap termination check.
4380-- 2) Use COARSE clock to get time with lower resolution but more efficiently.
4382-- 3) Use rdtscp/rdtsc to get time directly from the processor, compute the
4383-- termination value of rdtsc in the beginning and then in each iteration just
4384-- get rdtsc and check if we should terminate.
4388 | TakeByTimeCheck st s
4389 | TakeByTimeYield st s
4391{-# INLINE_NORMAL takeByTime #-}
4392takeByTime :: (MonadIO m, TimeUnit64 t) => t -> Stream m a -> Stream m a
4393takeByTime duration (Stream step1 state1) = Stream step (TakeByTimeInit state1)
4396 lim = toRelTime64 duration
4398 {-# INLINE_LATE step #-}
4399 step _ (TakeByTimeInit _) | lim == 0 = return Stop
4400 step _ (TakeByTimeInit st) = do
4401 t0 <- liftIO $ getTime Monotonic
4402 return $ Skip (TakeByTimeYield st t0)
4403 step _ (TakeByTimeCheck st t0) = do
4404 t <- liftIO $ getTime Monotonic
4406 if diffAbsTime64 t t0 > lim
4408 else Skip (TakeByTimeYield st t0)
4409 step gst (TakeByTimeYield st t0) = do
4412 Yield x s -> Yield x (TakeByTimeCheck s t0)
4413 Skip s -> Skip (TakeByTimeCheck s t0)
4416data DropByTime st s x
4418 | DropByTimeGen st s
4419 | DropByTimeCheck st s x
4420 | DropByTimeYield st
4422{-# INLINE_NORMAL dropByTime #-}
4423dropByTime :: (MonadIO m, TimeUnit64 t) => t -> Stream m a -> Stream m a
4424dropByTime duration (Stream step1 state1) = Stream step (DropByTimeInit state1)
4427 lim = toRelTime64 duration
4429 {-# INLINE_LATE step #-}
4430 step _ (DropByTimeInit st) = do
4431 t0 <- liftIO $ getTime Monotonic
4432 return $ Skip (DropByTimeGen st t0)
4433 step gst (DropByTimeGen st t0) = do
4436 Yield x s -> Skip (DropByTimeCheck s t0 x)
4437 Skip s -> Skip (DropByTimeGen s t0)
4439 step _ (DropByTimeCheck st t0 x) = do
4440 t <- liftIO $ getTime Monotonic
4441 if diffAbsTime64 t t0 <= lim
4442 then return $ Skip $ DropByTimeGen st t0
4443 else return $ Yield x $ DropByTimeYield st
4444 step gst (DropByTimeYield st) = do
4447 Yield x s -> Yield x (DropByTimeYield s)
4448 Skip s -> Skip (DropByTimeYield s)
4451-- XXX we should move this to stream generation section of this file. Also, the
4452-- take/drop combinators above should be moved to filtering section.
4453{-# INLINE_NORMAL currentTime #-}
4454currentTime :: MonadAsync m => Double -> Stream m AbsTime
4455currentTime g = Stream step Nothing
4459 g' = g * 10 ^ (6 :: Int)
4461 -- XXX should have a minimum granularity to avoid high CPU usage?
4462 {-# INLINE delayTime #-}
4464 if g' >= fromIntegral (maxBound :: Int)
4468 updateTimeVar timeVar = do
4469 threadDelay $ delayTime
4470 MicroSecond64 t <- fromAbsTime <$> getTime Monotonic
4471 modifyVar' timeVar (const t)
4473 {-# INLINE_LATE step #-}
4475 -- XXX note that this is safe only on a 64-bit machine. On a 32-bit
4476 -- machine a 64-bit 'Var' cannot be read consistently without a lock
4477 -- while another thread is writing to it.
4478 timeVar <- liftIO $ newVar (0 :: Int64)
4479 tid <- forkManaged $ liftIO $ forever (updateTimeVar timeVar)
4480 return $ Skip $ Just (timeVar, tid)
4482 step _ s@(Just (timeVar, _)) = do
4483 a <- liftIO $ readVar timeVar
4484 -- XXX we can perhaps use an AbsTime64 using a 64 bit Int for
4485 -- efficiency. or maybe we can use a representation using Double for
4486 -- floating precision time
4487 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