1{-# LANGUAGE BangPatterns #-}
3{-# LANGUAGE ConstraintKinds #-}
4{-# LANGUAGE ExistentialQuantification #-}
5{-# LANGUAGE FlexibleContexts #-}
6{-# LANGUAGE FlexibleInstances #-}
7{-# LANGUAGE MultiParamTypeClasses #-}
8{-# LANGUAGE PatternSynonyms #-}
9{-# LANGUAGE RecordWildCards #-}
10{-# LANGUAGE ScopedTypeVariables #-}
11{-# LANGUAGE ViewPatterns #-}
12{-# LANGUAGE RankNTypes #-}
13{-# LANGUAGE MagicHash #-}
15#if __GLASGOW_HASKELL__ >= 801
16{-# LANGUAGE TypeApplications #-}
22-- Module : Streamly.Internal.Data.Stream.StreamD
23-- Copyright : (c) 2018 Harendra Kumar
24-- Copyright : (c) Roman Leshchinskiy 2008-2010
25-- Copyright : (c) The University of Glasgow, 2009
26-- Copyright : (c) Bjoern Hoehrmann 2008-2009
29-- Maintainer : streamly@composewell.com
30-- Stability : experimental
33-- Direct style re-implementation of CPS style stream in StreamK module. The
34-- symbol or suffix 'D' in this module denotes the "Direct" style. GHC is able
35-- to INLINE and fuse direct style better, providing better performance than
39-- import qualified Streamly.Internal.Data.Stream.StreamD as D
42-- Some of the functions in this file have been adapted from the vector
43-- library, https://hackage.haskell.org/package/vector.
45module Streamly.Internal.Data.Stream.StreamD
50#if __GLASGOW_HASKELL__ >= 800
51 , Stream (Stream, UnStream)
71 -- ** Specialized Generation
72 -- | Generate a monadic stream from a seed.
83 , enumerateFromStepIntegral
84 , enumerateFromIntegral
85 , enumerateFromThenIntegral
86 , enumerateFromToIntegral
87 , enumerateFromThenToIntegral
89 , enumerateFromStepNum
92 , enumerateFromToFractional
93 , enumerateFromThenToFractional
96 -- | Transform an input structure into a stream.
97 -- | Direct style stream does not support @fromFoldable@.
126 -- ** Specialized Folds
150 -- ** Flattening nested streams
153 , ConcatMapUState (..)
155 , ConcatUnfoldInterleaveState (..)
156 , concatUnfoldInterleave
157 , concatUnfoldRoundrobin
160 , InterleaveState(..)
165 , roundRobin -- interleaveFair?/ParallelFair
198 -- | Transform a stream into another type.
215 -- ** By folding (scans)
258 , intersperseSuffixBySpan
284 -- * Transformation comprehensions
296 -- * Concurrent Application
303import Control.Concurrent (myThreadId, takeMVar)
304import Control.Exception (Exception, SomeException, fromException)
305import Control.Monad (void, when)
306import Control.Monad.Catch (MonadCatch, throwM)
307import Control.Monad.IO.Class (MonadIO(..))
308import Control.Monad.Reader (ReaderT)
309import Control.Monad.State.Strict (StateT)
310import Control.Monad.Trans (MonadTrans(lift))
311import Data.Bits (shiftR, shiftL, (.|.), (.&.))
312import Data.Functor.Identity (Identity(..))
313import Data.IORef (newIORef, readIORef, mkWeakIORef, writeIORef)
314import Data.Maybe (fromJust, isJust, isNothing)
315import Data.Word (Word32)
316import Foreign.Ptr (Ptr)
317import Foreign.Storable (Storable(..))
318import GHC.Types (SPEC(..))
319import System.Mem (performMajorGC)
321 hiding (map, mapM, mapM_, repeat, foldr, last, take, filter,
322 takeWhile, drop, dropWhile, all, any, maximum, minimum, elem,
323 notElem, null, head, tail, zipWith, lookup, foldr1, sequence,
324 (!!), scanl, scanl1, concatMap, replicate, enumFromTo, concat,
327import qualified Control.Monad.Catch as MC
328import qualified Control.Monad.Reader as Reader
329import qualified Control.Monad.State.Strict as State
330import qualified Prelude
332import Streamly.Internal.Data.Atomics (atomicModifyIORefCAS_)
333import Streamly.Internal.Memory.Array.Types (Array(..))
334import Streamly.Internal.Data.Fold.Types (Fold(..))
335import Streamly.Internal.Data.Pipe.Types (Pipe(..), PipeState(..))
336import Streamly.Internal.Data.Time.Clock (Clock(Monotonic), getTime)
337import Streamly.Internal.Data.Unfold.Types (Unfold(..))
339import Streamly.Internal.Data.Stream.StreamD.Type
340import Streamly.Internal.Data.SVar
341import Streamly.Internal.Data.Stream.SVar (fromConsumer, pushToFold)
343import qualified Streamly.Internal.Data.Pipe.Types as Pipe
344import qualified Streamly.Internal.Memory.Array.Types as A
345import qualified Streamly.Internal.Data.Fold as FL
346import qualified Streamly.Memory.Ring as RB
347import qualified Streamly.Internal.Data.Stream.StreamK as K
349------------------------------------------------------------------------------
351------------------------------------------------------------------------------
353-- | An empty 'Stream'.
354{-# INLINE_NORMAL nil #-}
355nil :: Monad m => Stream m a
356nil = Stream (\_ _ -> return Stop) ()
358-- | An empty 'Stream' with a side effect.
359{-# INLINE_NORMAL nilM #-}
360nilM :: Monad m => m b -> Stream m a
361nilM m = Stream (\_ _ -> m >> return Stop) ()
363{-# INLINE_NORMAL consM #-}
364consM :: Monad m => m a -> Stream m a -> Stream m a
365consM m (Stream step state) = Stream step1 Nothing
367 {-# INLINE_LATE step1 #-}
368 step1 _ Nothing = m >>= \x -> return $ Yield x (Just state)
369 step1 gst (Just st) = do
373 Yield a s -> Yield a (Just s)
374 Skip s -> Skip (Just s)
377-- XXX implement in terms of consM?
378-- cons x = consM (return x)
380-- | Can fuse but has O(n^2) complexity.
381{-# INLINE_NORMAL cons #-}
382cons :: Monad m => a -> Stream m a -> Stream m a
383cons x (Stream step state) = Stream step1 Nothing
385 {-# INLINE_LATE step1 #-}
386 step1 _ Nothing = return $ Yield x (Just state)
387 step1 gst (Just st) = do
391 Yield a s -> Yield a (Just s)
392 Skip s -> Skip (Just s)
395-------------------------------------------------------------------------------
397-------------------------------------------------------------------------------
399-- Does not fuse, has the same performance as the StreamK version.
400{-# INLINE_NORMAL uncons #-}
401uncons :: Monad m => Stream m a -> m (Maybe (a, Stream m a))
402uncons (UnStream step state) = go state
405 r <- step defState st
407 Yield x s -> return $ Just (x, Stream step s)
409 Stop -> return Nothing
411------------------------------------------------------------------------------
412-- Generation by unfold
413------------------------------------------------------------------------------
415{-# INLINE_NORMAL unfoldrM #-}
416unfoldrM :: Monad m => (s -> m (Maybe (a, s))) -> s -> Stream m a
417unfoldrM next state = Stream step state
419 {-# INLINE_LATE step #-}
423 Just (x, s) -> Yield x s
426{-# INLINE_LATE unfoldr #-}
427unfoldr :: Monad m => (s -> Maybe (a, s)) -> s -> Stream m a
428unfoldr f = unfoldrM (return . f)
430-- | Convert an 'Unfold' into a 'Stream' by supplying it a seed.
432{-# INLINE_NORMAL unfold #-}
433unfold :: Monad m => Unfold m a b -> a -> Stream m b
434unfold (Unfold ustep inject) seed = Stream step Nothing
436 {-# INLINE_LATE step #-}
437 step _ Nothing = inject seed >>= return . Skip . Just
438 step _ (Just st) = do
441 Yield x s -> Yield x (Just s)
442 Skip s -> Skip (Just s)
445------------------------------------------------------------------------------
446-- Specialized Generation
447------------------------------------------------------------------------------
449repeatM :: Monad m => m a -> Stream m a
450repeatM x = Stream (\_ _ -> x >>= \r -> return $ Yield r ()) ()
452repeat :: Monad m => a -> Stream m a
453repeat x = Stream (\_ _ -> return $ Yield x ()) ()
455{-# INLINE_NORMAL replicateM #-}
456replicateM :: forall m a. Monad m => Int -> m a -> Stream m a
457replicateM n p = Stream step n
459 {-# INLINE_LATE step #-}
461 | i <= 0 = return Stop
464 return $ Yield x (i - 1)
466{-# INLINE_NORMAL replicate #-}
467replicate :: Monad m => Int -> a -> Stream m a
468replicate n x = replicateM n (return x)
470-- This would not work properly for floats, therefore we put an Integral
472-- | Can be used to enumerate unbounded integrals. This does not check for
473-- overflow or underflow for bounded integrals.
474{-# INLINE_NORMAL enumerateFromStepIntegral #-}
475enumerateFromStepIntegral :: (Integral a, Monad m) => a -> a -> Stream m a
476enumerateFromStepIntegral from stride =
477 from `seq` stride `seq` Stream step from
479 {-# INLINE_LATE step #-}
480 step _ !x = return $ Yield x $! (x + stride)
482-- We are assuming that "to" is constrained by the type to be within
484{-# INLINE enumerateFromToIntegral #-}
485enumerateFromToIntegral :: (Monad m, Integral a) => a -> a -> Stream m a
486enumerateFromToIntegral from to =
487 takeWhile (<= to) $ enumerateFromStepIntegral from 1
489{-# INLINE enumerateFromIntegral #-}
490enumerateFromIntegral :: (Monad m, Integral a, Bounded a) => a -> Stream m a
491enumerateFromIntegral from = enumerateFromToIntegral from maxBound
493data EnumState a = EnumInit | EnumYield a a a | EnumStop
495{-# INLINE_NORMAL enumerateFromThenToIntegralUp #-}
496enumerateFromThenToIntegralUp
497 :: (Monad m, Integral a)
498 => a -> a -> a -> Stream m a
499enumerateFromThenToIntegralUp from next to = Stream step EnumInit
501 {-# INLINE_LATE step #-}
507 else Yield from EnumStop
508 else -- from <= next <= to
509 let stride = next - from
510 in Skip $ EnumYield from stride (to - stride)
512 step _ (EnumYield x stride toMinus) =
515 then Yield x EnumStop
516 else Yield x $ EnumYield (x + stride) stride toMinus
518 step _ EnumStop = return Stop
520{-# INLINE_NORMAL enumerateFromThenToIntegralDn #-}
521enumerateFromThenToIntegralDn
522 :: (Monad m, Integral a)
523 => a -> a -> a -> Stream m a
524enumerateFromThenToIntegralDn from next to = Stream step EnumInit
526 {-# INLINE_LATE step #-}
528 return $ if to > next
531 else Yield from EnumStop
532 else -- from >= next >= to
533 let stride = next - from
534 in Skip $ EnumYield from stride (to - stride)
536 step _ (EnumYield x stride toMinus) =
539 then Yield x EnumStop
540 else Yield x $ EnumYield (x + stride) stride toMinus
542 step _ EnumStop = return Stop
544{-# INLINE_NORMAL enumerateFromThenToIntegral #-}
545enumerateFromThenToIntegral
546 :: (Monad m, Integral a)
547 => a -> a -> a -> Stream m a
548enumerateFromThenToIntegral from next to
549 | next >= from = enumerateFromThenToIntegralUp from next to
550 | otherwise = enumerateFromThenToIntegralDn from next to
552{-# INLINE_NORMAL enumerateFromThenIntegral #-}
553enumerateFromThenIntegral
554 :: (Monad m, Integral a, Bounded a)
555 => a -> a -> Stream m a
556enumerateFromThenIntegral from next =
558 then enumerateFromThenToIntegralUp from next maxBound
559 else enumerateFromThenToIntegralDn from next minBound
561-- For floating point numbers if the increment is less than the precision then
562-- it just gets lost. Therefore we cannot always increment it correctly by just
564-- 9007199254740992 + 1 + 1 :: Double => 9.007199254740992e15
565-- 9007199254740992 + 2 :: Double => 9.007199254740994e15
567-- Instead we accumulate the increment counter and compute the increment
568-- every time before adding it to the starting number.
570-- This works for Integrals as well as floating point numbers, but
571-- enumerateFromStepIntegral is faster for integrals.
572{-# INLINE_NORMAL enumerateFromStepNum #-}
573enumerateFromStepNum :: (Monad m, Num a) => a -> a -> Stream m a
574enumerateFromStepNum from stride = Stream step 0
576 {-# INLINE_LATE step #-}
577 step _ !i = return $ (Yield $! (from + i * stride)) $! (i + 1)
579{-# INLINE_NORMAL numFrom #-}
580numFrom :: (Monad m, Num a) => a -> Stream m a
581numFrom from = enumerateFromStepNum from 1
583{-# INLINE_NORMAL numFromThen #-}
584numFromThen :: (Monad m, Num a) => a -> a -> Stream m a
585numFromThen from next = enumerateFromStepNum from (next - from)
587-- We cannot write a general function for Num. The only way to write code
588-- portable between the two is to use a 'Real' constraint and convert between
589-- Fractional and Integral using fromRational which is horribly slow.
590{-# INLINE_NORMAL enumerateFromToFractional #-}
591enumerateFromToFractional
592 :: (Monad m, Fractional a, Ord a)
593 => a -> a -> Stream m a
594enumerateFromToFractional from to =
595 takeWhile (<= to + 1 / 2) $ enumerateFromStepNum from 1
597{-# INLINE_NORMAL enumerateFromThenToFractional #-}
598enumerateFromThenToFractional
599 :: (Monad m, Fractional a, Ord a)
600 => a -> a -> a -> Stream m a
601enumerateFromThenToFractional from next to =
602 takeWhile predicate $ numFromThen from next
604 mid = (next - from) / 2
605 predicate | next >= from = (<= to + mid)
606 | otherwise = (>= to + mid)
608-------------------------------------------------------------------------------
609-- Generation by Conversion
610-------------------------------------------------------------------------------
612{-# INLINE_NORMAL fromIndicesM #-}
613fromIndicesM :: Monad m => (Int -> m a) -> Stream m a
614fromIndicesM gen = Stream step 0
616 {-# INLINE_LATE step #-}
619 return $ Yield x (i + 1)
621{-# INLINE fromIndices #-}
622fromIndices :: Monad m => (Int -> a) -> Stream m a
623fromIndices gen = fromIndicesM (return . gen)
625{-# INLINE_NORMAL generateM #-}
626generateM :: Monad m => Int -> (Int -> m a) -> Stream m a
627generateM n gen = n `seq` Stream step 0
629 {-# INLINE_LATE step #-}
630 step _ i | i < n = do
632 return $ Yield x (i + 1)
633 | otherwise = return Stop
635{-# INLINE generate #-}
636generate :: Monad m => Int -> (Int -> a) -> Stream m a
637generate n gen = generateM n (return . gen)
639-- XXX we need the MonadAsync constraint because of a rewrite rule.
640-- | Convert a list of monadic actions to a 'Stream'
641{-# INLINE_LATE fromListM #-}
642fromListM :: MonadAsync m => [m a] -> Stream m a
643fromListM = Stream step
645 {-# INLINE_LATE step #-}
646 step _ (m:ms) = m >>= \x -> return $ Yield x ms
647 step _ [] = return Stop
649{-# INLINE toStreamD #-}
650toStreamD :: (K.IsStream t, Monad m) => t m a -> Stream m a
651toStreamD = fromStreamK . K.toStream
653-------------------------------------------------------------------------------
654-- Generation from SVar
655-------------------------------------------------------------------------------
657data FromSVarState t m a =
659 | FromSVarRead (SVar t m a)
660 | FromSVarLoop (SVar t m a) [ChildEvent a]
661 | FromSVarDone (SVar t m a)
663{-# INLINE_NORMAL fromSVar #-}
664fromSVar :: (MonadAsync m) => SVar t m a -> Stream m a
665fromSVar svar = Stream step FromSVarInit
668 {-# INLINE_LATE step #-}
669 step _ FromSVarInit = do
670 ref <- liftIO $ newIORef ()
671 _ <- liftIO $ mkWeakIORef ref hook
672 -- when this copy of svar gets garbage collected "ref" will get
673 -- garbage collected and our GC hook will be called.
674 let sv = svar{svarRef = Just ref}
675 return $ Skip (FromSVarRead sv)
679 {-# NOINLINE hook #-}
681 when (svarInspectMode svar) $ do
682 r <- liftIO $ readIORef (svarStopTime (svarStats svar))
684 printSVar svar "SVar Garbage Collected"
686 -- If there are any SVars referenced by this SVar a GC will prompt
687 -- them to be cleaned up quickly.
688 when (svarInspectMode svar) performMajorGC
690 step _ (FromSVarRead sv) = do
691 list <- readOutputQ sv
692 -- Reversing the output is important to guarantee that we process the
693 -- outputs in the same order as they were generated by the constituent
695 return $ Skip $ FromSVarLoop sv (Prelude.reverse list)
697 step _ (FromSVarLoop sv []) = do
698 done <- postProcess sv
699 return $ Skip $ if done
700 then (FromSVarDone sv)
701 else (FromSVarRead sv)
703 step _ (FromSVarLoop sv (ev : es)) = do
705 ChildYield a -> return $ Yield a (FromSVarLoop sv es)
706 ChildStop tid e -> do
710 stop <- shouldStop tid
713 liftIO (cleanupSVar sv)
714 return $ Skip (FromSVarDone sv)
715 else return $ Skip (FromSVarLoop sv es)
717 case fromException ex of
719 return $ Skip (FromSVarLoop sv es)
720 Nothing -> liftIO (cleanupSVar sv) >> throwM ex
724 case svarStopStyle sv of
725 StopNone -> return False
726 StopAny -> return True
728 sid <- liftIO $ readIORef (svarStopBy sv)
729 return $ if tid == sid then True else False
731 step _ (FromSVarDone sv) = do
732 when (svarInspectMode sv) $ do
733 t <- liftIO $ getTime Monotonic
734 liftIO $ writeIORef (svarStopTime (svarStats sv)) (Just t)
735 liftIO $ printSVar sv "SVar Done"
738-------------------------------------------------------------------------------
739-- Process events received by a fold consumer from a stream producer
740-------------------------------------------------------------------------------
742{-# INLINE_NORMAL fromProducer #-}
743fromProducer :: (MonadAsync m) => SVar t m a -> Stream m a
744fromProducer svar = Stream step (FromSVarRead svar)
747 {-# INLINE_LATE step #-}
748 step _ (FromSVarRead sv) = do
749 list <- readOutputQ sv
750 -- Reversing the output is important to guarantee that we process the
751 -- outputs in the same order as they were generated by the constituent
753 return $ Skip $ FromSVarLoop sv (Prelude.reverse list)
755 step _ (FromSVarLoop sv []) = return $ Skip $ FromSVarRead sv
756 step _ (FromSVarLoop sv (ev : es)) = do
758 ChildYield a -> return $ Yield a (FromSVarLoop sv es)
759 ChildStop tid e -> do
763 sendStopToProducer sv
764 return $ Skip (FromSVarDone sv)
765 Just _ -> error "Bug: fromProducer: received exception"
767 step _ (FromSVarDone sv) = do
768 when (svarInspectMode sv) $ do
769 t <- liftIO $ getTime Monotonic
770 liftIO $ writeIORef (svarStopTime (svarStats sv)) (Just t)
771 liftIO $ printSVar sv "SVar Done"
774 step _ FromSVarInit = undefined
776-------------------------------------------------------------------------------
777-- Hoisting the inner monad
778-------------------------------------------------------------------------------
780{-# INLINE_NORMAL hoist #-}
781hoist :: Monad n => (forall x. m x -> n x) -> Stream m a -> Stream n a
782hoist f (Stream step state) = (Stream step' state)
784 {-# INLINE_LATE step' #-}
786 r <- f $ step (adaptState gst) st
788 Yield x s -> Yield x s
792{-# INLINE generally #-}
793generally :: Monad m => Stream Identity a -> Stream m a
794generally = hoist (return . runIdentity)
796{-# INLINE_NORMAL liftInner #-}
797liftInner :: (Monad m, MonadTrans t, Monad (t m))
798 => Stream m a -> Stream (t m) a
799liftInner (Stream step state) = Stream step' state
801 {-# INLINE_LATE step' #-}
803 r <- lift $ step (adaptState gst) st
805 Yield x s -> Yield x s
809{-# INLINE_NORMAL runReaderT #-}
810runReaderT :: Monad m => s -> Stream (ReaderT s m) a -> Stream m a
811runReaderT sval (Stream step state) = Stream step' state
813 {-# INLINE_LATE step' #-}
815 r <- Reader.runReaderT (step (adaptState gst) st) sval
817 Yield x s -> Yield x s
821{-# INLINE_NORMAL evalStateT #-}
822evalStateT :: Monad m => s -> Stream (StateT s m) a -> Stream m a
823evalStateT sval (Stream step state) = Stream step' (state, sval)
825 {-# INLINE_LATE step' #-}
826 step' gst (st, sv) = do
827 (r, sv') <- State.runStateT (step (adaptState gst) st) sv
829 Yield x s -> Yield x (s, sv')
830 Skip s -> Skip (s, sv')
833{-# INLINE_NORMAL runStateT #-}
834runStateT :: Monad m => s -> Stream (StateT s m) a -> Stream m (s, a)
835runStateT sval (Stream step state) = Stream step' (state, sval)
837 {-# INLINE_LATE step' #-}
838 step' gst (st, sv) = do
839 (r, sv') <- State.runStateT (step (adaptState gst) st) sv
841 Yield x s -> Yield (sv', x) (s, sv')
842 Skip s -> Skip (s, sv')
845------------------------------------------------------------------------------
846-- Elimination by Folds
847------------------------------------------------------------------------------
849------------------------------------------------------------------------------
851------------------------------------------------------------------------------
853{-# INLINE_NORMAL foldr1 #-}
854foldr1 :: Monad m => (a -> a -> a) -> Stream m a -> m (Maybe a)
858 Nothing -> return Nothing
859 Just (h, t) -> fmap Just (foldr f h t)
861------------------------------------------------------------------------------
863------------------------------------------------------------------------------
865{-# INLINE_NORMAL foldlT #-}
866foldlT :: (Monad m, Monad (s m), MonadTrans s)
867 => (s m b -> a -> s m b) -> s m b -> Stream m a -> s m b
868foldlT fstep begin (Stream step state) = go SPEC begin state
871 r <- lift $ step defState st
873 Yield x s -> go SPEC (fstep acc x) s
874 Skip s -> go SPEC acc s
877-- Note, this is going to have horrible performance, because of the nature of
878-- the stream type (i.e. direct stream vs CPS). Its only for reference, it is
879-- likely be practically unusable.
880{-# INLINE_NORMAL foldlS #-}
882 => (Stream m b -> a -> Stream m b) -> Stream m b -> Stream m a -> Stream m b
883foldlS fstep begin (Stream step state) = Stream step' (Left (state, begin))
885 step' gst (Left (st, acc)) = do
886 r <- step (adaptState gst) st
888 Yield x s -> Skip (Left (s, fstep acc x))
889 Skip s -> Skip (Left (s, acc))
890 Stop -> Skip (Right acc)
892 step' gst (Right (Stream stp stt)) = do
893 r <- stp (adaptState gst) stt
895 Yield x s -> Yield x (Right (Stream stp s))
896 Skip s -> Skip (Right (Stream stp s))
899------------------------------------------------------------------------------
901------------------------------------------------------------------------------
903-- | Run a streaming composition, discard the results.
904{-# INLINE_LATE drain #-}
905drain :: Monad m => Stream m a -> m ()
906-- drain = foldrM (\_ xs -> xs) (return ())
907drain (Stream step state) = go SPEC state
910 r <- step defState st
912 Yield _ s -> go SPEC s
916{-# INLINE_NORMAL null #-}
917null :: Monad m => Stream m a -> m Bool
918null m = foldrM (\_ _ -> return False) (return True) m
920{-# INLINE_NORMAL head #-}
921head :: Monad m => Stream m a -> m (Maybe a)
922head m = foldrM (\x _ -> return (Just x)) (return Nothing) m
924{-# INLINE_NORMAL headElse #-}
925headElse :: Monad m => a -> Stream m a -> m a
926headElse a m = foldrM (\x _ -> return x) (return a) m
928-- Does not fuse, has the same performance as the StreamK version.
929{-# INLINE_NORMAL tail #-}
930tail :: Monad m => Stream m a -> m (Maybe (Stream m a))
931tail (UnStream step state) = go state
934 r <- step defState st
936 Yield _ s -> return (Just $ Stream step s)
938 Stop -> return Nothing
940-- XXX will it fuse? need custom impl?
941{-# INLINE_NORMAL last #-}
942last :: Monad m => Stream m a -> m (Maybe a)
943last = foldl' (\_ y -> Just y) Nothing
945{-# INLINE_NORMAL elem #-}
946elem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
947-- elem e m = foldrM (\x xs -> if x == e then return True else xs) (return False) m
948elem e (Stream step state) = go state
951 r <- step defState st
954 | x == e -> return True
959{-# INLINE_NORMAL notElem #-}
960notElem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
961notElem e s = fmap not (elem e s)
963{-# INLINE_NORMAL all #-}
964all :: Monad m => (a -> Bool) -> Stream m a -> m Bool
965-- all p m = foldrM (\x xs -> if p x then xs else return False) (return True) m
966all p (Stream step state) = go state
969 r <- step defState st
973 | otherwise -> return False
977{-# INLINE_NORMAL any #-}
978any :: Monad m => (a -> Bool) -> Stream m a -> m Bool
979-- any p m = foldrM (\x xs -> if p x then return True else xs) (return False) m
980any p (Stream step state) = go state
983 r <- step defState st
991{-# INLINE_NORMAL maximum #-}
992maximum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
993maximum (Stream step state) = go Nothing state
996 r <- step defState st
998 Yield x s -> go (Just x) s
999 Skip s -> go Nothing s
1000 Stop -> return Nothing
1001 go (Just acc) st = do
1002 r <- step defState st
1005 | acc <= x -> go (Just x) s
1006 | otherwise -> go (Just acc) s
1007 Skip s -> go (Just acc) s
1008 Stop -> return (Just acc)
1010{-# INLINE_NORMAL maximumBy #-}
1011maximumBy :: Monad m => (a -> a -> Ordering) -> Stream m a -> m (Maybe a)
1012maximumBy cmp (Stream step state) = go Nothing state
1015 r <- step defState st
1017 Yield x s -> go (Just x) s
1018 Skip s -> go Nothing s
1019 Stop -> return Nothing
1020 go (Just acc) st = do
1021 r <- step defState st
1023 Yield x s -> case cmp acc x of
1024 GT -> go (Just acc) s
1026 Skip s -> go (Just acc) s
1027 Stop -> return (Just acc)
1029{-# INLINE_NORMAL minimum #-}
1030minimum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
1031minimum (Stream step state) = go Nothing state
1034 r <- step defState st
1036 Yield x s -> go (Just x) s
1037 Skip s -> go Nothing s
1038 Stop -> return Nothing
1039 go (Just acc) st = do
1040 r <- step defState st
1043 | acc <= x -> go (Just acc) s
1044 | otherwise -> go (Just x) s
1045 Skip s -> go (Just acc) s
1046 Stop -> return (Just acc)
1048{-# INLINE_NORMAL minimumBy #-}
1049minimumBy :: Monad m => (a -> a -> Ordering) -> Stream m a -> m (Maybe a)
1050minimumBy cmp (Stream step state) = go Nothing state
1053 r <- step defState st
1055 Yield x s -> go (Just x) s
1056 Skip s -> go Nothing s
1057 Stop -> return Nothing
1058 go (Just acc) st = do
1059 r <- step defState st
1061 Yield x s -> case cmp acc x of
1063 _ -> go (Just acc) s
1064 Skip s -> go (Just acc) s
1065 Stop -> return (Just acc)
1067{-# INLINE_NORMAL (!!) #-}
1068(!!) :: (Monad m) => Stream m a -> Int -> m (Maybe a)
1069(Stream step state) !! i = go i state
1072 r <- step defState st
1074 Yield x s | n < 0 -> return Nothing
1075 | n == 0 -> return $ Just x
1076 | otherwise -> go (n - 1) s
1078 Stop -> return Nothing
1080{-# INLINE_NORMAL lookup #-}
1081lookup :: (Monad m, Eq a) => a -> Stream m (a, b) -> m (Maybe b)
1082lookup e m = foldrM (\(a, b) xs -> if e == a then return (Just b) else xs)
1085{-# INLINE_NORMAL findM #-}
1086findM :: Monad m => (a -> m Bool) -> Stream m a -> m (Maybe a)
1087findM p m = foldrM (\x xs -> p x >>= \r -> if r then return (Just x) else xs)
1091find :: Monad m => (a -> Bool) -> Stream m a -> m (Maybe a)
1092find p = findM (return . p)
1094{-# INLINE_NORMAL findIndices #-}
1095findIndices :: Monad m => (a -> Bool) -> Stream m a -> Stream m Int
1096findIndices p (Stream step state) = Stream step' (state, 0)
1098 {-# INLINE_LATE step' #-}
1099 step' gst (st, i) = i `seq` do
1100 r <- step (adaptState gst) st
1102 Yield x s -> if p x then Yield i (s, i+1) else Skip (s, i+1)
1103 Skip s -> Skip (s, i)
1106{-# INLINE toListRev #-}
1107toListRev :: Monad m => Stream m a -> m [a]
1108toListRev = foldl' (flip (:)) []
1110-- We can implement reverse as:
1112-- > reverse = foldlS (flip cons) nil
1114-- However, this implementation is unusable because of the horrible performance
1115-- of cons. So we just convert it to a list first and then stream from the
1118-- XXX Maybe we can use an Array instead of a list here?
1119{-# INLINE_NORMAL reverse #-}
1120reverse :: Monad m => Stream m a -> Stream m a
1121reverse m = Stream step Nothing
1123 {-# INLINE_LATE step #-}
1126 return $ Skip (Just xs)
1127 step _ (Just (x:xs)) = return $ Yield x (Just xs)
1128 step _ (Just []) = return Stop
1130-- Much faster reverse for Storables
1131{-# INLINE_NORMAL reverse' #-}
1132reverse' :: forall m a. (MonadIO m, Storable a) => Stream m a -> Stream m a
1134-- This commented implementation copies the whole stream into one single array
1135-- and then streams from that array, this is 3-4x faster than the chunked code
1136-- that follows. Though this could be problematic due to unbounded large
1137-- allocations. We need to figure out why the chunked code is slower and if we
1138-- can optimize the chunked code to work as fast as this one. It may be a
1140import Foreign.ForeignPtr (touchForeignPtr)
1141import Foreign.ForeignPtr.Unsafe (unsafeForeignPtrToPtr)
1142import Foreign.Ptr (Ptr, plusPtr)
1143reverse' m = Stream step Nothing
1145 {-# INLINE_LATE step #-}
1147 arr <- A.fromStreamD m
1148 let p = aEnd arr `plusPtr` negate (sizeOf (undefined :: a))
1149 return $ Skip $ Just (aStart arr, p)
1151 step _ (Just (start, p)) | p < unsafeForeignPtrToPtr start = return Stop
1153 step _ (Just (start, p)) = do
1154 let !x = A.unsafeInlineIO $ do
1156 touchForeignPtr start
1158 next = p `plusPtr` negate (sizeOf (undefined :: a))
1159 return $ Yield x (Just (start, next))
1166 $ A.fromStreamDArraysOf A.defaultChunkSize m
1169------------------------------------------------------------------------------
1170-- Grouping/Splitting
1171------------------------------------------------------------------------------
1173{-# INLINE_NORMAL splitSuffixBy' #-}
1174splitSuffixBy' :: Monad m
1175 => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1176splitSuffixBy' predicate f (Stream step state) =
1177 Stream (stepOuter f) (Just state)
1181 {-# INLINE_LATE stepOuter #-}
1182 stepOuter (Fold fstep initial done) gst (Just st) = do
1183 res <- step (adaptState gst) st
1189 then done acc' >>= \val -> return $ Yield val (Just s)
1192 Skip s -> return $ Skip $ Just s
1198 res <- step (adaptState gst) stt
1203 then done acc' >>= \val -> return $ Yield val (Just s)
1205 Skip s -> go SPEC s acc
1206 Stop -> done acc >>= \val -> return $ Yield val Nothing
1208 stepOuter _ _ Nothing = return Stop
1210{-# INLINE_NORMAL groupsBy #-}
1216groupsBy cmp f (Stream step state) = Stream (stepOuter f) (Just state, Nothing)
1220 {-# INLINE_LATE stepOuter #-}
1221 stepOuter (Fold fstep initial done) gst (Just st, Nothing) = do
1222 res <- step (adaptState gst) st
1229 Skip s -> return $ Skip $ (Just s, Nothing)
1234 go !_ prev stt !acc = do
1235 res <- step (adaptState gst) stt
1242 else done acc >>= \r -> return $ Yield r (Just s, Just x)
1243 Skip s -> go SPEC prev s acc
1244 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1246 stepOuter (Fold fstep initial done) gst (Just st, Just prev) = do
1248 acc' <- fstep acc prev
1253 -- XXX code duplicated from the previous equation
1255 res <- step (adaptState gst) stt
1262 else done acc >>= \r -> return $ Yield r (Just s, Just x)
1263 Skip s -> go SPEC s acc
1264 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1266 stepOuter _ _ (Nothing,_) = return Stop
1268{-# INLINE_NORMAL groupsRollingBy #-}
1269groupsRollingBy :: Monad m
1274groupsRollingBy cmp f (Stream step state) =
1275 Stream (stepOuter f) (Just state, Nothing)
1278 {-# INLINE_LATE stepOuter #-}
1279 stepOuter (Fold fstep initial done) gst (Just st, Nothing) = do
1280 res <- step (adaptState gst) st
1287 Skip s -> return $ Skip $ (Just s, Nothing)
1291 go !_ prev stt !acc = do
1292 res <- step (adaptState gst) stt
1300 done acc >>= \r -> return $ Yield r (Just s, Just x)
1301 Skip s -> go SPEC prev s acc
1302 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1304 stepOuter (Fold fstep initial done) gst (Just st, Just prev') = do
1306 acc' <- fstep acc prev'
1307 go SPEC prev' st acc'
1310 go !_ prevv stt !acc = do
1311 res <- step (adaptState gst) stt
1318 else done acc >>= \r -> return $ Yield r (Just s, Just x)
1319 Skip s -> go SPEC prevv s acc
1320 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1322 stepOuter _ _ (Nothing, _) = return Stop
1324{-# INLINE_NORMAL splitBy #-}
1325splitBy :: Monad m => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1326splitBy predicate f (Stream step state) = Stream (step' f) (Just state)
1330 {-# INLINE_LATE step' #-}
1331 step' (Fold fstep initial done) gst (Just st) = initial >>= go SPEC st
1336 res <- step (adaptState gst) stt
1340 then done acc >>= \r -> return $ Yield r (Just s)
1344 Skip s -> go SPEC s acc
1345 Stop -> done acc >>= \r -> return $ Yield r Nothing
1347 step' _ _ Nothing = return Stop
1349-- XXX requires -funfolding-use-threshold=150 in lines-unlines benchmark
1350{-# INLINE_NORMAL splitSuffixBy #-}
1351splitSuffixBy :: Monad m
1352 => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1353splitSuffixBy predicate f (Stream step state) = Stream (step' f) (Just state)
1357 {-# INLINE_LATE step' #-}
1358 step' (Fold fstep initial done) gst (Just st) = do
1359 res <- step (adaptState gst) st
1364 then done acc >>= \val -> return $ Yield val (Just s)
1369 Skip s -> return $ Skip $ Just s
1375 res <- step (adaptState gst) stt
1379 then done acc >>= \r -> return $ Yield r (Just s)
1383 Skip s -> go SPEC s acc
1384 Stop -> done acc >>= \r -> return $ Yield r Nothing
1386 step' _ _ Nothing = return Stop
1388{-# INLINE_NORMAL wordsBy #-}
1389wordsBy :: Monad m => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1390wordsBy predicate f (Stream step state) = Stream (stepOuter f) (Just state)
1394 {-# INLINE_LATE stepOuter #-}
1395 stepOuter (Fold fstep initial done) gst (Just st) = do
1396 res <- step (adaptState gst) st
1400 then return $ Skip (Just s)
1406 Skip s -> return $ Skip $ Just s
1412 res <- step (adaptState gst) stt
1416 then done acc >>= \r -> return $ Yield r (Just s)
1420 Skip s -> go SPEC s acc
1421 Stop -> done acc >>= \r -> return $ Yield r Nothing
1423 stepOuter _ _ Nothing = return Stop
1425-- String search algorithms:
1426-- http://www-igm.univ-mlv.fr/~lecroq/string/index.html
1429-- TODO can we unify the splitting operations using a splitting configuration
1430-- like in the split package.
1432data SplitStyle = Infix | Suffix | Prefix deriving (Eq, Show)
1434data SplitOptions = SplitOptions
1435 { style :: SplitStyle
1436 , withSep :: Bool -- ^ keep the separators in output
1437 -- , compact :: Bool -- ^ treat multiple consecutive separators as one
1438 -- , trimHead :: Bool -- ^ drop blank at head
1439 -- , trimTail :: Bool -- ^ drop blank at tail
1443data SplitOnState s a =
1448 | GO_KARP_RABIN s !(RB.Ring a) !(Ptr a)
1451{-# INLINE_NORMAL splitOn #-}
1453 :: forall m a b. (MonadIO m, Storable a, Enum a, Eq a)
1458splitOn patArr@Array{..} (Fold fstep initial done) (Stream step state) =
1459 Stream stepOuter GO_START
1463 patLen = A.length patArr
1464 maxIndex = patLen - 1
1465 elemBits = sizeOf (undefined :: a) * 8
1467 {-# INLINE_LATE stepOuter #-}
1468 stepOuter _ GO_START =
1470 then return $ Skip $ GO_EMPTY_PAT state
1473 r <- liftIO $ (A.unsafeIndexIO patArr 0)
1474 return $ Skip $ GO_SINGLE_PAT state r
1475 else if sizeOf (undefined :: a) * patLen
1476 <= sizeOf (undefined :: Word)
1477 then return $ Skip $ GO_SHORT_PAT state
1479 (rb, rhead) <- liftIO $ RB.new patLen
1480 return $ Skip $ GO_KARP_RABIN state rb rhead
1482 stepOuter gst (GO_SINGLE_PAT stt pat) = initial >>= go SPEC stt
1487 res <- step (adaptState gst) st
1493 return $ Yield r (GO_SINGLE_PAT s pat)
1494 else fstep acc x >>= go SPEC s
1495 Skip s -> go SPEC s acc
1496 Stop -> done acc >>= \r -> return $ Yield r GO_DONE
1498 stepOuter gst (GO_SHORT_PAT stt) = initial >>= go0 SPEC 0 (0 :: Word) stt
1503 mask = (1 `shiftL` (elemBits * patLen)) - 1
1505 addToWord wrd a = (wrd `shiftL` elemBits) .|. fromIntegral (fromEnum a)
1508 patWord = mask .&. A.foldl' addToWord 0 patArr
1510 go0 !_ !idx wrd st !acc = do
1511 res <- step (adaptState gst) st
1514 let wrd' = addToWord wrd x
1517 if wrd' .&. mask == patWord
1520 return $ Yield r (GO_SHORT_PAT s)
1521 else go1 SPEC wrd' s acc
1522 else go0 SPEC (idx + 1) wrd' s acc
1523 Skip s -> go0 SPEC idx wrd s acc
1526 then go2 wrd idx acc
1528 done acc' >>= \r -> return $ Yield r GO_DONE
1531 go1 !_ wrd st !acc = do
1532 res <- step (adaptState gst) st
1535 let wrd' = addToWord wrd x
1536 old = (mask .&. wrd) `shiftR` (elemBits * (patLen - 1))
1537 acc' <- fstep acc (toEnum $ fromIntegral old)
1538 if wrd' .&. mask == patWord
1539 then done acc' >>= \r -> return $ Yield r (GO_SHORT_PAT s)
1540 else go1 SPEC wrd' s acc'
1541 Skip s -> go1 SPEC wrd s acc
1543 acc' <- go2 wrd patLen acc
1544 done acc' >>= \r -> return $ Yield r GO_DONE
1546 go2 !wrd !n !acc | n > 0 = do
1547 let old = (mask .&. wrd) `shiftR` (elemBits * (n - 1))
1548 fstep acc (toEnum $ fromIntegral old) >>= go2 wrd (n - 1)
1549 go2 _ _ acc = return acc
1551 stepOuter gst (GO_KARP_RABIN stt rb rhead) = do
1552 initial >>= go0 SPEC 0 rhead stt
1556 k = 2891336453 :: Word32
1558 addCksum cksum a = cksum * k + fromIntegral (fromEnum a)
1559 deltaCksum cksum old new =
1560 addCksum cksum new - coeff * fromIntegral (fromEnum old)
1562 -- XXX shall we use a random starting hash or 1 instead of 0?
1563 patHash = A.foldl' addCksum 0 patArr
1566 go0 !_ !idx !rh st !acc = do
1567 res <- step (adaptState gst) st
1570 rh' <- liftIO $ RB.unsafeInsert rb rh x
1573 let fold = RB.unsafeFoldRing (RB.ringBound rb)
1574 let !ringHash = fold addCksum 0 rb
1575 if ringHash == patHash
1576 then go2 SPEC ringHash rh' s acc
1577 else go1 SPEC ringHash rh' s acc
1578 else go0 SPEC (idx + 1) rh' s acc
1579 Skip s -> go0 SPEC idx rh s acc
1581 !acc' <- if idx /= 0
1582 then RB.unsafeFoldRingM rh fstep acc rb
1584 done acc' >>= \r -> return $ Yield r GO_DONE
1586 -- XXX Theoretically this code can do 4 times faster if GHC generates
1587 -- optimal code. If we use just "(cksum' == patHash)" condition it goes
1588 -- 4x faster, as soon as we add the "RB.unsafeEqArray rb v" condition
1589 -- the generated code changes drastically and becomes 4x slower. Need
1590 -- to investigate what is going on with GHC.
1592 go1 !_ !cksum !rh st !acc = do
1593 res <- step (adaptState gst) st
1596 old <- liftIO $ peek rh
1597 let cksum' = deltaCksum cksum old x
1598 acc' <- fstep acc old
1600 if (cksum' == patHash)
1602 rh' <- liftIO (RB.unsafeInsert rb rh x)
1603 go2 SPEC cksum' rh' s acc'
1605 rh' <- liftIO (RB.unsafeInsert rb rh x)
1606 go1 SPEC cksum' rh' s acc'
1607 Skip s -> go1 SPEC cksum rh s acc
1609 acc' <- RB.unsafeFoldRingFullM rh fstep acc rb
1610 done acc' >>= \r -> return $ Yield r GO_DONE
1612 go2 !_ !cksum' !rh' s !acc' = do
1613 if RB.unsafeEqArray rb rh' patArr
1616 return $ Yield r (GO_KARP_RABIN s rb rhead)
1617 else go1 SPEC cksum' rh' s acc'
1619 stepOuter gst (GO_EMPTY_PAT st) = do
1620 res <- step (adaptState gst) st
1625 done acc' >>= \r -> return $ Yield r (GO_EMPTY_PAT s)
1626 Skip s -> return $ Skip (GO_EMPTY_PAT s)
1629 stepOuter _ GO_DONE = return Stop
1631{-# INLINE_NORMAL splitSuffixOn #-}
1633 :: forall m a b. (MonadIO m, Storable a, Enum a, Eq a)
1639splitSuffixOn withSep patArr@Array{..} (Fold fstep initial done)
1640 (Stream step state) =
1641 Stream stepOuter GO_START
1645 patLen = A.length patArr
1646 maxIndex = patLen - 1
1647 elemBits = sizeOf (undefined :: a) * 8
1649 {-# INLINE_LATE stepOuter #-}
1650 stepOuter _ GO_START =
1652 then return $ Skip $ GO_EMPTY_PAT state
1655 r <- liftIO $ (A.unsafeIndexIO patArr 0)
1656 return $ Skip $ GO_SINGLE_PAT state r
1657 else if sizeOf (undefined :: a) * patLen
1658 <= sizeOf (undefined :: Word)
1659 then return $ Skip $ GO_SHORT_PAT state
1661 (rb, rhead) <- liftIO $ RB.new patLen
1662 return $ Skip $ GO_KARP_RABIN state rb rhead
1664 stepOuter gst (GO_SINGLE_PAT stt pat) = do
1665 -- This first part is the only difference between splitOn and
1667 -- If the last element is a separator do not issue a blank segment.
1668 res <- step (adaptState gst) stt
1674 acc' <- if withSep then fstep acc x else return acc
1675 done acc' >>= \r -> return $ Yield r (GO_SINGLE_PAT s pat)
1676 else fstep acc x >>= go SPEC s
1677 Skip s -> return $ Skip $ (GO_SINGLE_PAT s pat)
1682 -- This is identical for splitOn and splitSuffixOn
1684 res <- step (adaptState gst) st
1689 acc' <- if withSep then fstep acc x else return acc
1691 return $ Yield r (GO_SINGLE_PAT s pat)
1692 else fstep acc x >>= go SPEC s
1693 Skip s -> go SPEC s acc
1694 Stop -> done acc >>= \r -> return $ Yield r GO_DONE
1696 stepOuter gst (GO_SHORT_PAT stt) = do
1698 -- Call "initial" only if the stream yields an element, otherwise we
1699 -- may call "initial" but never yield anything. initial may produce a
1700 -- side effect, therefore we will end up doing and discard a side
1705 res <- step (adaptState gst) stt
1709 let wrd' = addToWord wrd x
1710 acc' <- if withSep then fstep acc x else return acc
1713 if wrd' .&. mask == patWord
1714 then done acc' >>= \r -> return $ Yield r (GO_SHORT_PAT s)
1715 else go0 SPEC (idx + 1) wrd' s acc'
1716 else go0 SPEC (idx + 1) wrd' s acc'
1717 Skip s -> return $ Skip (GO_SHORT_PAT s)
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
1735 acc' <- if withSep then fstep acc x else return acc
1738 if wrd' .&. mask == patWord
1741 return $ Yield r (GO_SHORT_PAT s)
1742 else go1 SPEC wrd' s acc'
1743 else go0 SPEC (idx + 1) wrd' s acc'
1744 Skip s -> go0 SPEC idx wrd s acc
1746 if (idx == maxIndex) && (wrd .&. mask == patWord)
1749 acc' <- if idx /= 0 && not withSep
1750 then go2 wrd idx acc
1752 done acc' >>= \r -> return $ Yield r GO_DONE
1755 go1 !_ wrd st !acc = do
1756 res <- step (adaptState gst) st
1759 let wrd' = addToWord wrd x
1760 old = (mask .&. wrd) `shiftR` (elemBits * (patLen - 1))
1763 else fstep acc (toEnum $ fromIntegral old)
1764 if wrd' .&. mask == patWord
1765 then done acc' >>= \r -> return $ Yield r (GO_SHORT_PAT s)
1766 else go1 SPEC wrd' s acc'
1767 Skip s -> go1 SPEC wrd s acc
1769 -- If the last sequence is a separator do not issue a blank
1771 if wrd .&. mask == patWord
1776 else go2 wrd patLen acc
1777 done acc' >>= \r -> return $ Yield r GO_DONE
1779 go2 !wrd !n !acc | n > 0 = do
1780 let old = (mask .&. wrd) `shiftR` (elemBits * (n - 1))
1781 fstep acc (toEnum $ fromIntegral old) >>= go2 wrd (n - 1)
1782 go2 _ _ acc = return acc
1784 stepOuter gst (GO_KARP_RABIN stt rb rhead) = do
1786 res <- step (adaptState gst) stt
1790 acc' <- if withSep then fstep acc x else return acc
1791 rh' <- liftIO (RB.unsafeInsert rb rhead x)
1794 let fold = RB.unsafeFoldRing (RB.ringBound rb)
1795 let !ringHash = fold addCksum 0 rb
1796 if ringHash == patHash
1797 then go2 SPEC ringHash rh' s acc'
1798 else go0 SPEC (idx + 1) rh' s acc'
1799 else go0 SPEC (idx + 1) rh' s acc'
1800 Skip s -> return $ Skip (GO_KARP_RABIN s rb rhead)
1805 k = 2891336453 :: Word32
1807 addCksum cksum a = cksum * k + fromIntegral (fromEnum a)
1808 deltaCksum cksum old new =
1809 addCksum cksum new - coeff * fromIntegral (fromEnum old)
1811 -- XXX shall we use a random starting hash or 1 instead of 0?
1812 patHash = A.foldl' addCksum 0 patArr
1815 go0 !_ !idx !rh st !acc = do
1816 res <- step (adaptState gst) st
1819 acc' <- if withSep then fstep acc x else return acc
1820 rh' <- liftIO (RB.unsafeInsert rb rh x)
1823 let fold = RB.unsafeFoldRing (RB.ringBound rb)
1824 let !ringHash = fold addCksum 0 rb
1825 if ringHash == patHash
1826 then go2 SPEC ringHash rh' s acc'
1827 else go1 SPEC ringHash rh' s acc'
1828 else go0 SPEC (idx + 1) rh' s acc'
1829 Skip s -> go0 SPEC idx rh s acc
1831 -- do not issue a blank segment when we end at pattern
1832 if (idx == maxIndex) && RB.unsafeEqArray rb rh patArr
1835 !acc' <- if idx /= 0 && not withSep
1836 then RB.unsafeFoldRingM rh fstep acc rb
1838 done acc' >>= \r -> return $ Yield r GO_DONE
1840 -- XXX Theoretically this code can do 4 times faster if GHC generates
1841 -- optimal code. If we use just "(cksum' == patHash)" condition it goes
1842 -- 4x faster, as soon as we add the "RB.unsafeEqArray rb v" condition
1843 -- the generated code changes drastically and becomes 4x slower. Need
1844 -- to investigate what is going on with GHC.
1846 go1 !_ !cksum !rh st !acc = do
1847 res <- step (adaptState gst) st
1850 old <- liftIO $ peek rh
1851 let cksum' = deltaCksum cksum old x
1856 if (cksum' == patHash)
1858 rh' <- liftIO (RB.unsafeInsert rb rh x)
1859 go2 SPEC cksum' rh' s acc'
1861 rh' <- liftIO (RB.unsafeInsert rb rh x)
1862 go1 SPEC cksum' rh' s acc'
1863 Skip s -> go1 SPEC cksum rh s acc
1865 if RB.unsafeEqArray rb rh patArr
1870 else RB.unsafeFoldRingFullM rh fstep acc rb
1871 done acc' >>= \r -> return $ Yield r GO_DONE
1873 go2 !_ !cksum' !rh' s !acc' = do
1874 if RB.unsafeEqArray rb rh' patArr
1877 return $ Yield r (GO_KARP_RABIN s rb rhead)
1878 else go1 SPEC cksum' rh' s acc'
1880 stepOuter gst (GO_EMPTY_PAT st) = do
1881 res <- step (adaptState gst) st
1886 done acc' >>= \r -> return $ Yield r (GO_EMPTY_PAT s)
1887 Skip s -> return $ Skip (GO_EMPTY_PAT s)
1890 stepOuter _ GO_DONE = return Stop
1892data SplitState s arr
1894 | SplitBuffering s arr
1895 | SplitSplitting s arr
1896 | SplitYielding arr (SplitState s arr)
1899-- XXX An alternative approach would be to use a partial fold (Fold m a b) to
1900-- split using a splitBy like combinator. The Fold would consume upto the
1901-- separator and return any leftover which can then be fed to the next fold.
1903-- We can revisit this once we have partial folds/parsers.
1905-- | Performs infix separator style splitting.
1906{-# INLINE_NORMAL splitInnerBy #-}
1909 => (f a -> m (f a, Maybe (f a))) -- splitter
1910 -> (f a -> f a -> m (f a)) -- joiner
1913splitInnerBy splitter joiner (Stream step1 state1) =
1914 (Stream step (SplitInitial state1))
1918 {-# INLINE_LATE step #-}
1919 step gst (SplitInitial st) = do
1923 (x1, mx2) <- splitter x
1924 return $ case mx2 of
1925 Nothing -> Skip (SplitBuffering s x1)
1926 Just x2 -> Skip (SplitYielding x1 (SplitSplitting s x2))
1927 Skip s -> return $ Skip (SplitInitial s)
1928 Stop -> return $ Stop
1930 step gst (SplitBuffering st buf) = do
1934 (x1, mx2) <- splitter x
1935 buf' <- joiner buf x1
1936 return $ case mx2 of
1937 Nothing -> Skip (SplitBuffering s buf')
1938 Just x2 -> Skip (SplitYielding buf' (SplitSplitting s x2))
1939 Skip s -> return $ Skip (SplitBuffering s buf)
1940 Stop -> return $ Skip (SplitYielding buf SplitFinishing)
1942 step _ (SplitSplitting st buf) = do
1943 (x1, mx2) <- splitter buf
1944 return $ case mx2 of
1945 Nothing -> Skip $ SplitBuffering st x1
1946 Just x2 -> Skip $ SplitYielding x1 (SplitSplitting st x2)
1948 step _ (SplitYielding x next) = return $ Yield x next
1949 step _ SplitFinishing = return $ Stop
1951-- | Performs infix separator style splitting.
1952{-# INLINE_NORMAL splitInnerBySuffix #-}
1954 :: (Monad m, Eq (f a), Monoid (f a))
1955 => (f a -> m (f a, Maybe (f a))) -- splitter
1956 -> (f a -> f a -> m (f a)) -- joiner
1959splitInnerBySuffix splitter joiner (Stream step1 state1) =
1960 (Stream step (SplitInitial state1))
1964 {-# INLINE_LATE step #-}
1965 step gst (SplitInitial st) = do
1969 (x1, mx2) <- splitter x
1970 return $ case mx2 of
1971 Nothing -> Skip (SplitBuffering s x1)
1972 Just x2 -> Skip (SplitYielding x1 (SplitSplitting s x2))
1973 Skip s -> return $ Skip (SplitInitial s)
1974 Stop -> return $ Stop
1976 step gst (SplitBuffering st buf) = do
1980 (x1, mx2) <- splitter x
1981 buf' <- joiner buf x1
1982 return $ case mx2 of
1983 Nothing -> Skip (SplitBuffering s buf')
1984 Just x2 -> Skip (SplitYielding buf' (SplitSplitting s x2))
1985 Skip s -> return $ Skip (SplitBuffering s buf)
1989 else Skip (SplitYielding buf SplitFinishing)
1991 step _ (SplitSplitting st buf) = do
1992 (x1, mx2) <- splitter buf
1993 return $ case mx2 of
1994 Nothing -> Skip $ SplitBuffering st x1
1995 Just x2 -> Skip $ SplitYielding x1 (SplitSplitting st x2)
1997 step _ (SplitYielding x next) = return $ Yield x next
1998 step _ SplitFinishing = return $ Stop
2000------------------------------------------------------------------------------
2002------------------------------------------------------------------------------
2004{-# INLINE_NORMAL isPrefixOf #-}
2005isPrefixOf :: (Eq a, Monad m) => Stream m a -> Stream m a -> m Bool
2006isPrefixOf (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
2008 go (sa, sb, Nothing) = do
2009 r <- stepa defState sa
2011 Yield x sa' -> go (sa', sb, Just x)
2012 Skip sa' -> go (sa', sb, Nothing)
2015 go (sa, sb, Just x) = do
2016 r <- stepb defState sb
2020 then go (sa, sb', Nothing)
2022 Skip sb' -> go (sa, sb', Just x)
2023 Stop -> return False
2025{-# INLINE_NORMAL isSubsequenceOf #-}
2026isSubsequenceOf :: (Eq a, Monad m) => Stream m a -> Stream m a -> m Bool
2027isSubsequenceOf (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
2029 go (sa, sb, Nothing) = do
2030 r <- stepa defState sa
2032 Yield x sa' -> go (sa', sb, Just x)
2033 Skip sa' -> go (sa', sb, Nothing)
2036 go (sa, sb, Just x) = do
2037 r <- stepb defState sb
2041 then go (sa, sb', Nothing)
2042 else go (sa, sb', Just x)
2043 Skip sb' -> go (sa, sb', Just x)
2044 Stop -> return False
2046{-# INLINE_NORMAL stripPrefix #-}
2049 => Stream m a -> Stream m a -> m (Maybe (Stream m a))
2050stripPrefix (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
2052 go (sa, sb, Nothing) = do
2053 r <- stepa defState sa
2055 Yield x sa' -> go (sa', sb, Just x)
2056 Skip sa' -> go (sa', sb, Nothing)
2057 Stop -> return $ Just (Stream stepb sb)
2059 go (sa, sb, Just x) = do
2060 r <- stepb defState sb
2064 then go (sa, sb', Nothing)
2066 Skip sb' -> go (sa, sb', Just x)
2067 Stop -> return Nothing
2069------------------------------------------------------------------------------
2071------------------------------------------------------------------------------
2073-- | Execute a monadic action for each element of the 'Stream'
2074{-# INLINE_NORMAL mapM_ #-}
2075mapM_ :: Monad m => (a -> m b) -> Stream m a -> m ()
2076mapM_ m = drain . mapM m
2078-------------------------------------------------------------------------------
2079-- Stream transformations using Unfolds
2080-------------------------------------------------------------------------------
2082-- Define a unique structure to use in inspection testing
2083data ConcatMapUState o i =
2085 | ConcatMapUInner o i
2087-- | @concatMapU unfold stream@ uses @unfold@ to map the input stream elements
2088-- to streams and then flattens the generated streams into a single output
2091-- This is like 'concatMap' but uses an unfold with an explicit state to
2092-- generate the stream instead of a 'Stream' type generator. This allows better
2093-- optimization via fusion. This can be many times more efficient than
2096{-# INLINE_NORMAL concatMapU #-}
2097concatMapU :: Monad m => Unfold m a b -> Stream m a -> Stream m b
2098concatMapU (Unfold istep inject) (Stream ostep ost) =
2099 Stream step (ConcatMapUOuter ost)
2101 {-# INLINE_LATE step #-}
2102 step gst (ConcatMapUOuter o) = do
2103 r <- ostep (adaptState gst) o
2107 i `seq` return (Skip (ConcatMapUInner o' i))
2108 Skip o' -> return $ Skip (ConcatMapUOuter o')
2109 Stop -> return $ Stop
2111 step _ (ConcatMapUInner o i) = do
2114 Yield x i' -> Yield x (ConcatMapUInner o i')
2115 Skip i' -> Skip (ConcatMapUInner o i')
2116 Stop -> Skip (ConcatMapUOuter o)
2118data ConcatUnfoldInterleaveState o i =
2119 ConcatUnfoldInterleaveOuter o [i]
2120 | ConcatUnfoldInterleaveInner o [i]
2121 | ConcatUnfoldInterleaveInnerL [i] [i]
2122 | ConcatUnfoldInterleaveInnerR [i] [i]
2124-- XXX use arrays to store state instead of lists.
2125-- XXX In general we can use different scheduling strategies e.g. how to
2126-- schedule the outer vs inner loop or assigning weights to different streams
2127-- or outer and inner loops.
2129-- After a yield, switch to the next stream. Do not switch streams on Skip.
2130-- Yield from outer stream switches to the inner stream.
2132-- There are two choices here, (1) exhaust the outer stream first and then
2133-- start yielding from the inner streams, this is much simpler to implement,
2134-- (2) yield at least one element from an inner stream before going back to
2135-- outer stream and opening the next stream from it.
2137-- Ideally, we need some scheduling bias to inner streams vs outer stream.
2138-- Maybe we can configure the behavior.
2140{-# INLINE_NORMAL concatUnfoldInterleave #-}
2141concatUnfoldInterleave :: Monad m => Unfold m a b -> Stream m a -> Stream m b
2142concatUnfoldInterleave (Unfold istep inject) (Stream ostep ost) =
2143 Stream step (ConcatUnfoldInterleaveOuter ost [])
2145 {-# INLINE_LATE step #-}
2146 step gst (ConcatUnfoldInterleaveOuter o ls) = do
2147 r <- ostep (adaptState gst) o
2151 i `seq` return (Skip (ConcatUnfoldInterleaveInner o' (i : ls)))
2152 Skip o' -> return $ Skip (ConcatUnfoldInterleaveOuter o' ls)
2153 Stop -> return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2155 step _ (ConcatUnfoldInterleaveInner _ []) = undefined
2156 step _ (ConcatUnfoldInterleaveInner o (st:ls)) = do
2159 Yield x s -> Yield x (ConcatUnfoldInterleaveOuter o (s:ls))
2160 Skip s -> Skip (ConcatUnfoldInterleaveInner o (s:ls))
2161 Stop -> Skip (ConcatUnfoldInterleaveOuter o ls)
2163 step _ (ConcatUnfoldInterleaveInnerL [] []) = return Stop
2164 step _ (ConcatUnfoldInterleaveInnerL [] rs) =
2165 return $ Skip (ConcatUnfoldInterleaveInnerR [] rs)
2167 step _ (ConcatUnfoldInterleaveInnerL (st:ls) rs) = do
2170 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerL ls (s:rs))
2171 Skip s -> Skip (ConcatUnfoldInterleaveInnerL (s:ls) rs)
2172 Stop -> Skip (ConcatUnfoldInterleaveInnerL ls rs)
2174 step _ (ConcatUnfoldInterleaveInnerR [] []) = return Stop
2175 step _ (ConcatUnfoldInterleaveInnerR ls []) =
2176 return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2178 step _ (ConcatUnfoldInterleaveInnerR ls (st:rs)) = do
2181 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerR (s:ls) rs)
2182 Skip s -> Skip (ConcatUnfoldInterleaveInnerR ls (s:rs))
2183 Stop -> Skip (ConcatUnfoldInterleaveInnerR ls rs)
2185-- XXX In general we can use different scheduling strategies e.g. how to
2186-- schedule the outer vs inner loop or assigning weights to different streams
2187-- or outer and inner loops.
2189-- This could be inefficient if the tasks are too small.
2191-- Compared to concatUnfoldInterleave this one switches streams on Skips.
2193{-# INLINE_NORMAL concatUnfoldRoundrobin #-}
2194concatUnfoldRoundrobin :: Monad m => Unfold m a b -> Stream m a -> Stream m b
2195concatUnfoldRoundrobin (Unfold istep inject) (Stream ostep ost) =
2196 Stream step (ConcatUnfoldInterleaveOuter ost [])
2198 {-# INLINE_LATE step #-}
2199 step gst (ConcatUnfoldInterleaveOuter o ls) = do
2200 r <- ostep (adaptState gst) o
2204 i `seq` return (Skip (ConcatUnfoldInterleaveInner o' (i : ls)))
2205 Skip o' -> return $ Skip (ConcatUnfoldInterleaveInner o' ls)
2206 Stop -> return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2208 step _ (ConcatUnfoldInterleaveInner o []) =
2209 return $ Skip (ConcatUnfoldInterleaveOuter o [])
2211 step _ (ConcatUnfoldInterleaveInner o (st:ls)) = do
2214 Yield x s -> Yield x (ConcatUnfoldInterleaveOuter o (s:ls))
2215 Skip s -> Skip (ConcatUnfoldInterleaveOuter o (s:ls))
2216 Stop -> Skip (ConcatUnfoldInterleaveOuter o ls)
2218 step _ (ConcatUnfoldInterleaveInnerL [] []) = return Stop
2219 step _ (ConcatUnfoldInterleaveInnerL [] rs) =
2220 return $ Skip (ConcatUnfoldInterleaveInnerR [] rs)
2222 step _ (ConcatUnfoldInterleaveInnerL (st:ls) rs) = do
2225 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerL ls (s:rs))
2226 Skip s -> Skip (ConcatUnfoldInterleaveInnerL ls (s:rs))
2227 Stop -> Skip (ConcatUnfoldInterleaveInnerL ls rs)
2229 step _ (ConcatUnfoldInterleaveInnerR [] []) = return Stop
2230 step _ (ConcatUnfoldInterleaveInnerR ls []) =
2231 return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2233 step _ (ConcatUnfoldInterleaveInnerR ls (st:rs)) = do
2236 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerR (s:ls) rs)
2237 Skip s -> Skip (ConcatUnfoldInterleaveInnerR (s:ls) rs)
2238 Stop -> Skip (ConcatUnfoldInterleaveInnerR ls rs)
2240data AppendState s1 s2 = AppendFirst s1 | AppendSecond s2
2242-- Note that this could be much faster compared to the CPS stream. However, as
2243-- the number of streams being composed increases this may become expensive.
2244-- Need to see where the breaking point is between the two.
2246{-# INLINE_NORMAL append #-}
2247append :: Monad m => Stream m a -> Stream m a -> Stream m a
2248append (Stream step1 state1) (Stream step2 state2) =
2249 Stream step (AppendFirst state1)
2253 {-# INLINE_LATE step #-}
2254 step gst (AppendFirst st) = do
2257 Yield a s -> Yield a (AppendFirst s)
2258 Skip s -> Skip (AppendFirst s)
2259 Stop -> Skip (AppendSecond state2)
2261 step gst (AppendSecond st) = do
2264 Yield a s -> Yield a (AppendSecond s)
2265 Skip s -> Skip (AppendSecond s)
2268data InterleaveState s1 s2 = InterleaveFirst s1 s2 | InterleaveSecond s1 s2
2269 | InterleaveSecondOnly s2 | InterleaveFirstOnly s1
2271{-# INLINE_NORMAL interleave #-}
2272interleave :: Monad m => Stream m a -> Stream m a -> Stream m a
2273interleave (Stream step1 state1) (Stream step2 state2) =
2274 Stream step (InterleaveFirst state1 state2)
2278 {-# INLINE_LATE step #-}
2279 step gst (InterleaveFirst st1 st2) = do
2282 Yield a s -> Yield a (InterleaveSecond s st2)
2283 Skip s -> Skip (InterleaveFirst s st2)
2284 Stop -> Skip (InterleaveSecondOnly st2)
2286 step gst (InterleaveSecond st1 st2) = do
2289 Yield a s -> Yield a (InterleaveFirst st1 s)
2290 Skip s -> Skip (InterleaveSecond st1 s)
2291 Stop -> Skip (InterleaveFirstOnly st1)
2293 step gst (InterleaveFirstOnly st1) = do
2296 Yield a s -> Yield a (InterleaveFirstOnly s)
2297 Skip s -> Skip (InterleaveFirstOnly s)
2300 step gst (InterleaveSecondOnly st2) = do
2303 Yield a s -> Yield a (InterleaveSecondOnly s)
2304 Skip s -> Skip (InterleaveSecondOnly s)
2307{-# INLINE_NORMAL interleaveMin #-}
2308interleaveMin :: Monad m => Stream m a -> Stream m a -> Stream m a
2309interleaveMin (Stream step1 state1) (Stream step2 state2) =
2310 Stream step (InterleaveFirst state1 state2)
2314 {-# INLINE_LATE step #-}
2315 step gst (InterleaveFirst st1 st2) = do
2318 Yield a s -> Yield a (InterleaveSecond s st2)
2319 Skip s -> Skip (InterleaveFirst s st2)
2322 step gst (InterleaveSecond st1 st2) = do
2325 Yield a s -> Yield a (InterleaveFirst st1 s)
2326 Skip s -> Skip (InterleaveSecond st1 s)
2329 step _ (InterleaveFirstOnly _) = undefined
2330 step _ (InterleaveSecondOnly _) = undefined
2332{-# INLINE_NORMAL interleaveSuffix #-}
2333interleaveSuffix :: Monad m => Stream m a -> Stream m a -> Stream m a
2334interleaveSuffix (Stream step1 state1) (Stream step2 state2) =
2335 Stream step (InterleaveFirst state1 state2)
2339 {-# INLINE_LATE step #-}
2340 step gst (InterleaveFirst st1 st2) = do
2343 Yield a s -> Yield a (InterleaveSecond s st2)
2344 Skip s -> Skip (InterleaveFirst s st2)
2347 step gst (InterleaveSecond st1 st2) = do
2350 Yield a s -> Yield a (InterleaveFirst st1 s)
2351 Skip s -> Skip (InterleaveSecond st1 s)
2352 Stop -> Skip (InterleaveFirstOnly st1)
2354 step gst (InterleaveFirstOnly st1) = do
2357 Yield a s -> Yield a (InterleaveFirstOnly s)
2358 Skip s -> Skip (InterleaveFirstOnly s)
2361 step _ (InterleaveSecondOnly _) = undefined
2363data InterleaveInfixState s1 s2 a
2364 = InterleaveInfixFirst s1 s2
2365 | InterleaveInfixSecondBuf s1 s2
2366 | InterleaveInfixSecondYield s1 s2 a
2367 | InterleaveInfixFirstYield s1 s2 a
2368 | InterleaveInfixFirstOnly s1
2370{-# INLINE_NORMAL interleaveInfix #-}
2371interleaveInfix :: Monad m => Stream m a -> Stream m a -> Stream m a
2372interleaveInfix (Stream step1 state1) (Stream step2 state2) =
2373 Stream step (InterleaveInfixFirst state1 state2)
2377 {-# INLINE_LATE step #-}
2378 step gst (InterleaveInfixFirst st1 st2) = do
2381 Yield a s -> Yield a (InterleaveInfixSecondBuf s st2)
2382 Skip s -> Skip (InterleaveInfixFirst s st2)
2385 step gst (InterleaveInfixSecondBuf st1 st2) = do
2388 Yield a s -> Skip (InterleaveInfixSecondYield st1 s a)
2389 Skip s -> Skip (InterleaveInfixSecondBuf st1 s)
2390 Stop -> Skip (InterleaveInfixFirstOnly st1)
2392 step gst (InterleaveInfixSecondYield st1 st2 x) = do
2395 Yield a s -> Yield x (InterleaveInfixFirstYield s st2 a)
2396 Skip s -> Skip (InterleaveInfixSecondYield s st2 x)
2399 step _ (InterleaveInfixFirstYield st1 st2 x) = do
2400 return $ Yield x (InterleaveInfixSecondBuf st1 st2)
2402 step gst (InterleaveInfixFirstOnly st1) = do
2405 Yield a s -> Yield a (InterleaveInfixFirstOnly s)
2406 Skip s -> Skip (InterleaveInfixFirstOnly s)
2409{-# INLINE_NORMAL roundRobin #-}
2410roundRobin :: Monad m => Stream m a -> Stream m a -> Stream m a
2411roundRobin (Stream step1 state1) (Stream step2 state2) =
2412 Stream step (InterleaveFirst state1 state2)
2416 {-# INLINE_LATE step #-}
2417 step gst (InterleaveFirst st1 st2) = do
2420 Yield a s -> Yield a (InterleaveSecond s st2)
2421 Skip s -> Skip (InterleaveSecond s st2)
2422 Stop -> Skip (InterleaveSecondOnly st2)
2424 step gst (InterleaveSecond st1 st2) = do
2427 Yield a s -> Yield a (InterleaveFirst st1 s)
2428 Skip s -> Skip (InterleaveFirst st1 s)
2429 Stop -> Skip (InterleaveFirstOnly st1)
2431 step gst (InterleaveSecondOnly st2) = do
2434 Yield a s -> Yield a (InterleaveSecondOnly s)
2435 Skip s -> Skip (InterleaveSecondOnly s)
2438 step gst (InterleaveFirstOnly st1) = do
2441 Yield a s -> Yield a (InterleaveFirstOnly s)
2442 Skip s -> Skip (InterleaveFirstOnly s)
2445data ICUState s1 s2 i1 i2 =
2450 | ICUFirstInner s1 s2 i1
2451 | ICUSecondInner s1 s2 i2
2452 | ICUFirstOnlyInner s1 i1
2453 | ICUSecondOnlyInner s2 i2
2455-- | Interleave streams (full streams, not the elements) unfolded from two
2456-- input streams and concat. Stop when the first stream stops. If the second
2457-- stream ends before the first one then first stream still keeps running alone
2458-- without any interleaving with the second stream.
2460-- [a1, a2, ... an] [b1, b2 ...]
2461-- => [streamA1, streamA2, ... streamAn] [streamB1, streamB2, ...]
2462-- => [streamA1, streamB1, streamA2...StreamAn, streamBn]
2463-- => [a11, a12, ...a1j, b11, b12, ...b1k, a21, a22, ...]
2465{-# INLINE_NORMAL gintercalateSuffix #-}
2468 => Unfold m a c -> Stream m a -> Unfold m b c -> Stream m b -> Stream m c
2470 (Unfold istep1 inject1) (Stream step1 state1)
2471 (Unfold istep2 inject2) (Stream step2 state2) =
2472 Stream step (ICUFirst state1 state2)
2476 {-# INLINE_LATE step #-}
2477 step gst (ICUFirst s1 s2) = do
2478 r <- step1 (adaptState gst) s1
2482 i `seq` return (Skip (ICUFirstInner s s2 i))
2483 Skip s -> return $ Skip (ICUFirst s s2)
2486 step gst (ICUFirstOnly s1) = do
2487 r <- step1 (adaptState gst) s1
2491 i `seq` return (Skip (ICUFirstOnlyInner s i))
2492 Skip s -> return $ Skip (ICUFirstOnly s)
2495 step _ (ICUFirstInner s1 s2 i1) = do
2498 Yield x i' -> Yield x (ICUFirstInner s1 s2 i')
2499 Skip i' -> Skip (ICUFirstInner s1 s2 i')
2500 Stop -> Skip (ICUSecond s1 s2)
2502 step _ (ICUFirstOnlyInner s1 i1) = do
2505 Yield x i' -> Yield x (ICUFirstOnlyInner s1 i')
2506 Skip i' -> Skip (ICUFirstOnlyInner s1 i')
2507 Stop -> Skip (ICUFirstOnly s1)
2509 step gst (ICUSecond s1 s2) = do
2510 r <- step2 (adaptState gst) s2
2514 i `seq` return (Skip (ICUSecondInner s1 s i))
2515 Skip s -> return $ Skip (ICUSecond s1 s)
2516 Stop -> return $ Skip (ICUFirstOnly s1)
2518 step _ (ICUSecondInner s1 s2 i2) = do
2521 Yield x i' -> Yield x (ICUSecondInner s1 s2 i')
2522 Skip i' -> Skip (ICUSecondInner s1 s2 i')
2523 Stop -> Skip (ICUFirst s1 s2)
2525 step _ (ICUSecondOnly _s2) = undefined
2526 step _ (ICUSecondOnlyInner _s2 _i2) = undefined
2528data InterposeSuffixState s1 i1 =
2529 InterposeSuffixFirst s1
2530 -- | InterposeSuffixFirstYield s1 i1
2531 | InterposeSuffixFirstInner s1 i1
2532 | InterposeSuffixSecond s1
2534-- Note that if an unfolded layer turns out to be nil we still emit the
2535-- separator effect. An alternate behavior could be to emit the separator
2536-- effect only if at least one element has been yielded by the unfolding.
2537-- However, that becomes a bit complicated, so we have chosen the former
2539{-# INLINE_NORMAL interposeSuffix #-}
2542 => m c -> Unfold m b c -> Stream m b -> Stream m c
2545 (Unfold istep1 inject1) (Stream step1 state1) =
2546 Stream step (InterposeSuffixFirst state1)
2550 {-# INLINE_LATE step #-}
2551 step gst (InterposeSuffixFirst s1) = do
2552 r <- step1 (adaptState gst) s1
2556 i `seq` return (Skip (InterposeSuffixFirstInner s i))
2557 -- i `seq` return (Skip (InterposeSuffixFirstYield s i))
2558 Skip s -> return $ Skip (InterposeSuffixFirst s)
2562 step _ (InterposeSuffixFirstYield s1 i1) = do
2565 Yield x i' -> Yield x (InterposeSuffixFirstInner s1 i')
2566 Skip i' -> Skip (InterposeSuffixFirstYield s1 i')
2567 Stop -> Skip (InterposeSuffixFirst s1)
2570 step _ (InterposeSuffixFirstInner s1 i1) = do
2573 Yield x i' -> Yield x (InterposeSuffixFirstInner s1 i')
2574 Skip i' -> Skip (InterposeSuffixFirstInner s1 i')
2575 Stop -> Skip (InterposeSuffixSecond s1)
2577 step _ (InterposeSuffixSecond s1) = do
2579 return $ Yield r (InterposeSuffixFirst s1)
2581data ICALState s1 s2 i1 i2 a =
2583 -- | ICALFirstYield s1 s2 i1
2584 | ICALFirstInner s1 s2 i1
2586 | ICALFirstOnlyInner s1 i1
2587 | ICALSecondInject s1 s2
2588 | ICALFirstInject s1 s2 i2
2589 -- | ICALFirstBuf s1 s2 i1 i2
2590 | ICALSecondInner s1 s2 i1 i2
2591 -- -- | ICALSecondInner s1 s2 i1 i2 a
2592 -- -- | ICALFirstResume s1 s2 i1 i2 a
2594-- | Interleave streams (full streams, not the elements) unfolded from two
2595-- input streams and concat. Stop when the first stream stops. If the second
2596-- stream ends before the first one then first stream still keeps running alone
2597-- without any interleaving with the second stream.
2599-- [a1, a2, ... an] [b1, b2 ...]
2600-- => [streamA1, streamA2, ... streamAn] [streamB1, streamB2, ...]
2601-- => [streamA1, streamB1, streamA2...StreamAn, streamBn]
2602-- => [a11, a12, ...a1j, b11, b12, ...b1k, a21, a22, ...]
2604{-# INLINE_NORMAL gintercalate #-}
2607 => Unfold m a c -> Stream m a -> Unfold m b c -> Stream m b -> Stream m c
2609 (Unfold istep1 inject1) (Stream step1 state1)
2610 (Unfold istep2 inject2) (Stream step2 state2) =
2611 Stream step (ICALFirst state1 state2)
2615 {-# INLINE_LATE step #-}
2616 step gst (ICALFirst s1 s2) = do
2617 r <- step1 (adaptState gst) s1
2621 i `seq` return (Skip (ICALFirstInner s s2 i))
2622 -- i `seq` return (Skip (ICALFirstYield s s2 i))
2623 Skip s -> return $ Skip (ICALFirst s s2)
2627 step _ (ICALFirstYield s1 s2 i1) = do
2630 Yield x i' -> Yield x (ICALFirstInner s1 s2 i')
2631 Skip i' -> Skip (ICALFirstYield s1 s2 i')
2632 Stop -> Skip (ICALFirst s1 s2)
2635 step _ (ICALFirstInner s1 s2 i1) = do
2638 Yield x i' -> Yield x (ICALFirstInner s1 s2 i')
2639 Skip i' -> Skip (ICALFirstInner s1 s2 i')
2640 Stop -> Skip (ICALSecondInject s1 s2)
2642 step gst (ICALFirstOnly s1) = do
2643 r <- step1 (adaptState gst) s1
2647 i `seq` return (Skip (ICALFirstOnlyInner s i))
2648 Skip s -> return $ Skip (ICALFirstOnly s)
2651 step _ (ICALFirstOnlyInner s1 i1) = do
2654 Yield x i' -> Yield x (ICALFirstOnlyInner s1 i')
2655 Skip i' -> Skip (ICALFirstOnlyInner s1 i')
2656 Stop -> Skip (ICALFirstOnly s1)
2658 -- We inject the second stream even before checking if the first stream
2659 -- would yield any more elements. There is no clear choice whether we
2660 -- should do this before or after that. Doing it after may make the state
2661 -- machine a bit simpler though.
2662 step gst (ICALSecondInject s1 s2) = do
2663 r <- step2 (adaptState gst) s2
2667 i `seq` return (Skip (ICALFirstInject s1 s i))
2668 Skip s -> return $ Skip (ICALSecondInject s1 s)
2669 Stop -> return $ Skip (ICALFirstOnly s1)
2671 step gst (ICALFirstInject s1 s2 i2) = do
2672 r <- step1 (adaptState gst) s1
2676 i `seq` return (Skip (ICALSecondInner s s2 i i2))
2677 -- i `seq` return (Skip (ICALFirstBuf s s2 i i2))
2678 Skip s -> return $ Skip (ICALFirstInject s s2 i2)
2682 step _ (ICALFirstBuf s1 s2 i1 i2) = do
2685 Yield x i' -> Skip (ICALSecondInner s1 s2 i' i2 x)
2686 Skip i' -> Skip (ICALFirstBuf s1 s2 i' i2)
2689 step _ (ICALSecondInner s1 s2 i1 i2 v) = do
2692 Yield x i' -> Yield x (ICALSecondInner s1 s2 i1 i' v)
2693 Skip i' -> Skip (ICALSecondInner s1 s2 i1 i' v)
2694 Stop -> Skip (ICALFirstResume s1 s2 i1 i2 v)
2697 step _ (ICALSecondInner s1 s2 i1 i2) = do
2700 Yield x i' -> Yield x (ICALSecondInner s1 s2 i1 i')
2701 Skip i' -> Skip (ICALSecondInner s1 s2 i1 i')
2702 Stop -> Skip (ICALFirstInner s1 s2 i1)
2703 -- Stop -> Skip (ICALFirstResume s1 s2 i1 i2)
2706 step _ (ICALFirstResume s1 s2 i1 i2 x) = do
2707 return $ Yield x (ICALFirstInner s1 s2 i1 i2)
2710data InterposeState s1 i1 a =
2712 -- | InterposeFirstYield s1 i1
2713 | InterposeFirstInner s1 i1
2714 | InterposeFirstInject s1
2715 -- | InterposeFirstBuf s1 i1
2716 | InterposeSecondYield s1 i1
2717 -- -- | InterposeSecondYield s1 i1 a
2718 -- -- | InterposeFirstResume s1 i1 a
2720-- Note that this only interposes the pure values, we may run many effects to
2721-- generate those values as some effects may not generate anything (Skip).
2722{-# INLINE_NORMAL interpose #-}
2723interpose :: Monad m => m c -> Unfold m b c -> Stream m b -> Stream m c
2726 (Unfold istep1 inject1) (Stream step1 state1) =
2727 Stream step (InterposeFirst state1)
2731 {-# INLINE_LATE step #-}
2732 step gst (InterposeFirst s1) = do
2733 r <- step1 (adaptState gst) s1
2737 i `seq` return (Skip (InterposeFirstInner s i))
2738 -- i `seq` return (Skip (InterposeFirstYield s i))
2739 Skip s -> return $ Skip (InterposeFirst s)
2743 step _ (InterposeFirstYield s1 i1) = do
2746 Yield x i' -> Yield x (InterposeFirstInner s1 i')
2747 Skip i' -> Skip (InterposeFirstYield s1 i')
2748 Stop -> Skip (InterposeFirst s1)
2751 step _ (InterposeFirstInner s1 i1) = do
2754 Yield x i' -> Yield x (InterposeFirstInner s1 i')
2755 Skip i' -> Skip (InterposeFirstInner s1 i')
2756 Stop -> Skip (InterposeFirstInject s1)
2758 step gst (InterposeFirstInject s1) = do
2759 r <- step1 (adaptState gst) s1
2763 -- i `seq` return (Skip (InterposeFirstBuf s i))
2764 i `seq` return (Skip (InterposeSecondYield s i))
2765 Skip s -> return $ Skip (InterposeFirstInject s)
2769 step _ (InterposeFirstBuf s1 i1) = do
2772 Yield x i' -> Skip (InterposeSecondYield s1 i' x)
2773 Skip i' -> Skip (InterposeFirstBuf s1 i')
2778 step _ (InterposeSecondYield s1 i1 v) = do
2780 return $ Yield r (InterposeFirstResume s1 i1 v)
2782 step _ (InterposeSecondYield s1 i1) = do
2784 return $ Yield r (InterposeFirstInner s1 i1)
2787 step _ (InterposeFirstResume s1 i1 v) = do
2788 return $ Yield v (InterposeFirstInner s1 i1)
2791------------------------------------------------------------------------------
2793------------------------------------------------------------------------------
2795data GbracketState s1 s2 v
2797 | GBracketNormal s1 v
2798 | GBracketException s2
2800-- | The most general bracketing and exception combinator. All other
2801-- combinators can be expressed in terms of this combinator. This can also be
2802-- used for cases which are not covered by the standard combinators.
2806{-# INLINE_NORMAL gbracket #-}
2810 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
2811 -> (c -> m d) -- ^ after, on normal stop
2812 -> (c -> e -> Stream m b) -- ^ on exception
2813 -> (c -> Stream m b) -- ^ stream generator
2815gbracket bef exc aft fexc fnormal =
2816 Stream step GBracketInit
2820 {-# INLINE_LATE step #-}
2821 step _ GBracketInit = do
2823 return $ Skip $ GBracketNormal (fnormal r) r
2825 step gst (GBracketNormal (UnStream step1 st) v) = do
2826 res <- exc $ step1 gst st
2828 Right r -> case r of
2830 return $ Yield x (GBracketNormal (Stream step1 s) v)
2831 Skip s -> return $ Skip (GBracketNormal (Stream step1 s) v)
2832 Stop -> aft v >> return Stop
2833 Left e -> return $ Skip (GBracketException (fexc v e))
2834 step gst (GBracketException (UnStream step1 st)) = do
2837 Yield x s -> return $ Yield x (GBracketException (Stream step1 s))
2838 Skip s -> return $ Skip (GBracketException (Stream step1 s))
2841-- | Run a side effect before the stream yields its first element.
2842{-# INLINE_NORMAL before #-}
2843before :: Monad m => m b -> Stream m a -> Stream m a
2844before action (Stream step state) = Stream step' Nothing
2848 {-# INLINE_LATE step' #-}
2849 step' _ Nothing = action >> return (Skip (Just state))
2851 step' gst (Just st) = do
2854 Yield x s -> return $ Yield x (Just s)
2855 Skip s -> return $ Skip (Just s)
2858-- | Run a side effect whenever the stream stops normally.
2859{-# INLINE_NORMAL after #-}
2860after :: Monad m => m b -> Stream m a -> Stream m a
2861after action (Stream step state) = Stream step' state
2865 {-# INLINE_LATE step' #-}
2869 Yield x s -> return $ Yield x s
2870 Skip s -> return $ Skip s
2871 Stop -> action >> return Stop
2873-- XXX These combinators are expensive due to the call to
2874-- onException/handle/try on each step. Therefore, when possible, they should
2875-- be called in an outer loop where we perform less iterations. For example, we
2876-- cannot call them on each iteration in a char stream, instead we can call
2877-- them when doing an IO on an array.
2879-- XXX For high performance error checks in busy streams we may need another
2880-- Error constructor in step.
2882-- | Run a side effect whenever the stream aborts due to an exception. The
2883-- exception is not caught, simply rethrown.
2884{-# INLINE_NORMAL onException #-}
2885onException :: MonadCatch m => m b -> Stream m a -> Stream m a
2886onException action str =
2887 gbracket (return ()) MC.try return
2888 (\_ (e :: MC.SomeException) -> nilM (action >> MC.throwM e))
2891{-# INLINE_NORMAL _onException #-}
2892_onException :: MonadCatch m => m b -> Stream m a -> Stream m a
2893_onException action (Stream step state) = Stream step' state
2897 {-# INLINE_LATE step' #-}
2899 res <- step gst st `MC.onException` action
2901 Yield x s -> return $ Yield x s
2902 Skip s -> return $ Skip s
2905-- XXX bracket is like concatMap, it generates a stream and then flattens it.
2906-- Like concatMap it has 10x worse performance compared to linear fused
2909-- | Run the first action before the stream starts and remember its output,
2910-- generate a stream using the output, run the second action providing the
2911-- remembered value as an argument whenever the stream ends normally or due to
2913{-# INLINE_NORMAL bracket #-}
2914bracket :: MonadCatch m => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
2915bracket bef aft bet =
2916 gbracket bef MC.try aft
2917 (\a (e :: SomeException) -> nilM (aft a >> MC.throwM e)) bet
2919data BracketState s v = BracketInit | BracketRun s v
2921{-# INLINE_NORMAL _bracket #-}
2922_bracket :: MonadCatch m => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
2923_bracket bef aft bet = Stream step' BracketInit
2927 {-# INLINE_LATE step' #-}
2928 step' _ BracketInit = bef >>= \x -> return (Skip (BracketRun (bet x) x))
2930 -- NOTE: It is important to use UnStream instead of the Stream pattern
2931 -- here, otherwise we get huge perf degradation, see note in concatMap.
2932 step' gst (BracketRun (UnStream step state) v) = do
2933 -- res <- step gst state `MC.onException` aft v
2934 res <- MC.try $ step gst state
2936 Left (e :: SomeException) -> aft v >> MC.throwM e >> return Stop
2937 Right r -> case r of
2938 Yield x s -> return $ Yield x (BracketRun (Stream step s) v)
2939 Skip s -> return $ Skip (BracketRun (Stream step s) v)
2940 Stop -> aft v >> return Stop
2942-- | Run a side effect whenever the stream stops normally or aborts due to an
2944{-# INLINE finally #-}
2945finally :: MonadCatch m => m b -> Stream m a -> Stream m a
2946-- finally action xs = after action $ onException action xs
2947finally action xs = bracket (return ()) (\_ -> action) (const xs)
2949-- | When evaluating a stream if an exception occurs, stream evaluation aborts
2950-- and the specified exception handler is run with the exception as argument.
2951{-# INLINE_NORMAL handle #-}
2952handle :: (MonadCatch m, Exception e)
2953 => (e -> Stream m a) -> Stream m a -> Stream m a
2955 gbracket (return ()) MC.try return (\_ e -> f e) (\_ -> str)
2957{-# INLINE_NORMAL _handle #-}
2958_handle :: (MonadCatch m, Exception e)
2959 => (e -> Stream m a) -> Stream m a -> Stream m a
2960_handle f (Stream step state) = Stream step' (Left state)
2964 {-# INLINE_LATE step' #-}
2965 step' gst (Left st) = do
2966 res <- MC.try $ step gst st
2968 Left e -> return $ Skip $ Right (f e)
2969 Right r -> case r of
2970 Yield x s -> return $ Yield x (Left s)
2971 Skip s -> return $ Skip (Left s)
2974 step' gst (Right (UnStream step1 st)) = do
2977 Yield x s -> return $ Yield x (Right (Stream step1 s))
2978 Skip s -> return $ Skip (Right (Stream step1 s))
2981-------------------------------------------------------------------------------
2982-- General transformation
2983-------------------------------------------------------------------------------
2985{-# INLINE_NORMAL transform #-}
2986transform :: Monad m => Pipe m a b -> Stream m a -> Stream m b
2987transform (Pipe pstep1 pstep2 pstate) (Stream step state) =
2988 Stream step' (Consume pstate, state)
2992 {-# INLINE_LATE step' #-}
2994 step' gst (Consume pst, st) = pst `seq` do
2995 r <- step (adaptState gst) st
3000 Pipe.Yield b pst' -> return $ Yield b (pst', s)
3001 Pipe.Continue pst' -> return $ Skip (pst', s)
3002 Skip s -> return $ Skip (Consume pst, s)
3005 step' _ (Produce pst, st) = pst `seq` do
3008 Pipe.Yield b pst' -> return $ Yield b (pst', st)
3009 Pipe.Continue pst' -> return $ Skip (pst', st)
3011------------------------------------------------------------------------------
3012-- Transformation by Folding (Scans)
3013------------------------------------------------------------------------------
3015------------------------------------------------------------------------------
3017------------------------------------------------------------------------------
3019-- XXX Is a prescan useful, discarding the last step does not sound useful? I
3020-- am not sure about the utility of this function, so this is implemented but
3021-- not exposed. We can expose it if someone provides good reasons why this is
3024-- XXX We have to execute the stream one step ahead to know that we are at the
3025-- last step. The vector implementation of prescan executes the last fold step
3026-- but does not yield the result. This means we have executed the effect but
3027-- discarded value. This does not sound right. In this implementation we are
3028-- not executing the last fold step.
3029{-# INLINE_NORMAL prescanlM' #-}
3030prescanlM' :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
3031prescanlM' f mz (Stream step state) = Stream step' (state, mz)
3033 {-# INLINE_LATE step' #-}
3034 step' gst (st, prev) = do
3035 r <- step (adaptState gst) st
3039 return $ Yield acc (s, f acc x)
3040 Skip s -> return $ Skip (s, prev)
3043{-# INLINE prescanl' #-}
3044prescanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
3045prescanl' f z = prescanlM' (\a b -> return (f a b)) (return z)
3047------------------------------------------------------------------------------
3048-- Monolithic postscans (postscan followed by a map)
3049------------------------------------------------------------------------------
3051-- The performance of a modular postscan followed by a map seems to be
3052-- equivalent to this monolithic scan followed by map therefore we may not need
3053-- this implementation. We just have it for performance comparison and in case
3054-- modular version does not perform well in some situation.
3056{-# INLINE_NORMAL postscanlMx' #-}
3057postscanlMx' :: Monad m
3058 => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> Stream m b
3059postscanlMx' fstep begin done (Stream step state) = do
3060 Stream step' (state, begin)
3062 {-# INLINE_LATE step' #-}
3063 step' gst (st, acc) = do
3064 r <- step (adaptState gst) st
3070 v `seq` y `seq` return (Yield v (s, return y))
3071 Skip s -> return $ Skip (s, acc)
3074{-# INLINE_NORMAL postscanlx' #-}
3075postscanlx' :: Monad m
3076 => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> Stream m b
3077postscanlx' fstep begin done s =
3078 postscanlMx' (\b a -> return (fstep b a)) (return begin) (return . done) s
3080-- XXX do we need consM strict to evaluate the begin value?
3081{-# INLINE scanlMx' #-}
3083 => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> Stream m b
3084scanlMx' fstep begin done s =
3085 (begin >>= \x -> x `seq` done x) `consM` postscanlMx' fstep begin done s
3087{-# INLINE scanlx' #-}
3089 => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> Stream m b
3090scanlx' fstep begin done s =
3091 scanlMx' (\b a -> return (fstep b a)) (return begin) (return . done) s
3093------------------------------------------------------------------------------
3095------------------------------------------------------------------------------
3097{-# INLINE_NORMAL postscanlM' #-}
3098postscanlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
3099postscanlM' fstep begin (Stream step state) =
3100 begin `seq` Stream step' (state, begin)
3102 {-# INLINE_LATE step' #-}
3103 step' gst (st, acc) = acc `seq` do
3104 r <- step (adaptState gst) st
3108 y `seq` return (Yield y (s, y))
3109 Skip s -> return $ Skip (s, acc)
3112{-# INLINE_NORMAL postscanl' #-}
3113postscanl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
3114postscanl' f = postscanlM' (\a b -> return (f a b))
3116{-# INLINE_NORMAL postscanlM #-}
3117postscanlM :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
3118postscanlM fstep begin (Stream step state) = Stream step' (state, begin)
3120 {-# INLINE_LATE step' #-}
3121 step' gst (st, acc) = do
3122 r <- step (adaptState gst) st
3126 return (Yield y (s, y))
3127 Skip s -> return $ Skip (s, acc)
3130{-# INLINE_NORMAL postscanl #-}
3131postscanl :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
3132postscanl f = postscanlM (\a b -> return (f a b))
3134{-# INLINE_NORMAL scanlM' #-}
3135scanlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
3136scanlM' fstep begin s = begin `seq` (begin `cons` postscanlM' fstep begin s)
3138{-# INLINE scanl' #-}
3139scanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
3140scanl' f = scanlM' (\a b -> return (f a b))
3142{-# INLINE_NORMAL scanlM #-}
3143scanlM :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
3144scanlM fstep begin s = begin `cons` postscanlM fstep begin s
3147scanl :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
3148scanl f = scanlM (\a b -> return (f a b))
3150{-# INLINE_NORMAL scanl1M #-}
3151scanl1M :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
3152scanl1M fstep (Stream step state) = Stream step' (state, Nothing)
3154 {-# INLINE_LATE step' #-}
3155 step' gst (st, Nothing) = do
3158 Yield x s -> return $ Yield x (s, Just x)
3159 Skip s -> return $ Skip (s, Nothing)
3162 step' gst (st, Just acc) = do
3167 return $ Yield z (s, Just z)
3168 Skip s -> return $ Skip (s, Just acc)
3171{-# INLINE scanl1 #-}
3172scanl1 :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
3173scanl1 f = scanl1M (\x y -> return (f x y))
3175{-# INLINE_NORMAL scanl1M' #-}
3176scanl1M' :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
3177scanl1M' fstep (Stream step state) = Stream step' (state, Nothing)
3179 {-# INLINE_LATE step' #-}
3180 step' gst (st, Nothing) = do
3183 Yield x s -> x `seq` return $ Yield x (s, Just x)
3184 Skip s -> return $ Skip (s, Nothing)
3187 step' gst (st, Just acc) = acc `seq` do
3192 z `seq` return $ Yield z (s, Just z)
3193 Skip s -> return $ Skip (s, Just acc)
3196{-# INLINE scanl1' #-}
3197scanl1' :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
3198scanl1' f = scanl1M' (\x y -> return (f x y))
3201tap :: Monad m => Fold m a b -> Stream m a -> Stream m a
3202tap (Fold fstep initial extract) (Stream step state) = Stream step' Nothing
3206 step' _ Nothing = do
3208 return $ Skip (Just (r, state))
3210 step' gst (Just (acc, st)) = acc `seq` do
3215 return $ Yield x (Just (acc', s))
3216 Skip s -> return $ Skip (Just (acc, s))
3221-------------------------------------------------------------------------------
3223-------------------------------------------------------------------------------
3225{-# INLINE_NORMAL takeWhileM #-}
3226takeWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
3227takeWhileM f (Stream step state) = Stream step' state
3229 {-# INLINE_LATE step' #-}
3235 return $ if b then Yield x s else Stop
3236 Skip s -> return $ Skip s
3239{-# INLINE takeWhile #-}
3240takeWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
3241takeWhile f = takeWhileM (return . f)
3243{-# INLINE_NORMAL drop #-}
3244drop :: Monad m => Int -> Stream m a -> Stream m a
3245drop n (Stream step state) = Stream step' (state, Just n)
3247 {-# INLINE_LATE step' #-}
3248 step' gst (st, Just i)
3253 Yield _ s -> Skip (s, Just (i - 1))
3254 Skip s -> Skip (s, Just i)
3256 | otherwise = return $ Skip (st, Nothing)
3258 step' gst (st, Nothing) = do
3262 Yield x s -> Yield x (s, Nothing)
3263 Skip s -> Skip (s, Nothing)
3266data DropWhileState s a
3268 | DropWhileYield a s
3271{-# INLINE_NORMAL dropWhileM #-}
3272dropWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
3273dropWhileM f (Stream step state) = Stream step' (DropWhileDrop state)
3275 {-# INLINE_LATE step' #-}
3276 step' gst (DropWhileDrop st) = do
3282 then return $ Skip (DropWhileDrop s)
3283 else return $ Skip (DropWhileYield x s)
3284 Skip s -> return $ Skip (DropWhileDrop s)
3287 step' gst (DropWhileNext st) = do
3290 Yield x s -> return $ Skip (DropWhileYield x s)
3291 Skip s -> return $ Skip (DropWhileNext s)
3294 step' _ (DropWhileYield x st) = return $ Yield x (DropWhileNext st)
3296{-# INLINE dropWhile #-}
3297dropWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
3298dropWhile f = dropWhileM (return . f)
3300{-# INLINE_NORMAL filterM #-}
3301filterM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
3302filterM f (Stream step state) = Stream step' state
3304 {-# INLINE_LATE step' #-}
3313 Skip s -> return $ Skip s
3316{-# INLINE filter #-}
3317filter :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
3318filter f = filterM (return . f)
3320{-# INLINE_NORMAL uniq #-}
3321uniq :: (Eq a, Monad m) => Stream m a -> Stream m a
3322uniq (Stream step state) = Stream step' (Nothing, state)
3324 {-# INLINE_LATE step' #-}
3325 step' gst (Nothing, st) = do
3328 Yield x s -> return $ Yield x (Just x, s)
3329 Skip s -> return $ Skip (Nothing, s)
3331 step' gst (Just x, st) = do
3334 Yield y s | x == y -> return $ Skip (Just x, s)
3335 | otherwise -> return $ Yield y (Just y, s)
3336 Skip s -> return $ Skip (Just x, s)
3339------------------------------------------------------------------------------
3340-- Transformation by Mapping
3341------------------------------------------------------------------------------
3343{-# INLINE_NORMAL sequence #-}
3344sequence :: Monad m => Stream m (m a) -> Stream m a
3345sequence (Stream step state) = Stream step' state
3347 {-# INLINE_LATE step' #-}
3349 r <- step (adaptState gst) st
3351 Yield x s -> x >>= \a -> return (Yield a s)
3352 Skip s -> return $ Skip s
3355------------------------------------------------------------------------------
3357------------------------------------------------------------------------------
3359data LoopState x s = FirstYield s
3360 | InterspersingYield s
3363{-# INLINE_NORMAL intersperseM #-}
3364intersperseM :: Monad m => m a -> Stream m a -> Stream m a
3365intersperseM m (Stream step state) = Stream step' (FirstYield state)
3367 {-# INLINE_LATE step' #-}
3368 step' gst (FirstYield st) = do
3372 Yield x s -> Skip (YieldAndCarry x s)
3373 Skip s -> Skip (FirstYield s)
3376 step' gst (InterspersingYield st) = do
3381 return $ Yield a (YieldAndCarry x s)
3382 Skip s -> return $ Skip $ InterspersingYield s
3385 step' _ (YieldAndCarry x st) = return $ Yield x (InterspersingYield st)
3390 | SuffixYield a (SuffixState s a)
3392{-# INLINE_NORMAL intersperseSuffix #-}
3393intersperseSuffix :: forall m a. Monad m => m a -> Stream m a -> Stream m a
3394intersperseSuffix action (Stream step state) = Stream step' (SuffixElem state)
3396 {-# INLINE_LATE step' #-}
3397 step' gst (SuffixElem st) = do
3400 Yield x s -> Skip (SuffixYield x (SuffixSuffix s))
3401 Skip s -> Skip (SuffixElem s)
3404 step' _ (SuffixSuffix st) = do
3405 action >>= \r -> return $ Skip (SuffixYield r (SuffixElem st))
3407 step' _ (SuffixYield x next) = return $ Yield x next
3409data SuffixSpanState s a
3410 = SuffixSpanElem s Int
3411 | SuffixSpanSuffix s
3412 | SuffixSpanYield a (SuffixSpanState s a)
3416-- | intersperse after every n items
3417{-# INLINE_NORMAL intersperseSuffixBySpan #-}
3418intersperseSuffixBySpan :: forall m a. Monad m
3419 => Int -> m a -> Stream m a -> Stream m a
3420intersperseSuffixBySpan n action (Stream step state) =
3421 Stream step' (SuffixSpanElem state n)
3423 {-# INLINE_LATE step' #-}
3424 step' gst (SuffixSpanElem st i) | i > 0 = do
3427 Yield x s -> Skip (SuffixSpanYield x (SuffixSpanElem s (i - 1)))
3428 Skip s -> Skip (SuffixSpanElem s i)
3429 Stop -> if i == n then Stop else Skip SuffixSpanLast
3430 step' _ (SuffixSpanElem st _) = return $ Skip (SuffixSpanSuffix st)
3432 step' _ (SuffixSpanSuffix st) = do
3433 action >>= \r -> return $ Skip (SuffixSpanYield r (SuffixSpanElem st n))
3435 step' _ (SuffixSpanLast) = do
3436 action >>= \r -> return $ Skip (SuffixSpanYield r SuffixSpanStop)
3438 step' _ (SuffixSpanYield x next) = return $ Yield x next
3440 step' _ (SuffixSpanStop) = return Stop
3442{-# INLINE intersperse #-}
3443intersperse :: Monad m => a -> Stream m a -> Stream m a
3444intersperse a = intersperseM (return a)
3446{-# INLINE_NORMAL insertBy #-}
3447insertBy :: Monad m => (a -> a -> Ordering) -> a -> Stream m a -> Stream m a
3448insertBy cmp a (Stream step state) = Stream step' (state, False, Nothing)
3450 {-# INLINE_LATE step' #-}
3451 step' gst (st, False, _) = do
3454 Yield x s -> case cmp a x of
3455 GT -> return $ Yield x (s, False, Nothing)
3456 _ -> return $ Yield a (s, True, Just x)
3457 Skip s -> return $ Skip (s, False, Nothing)
3458 Stop -> return $ Yield a (st, True, Nothing)
3460 step' _ (_, True, Nothing) = return Stop
3462 step' gst (st, True, Just prev) = do
3465 Yield x s -> return $ Yield prev (s, True, Just x)
3466 Skip s -> return $ Skip (s, True, Just prev)
3467 Stop -> return $ Yield prev (st, True, Nothing)
3469------------------------------------------------------------------------------
3471------------------------------------------------------------------------------
3473{-# INLINE_NORMAL deleteBy #-}
3474deleteBy :: Monad m => (a -> a -> Bool) -> a -> Stream m a -> Stream m a
3475deleteBy eq x (Stream step state) = Stream step' (state, False)
3477 {-# INLINE_LATE step' #-}
3478 step' gst (st, False) = do
3481 Yield y s -> return $
3482 if eq x y then Skip (s, True) else Yield y (s, False)
3483 Skip s -> return $ Skip (s, False)
3486 step' gst (st, True) = do
3489 Yield y s -> return $ Yield y (s, True)
3490 Skip s -> return $ Skip (s, True)
3493------------------------------------------------------------------------------
3494-- Transformation by Map and Filter
3495------------------------------------------------------------------------------
3497-- XXX Will this always fuse properly?
3498{-# INLINE_NORMAL mapMaybe #-}
3499mapMaybe :: Monad m => (a -> Maybe b) -> Stream m a -> Stream m b
3500mapMaybe f = fmap fromJust . filter isJust . map f
3502{-# INLINE_NORMAL mapMaybeM #-}
3503mapMaybeM :: Monad m => (a -> m (Maybe b)) -> Stream m a -> Stream m b
3504mapMaybeM f = fmap fromJust . filter isJust . mapM f
3506------------------------------------------------------------------------------
3508------------------------------------------------------------------------------
3510{-# INLINE_NORMAL indexed #-}
3511indexed :: Monad m => Stream m a -> Stream m (Int, a)
3512indexed (Stream step state) = Stream step' (state, 0)
3514 {-# INLINE_LATE step' #-}
3515 step' gst (st, i) = i `seq` do
3516 r <- step (adaptState gst) st
3518 Yield x s -> return $ Yield (i, x) (s, i+1)
3519 Skip s -> return $ Skip (s, i)
3522{-# INLINE_NORMAL indexedR #-}
3523indexedR :: Monad m => Int -> Stream m a -> Stream m (Int, a)
3524indexedR m (Stream step state) = Stream step' (state, m)
3526 {-# INLINE_LATE step' #-}
3527 step' gst (st, i) = i `seq` do
3528 r <- step (adaptState gst) st
3530 Yield x s -> let i' = i - 1
3531 in return $ Yield (i, x) (s, i')
3532 Skip s -> return $ Skip (s, i)
3535{-# INLINE_NORMAL zipWithM #-}
3537 => (a -> b -> m c) -> Stream m a -> Stream m b -> Stream m c
3538zipWithM f (Stream stepa ta) (Stream stepb tb) = Stream step (ta, tb, Nothing)
3540 {-# INLINE_LATE step #-}
3541 step gst (sa, sb, Nothing) = do
3542 r <- stepa (adaptState gst) sa
3545 Yield x sa' -> Skip (sa', sb, Just x)
3546 Skip sa' -> Skip (sa', sb, Nothing)
3549 step gst (sa, sb, Just x) = do
3550 r <- stepb (adaptState gst) sb
3554 return $ Yield z (sa, sb', Nothing)
3555 Skip sb' -> return $ Skip (sa, sb', Just x)
3558#if __GLASGOW_HASKELL__ >= 801
3559{-# RULES "zipWithM xs xs"
3560 forall f xs. zipWithM @Identity f xs xs = mapM (\x -> f x x) xs #-}
3563{-# INLINE zipWith #-}
3564zipWith :: Monad m => (a -> b -> c) -> Stream m a -> Stream m b -> Stream m c
3565zipWith f = zipWithM (\a b -> return (f a b))
3567-- | Like 'zipWithM' but zips concurrently i.e. both the streams being zipped
3568-- are generated concurrently.
3570{-# INLINE_NORMAL zipAsyncWithM #-}
3571zipAsyncWithM :: MonadAsync m
3572 => (a -> b -> m c) -> Stream m a -> Stream m b -> Stream m c
3573zipAsyncWithM f m1 m2 = Stream step Nothing
3577 {-# INLINE_LATE step #-}
3578 step gst Nothing = do
3579 ma <- mkParallel (adaptState gst) m1
3580 mb <- mkParallel (adaptState gst) m2
3581 return $ Skip (Just $ zipWithM f ma mb)
3583 step gst (Just (UnStream step1 st)) = do
3586 Yield a s -> Yield a (Just $ Stream step1 s)
3587 Skip s -> Skip (Just $ Stream step1 s)
3590-- | Like 'zipWith' but zips concurrently i.e. both the streams being zipped
3591-- are generated concurrently.
3593{-# INLINE zipAsyncWith #-}
3594zipAsyncWith :: MonadAsync m
3595 => (a -> b -> c) -> Stream m a -> Stream m b -> Stream m c
3596zipAsyncWith f = zipAsyncWithM (\a b -> return (f a b))
3598------------------------------------------------------------------------------
3600------------------------------------------------------------------------------
3602{-# INLINE_NORMAL mergeByM #-}
3605 => (a -> a -> m Ordering) -> Stream m a -> Stream m a -> Stream m a
3606mergeByM cmp (Stream stepa ta) (Stream stepb tb) =
3607 Stream step (Just ta, Just tb, Nothing, Nothing)
3609 {-# INLINE_LATE step #-}
3611 -- one of the values is missing, and the corresponding stream is running
3612 step gst (Just sa, sb, Nothing, b) = do
3615 Yield a sa' -> Skip (Just sa', sb, Just a, b)
3616 Skip sa' -> Skip (Just sa', sb, Nothing, b)
3617 Stop -> Skip (Nothing, sb, Nothing, b)
3619 step gst (sa, Just sb, a, Nothing) = do
3622 Yield b sb' -> Skip (sa, Just sb', a, Just b)
3623 Skip sb' -> Skip (sa, Just sb', a, Nothing)
3624 Stop -> Skip (sa, Nothing, a, Nothing)
3626 -- both the values are available
3627 step _ (sa, sb, Just a, Just b) = do
3629 return $ case res of
3630 GT -> Yield b (sa, sb, Just a, Nothing)
3631 _ -> Yield a (sa, sb, Nothing, Just b)
3633 -- one of the values is missing, corresponding stream is done
3634 step _ (Nothing, sb, Nothing, Just b) =
3635 return $ Yield b (Nothing, sb, Nothing, Nothing)
3637 step _ (sa, Nothing, Just a, Nothing) =
3638 return $ Yield a (sa, Nothing, Nothing, Nothing)
3640 step _ (Nothing, Nothing, Nothing, Nothing) = return Stop
3642{-# INLINE mergeBy #-}
3645 => (a -> a -> Ordering) -> Stream m a -> Stream m a -> Stream m a
3646mergeBy cmp = mergeByM (\a b -> return $ cmp a b)
3648------------------------------------------------------------------------------
3649-- Transformation comprehensions
3650------------------------------------------------------------------------------
3652{-# INLINE_NORMAL the #-}
3653the :: (Eq a, Monad m) => Stream m a -> m (Maybe a)
3654the (Stream step state) = go state
3657 r <- step defState st
3659 Yield x s -> go' x s
3661 Stop -> return Nothing
3663 r <- step defState st
3665 Yield x s | x == n -> go' n s
3666 | otherwise -> return Nothing
3668 Stop -> return (Just n)
3670{-# INLINE runFold #-}
3671runFold :: (Monad m) => Fold m a b -> Stream m a -> m b
3672runFold (Fold step begin done) = foldlMx' step begin done
3674-------------------------------------------------------------------------------
3675-- Concurrent application and fold
3676-------------------------------------------------------------------------------
3678-- Using StreamD the worker stream producing code can fuse with the code to
3679-- queue output to the SVar giving some perf boost.
3681-- Note that StreamD can only be used in limited situations, specifically, we
3682-- cannot implement joinStreamVarPar using this.
3684-- XXX make sure that the SVar passed is a Parallel style SVar.
3686-- | Fold the supplied stream to the SVar asynchronously using Parallel
3687-- concurrency style.
3688-- {-# INLINE_NORMAL toSVarParallel #-}
3689{-# INLINE toSVarParallel #-}
3690toSVarParallel :: MonadAsync m
3691 => State t m a -> SVar t m a -> Stream m a -> m ()
3692toSVarParallel st sv xs =
3693 if svarInspectMode sv
3697 case getYieldLimit st of
3698 Nothing -> doFork (work Nothing)
3700 (handleChildException sv)
3701 Just _ -> doFork (workLim Nothing)
3703 (handleChildException sv)
3708 {-# NOINLINE work #-}
3709 work info = (runFold (FL.toParallelSVar sv info) xs)
3711 {-# NOINLINE workLim #-}
3712 workLim info = runFold (FL.toParallelSVarLimited sv info) xs
3714 {-# NOINLINE forkWithDiag #-}
3716 -- We do not use workerCount in case of ParallelVar but still there is
3717 -- no harm in maintaining it correctly.
3718 liftIO $ atomicModifyIORefCAS_ (workerCount sv) $ \n -> n + 1
3720 -- This allocation matters when significant number of workers are being
3721 -- sent. We allocate it only when needed. The overhead increases by 4x.
3723 case yieldRateInfo sv of
3724 Nothing -> return Nothing
3725 Just _ -> liftIO $ do
3726 cntRef <- newIORef 0
3727 t <- getTime Monotonic
3728 lat <- newIORef (0, t)
3729 return $ Just WorkerInfo
3730 { workerYieldMax = 0
3731 , workerYieldCount = cntRef
3732 , workerLatencyStart = lat
3735 case getYieldLimit st of
3736 Nothing -> doFork (work winfo)
3738 (handleChildException sv)
3739 Just _ -> doFork (workLim winfo)
3741 (handleChildException sv)
3744{-# INLINE_NORMAL mkParallel #-}
3745mkParallel :: MonadAsync m => State t m a -> Stream m a -> m (Stream m a)
3747 sv <- newParallelVar StopNone defState
3748 toSVarParallel st sv m
3749 return $ fromSVar sv
3751{-# INLINE_NORMAL applyParallel #-}
3752applyParallel :: MonadAsync m
3753 => (Stream m a -> Stream m b) -> Stream m a -> Stream m b
3754applyParallel f m = Stream step Nothing
3757 {-# INLINE_LATE step #-}
3758 step gst Nothing = do
3759 sv <- newParallelVar StopNone (adaptState gst)
3760 toSVarParallel (adaptState gst) sv m
3761 return $ Skip $ Just $ f (fromSVar sv)
3763 step gst (Just (UnStream step1 st)) = do
3766 Yield a s -> Yield a (Just $ Stream step1 s)
3767 Skip s -> Skip (Just $ Stream step1 s)
3770{-# INLINE foldParallel #-}
3771foldParallel :: (K.IsStream t, MonadAsync m) => (t m a -> m b) -> t m a -> m b
3772foldParallel f m = do
3773 sv <- newParallelVar StopNone defState
3774 toSVarParallel defState sv (toStreamD m)
3775 f $ fromStreamD $ fromSVar sv
3777-------------------------------------------------------------------------------
3779-------------------------------------------------------------------------------
3781-- | Create an SVar with a fold consumer that will fold any elements sent to it
3782-- using the supplied fold function.
3783{-# INLINE newFoldSVar #-}
3784newFoldSVar :: MonadAsync m => State t m a -> Fold m a b -> m (SVar t m a)
3785newFoldSVar stt f = do
3786 -- Buffer size for the SVar is derived from the current state
3787 sv <- newParallelVar StopAny (adaptState stt)
3788 -- Add the producer thread-id to the SVar.
3789 liftIO myThreadId >>= modifyThread sv
3790 void $ doFork (work sv) (svarMrun sv) (handleFoldException sv)
3795 {-# NOINLINE work #-}
3796 work sv = void $ runFold f $ fromProducer sv
3798data TapState sv st = TapInit | Tapping sv st | TapDone st
3800{-# INLINE_NORMAL tapAsync #-}
3801tapAsync :: MonadAsync m => Fold m a b -> Stream m a -> Stream m a
3802tapAsync f (Stream step1 state1) = Stream step TapInit
3806 -- In general, a Stop event would come equipped with the result
3807 -- of the fold. It is not used here but it would be useful in
3808 -- applicative and distribute.
3809 done <- fromConsumer svr
3810 when (not done) $ do
3811 liftIO $ withDiagMVar svr "teeToSVar: waiting to drain"
3812 $ takeMVar (outputDoorBellFromConsumer svr)
3816 liftIO $ sendStop svr Nothing
3817 -- drain/wait until a stop event arrives from the fold.
3820 {-# INLINE_LATE step #-}
3821 step gst TapInit = do
3822 sv <- newFoldSVar gst f
3823 return $ Skip (Tapping sv state1)
3825 step gst (Tapping sv st) = do
3829 done <- pushToFold sv a
3832 -- XXX we do not need to wait synchronously here
3834 return $ Yield a (TapDone s)
3835 else return $ Yield a (Tapping sv s)
3836 Skip s -> return $ Skip (Tapping sv s)
3841 step gst (TapDone st) = do
3844 Yield a s -> Yield a (TapDone s)
3845 Skip s -> Skip (TapDone 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