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 splitParse 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 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 _ b -> return b
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 _ b -> return b
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 splitParse #-}
1043splitParse (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?
1081 assert (n <= length (x:buf)) (return ())
1082 let src = Prelude.reverse (Prelude.take n (x:buf))
1084 ParseChunksYield b (ParseChunksInit src s)
1085 Skip s -> return $ Skip $ ParseChunksStream s buf pst
1088 let src = Prelude.reverse buf
1090 ParseChunksYield b (ParseChunksInitLeftOver src)
1092 -- go back to stream processing mode
1093 stepOuter _ (ParseChunksBuf [] s buf pst) =
1094 return $ Skip $ ParseChunksStream s buf pst
1096 -- buffered processing loop
1097 stepOuter _ (ParseChunksBuf (x:xs) s buf pst) = do
1100 -- PR.Yield 0 pst1 ->
1101 PR.Yield n pst1 -> do
1102 assert (n <= length (x:buf)) (return ())
1103 let buf1 = Prelude.take n (x:buf)
1104 return $ Skip $ ParseChunksBuf xs s buf1 pst1
1105 -- PR.Skip 0 pst1 -> return $ Skip $ ParseChunksBuf xs s (x:buf) pst1
1106 PR.Skip n pst1 -> do
1107 assert (n <= length (x:buf)) (return ())
1108 let (src0, buf1) = splitAt n (x:buf)
1109 src = Prelude.reverse src0 ++ xs
1110 return $ Skip $ ParseChunksBuf src s buf1 pst1
1113 assert (n <= length (x:buf)) (return ())
1114 let src = Prelude.reverse (Prelude.take n (x:buf)) ++ xs
1115 return $ Skip $ ParseChunksYield b (ParseChunksInit src s)
1117 stepOuter _ (ParseChunksYield a next) = return $ Yield a next
1119------------------------------------------------------------------------------
1121------------------------------------------------------------------------------
1123-- | Run a streaming composition, discard the results.
1124{-# INLINE_LATE drain #-}
1125drain :: Monad m => Stream m a -> m ()
1126-- drain = foldrM (\_ xs -> xs) (return ())
1127drain (Stream step state) = go SPEC state
1130 r <- step defState st
1132 Yield _ s -> go SPEC s
1136{-# INLINE_NORMAL null #-}
1137null :: Monad m => Stream m a -> m Bool
1138null m = foldrM (\_ _ -> return False) (return True) m
1140{-# INLINE_NORMAL head #-}
1141head :: Monad m => Stream m a -> m (Maybe a)
1142head m = foldrM (\x _ -> return (Just x)) (return Nothing) m
1144{-# INLINE_NORMAL headElse #-}
1145headElse :: Monad m => a -> Stream m a -> m a
1146headElse a m = foldrM (\x _ -> return x) (return a) m
1148-- Does not fuse, has the same performance as the StreamK version.
1149{-# INLINE_NORMAL tail #-}
1150tail :: Monad m => Stream m a -> m (Maybe (Stream m a))
1151tail (UnStream step state) = go state
1154 r <- step defState st
1156 Yield _ s -> return (Just $ Stream step s)
1158 Stop -> return Nothing
1160-- XXX will it fuse? need custom impl?
1161{-# INLINE_NORMAL last #-}
1162last :: Monad m => Stream m a -> m (Maybe a)
1163last = foldl' (\_ y -> Just y) Nothing
1165{-# INLINE_NORMAL elem #-}
1166elem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
1167-- elem e m = foldrM (\x xs -> if x == e then return True else xs) (return False) m
1168elem e (Stream step state) = go state
1171 r <- step defState st
1174 | x == e -> return True
1177 Stop -> return False
1179{-# INLINE_NORMAL notElem #-}
1180notElem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
1181notElem e s = fmap not (elem e s)
1183{-# INLINE_NORMAL all #-}
1184all :: Monad m => (a -> Bool) -> Stream m a -> m Bool
1185-- all p m = foldrM (\x xs -> if p x then xs else return False) (return True) m
1186all p (Stream step state) = go state
1189 r <- step defState st
1193 | otherwise -> return False
1197{-# INLINE_NORMAL any #-}
1198any :: Monad m => (a -> Bool) -> Stream m a -> m Bool
1199-- any p m = foldrM (\x xs -> if p x then return True else xs) (return False) m
1200any p (Stream step state) = go state
1203 r <- step defState st
1206 | p x -> return True
1209 Stop -> return False
1211{-# INLINE_NORMAL maximum #-}
1212maximum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
1213maximum (Stream step state) = go Nothing state
1216 r <- step defState st
1218 Yield x s -> go (Just x) s
1219 Skip s -> go Nothing s
1220 Stop -> return Nothing
1221 go (Just acc) st = do
1222 r <- step defState st
1225 | acc <= x -> go (Just x) s
1226 | otherwise -> go (Just acc) s
1227 Skip s -> go (Just acc) s
1228 Stop -> return (Just acc)
1230{-# INLINE_NORMAL maximumBy #-}
1231maximumBy :: Monad m => (a -> a -> Ordering) -> Stream m a -> m (Maybe a)
1232maximumBy cmp (Stream step state) = go Nothing state
1235 r <- step defState st
1237 Yield x s -> go (Just x) s
1238 Skip s -> go Nothing s
1239 Stop -> return Nothing
1240 go (Just acc) st = do
1241 r <- step defState st
1243 Yield x s -> case cmp acc x of
1244 GT -> go (Just acc) s
1246 Skip s -> go (Just acc) s
1247 Stop -> return (Just acc)
1249{-# INLINE_NORMAL minimum #-}
1250minimum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
1251minimum (Stream step state) = go Nothing state
1254 r <- step defState st
1256 Yield x s -> go (Just x) s
1257 Skip s -> go Nothing s
1258 Stop -> return Nothing
1259 go (Just acc) st = do
1260 r <- step defState st
1263 | acc <= x -> go (Just acc) s
1264 | otherwise -> go (Just x) s
1265 Skip s -> go (Just acc) s
1266 Stop -> return (Just acc)
1268{-# INLINE_NORMAL minimumBy #-}
1269minimumBy :: Monad m => (a -> a -> Ordering) -> Stream m a -> m (Maybe a)
1270minimumBy cmp (Stream step state) = go Nothing state
1273 r <- step defState st
1275 Yield x s -> go (Just x) s
1276 Skip s -> go Nothing s
1277 Stop -> return Nothing
1278 go (Just acc) st = do
1279 r <- step defState st
1281 Yield x s -> case cmp acc x of
1283 _ -> go (Just acc) s
1284 Skip s -> go (Just acc) s
1285 Stop -> return (Just acc)
1287{-# INLINE_NORMAL (!!) #-}
1288(!!) :: (Monad m) => Stream m a -> Int -> m (Maybe a)
1289(Stream step state) !! i = go i state
1292 r <- step defState st
1294 Yield x s | n < 0 -> return Nothing
1295 | n == 0 -> return $ Just x
1296 | otherwise -> go (n - 1) s
1298 Stop -> return Nothing
1300{-# INLINE_NORMAL lookup #-}
1301lookup :: (Monad m, Eq a) => a -> Stream m (a, b) -> m (Maybe b)
1302lookup e m = foldrM (\(a, b) xs -> if e == a then return (Just b) else xs)
1305{-# INLINE_NORMAL findM #-}
1306findM :: Monad m => (a -> m Bool) -> Stream m a -> m (Maybe a)
1307findM p m = foldrM (\x xs -> p x >>= \r -> if r then return (Just x) else xs)
1311find :: Monad m => (a -> Bool) -> Stream m a -> m (Maybe a)
1312find p = findM (return . p)
1314{-# INLINE_NORMAL findIndices #-}
1315findIndices :: Monad m => (a -> Bool) -> Stream m a -> Stream m Int
1316findIndices p (Stream step state) = Stream step' (state, 0)
1318 {-# INLINE_LATE step' #-}
1319 step' gst (st, i) = i `seq` do
1320 r <- step (adaptState gst) st
1322 Yield x s -> if p x then Yield i (s, i+1) else Skip (s, i+1)
1323 Skip s -> Skip (s, i)
1326{-# INLINE toListRev #-}
1327toListRev :: Monad m => Stream m a -> m [a]
1328toListRev = foldl' (flip (:)) []
1330-- We can implement reverse as:
1332-- > reverse = foldlS (flip cons) nil
1334-- However, this implementation is unusable because of the horrible performance
1335-- of cons. So we just convert it to a list first and then stream from the
1338-- XXX Maybe we can use an Array instead of a list here?
1339{-# INLINE_NORMAL reverse #-}
1340reverse :: Monad m => Stream m a -> Stream m a
1341reverse m = Stream step Nothing
1343 {-# INLINE_LATE step #-}
1346 return $ Skip (Just xs)
1347 step _ (Just (x:xs)) = return $ Yield x (Just xs)
1348 step _ (Just []) = return Stop
1350-- Much faster reverse for Storables
1351{-# INLINE_NORMAL reverse' #-}
1352reverse' :: forall m a. (MonadIO m, Storable a) => Stream m a -> Stream m a
1354-- This commented implementation copies the whole stream into one single array
1355-- and then streams from that array, this is 3-4x faster than the chunked code
1356-- that follows. Though this could be problematic due to unbounded large
1357-- allocations. We need to figure out why the chunked code is slower and if we
1358-- can optimize the chunked code to work as fast as this one. It may be a
1360import Foreign.ForeignPtr (touchForeignPtr)
1361import Foreign.ForeignPtr.Unsafe (unsafeForeignPtrToPtr)
1362import Foreign.Ptr (Ptr, plusPtr)
1363reverse' m = Stream step Nothing
1365 {-# INLINE_LATE step #-}
1367 arr <- A.fromStreamD m
1368 let p = aEnd arr `plusPtr` negate (sizeOf (undefined :: a))
1369 return $ Skip $ Just (aStart arr, p)
1371 step _ (Just (start, p)) | p < unsafeForeignPtrToPtr start = return Stop
1373 step _ (Just (start, p)) = do
1374 let !x = A.unsafeInlineIO $ do
1376 touchForeignPtr start
1378 next = p `plusPtr` negate (sizeOf (undefined :: a))
1379 return $ Yield x (Just (start, next))
1386 $ A.fromStreamDArraysOf A.defaultChunkSize m
1389------------------------------------------------------------------------------
1390-- Grouping/Splitting
1391------------------------------------------------------------------------------
1393{-# INLINE_NORMAL splitSuffixBy' #-}
1394splitSuffixBy' :: Monad m
1395 => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1396splitSuffixBy' predicate f (Stream step state) =
1397 Stream (stepOuter f) (Just state)
1401 {-# INLINE_LATE stepOuter #-}
1402 stepOuter (Fold fstep initial done) gst (Just st) = do
1403 res <- step (adaptState gst) st
1409 then done acc' >>= \val -> return $ Yield val (Just s)
1412 Skip s -> return $ Skip $ Just s
1418 res <- step (adaptState gst) stt
1423 then done acc' >>= \val -> return $ Yield val (Just s)
1425 Skip s -> go SPEC s acc
1426 Stop -> done acc >>= \val -> return $ Yield val Nothing
1428 stepOuter _ _ Nothing = return Stop
1430{-# INLINE_NORMAL groupsBy #-}
1436groupsBy cmp f (Stream step state) = Stream (stepOuter f) (Just state, Nothing)
1440 {-# INLINE_LATE stepOuter #-}
1441 stepOuter (Fold fstep initial done) gst (Just st, Nothing) = do
1442 res <- step (adaptState gst) st
1449 Skip s -> return $ Skip $ (Just s, Nothing)
1454 go !_ prev stt !acc = do
1455 res <- step (adaptState gst) stt
1462 else done acc >>= \r -> return $ Yield r (Just s, Just x)
1463 Skip s -> go SPEC prev s acc
1464 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1466 stepOuter (Fold fstep initial done) gst (Just st, Just prev) = do
1468 acc' <- fstep acc prev
1473 -- XXX code duplicated from the previous equation
1475 res <- step (adaptState gst) stt
1482 else done acc >>= \r -> return $ Yield r (Just s, Just x)
1483 Skip s -> go SPEC s acc
1484 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1486 stepOuter _ _ (Nothing,_) = return Stop
1488{-# INLINE_NORMAL groupsRollingBy #-}
1489groupsRollingBy :: Monad m
1494groupsRollingBy cmp f (Stream step state) =
1495 Stream (stepOuter f) (Just state, Nothing)
1498 {-# INLINE_LATE stepOuter #-}
1499 stepOuter (Fold fstep initial done) gst (Just st, Nothing) = do
1500 res <- step (adaptState gst) st
1507 Skip s -> return $ Skip $ (Just s, Nothing)
1511 go !_ prev stt !acc = do
1512 res <- step (adaptState gst) stt
1520 done acc >>= \r -> return $ Yield r (Just s, Just x)
1521 Skip s -> go SPEC prev s acc
1522 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1524 stepOuter (Fold fstep initial done) gst (Just st, Just prev') = do
1526 acc' <- fstep acc prev'
1527 go SPEC prev' st acc'
1530 go !_ prevv stt !acc = do
1531 res <- step (adaptState gst) stt
1538 else done acc >>= \r -> return $ Yield r (Just s, Just x)
1539 Skip s -> go SPEC prevv s acc
1540 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1542 stepOuter _ _ (Nothing, _) = return Stop
1544{-# INLINE_NORMAL splitBy #-}
1545splitBy :: Monad m => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1546splitBy predicate f (Stream step state) = Stream (step' f) (Just state)
1550 {-# INLINE_LATE step' #-}
1551 step' (Fold fstep initial done) gst (Just st) = initial >>= go SPEC st
1556 res <- step (adaptState gst) stt
1560 then done acc >>= \r -> return $ Yield r (Just s)
1564 Skip s -> go SPEC s acc
1565 Stop -> done acc >>= \r -> return $ Yield r Nothing
1567 step' _ _ Nothing = return Stop
1569-- XXX requires -funfolding-use-threshold=150 in lines-unlines benchmark
1570{-# INLINE_NORMAL splitSuffixBy #-}
1571splitSuffixBy :: Monad m
1572 => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1573splitSuffixBy predicate f (Stream step state) = Stream (step' f) (Just state)
1577 {-# INLINE_LATE step' #-}
1578 step' (Fold fstep initial done) gst (Just st) = do
1579 res <- step (adaptState gst) st
1584 then done acc >>= \val -> return $ Yield val (Just s)
1589 Skip s -> return $ Skip $ Just s
1595 res <- step (adaptState gst) stt
1599 then done acc >>= \r -> return $ Yield r (Just s)
1603 Skip s -> go SPEC s acc
1604 Stop -> done acc >>= \r -> return $ Yield r Nothing
1606 step' _ _ Nothing = return Stop
1608{-# INLINE_NORMAL wordsBy #-}
1609wordsBy :: Monad m => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1610wordsBy predicate f (Stream step state) = Stream (stepOuter f) (Just state)
1614 {-# INLINE_LATE stepOuter #-}
1615 stepOuter (Fold fstep initial done) gst (Just st) = do
1616 res <- step (adaptState gst) st
1620 then return $ Skip (Just s)
1626 Skip s -> return $ Skip $ Just s
1632 res <- step (adaptState gst) stt
1636 then done acc >>= \r -> return $ Yield r (Just s)
1640 Skip s -> go SPEC s acc
1641 Stop -> done acc >>= \r -> return $ Yield r Nothing
1643 stepOuter _ _ Nothing = return Stop
1645-- String search algorithms:
1646-- http://www-igm.univ-mlv.fr/~lecroq/string/index.html
1649-- TODO can we unify the splitting operations using a splitting configuration
1650-- like in the split package.
1652data SplitStyle = Infix | Suffix | Prefix deriving (Eq, Show)
1654data SplitOptions = SplitOptions
1655 { style :: SplitStyle
1656 , withSep :: Bool -- ^ keep the separators in output
1657 -- , compact :: Bool -- ^ treat multiple consecutive separators as one
1658 -- , trimHead :: Bool -- ^ drop blank at head
1659 -- , trimTail :: Bool -- ^ drop blank at tail
1663data SplitOnState s a =
1668 | GO_KARP_RABIN s !(RB.Ring a) !(Ptr a)
1671{-# INLINE_NORMAL splitOn #-}
1673 :: forall m a b. (MonadIO m, Storable a, Enum a, Eq a)
1678splitOn patArr (Fold fstep initial done) (Stream step state) =
1679 Stream stepOuter GO_START
1683 patLen = A.length patArr
1684 maxIndex = patLen - 1
1685 elemBits = sizeOf (undefined :: a) * 8
1687 {-# INLINE_LATE stepOuter #-}
1688 stepOuter _ GO_START =
1690 then return $ Skip $ GO_EMPTY_PAT state
1693 r <- liftIO $ (A.unsafeIndexIO patArr 0)
1694 return $ Skip $ GO_SINGLE_PAT state r
1695 else if sizeOf (undefined :: a) * patLen
1696 <= sizeOf (undefined :: Word)
1697 then return $ Skip $ GO_SHORT_PAT state
1699 (rb, rhead) <- liftIO $ RB.new patLen
1700 return $ Skip $ GO_KARP_RABIN state rb rhead
1702 stepOuter gst (GO_SINGLE_PAT stt pat) = initial >>= go SPEC stt
1707 res <- step (adaptState gst) st
1713 return $ Yield r (GO_SINGLE_PAT s pat)
1714 else fstep acc x >>= go SPEC s
1715 Skip s -> go SPEC s acc
1716 Stop -> done acc >>= \r -> return $ Yield r GO_DONE
1718 stepOuter gst (GO_SHORT_PAT stt) = initial >>= go0 SPEC 0 (0 :: Word) stt
1723 mask = (1 `shiftL` (elemBits * patLen)) - 1
1725 addToWord wrd a = (wrd `shiftL` elemBits) .|. fromIntegral (fromEnum a)
1728 patWord = mask .&. A.foldl' addToWord 0 patArr
1730 go0 !_ !idx wrd st !acc = do
1731 res <- step (adaptState gst) st
1734 let wrd' = addToWord wrd x
1737 if wrd' .&. mask == patWord
1740 return $ Yield r (GO_SHORT_PAT s)
1741 else go1 SPEC wrd' s acc
1742 else go0 SPEC (idx + 1) wrd' s acc
1743 Skip s -> go0 SPEC idx wrd s acc
1746 then go2 wrd idx acc
1748 done acc' >>= \r -> return $ Yield r GO_DONE
1751 go1 !_ wrd st !acc = do
1752 res <- step (adaptState gst) st
1755 let wrd' = addToWord wrd x
1756 old = (mask .&. wrd) `shiftR` (elemBits * (patLen - 1))
1757 acc' <- fstep acc (toEnum $ fromIntegral old)
1758 if wrd' .&. mask == patWord
1759 then done acc' >>= \r -> return $ Yield r (GO_SHORT_PAT s)
1760 else go1 SPEC wrd' s acc'
1761 Skip s -> go1 SPEC wrd s acc
1763 acc' <- go2 wrd patLen acc
1764 done acc' >>= \r -> return $ Yield r GO_DONE
1766 go2 !wrd !n !acc | n > 0 = do
1767 let old = (mask .&. wrd) `shiftR` (elemBits * (n - 1))
1768 fstep acc (toEnum $ fromIntegral old) >>= go2 wrd (n - 1)
1769 go2 _ _ acc = return acc
1771 stepOuter gst (GO_KARP_RABIN stt rb rhead) = do
1772 initial >>= go0 SPEC 0 rhead stt
1776 k = 2891336453 :: Word32
1778 addCksum cksum a = cksum * k + fromIntegral (fromEnum a)
1779 deltaCksum cksum old new =
1780 addCksum cksum new - coeff * fromIntegral (fromEnum old)
1782 -- XXX shall we use a random starting hash or 1 instead of 0?
1783 patHash = A.foldl' addCksum 0 patArr
1786 go0 !_ !idx !rh st !acc = do
1787 res <- step (adaptState gst) st
1790 rh' <- liftIO $ RB.unsafeInsert rb rh x
1793 let fold = RB.unsafeFoldRing (RB.ringBound rb)
1794 let !ringHash = fold addCksum 0 rb
1795 if ringHash == patHash
1796 then go2 SPEC ringHash rh' s acc
1797 else go1 SPEC ringHash rh' s acc
1798 else go0 SPEC (idx + 1) rh' s acc
1799 Skip s -> go0 SPEC idx rh s acc
1801 !acc' <- if idx /= 0
1802 then RB.unsafeFoldRingM rh fstep acc rb
1804 done acc' >>= \r -> return $ Yield r GO_DONE
1806 -- XXX Theoretically this code can do 4 times faster if GHC generates
1807 -- optimal code. If we use just "(cksum' == patHash)" condition it goes
1808 -- 4x faster, as soon as we add the "RB.unsafeEqArray rb v" condition
1809 -- the generated code changes drastically and becomes 4x slower. Need
1810 -- to investigate what is going on with GHC.
1812 go1 !_ !cksum !rh st !acc = do
1813 res <- step (adaptState gst) st
1816 old <- liftIO $ peek rh
1817 let cksum' = deltaCksum cksum old x
1818 acc' <- fstep acc old
1820 if (cksum' == patHash)
1822 rh' <- liftIO (RB.unsafeInsert rb rh x)
1823 go2 SPEC cksum' rh' s acc'
1825 rh' <- liftIO (RB.unsafeInsert rb rh x)
1826 go1 SPEC cksum' rh' s acc'
1827 Skip s -> go1 SPEC cksum rh s acc
1829 acc' <- RB.unsafeFoldRingFullM rh fstep acc rb
1830 done acc' >>= \r -> return $ Yield r GO_DONE
1832 go2 !_ !cksum' !rh' s !acc' = do
1833 if RB.unsafeEqArray rb rh' patArr
1836 return $ Yield r (GO_KARP_RABIN s rb rhead)
1837 else go1 SPEC cksum' rh' s acc'
1839 stepOuter gst (GO_EMPTY_PAT st) = do
1840 res <- step (adaptState gst) st
1845 done acc' >>= \r -> return $ Yield r (GO_EMPTY_PAT s)
1846 Skip s -> return $ Skip (GO_EMPTY_PAT s)
1849 stepOuter _ GO_DONE = return Stop
1851{-# INLINE_NORMAL splitSuffixOn #-}
1853 :: forall m a b. (MonadIO m, Storable a, Enum a, Eq a)
1859splitSuffixOn withSep patArr (Fold fstep initial done)
1860 (Stream step state) =
1861 Stream stepOuter GO_START
1865 patLen = A.length patArr
1866 maxIndex = patLen - 1
1867 elemBits = sizeOf (undefined :: a) * 8
1869 {-# INLINE_LATE stepOuter #-}
1870 stepOuter _ GO_START =
1872 then return $ Skip $ GO_EMPTY_PAT state
1875 r <- liftIO $ (A.unsafeIndexIO patArr 0)
1876 return $ Skip $ GO_SINGLE_PAT state r
1877 else if sizeOf (undefined :: a) * patLen
1878 <= sizeOf (undefined :: Word)
1879 then return $ Skip $ GO_SHORT_PAT state
1881 (rb, rhead) <- liftIO $ RB.new patLen
1882 return $ Skip $ GO_KARP_RABIN state rb rhead
1884 stepOuter gst (GO_SINGLE_PAT stt pat) = do
1885 -- This first part is the only difference between splitOn and
1887 -- If the last element is a separator do not issue a blank segment.
1888 res <- step (adaptState gst) stt
1894 acc' <- if withSep then fstep acc x else return acc
1895 done acc' >>= \r -> return $ Yield r (GO_SINGLE_PAT s pat)
1896 else fstep acc x >>= go SPEC s
1897 Skip s -> return $ Skip $ (GO_SINGLE_PAT s pat)
1902 -- This is identical for splitOn and splitSuffixOn
1904 res <- step (adaptState gst) st
1909 acc' <- if withSep then fstep acc x else return acc
1911 return $ Yield r (GO_SINGLE_PAT s pat)
1912 else fstep acc x >>= go SPEC s
1913 Skip s -> go SPEC s acc
1914 Stop -> done acc >>= \r -> return $ Yield r GO_DONE
1916 stepOuter gst (GO_SHORT_PAT stt) = do
1918 -- Call "initial" only if the stream yields an element, otherwise we
1919 -- may call "initial" but never yield anything. initial may produce a
1920 -- side effect, therefore we will end up doing and discard a side
1925 res <- step (adaptState gst) stt
1929 let wrd' = addToWord wrd x
1930 acc' <- if withSep then fstep acc x else return acc
1933 if wrd' .&. mask == patWord
1934 then done acc' >>= \r -> return $ Yield r (GO_SHORT_PAT s)
1935 else go0 SPEC (idx + 1) wrd' s acc'
1936 else go0 SPEC (idx + 1) wrd' s acc'
1937 Skip s -> return $ Skip (GO_SHORT_PAT s)
1943 mask = (1 `shiftL` (elemBits * patLen)) - 1
1945 addToWord wrd a = (wrd `shiftL` elemBits) .|. fromIntegral (fromEnum a)
1948 patWord = mask .&. A.foldl' addToWord 0 patArr
1950 go0 !_ !idx wrd st !acc = do
1951 res <- step (adaptState gst) st
1954 let wrd' = addToWord wrd x
1955 acc' <- if withSep then fstep acc x else return acc
1958 if wrd' .&. mask == patWord
1961 return $ Yield r (GO_SHORT_PAT s)
1962 else go1 SPEC wrd' s acc'
1963 else go0 SPEC (idx + 1) wrd' s acc'
1964 Skip s -> go0 SPEC idx wrd s acc
1966 if (idx == maxIndex) && (wrd .&. mask == patWord)
1969 acc' <- if idx /= 0 && not withSep
1970 then go2 wrd idx acc
1972 done acc' >>= \r -> return $ Yield r GO_DONE
1975 go1 !_ wrd st !acc = do
1976 res <- step (adaptState gst) st
1979 let wrd' = addToWord wrd x
1980 old = (mask .&. wrd) `shiftR` (elemBits * (patLen - 1))
1983 else fstep acc (toEnum $ fromIntegral old)
1984 if wrd' .&. mask == patWord
1985 then done acc' >>= \r -> return $ Yield r (GO_SHORT_PAT s)
1986 else go1 SPEC wrd' s acc'
1987 Skip s -> go1 SPEC wrd s acc
1989 -- If the last sequence is a separator do not issue a blank
1991 if wrd .&. mask == patWord
1996 else go2 wrd patLen acc
1997 done acc' >>= \r -> return $ Yield r GO_DONE
1999 go2 !wrd !n !acc | n > 0 = do
2000 let old = (mask .&. wrd) `shiftR` (elemBits * (n - 1))
2001 fstep acc (toEnum $ fromIntegral old) >>= go2 wrd (n - 1)
2002 go2 _ _ acc = return acc
2004 stepOuter gst (GO_KARP_RABIN stt rb rhead) = do
2006 res <- step (adaptState gst) stt
2010 acc' <- if withSep then fstep acc x else return acc
2011 rh' <- liftIO (RB.unsafeInsert rb rhead x)
2014 let fold = RB.unsafeFoldRing (RB.ringBound rb)
2015 let !ringHash = fold addCksum 0 rb
2016 if ringHash == patHash
2017 then go2 SPEC ringHash rh' s acc'
2018 else go0 SPEC (idx + 1) rh' s acc'
2019 else go0 SPEC (idx + 1) rh' s acc'
2020 Skip s -> return $ Skip (GO_KARP_RABIN s rb rhead)
2025 k = 2891336453 :: Word32
2027 addCksum cksum a = cksum * k + fromIntegral (fromEnum a)
2028 deltaCksum cksum old new =
2029 addCksum cksum new - coeff * fromIntegral (fromEnum old)
2031 -- XXX shall we use a random starting hash or 1 instead of 0?
2032 patHash = A.foldl' addCksum 0 patArr
2035 go0 !_ !idx !rh st !acc = do
2036 res <- step (adaptState gst) st
2039 acc' <- if withSep then fstep acc x else return acc
2040 rh' <- liftIO (RB.unsafeInsert rb rh x)
2043 let fold = RB.unsafeFoldRing (RB.ringBound rb)
2044 let !ringHash = fold addCksum 0 rb
2045 if ringHash == patHash
2046 then go2 SPEC ringHash rh' s acc'
2047 else go1 SPEC ringHash rh' s acc'
2048 else go0 SPEC (idx + 1) rh' s acc'
2049 Skip s -> go0 SPEC idx rh s acc
2051 -- do not issue a blank segment when we end at pattern
2052 if (idx == maxIndex) && RB.unsafeEqArray rb rh patArr
2055 !acc' <- if idx /= 0 && not withSep
2056 then RB.unsafeFoldRingM rh fstep acc rb
2058 done acc' >>= \r -> return $ Yield r GO_DONE
2060 -- XXX Theoretically this code can do 4 times faster if GHC generates
2061 -- optimal code. If we use just "(cksum' == patHash)" condition it goes
2062 -- 4x faster, as soon as we add the "RB.unsafeEqArray rb v" condition
2063 -- the generated code changes drastically and becomes 4x slower. Need
2064 -- to investigate what is going on with GHC.
2066 go1 !_ !cksum !rh st !acc = do
2067 res <- step (adaptState gst) st
2070 old <- liftIO $ peek rh
2071 let cksum' = deltaCksum cksum old x
2076 if (cksum' == patHash)
2078 rh' <- liftIO (RB.unsafeInsert rb rh x)
2079 go2 SPEC cksum' rh' s acc'
2081 rh' <- liftIO (RB.unsafeInsert rb rh x)
2082 go1 SPEC cksum' rh' s acc'
2083 Skip s -> go1 SPEC cksum rh s acc
2085 if RB.unsafeEqArray rb rh patArr
2090 else RB.unsafeFoldRingFullM rh fstep acc rb
2091 done acc' >>= \r -> return $ Yield r GO_DONE
2093 go2 !_ !cksum' !rh' s !acc' = do
2094 if RB.unsafeEqArray rb rh' patArr
2097 return $ Yield r (GO_KARP_RABIN s rb rhead)
2098 else go1 SPEC cksum' rh' s acc'
2100 stepOuter gst (GO_EMPTY_PAT st) = do
2101 res <- step (adaptState gst) st
2106 done acc' >>= \r -> return $ Yield r (GO_EMPTY_PAT s)
2107 Skip s -> return $ Skip (GO_EMPTY_PAT s)
2110 stepOuter _ GO_DONE = return Stop
2112data SplitState s arr
2114 | SplitBuffering s arr
2115 | SplitSplitting s arr
2116 | SplitYielding arr (SplitState s arr)
2119-- XXX An alternative approach would be to use a partial fold (Fold m a b) to
2120-- split using a splitBy like combinator. The Fold would consume upto the
2121-- separator and return any leftover which can then be fed to the next fold.
2123-- We can revisit this once we have partial folds/parsers.
2125-- | Performs infix separator style splitting.
2126{-# INLINE_NORMAL splitInnerBy #-}
2129 => (f a -> m (f a, Maybe (f a))) -- splitter
2130 -> (f a -> f a -> m (f a)) -- joiner
2133splitInnerBy splitter joiner (Stream step1 state1) =
2134 (Stream step (SplitInitial state1))
2138 {-# INLINE_LATE step #-}
2139 step gst (SplitInitial st) = do
2143 (x1, mx2) <- splitter x
2144 return $ case mx2 of
2145 Nothing -> Skip (SplitBuffering s x1)
2146 Just x2 -> Skip (SplitYielding x1 (SplitSplitting s x2))
2147 Skip s -> return $ Skip (SplitInitial s)
2148 Stop -> return $ Stop
2150 step gst (SplitBuffering st buf) = do
2154 (x1, mx2) <- splitter x
2155 buf' <- joiner buf x1
2156 return $ case mx2 of
2157 Nothing -> Skip (SplitBuffering s buf')
2158 Just x2 -> Skip (SplitYielding buf' (SplitSplitting s x2))
2159 Skip s -> return $ Skip (SplitBuffering s buf)
2160 Stop -> return $ Skip (SplitYielding buf SplitFinishing)
2162 step _ (SplitSplitting st buf) = do
2163 (x1, mx2) <- splitter buf
2164 return $ case mx2 of
2165 Nothing -> Skip $ SplitBuffering st x1
2166 Just x2 -> Skip $ SplitYielding x1 (SplitSplitting st x2)
2168 step _ (SplitYielding x next) = return $ Yield x next
2169 step _ SplitFinishing = return $ Stop
2171-- | Performs infix separator style splitting.
2172{-# INLINE_NORMAL splitInnerBySuffix #-}
2174 :: (Monad m, Eq (f a), Monoid (f a))
2175 => (f a -> m (f a, Maybe (f a))) -- splitter
2176 -> (f a -> f a -> m (f a)) -- joiner
2179splitInnerBySuffix splitter joiner (Stream step1 state1) =
2180 (Stream step (SplitInitial state1))
2184 {-# INLINE_LATE step #-}
2185 step gst (SplitInitial st) = do
2189 (x1, mx2) <- splitter x
2190 return $ case mx2 of
2191 Nothing -> Skip (SplitBuffering s x1)
2192 Just x2 -> Skip (SplitYielding x1 (SplitSplitting s x2))
2193 Skip s -> return $ Skip (SplitInitial s)
2194 Stop -> return $ Stop
2196 step gst (SplitBuffering st buf) = do
2200 (x1, mx2) <- splitter x
2201 buf' <- joiner buf x1
2202 return $ case mx2 of
2203 Nothing -> Skip (SplitBuffering s buf')
2204 Just x2 -> Skip (SplitYielding buf' (SplitSplitting s x2))
2205 Skip s -> return $ Skip (SplitBuffering s buf)
2209 else Skip (SplitYielding buf SplitFinishing)
2211 step _ (SplitSplitting st buf) = do
2212 (x1, mx2) <- splitter buf
2213 return $ case mx2 of
2214 Nothing -> Skip $ SplitBuffering st x1
2215 Just x2 -> Skip $ SplitYielding x1 (SplitSplitting st x2)
2217 step _ (SplitYielding x next) = return $ Yield x next
2218 step _ SplitFinishing = return $ Stop
2220------------------------------------------------------------------------------
2222------------------------------------------------------------------------------
2224{-# INLINE_NORMAL isPrefixOf #-}
2225isPrefixOf :: (Eq a, Monad m) => Stream m a -> Stream m a -> m Bool
2226isPrefixOf (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
2228 go (sa, sb, Nothing) = do
2229 r <- stepa defState sa
2231 Yield x sa' -> go (sa', sb, Just x)
2232 Skip sa' -> go (sa', sb, Nothing)
2235 go (sa, sb, Just x) = do
2236 r <- stepb defState sb
2240 then go (sa, sb', Nothing)
2242 Skip sb' -> go (sa, sb', Just x)
2243 Stop -> return False
2245{-# INLINE_NORMAL isSubsequenceOf #-}
2246isSubsequenceOf :: (Eq a, Monad m) => Stream m a -> Stream m a -> m Bool
2247isSubsequenceOf (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
2249 go (sa, sb, Nothing) = do
2250 r <- stepa defState sa
2252 Yield x sa' -> go (sa', sb, Just x)
2253 Skip sa' -> go (sa', sb, Nothing)
2256 go (sa, sb, Just x) = do
2257 r <- stepb defState sb
2261 then go (sa, sb', Nothing)
2262 else go (sa, sb', Just x)
2263 Skip sb' -> go (sa, sb', Just x)
2264 Stop -> return False
2266{-# INLINE_NORMAL stripPrefix #-}
2269 => Stream m a -> Stream m a -> m (Maybe (Stream m a))
2270stripPrefix (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
2272 go (sa, sb, Nothing) = do
2273 r <- stepa defState sa
2275 Yield x sa' -> go (sa', sb, Just x)
2276 Skip sa' -> go (sa', sb, Nothing)
2277 Stop -> return $ Just (Stream stepb sb)
2279 go (sa, sb, Just x) = do
2280 r <- stepb defState sb
2284 then go (sa, sb', Nothing)
2286 Skip sb' -> go (sa, sb', Just x)
2287 Stop -> return Nothing
2289------------------------------------------------------------------------------
2291------------------------------------------------------------------------------
2293-- | Execute a monadic action for each element of the 'Stream'
2294{-# INLINE_NORMAL mapM_ #-}
2295mapM_ :: Monad m => (a -> m b) -> Stream m a -> m ()
2296mapM_ m = drain . mapM m
2298-------------------------------------------------------------------------------
2299-- Stream transformations using Unfolds
2300-------------------------------------------------------------------------------
2302-- Define a unique structure to use in inspection testing
2303data ConcatMapUState o i =
2305 | ConcatMapUInner o i
2307-- | @concatMapU unfold stream@ uses @unfold@ to map the input stream elements
2308-- to streams and then flattens the generated streams into a single output
2311-- This is like 'concatMap' but uses an unfold with an explicit state to
2312-- generate the stream instead of a 'Stream' type generator. This allows better
2313-- optimization via fusion. This can be many times more efficient than
2316{-# INLINE_NORMAL concatMapU #-}
2317concatMapU :: Monad m => Unfold m a b -> Stream m a -> Stream m b
2318concatMapU (Unfold istep inject) (Stream ostep ost) =
2319 Stream step (ConcatMapUOuter ost)
2321 {-# INLINE_LATE step #-}
2322 step gst (ConcatMapUOuter o) = do
2323 r <- ostep (adaptState gst) o
2327 i `seq` return (Skip (ConcatMapUInner o' i))
2328 Skip o' -> return $ Skip (ConcatMapUOuter o')
2329 Stop -> return $ Stop
2331 step _ (ConcatMapUInner o i) = do
2334 Yield x i' -> Yield x (ConcatMapUInner o i')
2335 Skip i' -> Skip (ConcatMapUInner o i')
2336 Stop -> Skip (ConcatMapUOuter o)
2338data ConcatUnfoldInterleaveState o i =
2339 ConcatUnfoldInterleaveOuter o [i]
2340 | ConcatUnfoldInterleaveInner o [i]
2341 | ConcatUnfoldInterleaveInnerL [i] [i]
2342 | ConcatUnfoldInterleaveInnerR [i] [i]
2344-- XXX use arrays to store state instead of lists.
2345-- XXX In general we can use different scheduling strategies e.g. how to
2346-- schedule the outer vs inner loop or assigning weights to different streams
2347-- or outer and inner loops.
2349-- After a yield, switch to the next stream. Do not switch streams on Skip.
2350-- Yield from outer stream switches to the inner stream.
2352-- There are two choices here, (1) exhaust the outer stream first and then
2353-- start yielding from the inner streams, this is much simpler to implement,
2354-- (2) yield at least one element from an inner stream before going back to
2355-- outer stream and opening the next stream from it.
2357-- Ideally, we need some scheduling bias to inner streams vs outer stream.
2358-- Maybe we can configure the behavior.
2360{-# INLINE_NORMAL concatUnfoldInterleave #-}
2361concatUnfoldInterleave :: Monad m => Unfold m a b -> Stream m a -> Stream m b
2362concatUnfoldInterleave (Unfold istep inject) (Stream ostep ost) =
2363 Stream step (ConcatUnfoldInterleaveOuter ost [])
2365 {-# INLINE_LATE step #-}
2366 step gst (ConcatUnfoldInterleaveOuter o ls) = do
2367 r <- ostep (adaptState gst) o
2371 i `seq` return (Skip (ConcatUnfoldInterleaveInner o' (i : ls)))
2372 Skip o' -> return $ Skip (ConcatUnfoldInterleaveOuter o' ls)
2373 Stop -> return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2375 step _ (ConcatUnfoldInterleaveInner _ []) = undefined
2376 step _ (ConcatUnfoldInterleaveInner o (st:ls)) = do
2379 Yield x s -> Yield x (ConcatUnfoldInterleaveOuter o (s:ls))
2380 Skip s -> Skip (ConcatUnfoldInterleaveInner o (s:ls))
2381 Stop -> Skip (ConcatUnfoldInterleaveOuter o ls)
2383 step _ (ConcatUnfoldInterleaveInnerL [] []) = return Stop
2384 step _ (ConcatUnfoldInterleaveInnerL [] rs) =
2385 return $ Skip (ConcatUnfoldInterleaveInnerR [] rs)
2387 step _ (ConcatUnfoldInterleaveInnerL (st:ls) rs) = do
2390 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerL ls (s:rs))
2391 Skip s -> Skip (ConcatUnfoldInterleaveInnerL (s:ls) rs)
2392 Stop -> Skip (ConcatUnfoldInterleaveInnerL ls rs)
2394 step _ (ConcatUnfoldInterleaveInnerR [] []) = return Stop
2395 step _ (ConcatUnfoldInterleaveInnerR ls []) =
2396 return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2398 step _ (ConcatUnfoldInterleaveInnerR ls (st:rs)) = do
2401 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerR (s:ls) rs)
2402 Skip s -> Skip (ConcatUnfoldInterleaveInnerR ls (s:rs))
2403 Stop -> Skip (ConcatUnfoldInterleaveInnerR ls rs)
2405-- XXX In general we can use different scheduling strategies e.g. how to
2406-- schedule the outer vs inner loop or assigning weights to different streams
2407-- or outer and inner loops.
2409-- This could be inefficient if the tasks are too small.
2411-- Compared to concatUnfoldInterleave this one switches streams on Skips.
2413{-# INLINE_NORMAL concatUnfoldRoundrobin #-}
2414concatUnfoldRoundrobin :: Monad m => Unfold m a b -> Stream m a -> Stream m b
2415concatUnfoldRoundrobin (Unfold istep inject) (Stream ostep ost) =
2416 Stream step (ConcatUnfoldInterleaveOuter ost [])
2418 {-# INLINE_LATE step #-}
2419 step gst (ConcatUnfoldInterleaveOuter o ls) = do
2420 r <- ostep (adaptState gst) o
2424 i `seq` return (Skip (ConcatUnfoldInterleaveInner o' (i : ls)))
2425 Skip o' -> return $ Skip (ConcatUnfoldInterleaveInner o' ls)
2426 Stop -> return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2428 step _ (ConcatUnfoldInterleaveInner o []) =
2429 return $ Skip (ConcatUnfoldInterleaveOuter o [])
2431 step _ (ConcatUnfoldInterleaveInner o (st:ls)) = do
2434 Yield x s -> Yield x (ConcatUnfoldInterleaveOuter o (s:ls))
2435 Skip s -> Skip (ConcatUnfoldInterleaveOuter o (s:ls))
2436 Stop -> Skip (ConcatUnfoldInterleaveOuter o ls)
2438 step _ (ConcatUnfoldInterleaveInnerL [] []) = return Stop
2439 step _ (ConcatUnfoldInterleaveInnerL [] rs) =
2440 return $ Skip (ConcatUnfoldInterleaveInnerR [] rs)
2442 step _ (ConcatUnfoldInterleaveInnerL (st:ls) rs) = do
2445 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerL ls (s:rs))
2446 Skip s -> Skip (ConcatUnfoldInterleaveInnerL ls (s:rs))
2447 Stop -> Skip (ConcatUnfoldInterleaveInnerL ls rs)
2449 step _ (ConcatUnfoldInterleaveInnerR [] []) = return Stop
2450 step _ (ConcatUnfoldInterleaveInnerR ls []) =
2451 return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2453 step _ (ConcatUnfoldInterleaveInnerR ls (st:rs)) = do
2456 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerR (s:ls) rs)
2457 Skip s -> Skip (ConcatUnfoldInterleaveInnerR (s:ls) rs)
2458 Stop -> Skip (ConcatUnfoldInterleaveInnerR ls rs)
2460data AppendState s1 s2 = AppendFirst s1 | AppendSecond s2
2462-- Note that this could be much faster compared to the CPS stream. However, as
2463-- the number of streams being composed increases this may become expensive.
2464-- Need to see where the breaking point is between the two.
2466{-# INLINE_NORMAL append #-}
2467append :: Monad m => Stream m a -> Stream m a -> Stream m a
2468append (Stream step1 state1) (Stream step2 state2) =
2469 Stream step (AppendFirst state1)
2473 {-# INLINE_LATE step #-}
2474 step gst (AppendFirst st) = do
2477 Yield a s -> Yield a (AppendFirst s)
2478 Skip s -> Skip (AppendFirst s)
2479 Stop -> Skip (AppendSecond state2)
2481 step gst (AppendSecond st) = do
2484 Yield a s -> Yield a (AppendSecond s)
2485 Skip s -> Skip (AppendSecond s)
2488data InterleaveState s1 s2 = InterleaveFirst s1 s2 | InterleaveSecond s1 s2
2489 | InterleaveSecondOnly s2 | InterleaveFirstOnly s1
2491{-# INLINE_NORMAL interleave #-}
2492interleave :: Monad m => Stream m a -> Stream m a -> Stream m a
2493interleave (Stream step1 state1) (Stream step2 state2) =
2494 Stream step (InterleaveFirst state1 state2)
2498 {-# INLINE_LATE step #-}
2499 step gst (InterleaveFirst st1 st2) = do
2502 Yield a s -> Yield a (InterleaveSecond s st2)
2503 Skip s -> Skip (InterleaveFirst s st2)
2504 Stop -> Skip (InterleaveSecondOnly st2)
2506 step gst (InterleaveSecond st1 st2) = do
2509 Yield a s -> Yield a (InterleaveFirst st1 s)
2510 Skip s -> Skip (InterleaveSecond st1 s)
2511 Stop -> Skip (InterleaveFirstOnly st1)
2513 step gst (InterleaveFirstOnly st1) = do
2516 Yield a s -> Yield a (InterleaveFirstOnly s)
2517 Skip s -> Skip (InterleaveFirstOnly s)
2520 step gst (InterleaveSecondOnly st2) = do
2523 Yield a s -> Yield a (InterleaveSecondOnly s)
2524 Skip s -> Skip (InterleaveSecondOnly s)
2527{-# INLINE_NORMAL interleaveMin #-}
2528interleaveMin :: Monad m => Stream m a -> Stream m a -> Stream m a
2529interleaveMin (Stream step1 state1) (Stream step2 state2) =
2530 Stream step (InterleaveFirst state1 state2)
2534 {-# INLINE_LATE step #-}
2535 step gst (InterleaveFirst st1 st2) = do
2538 Yield a s -> Yield a (InterleaveSecond s st2)
2539 Skip s -> Skip (InterleaveFirst s st2)
2542 step gst (InterleaveSecond st1 st2) = do
2545 Yield a s -> Yield a (InterleaveFirst st1 s)
2546 Skip s -> Skip (InterleaveSecond st1 s)
2549 step _ (InterleaveFirstOnly _) = undefined
2550 step _ (InterleaveSecondOnly _) = undefined
2552{-# INLINE_NORMAL interleaveSuffix #-}
2553interleaveSuffix :: Monad m => Stream m a -> Stream m a -> Stream m a
2554interleaveSuffix (Stream step1 state1) (Stream step2 state2) =
2555 Stream step (InterleaveFirst state1 state2)
2559 {-# INLINE_LATE step #-}
2560 step gst (InterleaveFirst st1 st2) = do
2563 Yield a s -> Yield a (InterleaveSecond s st2)
2564 Skip s -> Skip (InterleaveFirst s st2)
2567 step gst (InterleaveSecond st1 st2) = do
2570 Yield a s -> Yield a (InterleaveFirst st1 s)
2571 Skip s -> Skip (InterleaveSecond st1 s)
2572 Stop -> Skip (InterleaveFirstOnly st1)
2574 step gst (InterleaveFirstOnly st1) = do
2577 Yield a s -> Yield a (InterleaveFirstOnly s)
2578 Skip s -> Skip (InterleaveFirstOnly s)
2581 step _ (InterleaveSecondOnly _) = undefined
2583data InterleaveInfixState s1 s2 a
2584 = InterleaveInfixFirst s1 s2
2585 | InterleaveInfixSecondBuf s1 s2
2586 | InterleaveInfixSecondYield s1 s2 a
2587 | InterleaveInfixFirstYield s1 s2 a
2588 | InterleaveInfixFirstOnly s1
2590{-# INLINE_NORMAL interleaveInfix #-}
2591interleaveInfix :: Monad m => Stream m a -> Stream m a -> Stream m a
2592interleaveInfix (Stream step1 state1) (Stream step2 state2) =
2593 Stream step (InterleaveInfixFirst state1 state2)
2597 {-# INLINE_LATE step #-}
2598 step gst (InterleaveInfixFirst st1 st2) = do
2601 Yield a s -> Yield a (InterleaveInfixSecondBuf s st2)
2602 Skip s -> Skip (InterleaveInfixFirst s st2)
2605 step gst (InterleaveInfixSecondBuf st1 st2) = do
2608 Yield a s -> Skip (InterleaveInfixSecondYield st1 s a)
2609 Skip s -> Skip (InterleaveInfixSecondBuf st1 s)
2610 Stop -> Skip (InterleaveInfixFirstOnly st1)
2612 step gst (InterleaveInfixSecondYield st1 st2 x) = do
2615 Yield a s -> Yield x (InterleaveInfixFirstYield s st2 a)
2616 Skip s -> Skip (InterleaveInfixSecondYield s st2 x)
2619 step _ (InterleaveInfixFirstYield st1 st2 x) = do
2620 return $ Yield x (InterleaveInfixSecondBuf st1 st2)
2622 step gst (InterleaveInfixFirstOnly st1) = do
2625 Yield a s -> Yield a (InterleaveInfixFirstOnly s)
2626 Skip s -> Skip (InterleaveInfixFirstOnly s)
2629{-# INLINE_NORMAL roundRobin #-}
2630roundRobin :: Monad m => Stream m a -> Stream m a -> Stream m a
2631roundRobin (Stream step1 state1) (Stream step2 state2) =
2632 Stream step (InterleaveFirst state1 state2)
2636 {-# INLINE_LATE step #-}
2637 step gst (InterleaveFirst st1 st2) = do
2640 Yield a s -> Yield a (InterleaveSecond s st2)
2641 Skip s -> Skip (InterleaveSecond s st2)
2642 Stop -> Skip (InterleaveSecondOnly st2)
2644 step gst (InterleaveSecond st1 st2) = do
2647 Yield a s -> Yield a (InterleaveFirst st1 s)
2648 Skip s -> Skip (InterleaveFirst st1 s)
2649 Stop -> Skip (InterleaveFirstOnly st1)
2651 step gst (InterleaveSecondOnly st2) = do
2654 Yield a s -> Yield a (InterleaveSecondOnly s)
2655 Skip s -> Skip (InterleaveSecondOnly s)
2658 step gst (InterleaveFirstOnly st1) = do
2661 Yield a s -> Yield a (InterleaveFirstOnly s)
2662 Skip s -> Skip (InterleaveFirstOnly s)
2665data ICUState s1 s2 i1 i2 =
2670 | ICUFirstInner s1 s2 i1
2671 | ICUSecondInner s1 s2 i2
2672 | ICUFirstOnlyInner s1 i1
2673 | ICUSecondOnlyInner s2 i2
2675-- | Interleave streams (full streams, not the elements) unfolded from two
2676-- input streams and concat. Stop when the first stream stops. If the second
2677-- stream ends before the first one then first stream still keeps running alone
2678-- without any interleaving with the second stream.
2680-- [a1, a2, ... an] [b1, b2 ...]
2681-- => [streamA1, streamA2, ... streamAn] [streamB1, streamB2, ...]
2682-- => [streamA1, streamB1, streamA2...StreamAn, streamBn]
2683-- => [a11, a12, ...a1j, b11, b12, ...b1k, a21, a22, ...]
2685{-# INLINE_NORMAL gintercalateSuffix #-}
2688 => Unfold m a c -> Stream m a -> Unfold m b c -> Stream m b -> Stream m c
2690 (Unfold istep1 inject1) (Stream step1 state1)
2691 (Unfold istep2 inject2) (Stream step2 state2) =
2692 Stream step (ICUFirst state1 state2)
2696 {-# INLINE_LATE step #-}
2697 step gst (ICUFirst s1 s2) = do
2698 r <- step1 (adaptState gst) s1
2702 i `seq` return (Skip (ICUFirstInner s s2 i))
2703 Skip s -> return $ Skip (ICUFirst s s2)
2706 step gst (ICUFirstOnly s1) = do
2707 r <- step1 (adaptState gst) s1
2711 i `seq` return (Skip (ICUFirstOnlyInner s i))
2712 Skip s -> return $ Skip (ICUFirstOnly s)
2715 step _ (ICUFirstInner s1 s2 i1) = do
2718 Yield x i' -> Yield x (ICUFirstInner s1 s2 i')
2719 Skip i' -> Skip (ICUFirstInner s1 s2 i')
2720 Stop -> Skip (ICUSecond s1 s2)
2722 step _ (ICUFirstOnlyInner s1 i1) = do
2725 Yield x i' -> Yield x (ICUFirstOnlyInner s1 i')
2726 Skip i' -> Skip (ICUFirstOnlyInner s1 i')
2727 Stop -> Skip (ICUFirstOnly s1)
2729 step gst (ICUSecond s1 s2) = do
2730 r <- step2 (adaptState gst) s2
2734 i `seq` return (Skip (ICUSecondInner s1 s i))
2735 Skip s -> return $ Skip (ICUSecond s1 s)
2736 Stop -> return $ Skip (ICUFirstOnly s1)
2738 step _ (ICUSecondInner s1 s2 i2) = do
2741 Yield x i' -> Yield x (ICUSecondInner s1 s2 i')
2742 Skip i' -> Skip (ICUSecondInner s1 s2 i')
2743 Stop -> Skip (ICUFirst s1 s2)
2745 step _ (ICUSecondOnly _s2) = undefined
2746 step _ (ICUSecondOnlyInner _s2 _i2) = undefined
2748data InterposeSuffixState s1 i1 =
2749 InterposeSuffixFirst s1
2750 -- | InterposeSuffixFirstYield s1 i1
2751 | InterposeSuffixFirstInner s1 i1
2752 | InterposeSuffixSecond s1
2754-- Note that if an unfolded layer turns out to be nil we still emit the
2755-- separator effect. An alternate behavior could be to emit the separator
2756-- effect only if at least one element has been yielded by the unfolding.
2757-- However, that becomes a bit complicated, so we have chosen the former
2759{-# INLINE_NORMAL interposeSuffix #-}
2762 => m c -> Unfold m b c -> Stream m b -> Stream m c
2765 (Unfold istep1 inject1) (Stream step1 state1) =
2766 Stream step (InterposeSuffixFirst state1)
2770 {-# INLINE_LATE step #-}
2771 step gst (InterposeSuffixFirst s1) = do
2772 r <- step1 (adaptState gst) s1
2776 i `seq` return (Skip (InterposeSuffixFirstInner s i))
2777 -- i `seq` return (Skip (InterposeSuffixFirstYield s i))
2778 Skip s -> return $ Skip (InterposeSuffixFirst s)
2782 step _ (InterposeSuffixFirstYield s1 i1) = do
2785 Yield x i' -> Yield x (InterposeSuffixFirstInner s1 i')
2786 Skip i' -> Skip (InterposeSuffixFirstYield s1 i')
2787 Stop -> Skip (InterposeSuffixFirst s1)
2790 step _ (InterposeSuffixFirstInner s1 i1) = do
2793 Yield x i' -> Yield x (InterposeSuffixFirstInner s1 i')
2794 Skip i' -> Skip (InterposeSuffixFirstInner s1 i')
2795 Stop -> Skip (InterposeSuffixSecond s1)
2797 step _ (InterposeSuffixSecond s1) = do
2799 return $ Yield r (InterposeSuffixFirst s1)
2801data ICALState s1 s2 i1 i2 a =
2803 -- | ICALFirstYield s1 s2 i1
2804 | ICALFirstInner s1 s2 i1
2806 | ICALFirstOnlyInner s1 i1
2807 | ICALSecondInject s1 s2
2808 | ICALFirstInject s1 s2 i2
2809 -- | ICALFirstBuf s1 s2 i1 i2
2810 | ICALSecondInner s1 s2 i1 i2
2811 -- -- | ICALSecondInner s1 s2 i1 i2 a
2812 -- -- | ICALFirstResume s1 s2 i1 i2 a
2814-- | Interleave streams (full streams, not the elements) unfolded from two
2815-- input streams and concat. Stop when the first stream stops. If the second
2816-- stream ends before the first one then first stream still keeps running alone
2817-- without any interleaving with the second stream.
2819-- [a1, a2, ... an] [b1, b2 ...]
2820-- => [streamA1, streamA2, ... streamAn] [streamB1, streamB2, ...]
2821-- => [streamA1, streamB1, streamA2...StreamAn, streamBn]
2822-- => [a11, a12, ...a1j, b11, b12, ...b1k, a21, a22, ...]
2824{-# INLINE_NORMAL gintercalate #-}
2827 => Unfold m a c -> Stream m a -> Unfold m b c -> Stream m b -> Stream m c
2829 (Unfold istep1 inject1) (Stream step1 state1)
2830 (Unfold istep2 inject2) (Stream step2 state2) =
2831 Stream step (ICALFirst state1 state2)
2835 {-# INLINE_LATE step #-}
2836 step gst (ICALFirst s1 s2) = do
2837 r <- step1 (adaptState gst) s1
2841 i `seq` return (Skip (ICALFirstInner s s2 i))
2842 -- i `seq` return (Skip (ICALFirstYield s s2 i))
2843 Skip s -> return $ Skip (ICALFirst s s2)
2847 step _ (ICALFirstYield s1 s2 i1) = do
2850 Yield x i' -> Yield x (ICALFirstInner s1 s2 i')
2851 Skip i' -> Skip (ICALFirstYield s1 s2 i')
2852 Stop -> Skip (ICALFirst s1 s2)
2855 step _ (ICALFirstInner s1 s2 i1) = do
2858 Yield x i' -> Yield x (ICALFirstInner s1 s2 i')
2859 Skip i' -> Skip (ICALFirstInner s1 s2 i')
2860 Stop -> Skip (ICALSecondInject s1 s2)
2862 step gst (ICALFirstOnly s1) = do
2863 r <- step1 (adaptState gst) s1
2867 i `seq` return (Skip (ICALFirstOnlyInner s i))
2868 Skip s -> return $ Skip (ICALFirstOnly s)
2871 step _ (ICALFirstOnlyInner s1 i1) = do
2874 Yield x i' -> Yield x (ICALFirstOnlyInner s1 i')
2875 Skip i' -> Skip (ICALFirstOnlyInner s1 i')
2876 Stop -> Skip (ICALFirstOnly s1)
2878 -- We inject the second stream even before checking if the first stream
2879 -- would yield any more elements. There is no clear choice whether we
2880 -- should do this before or after that. Doing it after may make the state
2881 -- machine a bit simpler though.
2882 step gst (ICALSecondInject s1 s2) = do
2883 r <- step2 (adaptState gst) s2
2887 i `seq` return (Skip (ICALFirstInject s1 s i))
2888 Skip s -> return $ Skip (ICALSecondInject s1 s)
2889 Stop -> return $ Skip (ICALFirstOnly s1)
2891 step gst (ICALFirstInject s1 s2 i2) = do
2892 r <- step1 (adaptState gst) s1
2896 i `seq` return (Skip (ICALSecondInner s s2 i i2))
2897 -- i `seq` return (Skip (ICALFirstBuf s s2 i i2))
2898 Skip s -> return $ Skip (ICALFirstInject s s2 i2)
2902 step _ (ICALFirstBuf s1 s2 i1 i2) = do
2905 Yield x i' -> Skip (ICALSecondInner s1 s2 i' i2 x)
2906 Skip i' -> Skip (ICALFirstBuf s1 s2 i' i2)
2909 step _ (ICALSecondInner s1 s2 i1 i2 v) = do
2912 Yield x i' -> Yield x (ICALSecondInner s1 s2 i1 i' v)
2913 Skip i' -> Skip (ICALSecondInner s1 s2 i1 i' v)
2914 Stop -> Skip (ICALFirstResume s1 s2 i1 i2 v)
2917 step _ (ICALSecondInner s1 s2 i1 i2) = do
2920 Yield x i' -> Yield x (ICALSecondInner s1 s2 i1 i')
2921 Skip i' -> Skip (ICALSecondInner s1 s2 i1 i')
2922 Stop -> Skip (ICALFirstInner s1 s2 i1)
2923 -- Stop -> Skip (ICALFirstResume s1 s2 i1 i2)
2926 step _ (ICALFirstResume s1 s2 i1 i2 x) = do
2927 return $ Yield x (ICALFirstInner s1 s2 i1 i2)
2930data InterposeState s1 i1 a =
2932 -- | InterposeFirstYield s1 i1
2933 | InterposeFirstInner s1 i1
2934 | InterposeFirstInject s1
2935 -- | InterposeFirstBuf s1 i1
2936 | InterposeSecondYield s1 i1
2937 -- -- | InterposeSecondYield s1 i1 a
2938 -- -- | InterposeFirstResume s1 i1 a
2940-- Note that this only interposes the pure values, we may run many effects to
2941-- generate those values as some effects may not generate anything (Skip).
2942{-# INLINE_NORMAL interpose #-}
2943interpose :: Monad m => m c -> Unfold m b c -> Stream m b -> Stream m c
2946 (Unfold istep1 inject1) (Stream step1 state1) =
2947 Stream step (InterposeFirst state1)
2951 {-# INLINE_LATE step #-}
2952 step gst (InterposeFirst s1) = do
2953 r <- step1 (adaptState gst) s1
2957 i `seq` return (Skip (InterposeFirstInner s i))
2958 -- i `seq` return (Skip (InterposeFirstYield s i))
2959 Skip s -> return $ Skip (InterposeFirst s)
2963 step _ (InterposeFirstYield s1 i1) = do
2966 Yield x i' -> Yield x (InterposeFirstInner s1 i')
2967 Skip i' -> Skip (InterposeFirstYield s1 i')
2968 Stop -> Skip (InterposeFirst s1)
2971 step _ (InterposeFirstInner s1 i1) = do
2974 Yield x i' -> Yield x (InterposeFirstInner s1 i')
2975 Skip i' -> Skip (InterposeFirstInner s1 i')
2976 Stop -> Skip (InterposeFirstInject s1)
2978 step gst (InterposeFirstInject s1) = do
2979 r <- step1 (adaptState gst) s1
2983 -- i `seq` return (Skip (InterposeFirstBuf s i))
2984 i `seq` return (Skip (InterposeSecondYield s i))
2985 Skip s -> return $ Skip (InterposeFirstInject s)
2989 step _ (InterposeFirstBuf s1 i1) = do
2992 Yield x i' -> Skip (InterposeSecondYield s1 i' x)
2993 Skip i' -> Skip (InterposeFirstBuf s1 i')
2998 step _ (InterposeSecondYield s1 i1 v) = do
3000 return $ Yield r (InterposeFirstResume s1 i1 v)
3002 step _ (InterposeSecondYield s1 i1) = do
3004 return $ Yield r (InterposeFirstInner s1 i1)
3007 step _ (InterposeFirstResume s1 i1 v) = do
3008 return $ Yield v (InterposeFirstInner s1 i1)
3011------------------------------------------------------------------------------
3013------------------------------------------------------------------------------
3015data GbracketState s1 s2 v
3017 | GBracketNormal s1 v
3018 | GBracketException s2
3020-- | The most general bracketing and exception combinator. All other
3021-- combinators can be expressed in terms of this combinator. This can also be
3022-- used for cases which are not covered by the standard combinators.
3026{-# INLINE_NORMAL gbracket #-}
3030 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
3031 -> (c -> m d) -- ^ after, on normal stop
3032 -> (c -> e -> Stream m b) -- ^ on exception
3033 -> (c -> Stream m b) -- ^ stream generator
3035gbracket bef exc aft fexc fnormal =
3036 Stream step GBracketInit
3040 {-# INLINE_LATE step #-}
3041 step _ GBracketInit = do
3043 return $ Skip $ GBracketNormal (fnormal r) r
3045 step gst (GBracketNormal (UnStream step1 st) v) = do
3046 res <- exc $ step1 gst st
3048 Right r -> case r of
3050 return $ Yield x (GBracketNormal (Stream step1 s) v)
3051 Skip s -> return $ Skip (GBracketNormal (Stream step1 s) v)
3052 Stop -> aft v >> return Stop
3053 Left e -> return $ Skip (GBracketException (fexc v e))
3054 step gst (GBracketException (UnStream step1 st)) = do
3057 Yield x s -> return $ Yield x (GBracketException (Stream step1 s))
3058 Skip s -> return $ Skip (GBracketException (Stream step1 s))
3061-- | Create an IORef holding a finalizer that is called automatically when the
3062-- IORef is garbage collected. The IORef can be written to with a 'Nothing'
3063-- value to deactivate the finalizer.
3064newFinalizedIORef :: (MonadIO m, MonadBaseControl IO m)
3065 => m a -> m (IORef (Maybe (IO ())))
3066newFinalizedIORef finalizer = do
3067 mrun <- captureMonadState
3068 ref <- liftIO $ newIORef $ Just $ liftIO $ void $ do
3069 _ <- runInIO mrun finalizer
3072 res <- readIORef ref
3074 Nothing -> return ()
3076 _ <- liftIO $ mkWeakIORef ref finalizer1
3079-- | Run the finalizer stored in an IORef and deactivate it so that it is run
3082runIORefFinalizer :: MonadIO m => IORef (Maybe (IO ())) -> m ()
3083runIORefFinalizer ref = liftIO $ do
3084 res <- readIORef ref
3086 Nothing -> return ()
3087 Just f -> writeIORef ref Nothing >> f
3089-- | Deactivate the finalizer stored in an IORef without running it.
3091clearIORefFinalizer :: MonadIO m => IORef (Maybe (IO ())) -> m ()
3092clearIORefFinalizer ref = liftIO $ writeIORef ref Nothing
3094data GbracketIOState s1 s2 v wref
3096 | GBracketIONormal s1 v wref
3097 | GBracketIOException s2
3099-- | Like gbracket but also uses a finalizer to make sure when the stream is
3100-- garbage collected we run the finalizing action. This requires a MonadIO and
3101-- MonadBaseControl IO constraint.
3103-- | The most general bracketing and exception combinator. All other
3104-- combinators can be expressed in terms of this combinator. This can also be
3105-- used for cases which are not covered by the standard combinators.
3109{-# INLINE_NORMAL gbracketIO #-}
3111 :: (MonadIO m, MonadBaseControl IO m)
3113 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
3114 -> (c -> m d) -- ^ after, on normal stop or GC
3115 -> (c -> e -> Stream m b) -- ^ on exception
3116 -> (c -> Stream m b) -- ^ stream generator
3118gbracketIO bef exc aft fexc fnormal =
3119 Stream step GBracketIOInit
3123 -- If the stream is never evaluated the "aft" action will never be
3124 -- called. For that to occur we will need the user of this API to pass a
3125 -- weak pointer to us.
3126 {-# INLINE_LATE step #-}
3127 step _ GBracketIOInit = do
3128 -- We mask asynchronous exceptions to make the execution
3129 -- of 'bef' and the registration of 'aft' atomic.
3130 -- A similar thing is done in the resourcet package: https://git.io/JvKV3
3131 -- Tutorial: https://markkarpov.com/tutorial/exceptions.html
3132 (r, ref) <- liftBaseOp_ mask_ $ do
3134 ref <- newFinalizedIORef (aft r)
3136 return $ Skip $ GBracketIONormal (fnormal r) r ref
3138 step gst (GBracketIONormal (UnStream step1 st) v ref) = do
3139 res <- exc $ step1 gst st
3141 Right r -> case r of
3143 return $ Yield x (GBracketIONormal (Stream step1 s) v ref)
3145 return $ Skip (GBracketIONormal (Stream step1 s) v ref)
3147 runIORefFinalizer ref
3150 clearIORefFinalizer ref
3151 return $ Skip (GBracketIOException (fexc v e))
3152 step gst (GBracketIOException (UnStream step1 st)) = do
3156 return $ Yield x (GBracketIOException (Stream step1 s))
3157 Skip s -> return $ Skip (GBracketIOException (Stream step1 s))
3160-- | Run a side effect before the stream yields its first element.
3161{-# INLINE_NORMAL before #-}
3162before :: Monad m => m b -> Stream m a -> Stream m a
3163before action (Stream step state) = Stream step' Nothing
3167 {-# INLINE_LATE step' #-}
3168 step' _ Nothing = action >> return (Skip (Just state))
3170 step' gst (Just st) = do
3173 Yield x s -> return $ Yield x (Just s)
3174 Skip s -> return $ Skip (Just s)
3177-- | Run a side effect whenever the stream stops normally.
3178{-# INLINE_NORMAL after #-}
3179after :: Monad m => m b -> Stream m a -> Stream m a
3180after action (Stream step state) = Stream step' state
3184 {-# INLINE_LATE step' #-}
3188 Yield x s -> return $ Yield x s
3189 Skip s -> return $ Skip s
3190 Stop -> action >> return Stop
3192{-# INLINE_NORMAL afterIO #-}
3193afterIO :: (MonadIO m, MonadBaseControl IO m)
3194 => m b -> Stream m a -> Stream m a
3195afterIO action (Stream step state) = Stream step' Nothing
3199 {-# INLINE_LATE step' #-}
3200 step' _ Nothing = do
3201 ref <- newFinalizedIORef action
3202 return $ Skip $ Just (state, ref)
3203 step' gst (Just (st, ref)) = do
3206 Yield x s -> return $ Yield x (Just (s, ref))
3207 Skip s -> return $ Skip (Just (s, ref))
3209 runIORefFinalizer ref
3212-- XXX These combinators are expensive due to the call to
3213-- onException/handle/try on each step. Therefore, when possible, they should
3214-- be called in an outer loop where we perform less iterations. For example, we
3215-- cannot call them on each iteration in a char stream, instead we can call
3216-- them when doing an IO on an array.
3218-- XXX For high performance error checks in busy streams we may need another
3219-- Error constructor in step.
3221-- | Run a side effect whenever the stream aborts due to an exception. The
3222-- exception is not caught, simply rethrown.
3223{-# INLINE_NORMAL onException #-}
3224onException :: MonadCatch m => m b -> Stream m a -> Stream m a
3225onException action str =
3226 gbracket (return ()) MC.try return
3227 (\_ (e :: MC.SomeException) -> nilM (action >> MC.throwM e))
3230{-# INLINE_NORMAL _onException #-}
3231_onException :: MonadCatch m => m b -> Stream m a -> Stream m a
3232_onException action (Stream step state) = Stream step' state
3236 {-# INLINE_LATE step' #-}
3238 res <- step gst st `MC.onException` action
3240 Yield x s -> return $ Yield x s
3241 Skip s -> return $ Skip s
3244-- XXX bracket is like concatMap, it generates a stream and then flattens it.
3245-- Like concatMap it has 10x worse performance compared to linear fused
3248-- | Run the first action before the stream starts and remember its output,
3249-- generate a stream using the output, run the second action providing the
3250-- remembered value as an argument whenever the stream ends normally or due to
3252{-# INLINE_NORMAL bracket #-}
3253bracket :: MonadCatch m => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
3254bracket bef aft bet =
3255 gbracket bef MC.try aft
3256 (\a (e :: SomeException) -> nilM (aft a >> MC.throwM e)) bet
3258{-# INLINE_NORMAL bracketIO #-}
3259bracketIO :: (MonadAsync m, MonadCatch m)
3260 => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
3261bracketIO bef aft bet =
3262 gbracketIO bef MC.try aft
3263 (\a (e :: SomeException) -> nilM (aft a >> MC.throwM e)) bet
3265data BracketState s v = BracketInit | BracketRun s v
3267{-# INLINE_NORMAL _bracket #-}
3268_bracket :: MonadCatch m => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
3269_bracket bef aft bet = Stream step' BracketInit
3273 {-# INLINE_LATE step' #-}
3274 step' _ BracketInit = bef >>= \x -> return (Skip (BracketRun (bet x) x))
3276 -- NOTE: It is important to use UnStream instead of the Stream pattern
3277 -- here, otherwise we get huge perf degradation, see note in concatMap.
3278 step' gst (BracketRun (UnStream step state) v) = do
3279 -- res <- step gst state `MC.onException` aft v
3280 res <- MC.try $ step gst state
3282 Left (e :: SomeException) -> aft v >> MC.throwM e >> return Stop
3283 Right r -> case r of
3284 Yield x s -> return $ Yield x (BracketRun (Stream step s) v)
3285 Skip s -> return $ Skip (BracketRun (Stream step s) v)
3286 Stop -> aft v >> return Stop
3288-- | Run a side effect whenever the stream stops normally or aborts due to an
3290{-# INLINE finally #-}
3291finally :: MonadCatch m => m b -> Stream m a -> Stream m a
3292-- finally action xs = after action $ onException action xs
3293finally action xs = bracket (return ()) (\_ -> action) (const xs)
3295{-# INLINE finallyIO #-}
3296finallyIO :: (MonadAsync m, MonadCatch m) => m b -> Stream m a -> Stream m a
3297finallyIO action xs = bracketIO (return ()) (\_ -> action) (const xs)
3299-- | When evaluating a stream if an exception occurs, stream evaluation aborts
3300-- and the specified exception handler is run with the exception as argument.
3301{-# INLINE_NORMAL handle #-}
3302handle :: (MonadCatch m, Exception e)
3303 => (e -> Stream m a) -> Stream m a -> Stream m a
3305 gbracket (return ()) MC.try return (\_ e -> f e) (\_ -> str)
3307{-# INLINE_NORMAL _handle #-}
3308_handle :: (MonadCatch m, Exception e)
3309 => (e -> Stream m a) -> Stream m a -> Stream m a
3310_handle f (Stream step state) = Stream step' (Left state)
3314 {-# INLINE_LATE step' #-}
3315 step' gst (Left st) = do
3316 res <- MC.try $ step gst st
3318 Left e -> return $ Skip $ Right (f e)
3319 Right r -> case r of
3320 Yield x s -> return $ Yield x (Left s)
3321 Skip s -> return $ Skip (Left s)
3324 step' gst (Right (UnStream step1 st)) = do
3327 Yield x s -> return $ Yield x (Right (Stream step1 s))
3328 Skip s -> return $ Skip (Right (Stream step1 s))
3331-------------------------------------------------------------------------------
3332-- General transformation
3333-------------------------------------------------------------------------------
3335{-# INLINE_NORMAL transform #-}
3336transform :: Monad m => Pipe m a b -> Stream m a -> Stream m b
3337transform (Pipe pstep1 pstep2 pstate) (Stream step state) =
3338 Stream step' (Consume pstate, state)
3342 {-# INLINE_LATE step' #-}
3344 step' gst (Consume pst, st) = pst `seq` do
3345 r <- step (adaptState gst) st
3350 Pipe.Yield b pst' -> return $ Yield b (pst', s)
3351 Pipe.Continue pst' -> return $ Skip (pst', s)
3352 Skip s -> return $ Skip (Consume pst, s)
3355 step' _ (Produce pst, st) = pst `seq` do
3358 Pipe.Yield b pst' -> return $ Yield b (pst', st)
3359 Pipe.Continue pst' -> return $ Skip (pst', st)
3361------------------------------------------------------------------------------
3362-- Transformation by Folding (Scans)
3363------------------------------------------------------------------------------
3365------------------------------------------------------------------------------
3367------------------------------------------------------------------------------
3369-- XXX Is a prescan useful, discarding the last step does not sound useful? I
3370-- am not sure about the utility of this function, so this is implemented but
3371-- not exposed. We can expose it if someone provides good reasons why this is
3374-- XXX We have to execute the stream one step ahead to know that we are at the
3375-- last step. The vector implementation of prescan executes the last fold step
3376-- but does not yield the result. This means we have executed the effect but
3377-- discarded value. This does not sound right. In this implementation we are
3378-- not executing the last fold step.
3379{-# INLINE_NORMAL prescanlM' #-}
3380prescanlM' :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
3381prescanlM' f mz (Stream step state) = Stream step' (state, mz)
3383 {-# INLINE_LATE step' #-}
3384 step' gst (st, prev) = do
3385 r <- step (adaptState gst) st
3389 return $ Yield acc (s, f acc x)
3390 Skip s -> return $ Skip (s, prev)
3393{-# INLINE prescanl' #-}
3394prescanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
3395prescanl' f z = prescanlM' (\a b -> return (f a b)) (return z)
3397------------------------------------------------------------------------------
3398-- Monolithic postscans (postscan followed by a map)
3399------------------------------------------------------------------------------
3401-- The performance of a modular postscan followed by a map seems to be
3402-- equivalent to this monolithic scan followed by map therefore we may not need
3403-- this implementation. We just have it for performance comparison and in case
3404-- modular version does not perform well in some situation.
3406{-# INLINE_NORMAL postscanlMx' #-}
3407postscanlMx' :: Monad m
3408 => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> Stream m b
3409postscanlMx' fstep begin done (Stream step state) = do
3410 Stream step' (state, begin)
3412 {-# INLINE_LATE step' #-}
3413 step' gst (st, acc) = do
3414 r <- step (adaptState gst) st
3420 v `seq` y `seq` return (Yield v (s, return y))
3421 Skip s -> return $ Skip (s, acc)
3424{-# INLINE_NORMAL postscanlx' #-}
3425postscanlx' :: Monad m
3426 => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> Stream m b
3427postscanlx' fstep begin done s =
3428 postscanlMx' (\b a -> return (fstep b a)) (return begin) (return . done) s
3430-- XXX do we need consM strict to evaluate the begin value?
3431{-# INLINE scanlMx' #-}
3433 => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> Stream m b
3434scanlMx' fstep begin done s =
3435 (begin >>= \x -> x `seq` done x) `consM` postscanlMx' fstep begin done s
3437{-# INLINE scanlx' #-}
3439 => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> Stream m b
3440scanlx' fstep begin done s =
3441 scanlMx' (\b a -> return (fstep b a)) (return begin) (return . done) s
3443------------------------------------------------------------------------------
3445------------------------------------------------------------------------------
3447{-# INLINE_NORMAL postscanlM' #-}
3448postscanlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
3449postscanlM' fstep begin (Stream step state) =
3450 begin `seq` Stream step' (state, begin)
3452 {-# INLINE_LATE step' #-}
3453 step' gst (st, acc) = acc `seq` do
3454 r <- step (adaptState gst) st
3458 y `seq` return (Yield y (s, y))
3459 Skip s -> return $ Skip (s, acc)
3462{-# INLINE_NORMAL postscanl' #-}
3463postscanl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
3464postscanl' f = postscanlM' (\a b -> return (f a b))
3466{-# INLINE_NORMAL postscanlM #-}
3467postscanlM :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
3468postscanlM fstep begin (Stream step state) = Stream step' (state, begin)
3470 {-# INLINE_LATE step' #-}
3471 step' gst (st, acc) = do
3472 r <- step (adaptState gst) st
3476 return (Yield y (s, y))
3477 Skip s -> return $ Skip (s, acc)
3480{-# INLINE_NORMAL postscanl #-}
3481postscanl :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
3482postscanl f = postscanlM (\a b -> return (f a b))
3484{-# INLINE_NORMAL scanlM' #-}
3485scanlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
3486scanlM' fstep begin s = begin `seq` (begin `cons` postscanlM' fstep begin s)
3488{-# INLINE scanl' #-}
3489scanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
3490scanl' f = scanlM' (\a b -> return (f a b))
3492{-# INLINE_NORMAL scanlM #-}
3493scanlM :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
3494scanlM fstep begin s = begin `cons` postscanlM fstep begin s
3497scanl :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
3498scanl f = scanlM (\a b -> return (f a b))
3500{-# INLINE_NORMAL scanl1M #-}
3501scanl1M :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
3502scanl1M fstep (Stream step state) = Stream step' (state, Nothing)
3504 {-# INLINE_LATE step' #-}
3505 step' gst (st, Nothing) = do
3508 Yield x s -> return $ Yield x (s, Just x)
3509 Skip s -> return $ Skip (s, Nothing)
3512 step' gst (st, Just acc) = do
3517 return $ Yield z (s, Just z)
3518 Skip s -> return $ Skip (s, Just acc)
3521{-# INLINE scanl1 #-}
3522scanl1 :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
3523scanl1 f = scanl1M (\x y -> return (f x y))
3525{-# INLINE_NORMAL scanl1M' #-}
3526scanl1M' :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
3527scanl1M' fstep (Stream step state) = Stream step' (state, Nothing)
3529 {-# INLINE_LATE step' #-}
3530 step' gst (st, Nothing) = do
3533 Yield x s -> x `seq` return $ Yield x (s, Just x)
3534 Skip s -> return $ Skip (s, Nothing)
3537 step' gst (st, Just acc) = acc `seq` do
3542 z `seq` return $ Yield z (s, Just z)
3543 Skip s -> return $ Skip (s, Just acc)
3546{-# INLINE scanl1' #-}
3547scanl1' :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
3548scanl1' f = scanl1M' (\x y -> return (f x y))
3550------------------------------------------------------------------------------
3552------------------------------------------------------------------------------
3554data RollingMapState s a = RollingMapInit s | RollingMapGo s a
3556{-# INLINE rollingMapM #-}
3557rollingMapM :: Monad m => (a -> a -> m b) -> Stream m a -> Stream m b
3558rollingMapM f (Stream step1 state1) = Stream step (RollingMapInit state1)
3560 step gst (RollingMapInit st) = do
3561 r <- step1 (adaptState gst) st
3563 Yield x s -> Skip $ RollingMapGo s x
3564 Skip s -> Skip $ RollingMapInit s
3567 step gst (RollingMapGo s1 x1) = do
3568 r <- step1 (adaptState gst) s1
3572 return $ Yield res $ RollingMapGo s x
3573 Skip s -> return $ Skip $ RollingMapGo s x1
3574 Stop -> return $ Stop
3576{-# INLINE rollingMap #-}
3577rollingMap :: Monad m => (a -> a -> b) -> Stream m a -> Stream m b
3578rollingMap f = rollingMapM (\x y -> return $ f x y)
3580------------------------------------------------------------------------------
3581-- Tapping/Distributing
3582------------------------------------------------------------------------------
3585tap :: Monad m => Fold m a b -> Stream m a -> Stream m a
3586tap (Fold fstep initial extract) (Stream step state) = Stream step' Nothing
3590 step' _ Nothing = do
3592 return $ Skip (Just (r, state))
3594 step' gst (Just (acc, st)) = acc `seq` do
3599 return $ Yield x (Just (acc', s))
3600 Skip s -> return $ Skip (Just (acc, s))
3605{-# INLINE_NORMAL tapOffsetEvery #-}
3606tapOffsetEvery :: Monad m
3607 => Int -> Int -> Fold m a b -> Stream m a -> Stream m a
3608tapOffsetEvery offset n (Fold fstep initial extract) (Stream step state) =
3609 Stream step' Nothing
3613 {-# INLINE_LATE step' #-}
3614 step' _ Nothing = do
3616 return $ Skip (Just (r, state, offset `mod` n))
3618 step' gst (Just (acc, st, count)) | count <= 0 = do
3622 !acc' <- fstep acc x
3623 return $ Yield x (Just (acc', s, n - 1))
3624 Skip s -> return $ Skip (Just (acc, s, count))
3629 step' gst (Just (acc, st, count)) = do
3632 Yield x s -> return $ Yield x (Just (acc, s, count - 1))
3633 Skip s -> return $ Skip (Just (acc, s, count))
3638{-# INLINE_NORMAL pollCounts #-}
3642 -> (Stream m Int -> Stream m Int)
3646pollCounts predicate transf fld (Stream step state) = Stream step' Nothing
3649 {-# INLINE_LATE step' #-}
3650 step' _ Nothing = do
3651 -- As long as we are using an "Int" for counts lockfree reads from
3652 -- Var should work correctly on both 32-bit and 64-bit machines.
3653 -- However, an Int on a 32-bit machine may overflow quickly.
3654 countVar <- liftIO $ newVar (0 :: Int)
3656 $ void $ runFold fld
3657 $ transf $ fromPrimVar countVar
3658 return $ Skip (Just (countVar, tid, state))
3660 step' gst (Just (countVar, tid, st)) = do
3664 when (predicate x) $ liftIO $ modifyVar' countVar (+ 1)
3665 return $ Yield x (Just (countVar, tid, s))
3666 Skip s -> return $ Skip (Just (countVar, tid, s))
3668 liftIO $ killThread tid
3671{-# INLINE_NORMAL tapRate #-}
3673 (MonadAsync m, MonadCatch m)
3678tapRate samplingRate action (Stream step state) = Stream step' Nothing
3680 {-# NOINLINE loop #-}
3681 loop countVar prev = do
3684 (do liftIO $ threadDelay (round $ samplingRate * 1000000)
3685 i <- liftIO $ readVar countVar
3686 let !diff = i - prev
3689 (\(e :: AsyncException) -> do
3690 i <- liftIO $ readVar countVar
3691 let !diff = i - prev
3693 throwM (MC.toException e))
3696 {-# INLINE_LATE step' #-}
3697 step' _ Nothing = do
3698 countVar <- liftIO $ newVar 0
3699 tid <- fork $ loop countVar 0
3700 ref <- liftIO $ newIORef ()
3701 _ <- liftIO $ mkWeakIORef ref (killThread tid)
3702 return $ Skip (Just (countVar, tid, state, ref))
3704 step' gst (Just (countVar, tid, st, ref)) = do
3708 liftIO $ modifyVar' countVar (+ 1)
3709 return $ Yield x (Just (countVar, tid, s, ref))
3710 Skip s -> return $ Skip (Just (countVar, tid, s, ref))
3712 liftIO $ killThread tid
3716-------------------------------------------------------------------------------
3718-------------------------------------------------------------------------------
3720{-# INLINE_NORMAL takeWhileM #-}
3721takeWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
3722takeWhileM f (Stream step state) = Stream step' state
3724 {-# INLINE_LATE step' #-}
3730 return $ if b then Yield x s else Stop
3731 Skip s -> return $ Skip s
3734{-# INLINE takeWhile #-}
3735takeWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
3736takeWhile f = takeWhileM (return . f)
3738{-# INLINE_NORMAL drop #-}
3739drop :: Monad m => Int -> Stream m a -> Stream m a
3740drop n (Stream step state) = Stream step' (state, Just n)
3742 {-# INLINE_LATE step' #-}
3743 step' gst (st, Just i)
3748 Yield _ s -> Skip (s, Just (i - 1))
3749 Skip s -> Skip (s, Just i)
3751 | otherwise = return $ Skip (st, Nothing)
3753 step' gst (st, Nothing) = do
3757 Yield x s -> Yield x (s, Nothing)
3758 Skip s -> Skip (s, Nothing)
3761data DropWhileState s a
3763 | DropWhileYield a s
3766{-# INLINE_NORMAL dropWhileM #-}
3767dropWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
3768dropWhileM f (Stream step state) = Stream step' (DropWhileDrop state)
3770 {-# INLINE_LATE step' #-}
3771 step' gst (DropWhileDrop st) = do
3777 then return $ Skip (DropWhileDrop s)
3778 else return $ Skip (DropWhileYield x s)
3779 Skip s -> return $ Skip (DropWhileDrop s)
3782 step' gst (DropWhileNext st) = do
3785 Yield x s -> return $ Skip (DropWhileYield x s)
3786 Skip s -> return $ Skip (DropWhileNext s)
3789 step' _ (DropWhileYield x st) = return $ Yield x (DropWhileNext st)
3791{-# INLINE dropWhile #-}
3792dropWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
3793dropWhile f = dropWhileM (return . f)
3795{-# INLINE_NORMAL filterM #-}
3796filterM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
3797filterM f (Stream step state) = Stream step' state
3799 {-# INLINE_LATE step' #-}
3808 Skip s -> return $ Skip s
3811{-# INLINE filter #-}
3812filter :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
3813filter f = filterM (return . f)
3815{-# INLINE_NORMAL uniq #-}
3816uniq :: (Eq a, Monad m) => Stream m a -> Stream m a
3817uniq (Stream step state) = Stream step' (Nothing, state)
3819 {-# INLINE_LATE step' #-}
3820 step' gst (Nothing, st) = do
3823 Yield x s -> return $ Yield x (Just x, s)
3824 Skip s -> return $ Skip (Nothing, s)
3826 step' gst (Just x, st) = do
3829 Yield y s | x == y -> return $ Skip (Just x, s)
3830 | otherwise -> return $ Yield y (Just y, s)
3831 Skip s -> return $ Skip (Just x, s)
3834------------------------------------------------------------------------------
3835-- Transformation by Mapping
3836------------------------------------------------------------------------------
3838{-# INLINE_NORMAL sequence #-}
3839sequence :: Monad m => Stream m (m a) -> Stream m a
3840sequence (Stream step state) = Stream step' state
3842 {-# INLINE_LATE step' #-}
3844 r <- step (adaptState gst) st
3846 Yield x s -> x >>= \a -> return (Yield a s)
3847 Skip s -> return $ Skip s
3850------------------------------------------------------------------------------
3852------------------------------------------------------------------------------
3854data LoopState x s = FirstYield s
3855 | InterspersingYield s
3858{-# INLINE_NORMAL intersperseM #-}
3859intersperseM :: Monad m => m a -> Stream m a -> Stream m a
3860intersperseM m (Stream step state) = Stream step' (FirstYield state)
3862 {-# INLINE_LATE step' #-}
3863 step' gst (FirstYield st) = do
3867 Yield x s -> Skip (YieldAndCarry x s)
3868 Skip s -> Skip (FirstYield s)
3871 step' gst (InterspersingYield st) = do
3876 return $ Yield a (YieldAndCarry x s)
3877 Skip s -> return $ Skip $ InterspersingYield s
3880 step' _ (YieldAndCarry x st) = return $ Yield x (InterspersingYield st)
3885 | SuffixYield a (SuffixState s a)
3887{-# INLINE_NORMAL intersperseSuffix #-}
3888intersperseSuffix :: forall m a. Monad m => m a -> Stream m a -> Stream m a
3889intersperseSuffix action (Stream step state) = Stream step' (SuffixElem state)
3891 {-# INLINE_LATE step' #-}
3892 step' gst (SuffixElem st) = do
3895 Yield x s -> Skip (SuffixYield x (SuffixSuffix s))
3896 Skip s -> Skip (SuffixElem s)
3899 step' _ (SuffixSuffix st) = do
3900 action >>= \r -> return $ Skip (SuffixYield r (SuffixElem st))
3902 step' _ (SuffixYield x next) = return $ Yield x next
3904data SuffixSpanState s a
3905 = SuffixSpanElem s Int
3906 | SuffixSpanSuffix s
3907 | SuffixSpanYield a (SuffixSpanState s a)
3911-- | intersperse after every n items
3912{-# INLINE_NORMAL intersperseSuffixBySpan #-}
3913intersperseSuffixBySpan :: forall m a. Monad m
3914 => Int -> m a -> Stream m a -> Stream m a
3915intersperseSuffixBySpan n action (Stream step state) =
3916 Stream step' (SuffixSpanElem state n)
3918 {-# INLINE_LATE step' #-}
3919 step' gst (SuffixSpanElem st i) | i > 0 = do
3922 Yield x s -> Skip (SuffixSpanYield x (SuffixSpanElem s (i - 1)))
3923 Skip s -> Skip (SuffixSpanElem s i)
3924 Stop -> if i == n then Stop else Skip SuffixSpanLast
3925 step' _ (SuffixSpanElem st _) = return $ Skip (SuffixSpanSuffix st)
3927 step' _ (SuffixSpanSuffix st) = do
3928 action >>= \r -> return $ Skip (SuffixSpanYield r (SuffixSpanElem st n))
3930 step' _ (SuffixSpanLast) = do
3931 action >>= \r -> return $ Skip (SuffixSpanYield r SuffixSpanStop)
3933 step' _ (SuffixSpanYield x next) = return $ Yield x next
3935 step' _ (SuffixSpanStop) = return Stop
3937{-# INLINE intersperse #-}
3938intersperse :: Monad m => a -> Stream m a -> Stream m a
3939intersperse a = intersperseM (return a)
3941{-# INLINE_NORMAL insertBy #-}
3942insertBy :: Monad m => (a -> a -> Ordering) -> a -> Stream m a -> Stream m a
3943insertBy cmp a (Stream step state) = Stream step' (state, False, Nothing)
3945 {-# INLINE_LATE step' #-}
3946 step' gst (st, False, _) = do
3949 Yield x s -> case cmp a x of
3950 GT -> return $ Yield x (s, False, Nothing)
3951 _ -> return $ Yield a (s, True, Just x)
3952 Skip s -> return $ Skip (s, False, Nothing)
3953 Stop -> return $ Yield a (st, True, Nothing)
3955 step' _ (_, True, Nothing) = return Stop
3957 step' gst (st, True, Just prev) = do
3960 Yield x s -> return $ Yield prev (s, True, Just x)
3961 Skip s -> return $ Skip (s, True, Just prev)
3962 Stop -> return $ Yield prev (st, True, Nothing)
3964------------------------------------------------------------------------------
3966------------------------------------------------------------------------------
3968{-# INLINE_NORMAL deleteBy #-}
3969deleteBy :: Monad m => (a -> a -> Bool) -> a -> Stream m a -> Stream m a
3970deleteBy eq x (Stream step state) = Stream step' (state, False)
3972 {-# INLINE_LATE step' #-}
3973 step' gst (st, False) = do
3976 Yield y s -> return $
3977 if eq x y then Skip (s, True) else Yield y (s, False)
3978 Skip s -> return $ Skip (s, False)
3981 step' gst (st, True) = do
3984 Yield y s -> return $ Yield y (s, True)
3985 Skip s -> return $ Skip (s, True)
3988------------------------------------------------------------------------------
3989-- Transformation by Map and Filter
3990------------------------------------------------------------------------------
3992-- XXX Will this always fuse properly?
3993{-# INLINE_NORMAL mapMaybe #-}
3994mapMaybe :: Monad m => (a -> Maybe b) -> Stream m a -> Stream m b
3995mapMaybe f = fmap fromJust . filter isJust . map f
3997{-# INLINE_NORMAL mapMaybeM #-}
3998mapMaybeM :: Monad m => (a -> m (Maybe b)) -> Stream m a -> Stream m b
3999mapMaybeM f = fmap fromJust . filter isJust . mapM f
4001------------------------------------------------------------------------------
4003------------------------------------------------------------------------------
4005{-# INLINE_NORMAL indexed #-}
4006indexed :: Monad m => Stream m a -> Stream m (Int, a)
4007indexed (Stream step state) = Stream step' (state, 0)
4009 {-# INLINE_LATE step' #-}
4010 step' gst (st, i) = i `seq` do
4011 r <- step (adaptState gst) st
4013 Yield x s -> return $ Yield (i, x) (s, i+1)
4014 Skip s -> return $ Skip (s, i)
4017{-# INLINE_NORMAL indexedR #-}
4018indexedR :: Monad m => Int -> Stream m a -> Stream m (Int, a)
4019indexedR m (Stream step state) = Stream step' (state, m)
4021 {-# INLINE_LATE step' #-}
4022 step' gst (st, i) = i `seq` do
4023 r <- step (adaptState gst) st
4025 Yield x s -> let i' = i - 1
4026 in return $ Yield (i, x) (s, i')
4027 Skip s -> return $ Skip (s, i)
4030{-# INLINE_NORMAL zipWithM #-}
4032 => (a -> b -> m c) -> Stream m a -> Stream m b -> Stream m c
4033zipWithM f (Stream stepa ta) (Stream stepb tb) = Stream step (ta, tb, Nothing)
4035 {-# INLINE_LATE step #-}
4036 step gst (sa, sb, Nothing) = do
4037 r <- stepa (adaptState gst) sa
4040 Yield x sa' -> Skip (sa', sb, Just x)
4041 Skip sa' -> Skip (sa', sb, Nothing)
4044 step gst (sa, sb, Just x) = do
4045 r <- stepb (adaptState gst) sb
4049 return $ Yield z (sa, sb', Nothing)
4050 Skip sb' -> return $ Skip (sa, sb', Just x)
4053#if __GLASGOW_HASKELL__ >= 801
4054{-# RULES "zipWithM xs xs"
4055 forall f xs. zipWithM @Identity f xs xs = mapM (\x -> f x x) xs #-}
4058{-# INLINE zipWith #-}
4059zipWith :: Monad m => (a -> b -> c) -> Stream m a -> Stream m b -> Stream m c
4060zipWith f = zipWithM (\a b -> return (f a b))
4062------------------------------------------------------------------------------
4064------------------------------------------------------------------------------
4066{-# INLINE_NORMAL mergeByM #-}
4069 => (a -> a -> m Ordering) -> Stream m a -> Stream m a -> Stream m a
4070mergeByM cmp (Stream stepa ta) (Stream stepb tb) =
4071 Stream step (Just ta, Just tb, Nothing, Nothing)
4073 {-# INLINE_LATE step #-}
4075 -- one of the values is missing, and the corresponding stream is running
4076 step gst (Just sa, sb, Nothing, b) = do
4079 Yield a sa' -> Skip (Just sa', sb, Just a, b)
4080 Skip sa' -> Skip (Just sa', sb, Nothing, b)
4081 Stop -> Skip (Nothing, sb, Nothing, b)
4083 step gst (sa, Just sb, a, Nothing) = do
4086 Yield b sb' -> Skip (sa, Just sb', a, Just b)
4087 Skip sb' -> Skip (sa, Just sb', a, Nothing)
4088 Stop -> Skip (sa, Nothing, a, Nothing)
4090 -- both the values are available
4091 step _ (sa, sb, Just a, Just b) = do
4093 return $ case res of
4094 GT -> Yield b (sa, sb, Just a, Nothing)
4095 _ -> Yield a (sa, sb, Nothing, Just b)
4097 -- one of the values is missing, corresponding stream is done
4098 step _ (Nothing, sb, Nothing, Just b) =
4099 return $ Yield b (Nothing, sb, Nothing, Nothing)
4101 step _ (sa, Nothing, Just a, Nothing) =
4102 return $ Yield a (sa, Nothing, Nothing, Nothing)
4104 step _ (Nothing, Nothing, Nothing, Nothing) = return Stop
4106{-# INLINE mergeBy #-}
4109 => (a -> a -> Ordering) -> Stream m a -> Stream m a -> Stream m a
4110mergeBy cmp = mergeByM (\a b -> return $ cmp a b)
4112------------------------------------------------------------------------------
4113-- Transformation comprehensions
4114------------------------------------------------------------------------------
4116{-# INLINE_NORMAL the #-}
4117the :: (Eq a, Monad m) => Stream m a -> m (Maybe a)
4118the (Stream step state) = go state
4121 r <- step defState st
4123 Yield x s -> go' x s
4125 Stop -> return Nothing
4127 r <- step defState st
4129 Yield x s | x == n -> go' n s
4130 | otherwise -> return Nothing
4132 Stop -> return (Just n)
4134{-# INLINE runFold #-}
4135runFold :: (Monad m) => Fold m a b -> Stream m a -> m b
4136runFold (Fold step begin done) = foldlMx' step begin done
4138-------------------------------------------------------------------------------
4139-- Concurrent application and fold
4140-------------------------------------------------------------------------------
4142-- XXX These functions should be moved to Stream/Parallel.hs
4144-- Using StreamD the worker stream producing code can fuse with the code to
4145-- queue output to the SVar giving some perf boost.
4147-- Note that StreamD can only be used in limited situations, specifically, we
4148-- cannot implement joinStreamVarPar using this.
4150-- XXX make sure that the SVar passed is a Parallel style SVar.
4152-- | Fold the supplied stream to the SVar asynchronously using Parallel
4153-- concurrency style.
4154-- {-# INLINE_NORMAL toSVarParallel #-}
4155{-# INLINE toSVarParallel #-}
4156toSVarParallel :: MonadAsync m
4157 => State t m a -> SVar t m a -> Stream m a -> m ()
4158toSVarParallel st sv xs =
4159 if svarInspectMode sv
4163 case getYieldLimit st of
4164 Nothing -> doFork (work Nothing)
4166 (handleChildException sv)
4167 Just _ -> doFork (workLim Nothing)
4169 (handleChildException sv)
4174 {-# NOINLINE work #-}
4175 work info = (runFold (FL.toParallelSVar sv info) xs)
4177 {-# NOINLINE workLim #-}
4178 workLim info = runFold (FL.toParallelSVarLimited sv info) xs
4180 {-# NOINLINE forkWithDiag #-}
4182 -- We do not use workerCount in case of ParallelVar but still there is
4183 -- no harm in maintaining it correctly.
4184 liftIO $ atomicModifyIORefCAS_ (workerCount sv) $ \n -> n + 1
4186 -- This allocation matters when significant number of workers are being
4187 -- sent. We allocate it only when needed. The overhead increases by 4x.
4189 case yieldRateInfo sv of
4190 Nothing -> return Nothing
4191 Just _ -> liftIO $ do
4192 cntRef <- newIORef 0
4193 t <- getTime Monotonic
4194 lat <- newIORef (0, t)
4195 return $ Just WorkerInfo
4196 { workerYieldMax = 0
4197 , workerYieldCount = cntRef
4198 , workerLatencyStart = lat
4201 case getYieldLimit st of
4202 Nothing -> doFork (work winfo)
4204 (handleChildException sv)
4205 Just _ -> doFork (workLim winfo)
4207 (handleChildException sv)
4210{-# INLINE_NORMAL mkParallelD #-}
4211mkParallelD :: MonadAsync m => Stream m a -> Stream m a
4212mkParallelD m = Stream step Nothing
4215 step gst Nothing = do
4216 sv <- newParallelVar StopNone gst
4217 toSVarParallel gst sv m
4218 -- XXX use unfold instead?
4219 return $ Skip $ Just $ fromSVar sv
4221 step gst (Just (UnStream step1 st)) = do
4224 Yield a s -> Yield a (Just $ Stream step1 s)
4225 Skip s -> Skip (Just $ Stream step1 s)
4228-- Compare with mkAsync. mkAsync uses an Async style SVar whereas this uses a
4229-- parallel style SVar for evaluation. Currently, parallel style cannot use
4230-- rate control whereas Async style can use rate control. In async style SVar
4231-- the worker thread terminates when the buffer is full whereas in Parallel
4234-- | Make the stream producer and consumer run concurrently by introducing a
4235-- buffer between them. The producer thread evaluates the input stream until
4236-- the buffer fills, it blocks if the buffer is full until there is space in
4237-- the buffer. The consumer consumes the stream lazily from the buffer.
4241{-# INLINE_NORMAL mkParallel #-}
4242mkParallel :: (K.IsStream t, MonadAsync m) => t m a -> t m a
4243mkParallel = fromStreamD . mkParallelD . toStreamD
4245-- Note: we can use another API with two callbacks stop and yield if we want
4246-- the callback to be able to indicate end of stream.
4248-- | Generates a callback and a stream pair. The callback returned is used to
4249-- queue values to the stream. The stream is infinite, there is no way for the
4250-- callback to indicate that it is done now.
4254{-# INLINE_NORMAL newCallbackStream #-}
4255newCallbackStream :: (K.IsStream t, MonadAsync m) => m ((a -> m ()), t m a)
4256newCallbackStream = do
4257 sv <- newParallelVar StopNone defState
4259 -- XXX Add our own thread-id to the SVar as we can not know the callback's
4260 -- thread-id and the callback is not run in a managed worker. We need to
4261 -- handle this better.
4262 liftIO myThreadId >>= modifyThread sv
4264 let callback a = liftIO $ void $ send sv (ChildYield a)
4265 -- XXX we can return an SVar and then the consumer can unfold from the
4267 return (callback, fromStreamD (fromSVar sv))
4269-------------------------------------------------------------------------------
4271-------------------------------------------------------------------------------
4273-- | Create an SVar with a fold consumer that will fold any elements sent to it
4274-- using the supplied fold function.
4275{-# INLINE newFoldSVar #-}
4276newFoldSVar :: MonadAsync m => State t m a -> Fold m a b -> m (SVar t m a)
4277newFoldSVar stt f = do
4278 -- Buffer size for the SVar is derived from the current state
4279 sv <- newParallelVar StopAny (adaptState stt)
4280 -- Add the producer thread-id to the SVar.
4281 liftIO myThreadId >>= modifyThread sv
4282 void $ doFork (work sv) (svarMrun sv) (handleFoldException sv)
4287 {-# NOINLINE work #-}
4288 work sv = void $ runFold f $ fromProducer sv
4290data TapState sv st = TapInit | Tapping sv st | TapDone st
4292{-# INLINE_NORMAL tapAsync #-}
4293tapAsync :: MonadAsync m => Fold m a b -> Stream m a -> Stream m a
4294tapAsync f (Stream step1 state1) = Stream step TapInit
4298 -- In general, a Stop event would come equipped with the result
4299 -- of the fold. It is not used here but it would be useful in
4300 -- applicative and distribute.
4301 done <- fromConsumer svr
4302 when (not done) $ do
4303 liftIO $ withDiagMVar svr "teeToSVar: waiting to drain"
4304 $ takeMVar (outputDoorBellFromConsumer svr)
4308 liftIO $ sendStop svr Nothing
4309 -- drain/wait until a stop event arrives from the fold.
4312 {-# INLINE_LATE step #-}
4313 step gst TapInit = do
4314 sv <- newFoldSVar gst f
4315 return $ Skip (Tapping sv state1)
4317 step gst (Tapping sv st) = do
4321 done <- pushToFold sv a
4324 -- XXX we do not need to wait synchronously here
4326 return $ Yield a (TapDone s)
4327 else return $ Yield a (Tapping sv s)
4328 Skip s -> return $ Skip (Tapping sv s)
4333 step gst (TapDone st) = do
4336 Yield a s -> Yield a (TapDone s)
4337 Skip s -> Skip (TapDone s)
4340-- XXX Exported from Array again as this fold is specific to Array
4341-- | Take last 'n' elements from the stream and discard the rest.
4343lastN :: (Storable a, MonadIO m) => Int -> Fold m a (Array a)
4345 | n <= 0 = fmap (const mempty) FL.drain
4346 | otherwise = Fold step initial done
4348 step (Tuple3' rb rh i) a = do
4349 rh1 <- liftIO $ RB.unsafeInsert rb rh a
4350 return $ Tuple3' rb rh1 (i + 1)
4351 initial = fmap (\(a, b) -> Tuple3' a b (0 :: Int)) $ liftIO $ RB.new n
4352 done (Tuple3' rb rh i) = do
4353 arr <- liftIO $ A.newArray n
4354 foldFunc i rh snoc' arr rb
4355 snoc' b a = liftIO $ A.unsafeSnoc b a
4357 | i < n = RB.unsafeFoldRingM
4358 | otherwise = RB.unsafeFoldRingFullM
4360------------------------------------------------------------------------------
4362------------------------------------------------------------------------------
4364-- XXX using getTime in the loop can be pretty expensive especially for
4365-- computations where iterations are lightweight. We have the following
4368-- 1) Run a timeout thread updating a flag asynchronously and check that
4369-- flag here, that way we can have a cheap termination check.
4371-- 2) Use COARSE clock to get time with lower resolution but more efficiently.
4373-- 3) Use rdtscp/rdtsc to get time directly from the processor, compute the
4374-- termination value of rdtsc in the beginning and then in each iteration just
4375-- get rdtsc and check if we should terminate.
4379 | TakeByTimeCheck st s
4380 | TakeByTimeYield st s
4382{-# INLINE_NORMAL takeByTime #-}
4383takeByTime :: (MonadIO m, TimeUnit64 t) => t -> Stream m a -> Stream m a
4384takeByTime duration (Stream step1 state1) = Stream step (TakeByTimeInit state1)
4387 lim = toRelTime64 duration
4389 {-# INLINE_LATE step #-}
4390 step _ (TakeByTimeInit _) | lim == 0 = return Stop
4391 step _ (TakeByTimeInit st) = do
4392 t0 <- liftIO $ getTime Monotonic
4393 return $ Skip (TakeByTimeYield st t0)
4394 step _ (TakeByTimeCheck st t0) = do
4395 t <- liftIO $ getTime Monotonic
4397 if diffAbsTime64 t t0 > lim
4399 else Skip (TakeByTimeYield st t0)
4400 step gst (TakeByTimeYield st t0) = do
4403 Yield x s -> Yield x (TakeByTimeCheck s t0)
4404 Skip s -> Skip (TakeByTimeCheck s t0)
4407data DropByTime st s x
4409 | DropByTimeGen st s
4410 | DropByTimeCheck st s x
4411 | DropByTimeYield st
4413{-# INLINE_NORMAL dropByTime #-}
4414dropByTime :: (MonadIO m, TimeUnit64 t) => t -> Stream m a -> Stream m a
4415dropByTime duration (Stream step1 state1) = Stream step (DropByTimeInit state1)
4418 lim = toRelTime64 duration
4420 {-# INLINE_LATE step #-}
4421 step _ (DropByTimeInit st) = do
4422 t0 <- liftIO $ getTime Monotonic
4423 return $ Skip (DropByTimeGen st t0)
4424 step gst (DropByTimeGen st t0) = do
4427 Yield x s -> Skip (DropByTimeCheck s t0 x)
4428 Skip s -> Skip (DropByTimeGen s t0)
4430 step _ (DropByTimeCheck st t0 x) = do
4431 t <- liftIO $ getTime Monotonic
4432 if diffAbsTime64 t t0 <= lim
4433 then return $ Skip $ DropByTimeGen st t0
4434 else return $ Yield x $ DropByTimeYield st
4435 step gst (DropByTimeYield st) = do
4438 Yield x s -> Yield x (DropByTimeYield s)
4439 Skip s -> Skip (DropByTimeYield s)
4442-- XXX we should move this to stream generation section of this file. Also, the
4443-- take/drop combinators above should be moved to filtering section.
4444{-# INLINE_NORMAL currentTime #-}
4445currentTime :: MonadAsync m => Double -> Stream m AbsTime
4446currentTime g = Stream step Nothing
4450 g' = g * 10 ^ (6 :: Int)
4452 -- XXX should have a minimum granularity to avoid high CPU usage?
4453 {-# INLINE delayTime #-}
4455 if g' >= fromIntegral (maxBound :: Int)
4459 updateTimeVar timeVar = do
4460 threadDelay $ delayTime
4461 MicroSecond64 t <- fromAbsTime <$> getTime Monotonic
4462 modifyVar' timeVar (const t)
4464 {-# INLINE_LATE step #-}
4466 -- XXX note that this is safe only on a 64-bit machine. On a 32-bit
4467 -- machine a 64-bit 'Var' cannot be read consistently without a lock
4468 -- while another thread is writing to it.
4469 timeVar <- liftIO $ newVar (0 :: Int64)
4470 tid <- forkManaged $ liftIO $ forever (updateTimeVar timeVar)
4471 return $ Skip $ Just (timeVar, tid)
4473 step _ s@(Just (timeVar, _)) = do
4474 a <- liftIO $ readVar timeVar
4475 -- XXX we can perhaps use an AbsTime64 using a 64 bit Int for
4476 -- efficiency. or maybe we can use a representation using Double for
4477 -- floating precision time
4478 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