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
151 -- ** Flattening nested streams
154 , ConcatMapUState (..)
156 , ConcatUnfoldInterleaveState (..)
157 , concatUnfoldInterleave
158 , concatUnfoldRoundrobin
161 , InterleaveState(..)
166 , roundRobin -- interleaveFair?/ParallelFair
199 -- | Transform a stream into another type.
216 -- ** By folding (scans)
259 , intersperseSuffixBySpan
285 -- * Transformation comprehensions
297 -- * Concurrent Application
304import Control.Concurrent (killThread, myThreadId, takeMVar, threadDelay)
305import Control.Exception
306 (Exception, SomeException, AsyncException, fromException)
307import Control.Monad (void, when)
308import Control.Monad.Catch (MonadCatch, throwM)
309import Control.Monad.IO.Class (MonadIO(..))
310import Control.Monad.Reader (ReaderT)
311import Control.Monad.State.Strict (StateT)
312import Control.Monad.Trans (MonadTrans(lift))
313import Data.Bits (shiftR, shiftL, (.|.), (.&.))
314import Data.Functor.Identity (Identity(..))
315import Data.IORef (newIORef, readIORef, mkWeakIORef, writeIORef, modifyIORef')
316import Data.Maybe (fromJust, isJust, isNothing)
317import Data.Word (Word32)
318import Foreign.Ptr (Ptr)
319import Foreign.Storable (Storable(..))
320import GHC.Types (SPEC(..))
321import System.Mem (performMajorGC)
323 hiding (map, mapM, mapM_, repeat, foldr, last, take, filter,
324 takeWhile, drop, dropWhile, all, any, maximum, minimum, elem,
325 notElem, null, head, tail, zipWith, lookup, foldr1, sequence,
326 (!!), scanl, scanl1, concatMap, replicate, enumFromTo, concat,
329import qualified Control.Monad.Catch as MC
330import qualified Control.Monad.Reader as Reader
331import qualified Control.Monad.State.Strict as State
332import qualified Prelude
334import Streamly.Internal.Data.Atomics (atomicModifyIORefCAS_)
335import Streamly.Internal.Memory.Array.Types (Array(..))
336import Streamly.Internal.Data.Fold.Types (Fold(..))
337import Streamly.Internal.Data.Pipe.Types (Pipe(..), PipeState(..))
338import Streamly.Internal.Data.Time.Clock (Clock(Monotonic), getTime)
339import Streamly.Internal.Data.Unfold.Types (Unfold(..))
341import Streamly.Internal.Data.Stream.StreamD.Type
342import Streamly.Internal.Data.SVar
343import Streamly.Internal.Data.Stream.SVar (fromConsumer, pushToFold)
345import qualified Streamly.Internal.Data.Pipe.Types as Pipe
346import qualified Streamly.Internal.Memory.Array.Types as A
347import qualified Streamly.Internal.Data.Fold as FL
348import qualified Streamly.Memory.Ring as RB
349import qualified Streamly.Internal.Data.Stream.StreamK as K
351------------------------------------------------------------------------------
353------------------------------------------------------------------------------
355-- | An empty 'Stream'.
356{-# INLINE_NORMAL nil #-}
357nil :: Monad m => Stream m a
358nil = Stream (\_ _ -> return Stop) ()
360-- | An empty 'Stream' with a side effect.
361{-# INLINE_NORMAL nilM #-}
362nilM :: Monad m => m b -> Stream m a
363nilM m = Stream (\_ _ -> m >> return Stop) ()
365{-# INLINE_NORMAL consM #-}
366consM :: Monad m => m a -> Stream m a -> Stream m a
367consM m (Stream step state) = Stream step1 Nothing
369 {-# INLINE_LATE step1 #-}
370 step1 _ Nothing = m >>= \x -> return $ Yield x (Just state)
371 step1 gst (Just st) = do
375 Yield a s -> Yield a (Just s)
376 Skip s -> Skip (Just s)
379-- XXX implement in terms of consM?
380-- cons x = consM (return x)
382-- | Can fuse but has O(n^2) complexity.
383{-# INLINE_NORMAL cons #-}
384cons :: Monad m => a -> Stream m a -> Stream m a
385cons x (Stream step state) = Stream step1 Nothing
387 {-# INLINE_LATE step1 #-}
388 step1 _ Nothing = return $ Yield x (Just state)
389 step1 gst (Just st) = do
393 Yield a s -> Yield a (Just s)
394 Skip s -> Skip (Just s)
397-------------------------------------------------------------------------------
399-------------------------------------------------------------------------------
401-- Does not fuse, has the same performance as the StreamK version.
402{-# INLINE_NORMAL uncons #-}
403uncons :: Monad m => Stream m a -> m (Maybe (a, Stream m a))
404uncons (UnStream step state) = go state
407 r <- step defState st
409 Yield x s -> return $ Just (x, Stream step s)
411 Stop -> return Nothing
413------------------------------------------------------------------------------
414-- Generation by unfold
415------------------------------------------------------------------------------
417{-# INLINE_NORMAL unfoldrM #-}
418unfoldrM :: Monad m => (s -> m (Maybe (a, s))) -> s -> Stream m a
419unfoldrM next state = Stream step state
421 {-# INLINE_LATE step #-}
425 Just (x, s) -> Yield x s
428{-# INLINE_LATE unfoldr #-}
429unfoldr :: Monad m => (s -> Maybe (a, s)) -> s -> Stream m a
430unfoldr f = unfoldrM (return . f)
432-- | Convert an 'Unfold' into a 'Stream' by supplying it a seed.
434{-# INLINE_NORMAL unfold #-}
435unfold :: Monad m => Unfold m a b -> a -> Stream m b
436unfold (Unfold ustep inject) seed = Stream step Nothing
438 {-# INLINE_LATE step #-}
439 step _ Nothing = inject seed >>= return . Skip . Just
440 step _ (Just st) = do
443 Yield x s -> Yield x (Just s)
444 Skip s -> Skip (Just s)
447------------------------------------------------------------------------------
448-- Specialized Generation
449------------------------------------------------------------------------------
451repeatM :: Monad m => m a -> Stream m a
452repeatM x = Stream (\_ _ -> x >>= \r -> return $ Yield r ()) ()
454repeat :: Monad m => a -> Stream m a
455repeat x = Stream (\_ _ -> return $ Yield x ()) ()
457{-# INLINE_NORMAL replicateM #-}
458replicateM :: forall m a. Monad m => Int -> m a -> Stream m a
459replicateM n p = Stream step n
461 {-# INLINE_LATE step #-}
463 | i <= 0 = return Stop
466 return $ Yield x (i - 1)
468{-# INLINE_NORMAL replicate #-}
469replicate :: Monad m => Int -> a -> Stream m a
470replicate n x = replicateM n (return x)
472-- This would not work properly for floats, therefore we put an Integral
474-- | Can be used to enumerate unbounded integrals. This does not check for
475-- overflow or underflow for bounded integrals.
476{-# INLINE_NORMAL enumerateFromStepIntegral #-}
477enumerateFromStepIntegral :: (Integral a, Monad m) => a -> a -> Stream m a
478enumerateFromStepIntegral from stride =
479 from `seq` stride `seq` Stream step from
481 {-# INLINE_LATE step #-}
482 step _ !x = return $ Yield x $! (x + stride)
484-- We are assuming that "to" is constrained by the type to be within
486{-# INLINE enumerateFromToIntegral #-}
487enumerateFromToIntegral :: (Monad m, Integral a) => a -> a -> Stream m a
488enumerateFromToIntegral from to =
489 takeWhile (<= to) $ enumerateFromStepIntegral from 1
491{-# INLINE enumerateFromIntegral #-}
492enumerateFromIntegral :: (Monad m, Integral a, Bounded a) => a -> Stream m a
493enumerateFromIntegral from = enumerateFromToIntegral from maxBound
495data EnumState a = EnumInit | EnumYield a a a | EnumStop
497{-# INLINE_NORMAL enumerateFromThenToIntegralUp #-}
498enumerateFromThenToIntegralUp
499 :: (Monad m, Integral a)
500 => a -> a -> a -> Stream m a
501enumerateFromThenToIntegralUp from next to = Stream step EnumInit
503 {-# INLINE_LATE step #-}
509 else Yield from EnumStop
510 else -- from <= next <= to
511 let stride = next - from
512 in Skip $ EnumYield from stride (to - stride)
514 step _ (EnumYield x stride toMinus) =
517 then Yield x EnumStop
518 else Yield x $ EnumYield (x + stride) stride toMinus
520 step _ EnumStop = return Stop
522{-# INLINE_NORMAL enumerateFromThenToIntegralDn #-}
523enumerateFromThenToIntegralDn
524 :: (Monad m, Integral a)
525 => a -> a -> a -> Stream m a
526enumerateFromThenToIntegralDn from next to = Stream step EnumInit
528 {-# INLINE_LATE step #-}
530 return $ if to > next
533 else Yield from EnumStop
534 else -- from >= next >= to
535 let stride = next - from
536 in Skip $ EnumYield from stride (to - stride)
538 step _ (EnumYield x stride toMinus) =
541 then Yield x EnumStop
542 else Yield x $ EnumYield (x + stride) stride toMinus
544 step _ EnumStop = return Stop
546{-# INLINE_NORMAL enumerateFromThenToIntegral #-}
547enumerateFromThenToIntegral
548 :: (Monad m, Integral a)
549 => a -> a -> a -> Stream m a
550enumerateFromThenToIntegral from next to
551 | next >= from = enumerateFromThenToIntegralUp from next to
552 | otherwise = enumerateFromThenToIntegralDn from next to
554{-# INLINE_NORMAL enumerateFromThenIntegral #-}
555enumerateFromThenIntegral
556 :: (Monad m, Integral a, Bounded a)
557 => a -> a -> Stream m a
558enumerateFromThenIntegral from next =
560 then enumerateFromThenToIntegralUp from next maxBound
561 else enumerateFromThenToIntegralDn from next minBound
563-- For floating point numbers if the increment is less than the precision then
564-- it just gets lost. Therefore we cannot always increment it correctly by just
566-- 9007199254740992 + 1 + 1 :: Double => 9.007199254740992e15
567-- 9007199254740992 + 2 :: Double => 9.007199254740994e15
569-- Instead we accumulate the increment counter and compute the increment
570-- every time before adding it to the starting number.
572-- This works for Integrals as well as floating point numbers, but
573-- enumerateFromStepIntegral is faster for integrals.
574{-# INLINE_NORMAL enumerateFromStepNum #-}
575enumerateFromStepNum :: (Monad m, Num a) => a -> a -> Stream m a
576enumerateFromStepNum from stride = Stream step 0
578 {-# INLINE_LATE step #-}
579 step _ !i = return $ (Yield $! (from + i * stride)) $! (i + 1)
581{-# INLINE_NORMAL numFrom #-}
582numFrom :: (Monad m, Num a) => a -> Stream m a
583numFrom from = enumerateFromStepNum from 1
585{-# INLINE_NORMAL numFromThen #-}
586numFromThen :: (Monad m, Num a) => a -> a -> Stream m a
587numFromThen from next = enumerateFromStepNum from (next - from)
589-- We cannot write a general function for Num. The only way to write code
590-- portable between the two is to use a 'Real' constraint and convert between
591-- Fractional and Integral using fromRational which is horribly slow.
592{-# INLINE_NORMAL enumerateFromToFractional #-}
593enumerateFromToFractional
594 :: (Monad m, Fractional a, Ord a)
595 => a -> a -> Stream m a
596enumerateFromToFractional from to =
597 takeWhile (<= to + 1 / 2) $ enumerateFromStepNum from 1
599{-# INLINE_NORMAL enumerateFromThenToFractional #-}
600enumerateFromThenToFractional
601 :: (Monad m, Fractional a, Ord a)
602 => a -> a -> a -> Stream m a
603enumerateFromThenToFractional from next to =
604 takeWhile predicate $ numFromThen from next
606 mid = (next - from) / 2
607 predicate | next >= from = (<= to + mid)
608 | otherwise = (>= to + mid)
610-------------------------------------------------------------------------------
611-- Generation by Conversion
612-------------------------------------------------------------------------------
614{-# INLINE_NORMAL fromIndicesM #-}
615fromIndicesM :: Monad m => (Int -> m a) -> Stream m a
616fromIndicesM gen = Stream step 0
618 {-# INLINE_LATE step #-}
621 return $ Yield x (i + 1)
623{-# INLINE fromIndices #-}
624fromIndices :: Monad m => (Int -> a) -> Stream m a
625fromIndices gen = fromIndicesM (return . gen)
627{-# INLINE_NORMAL generateM #-}
628generateM :: Monad m => Int -> (Int -> m a) -> Stream m a
629generateM n gen = n `seq` Stream step 0
631 {-# INLINE_LATE step #-}
632 step _ i | i < n = do
634 return $ Yield x (i + 1)
635 | otherwise = return Stop
637{-# INLINE generate #-}
638generate :: Monad m => Int -> (Int -> a) -> Stream m a
639generate n gen = generateM n (return . gen)
641-- XXX we need the MonadAsync constraint because of a rewrite rule.
642-- | Convert a list of monadic actions to a 'Stream'
643{-# INLINE_LATE fromListM #-}
644fromListM :: MonadAsync m => [m a] -> Stream m a
645fromListM = Stream step
647 {-# INLINE_LATE step #-}
648 step _ (m:ms) = m >>= \x -> return $ Yield x ms
649 step _ [] = return Stop
651{-# INLINE toStreamD #-}
652toStreamD :: (K.IsStream t, Monad m) => t m a -> Stream m a
653toStreamD = fromStreamK . K.toStream
655-------------------------------------------------------------------------------
656-- Generation from SVar
657-------------------------------------------------------------------------------
659data FromSVarState t m a =
661 | FromSVarRead (SVar t m a)
662 | FromSVarLoop (SVar t m a) [ChildEvent a]
663 | FromSVarDone (SVar t m a)
665{-# INLINE_NORMAL fromSVar #-}
666fromSVar :: (MonadAsync m) => SVar t m a -> Stream m a
667fromSVar svar = Stream step FromSVarInit
670 {-# INLINE_LATE step #-}
671 step _ FromSVarInit = do
672 ref <- liftIO $ newIORef ()
673 _ <- liftIO $ mkWeakIORef ref hook
674 -- when this copy of svar gets garbage collected "ref" will get
675 -- garbage collected and our GC hook will be called.
676 let sv = svar{svarRef = Just ref}
677 return $ Skip (FromSVarRead sv)
681 {-# NOINLINE hook #-}
683 when (svarInspectMode svar) $ do
684 r <- liftIO $ readIORef (svarStopTime (svarStats svar))
686 printSVar svar "SVar Garbage Collected"
688 -- If there are any SVars referenced by this SVar a GC will prompt
689 -- them to be cleaned up quickly.
690 when (svarInspectMode svar) performMajorGC
692 step _ (FromSVarRead sv) = do
693 list <- readOutputQ sv
694 -- Reversing the output is important to guarantee that we process the
695 -- outputs in the same order as they were generated by the constituent
697 return $ Skip $ FromSVarLoop sv (Prelude.reverse list)
699 step _ (FromSVarLoop sv []) = do
700 done <- postProcess sv
701 return $ Skip $ if done
702 then (FromSVarDone sv)
703 else (FromSVarRead sv)
705 step _ (FromSVarLoop sv (ev : es)) = do
707 ChildYield a -> return $ Yield a (FromSVarLoop sv es)
708 ChildStop tid e -> do
712 stop <- shouldStop tid
715 liftIO (cleanupSVar sv)
716 return $ Skip (FromSVarDone sv)
717 else return $ Skip (FromSVarLoop sv es)
719 case fromException ex of
721 return $ Skip (FromSVarLoop sv es)
722 Nothing -> liftIO (cleanupSVar sv) >> throwM ex
726 case svarStopStyle sv of
727 StopNone -> return False
728 StopAny -> return True
730 sid <- liftIO $ readIORef (svarStopBy sv)
731 return $ if tid == sid then True else False
733 step _ (FromSVarDone sv) = do
734 when (svarInspectMode sv) $ do
735 t <- liftIO $ getTime Monotonic
736 liftIO $ writeIORef (svarStopTime (svarStats sv)) (Just t)
737 liftIO $ printSVar sv "SVar Done"
740-------------------------------------------------------------------------------
741-- Process events received by a fold consumer from a stream producer
742-------------------------------------------------------------------------------
744{-# INLINE_NORMAL fromProducer #-}
745fromProducer :: (MonadAsync m) => SVar t m a -> Stream m a
746fromProducer svar = Stream step (FromSVarRead svar)
749 {-# INLINE_LATE step #-}
750 step _ (FromSVarRead sv) = do
751 list <- readOutputQ sv
752 -- Reversing the output is important to guarantee that we process the
753 -- outputs in the same order as they were generated by the constituent
755 return $ Skip $ FromSVarLoop sv (Prelude.reverse list)
757 step _ (FromSVarLoop sv []) = return $ Skip $ FromSVarRead sv
758 step _ (FromSVarLoop sv (ev : es)) = do
760 ChildYield a -> return $ Yield a (FromSVarLoop sv es)
761 ChildStop tid e -> do
765 sendStopToProducer sv
766 return $ Skip (FromSVarDone sv)
767 Just _ -> error "Bug: fromProducer: received exception"
769 step _ (FromSVarDone sv) = do
770 when (svarInspectMode sv) $ do
771 t <- liftIO $ getTime Monotonic
772 liftIO $ writeIORef (svarStopTime (svarStats sv)) (Just t)
773 liftIO $ printSVar sv "SVar Done"
776 step _ FromSVarInit = undefined
778-------------------------------------------------------------------------------
779-- Hoisting the inner monad
780-------------------------------------------------------------------------------
782{-# INLINE_NORMAL hoist #-}
783hoist :: Monad n => (forall x. m x -> n x) -> Stream m a -> Stream n a
784hoist f (Stream step state) = (Stream step' state)
786 {-# INLINE_LATE step' #-}
788 r <- f $ step (adaptState gst) st
790 Yield x s -> Yield x s
794{-# INLINE generally #-}
795generally :: Monad m => Stream Identity a -> Stream m a
796generally = hoist (return . runIdentity)
798{-# INLINE_NORMAL liftInner #-}
799liftInner :: (Monad m, MonadTrans t, Monad (t m))
800 => Stream m a -> Stream (t m) a
801liftInner (Stream step state) = Stream step' state
803 {-# INLINE_LATE step' #-}
805 r <- lift $ step (adaptState gst) st
807 Yield x s -> Yield x s
811{-# INLINE_NORMAL runReaderT #-}
812runReaderT :: Monad m => s -> Stream (ReaderT s m) a -> Stream m a
813runReaderT sval (Stream step state) = Stream step' state
815 {-# INLINE_LATE step' #-}
817 r <- Reader.runReaderT (step (adaptState gst) st) sval
819 Yield x s -> Yield x s
823{-# INLINE_NORMAL evalStateT #-}
824evalStateT :: Monad m => s -> Stream (StateT s m) a -> Stream m a
825evalStateT sval (Stream step state) = Stream step' (state, sval)
827 {-# INLINE_LATE step' #-}
828 step' gst (st, sv) = do
829 (r, sv') <- State.runStateT (step (adaptState gst) st) sv
831 Yield x s -> Yield x (s, sv')
832 Skip s -> Skip (s, sv')
835{-# INLINE_NORMAL runStateT #-}
836runStateT :: Monad m => s -> Stream (StateT s m) a -> Stream m (s, a)
837runStateT sval (Stream step state) = Stream step' (state, sval)
839 {-# INLINE_LATE step' #-}
840 step' gst (st, sv) = do
841 (r, sv') <- State.runStateT (step (adaptState gst) st) sv
843 Yield x s -> Yield (sv', x) (s, sv')
844 Skip s -> Skip (s, sv')
847------------------------------------------------------------------------------
848-- Elimination by Folds
849------------------------------------------------------------------------------
851------------------------------------------------------------------------------
853------------------------------------------------------------------------------
855{-# INLINE_NORMAL foldr1 #-}
856foldr1 :: Monad m => (a -> a -> a) -> Stream m a -> m (Maybe a)
860 Nothing -> return Nothing
861 Just (h, t) -> fmap Just (foldr f h t)
863------------------------------------------------------------------------------
865------------------------------------------------------------------------------
867{-# INLINE_NORMAL foldlT #-}
868foldlT :: (Monad m, Monad (s m), MonadTrans s)
869 => (s m b -> a -> s m b) -> s m b -> Stream m a -> s m b
870foldlT fstep begin (Stream step state) = go SPEC begin state
873 r <- lift $ step defState st
875 Yield x s -> go SPEC (fstep acc x) s
876 Skip s -> go SPEC acc s
879-- Note, this is going to have horrible performance, because of the nature of
880-- the stream type (i.e. direct stream vs CPS). Its only for reference, it is
881-- likely be practically unusable.
882{-# INLINE_NORMAL foldlS #-}
884 => (Stream m b -> a -> Stream m b) -> Stream m b -> Stream m a -> Stream m b
885foldlS fstep begin (Stream step state) = Stream step' (Left (state, begin))
887 step' gst (Left (st, acc)) = do
888 r <- step (adaptState gst) st
890 Yield x s -> Skip (Left (s, fstep acc x))
891 Skip s -> Skip (Left (s, acc))
892 Stop -> Skip (Right acc)
894 step' gst (Right (Stream stp stt)) = do
895 r <- stp (adaptState gst) stt
897 Yield x s -> Yield x (Right (Stream stp s))
898 Skip s -> Skip (Right (Stream stp s))
901------------------------------------------------------------------------------
903------------------------------------------------------------------------------
905-- | Run a streaming composition, discard the results.
906{-# INLINE_LATE drain #-}
907drain :: Monad m => Stream m a -> m ()
908-- drain = foldrM (\_ xs -> xs) (return ())
909drain (Stream step state) = go SPEC state
912 r <- step defState st
914 Yield _ s -> go SPEC s
918{-# INLINE_NORMAL null #-}
919null :: Monad m => Stream m a -> m Bool
920null m = foldrM (\_ _ -> return False) (return True) m
922{-# INLINE_NORMAL head #-}
923head :: Monad m => Stream m a -> m (Maybe a)
924head m = foldrM (\x _ -> return (Just x)) (return Nothing) m
926{-# INLINE_NORMAL headElse #-}
927headElse :: Monad m => a -> Stream m a -> m a
928headElse a m = foldrM (\x _ -> return x) (return a) m
930-- Does not fuse, has the same performance as the StreamK version.
931{-# INLINE_NORMAL tail #-}
932tail :: Monad m => Stream m a -> m (Maybe (Stream m a))
933tail (UnStream step state) = go state
936 r <- step defState st
938 Yield _ s -> return (Just $ Stream step s)
940 Stop -> return Nothing
942-- XXX will it fuse? need custom impl?
943{-# INLINE_NORMAL last #-}
944last :: Monad m => Stream m a -> m (Maybe a)
945last = foldl' (\_ y -> Just y) Nothing
947{-# INLINE_NORMAL elem #-}
948elem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
949-- elem e m = foldrM (\x xs -> if x == e then return True else xs) (return False) m
950elem e (Stream step state) = go state
953 r <- step defState st
956 | x == e -> return True
961{-# INLINE_NORMAL notElem #-}
962notElem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
963notElem e s = fmap not (elem e s)
965{-# INLINE_NORMAL all #-}
966all :: Monad m => (a -> Bool) -> Stream m a -> m Bool
967-- all p m = foldrM (\x xs -> if p x then xs else return False) (return True) m
968all p (Stream step state) = go state
971 r <- step defState st
975 | otherwise -> return False
979{-# INLINE_NORMAL any #-}
980any :: Monad m => (a -> Bool) -> Stream m a -> m Bool
981-- any p m = foldrM (\x xs -> if p x then return True else xs) (return False) m
982any p (Stream step state) = go state
985 r <- step defState st
993{-# INLINE_NORMAL maximum #-}
994maximum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
995maximum (Stream step state) = go Nothing state
998 r <- step defState st
1000 Yield x s -> go (Just x) s
1001 Skip s -> go Nothing s
1002 Stop -> return Nothing
1003 go (Just acc) st = do
1004 r <- step defState st
1007 | acc <= x -> go (Just x) s
1008 | otherwise -> go (Just acc) s
1009 Skip s -> go (Just acc) s
1010 Stop -> return (Just acc)
1012{-# INLINE_NORMAL maximumBy #-}
1013maximumBy :: Monad m => (a -> a -> Ordering) -> Stream m a -> m (Maybe a)
1014maximumBy cmp (Stream step state) = go Nothing state
1017 r <- step defState st
1019 Yield x s -> go (Just x) s
1020 Skip s -> go Nothing s
1021 Stop -> return Nothing
1022 go (Just acc) st = do
1023 r <- step defState st
1025 Yield x s -> case cmp acc x of
1026 GT -> go (Just acc) s
1028 Skip s -> go (Just acc) s
1029 Stop -> return (Just acc)
1031{-# INLINE_NORMAL minimum #-}
1032minimum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
1033minimum (Stream step state) = go Nothing state
1036 r <- step defState st
1038 Yield x s -> go (Just x) s
1039 Skip s -> go Nothing s
1040 Stop -> return Nothing
1041 go (Just acc) st = do
1042 r <- step defState st
1045 | acc <= x -> go (Just acc) s
1046 | otherwise -> go (Just x) s
1047 Skip s -> go (Just acc) s
1048 Stop -> return (Just acc)
1050{-# INLINE_NORMAL minimumBy #-}
1051minimumBy :: Monad m => (a -> a -> Ordering) -> Stream m a -> m (Maybe a)
1052minimumBy cmp (Stream step state) = go Nothing state
1055 r <- step defState st
1057 Yield x s -> go (Just x) s
1058 Skip s -> go Nothing s
1059 Stop -> return Nothing
1060 go (Just acc) st = do
1061 r <- step defState st
1063 Yield x s -> case cmp acc x of
1065 _ -> go (Just acc) s
1066 Skip s -> go (Just acc) s
1067 Stop -> return (Just acc)
1069{-# INLINE_NORMAL (!!) #-}
1070(!!) :: (Monad m) => Stream m a -> Int -> m (Maybe a)
1071(Stream step state) !! i = go i state
1074 r <- step defState st
1076 Yield x s | n < 0 -> return Nothing
1077 | n == 0 -> return $ Just x
1078 | otherwise -> go (n - 1) s
1080 Stop -> return Nothing
1082{-# INLINE_NORMAL lookup #-}
1083lookup :: (Monad m, Eq a) => a -> Stream m (a, b) -> m (Maybe b)
1084lookup e m = foldrM (\(a, b) xs -> if e == a then return (Just b) else xs)
1087{-# INLINE_NORMAL findM #-}
1088findM :: Monad m => (a -> m Bool) -> Stream m a -> m (Maybe a)
1089findM p m = foldrM (\x xs -> p x >>= \r -> if r then return (Just x) else xs)
1093find :: Monad m => (a -> Bool) -> Stream m a -> m (Maybe a)
1094find p = findM (return . p)
1096{-# INLINE_NORMAL findIndices #-}
1097findIndices :: Monad m => (a -> Bool) -> Stream m a -> Stream m Int
1098findIndices p (Stream step state) = Stream step' (state, 0)
1100 {-# INLINE_LATE step' #-}
1101 step' gst (st, i) = i `seq` do
1102 r <- step (adaptState gst) st
1104 Yield x s -> if p x then Yield i (s, i+1) else Skip (s, i+1)
1105 Skip s -> Skip (s, i)
1108{-# INLINE toListRev #-}
1109toListRev :: Monad m => Stream m a -> m [a]
1110toListRev = foldl' (flip (:)) []
1112-- We can implement reverse as:
1114-- > reverse = foldlS (flip cons) nil
1116-- However, this implementation is unusable because of the horrible performance
1117-- of cons. So we just convert it to a list first and then stream from the
1120-- XXX Maybe we can use an Array instead of a list here?
1121{-# INLINE_NORMAL reverse #-}
1122reverse :: Monad m => Stream m a -> Stream m a
1123reverse m = Stream step Nothing
1125 {-# INLINE_LATE step #-}
1128 return $ Skip (Just xs)
1129 step _ (Just (x:xs)) = return $ Yield x (Just xs)
1130 step _ (Just []) = return Stop
1132-- Much faster reverse for Storables
1133{-# INLINE_NORMAL reverse' #-}
1134reverse' :: forall m a. (MonadIO m, Storable a) => Stream m a -> Stream m a
1136-- This commented implementation copies the whole stream into one single array
1137-- and then streams from that array, this is 3-4x faster than the chunked code
1138-- that follows. Though this could be problematic due to unbounded large
1139-- allocations. We need to figure out why the chunked code is slower and if we
1140-- can optimize the chunked code to work as fast as this one. It may be a
1142import Foreign.ForeignPtr (touchForeignPtr)
1143import Foreign.ForeignPtr.Unsafe (unsafeForeignPtrToPtr)
1144import Foreign.Ptr (Ptr, plusPtr)
1145reverse' m = Stream step Nothing
1147 {-# INLINE_LATE step #-}
1149 arr <- A.fromStreamD m
1150 let p = aEnd arr `plusPtr` negate (sizeOf (undefined :: a))
1151 return $ Skip $ Just (aStart arr, p)
1153 step _ (Just (start, p)) | p < unsafeForeignPtrToPtr start = return Stop
1155 step _ (Just (start, p)) = do
1156 let !x = A.unsafeInlineIO $ do
1158 touchForeignPtr start
1160 next = p `plusPtr` negate (sizeOf (undefined :: a))
1161 return $ Yield x (Just (start, next))
1168 $ A.fromStreamDArraysOf A.defaultChunkSize m
1171------------------------------------------------------------------------------
1172-- Grouping/Splitting
1173------------------------------------------------------------------------------
1175{-# INLINE_NORMAL splitSuffixBy' #-}
1176splitSuffixBy' :: Monad m
1177 => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1178splitSuffixBy' predicate f (Stream step state) =
1179 Stream (stepOuter f) (Just state)
1183 {-# INLINE_LATE stepOuter #-}
1184 stepOuter (Fold fstep initial done) gst (Just st) = do
1185 res <- step (adaptState gst) st
1191 then done acc' >>= \val -> return $ Yield val (Just s)
1194 Skip s -> return $ Skip $ Just s
1200 res <- step (adaptState gst) stt
1205 then done acc' >>= \val -> return $ Yield val (Just s)
1207 Skip s -> go SPEC s acc
1208 Stop -> done acc >>= \val -> return $ Yield val Nothing
1210 stepOuter _ _ Nothing = return Stop
1212{-# INLINE_NORMAL groupsBy #-}
1218groupsBy cmp f (Stream step state) = Stream (stepOuter f) (Just state, Nothing)
1222 {-# INLINE_LATE stepOuter #-}
1223 stepOuter (Fold fstep initial done) gst (Just st, Nothing) = do
1224 res <- step (adaptState gst) st
1231 Skip s -> return $ Skip $ (Just s, Nothing)
1236 go !_ prev stt !acc = do
1237 res <- step (adaptState gst) stt
1244 else done acc >>= \r -> return $ Yield r (Just s, Just x)
1245 Skip s -> go SPEC prev s acc
1246 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1248 stepOuter (Fold fstep initial done) gst (Just st, Just prev) = do
1250 acc' <- fstep acc prev
1255 -- XXX code duplicated from the previous equation
1257 res <- step (adaptState gst) stt
1264 else done acc >>= \r -> return $ Yield r (Just s, Just x)
1265 Skip s -> go SPEC s acc
1266 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1268 stepOuter _ _ (Nothing,_) = return Stop
1270{-# INLINE_NORMAL groupsRollingBy #-}
1271groupsRollingBy :: Monad m
1276groupsRollingBy cmp f (Stream step state) =
1277 Stream (stepOuter f) (Just state, Nothing)
1280 {-# INLINE_LATE stepOuter #-}
1281 stepOuter (Fold fstep initial done) gst (Just st, Nothing) = do
1282 res <- step (adaptState gst) st
1289 Skip s -> return $ Skip $ (Just s, Nothing)
1293 go !_ prev stt !acc = do
1294 res <- step (adaptState gst) stt
1302 done acc >>= \r -> return $ Yield r (Just s, Just x)
1303 Skip s -> go SPEC prev s acc
1304 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1306 stepOuter (Fold fstep initial done) gst (Just st, Just prev') = do
1308 acc' <- fstep acc prev'
1309 go SPEC prev' st acc'
1312 go !_ prevv stt !acc = do
1313 res <- step (adaptState gst) stt
1320 else done acc >>= \r -> return $ Yield r (Just s, Just x)
1321 Skip s -> go SPEC prevv s acc
1322 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1324 stepOuter _ _ (Nothing, _) = return Stop
1326{-# INLINE_NORMAL splitBy #-}
1327splitBy :: Monad m => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1328splitBy predicate f (Stream step state) = Stream (step' f) (Just state)
1332 {-# INLINE_LATE step' #-}
1333 step' (Fold fstep initial done) gst (Just st) = initial >>= go SPEC st
1338 res <- step (adaptState gst) stt
1342 then done acc >>= \r -> return $ Yield r (Just s)
1346 Skip s -> go SPEC s acc
1347 Stop -> done acc >>= \r -> return $ Yield r Nothing
1349 step' _ _ Nothing = return Stop
1351-- XXX requires -funfolding-use-threshold=150 in lines-unlines benchmark
1352{-# INLINE_NORMAL splitSuffixBy #-}
1353splitSuffixBy :: Monad m
1354 => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1355splitSuffixBy predicate f (Stream step state) = Stream (step' f) (Just state)
1359 {-# INLINE_LATE step' #-}
1360 step' (Fold fstep initial done) gst (Just st) = do
1361 res <- step (adaptState gst) st
1366 then done acc >>= \val -> return $ Yield val (Just s)
1371 Skip s -> return $ Skip $ Just s
1377 res <- step (adaptState gst) stt
1381 then done acc >>= \r -> return $ Yield r (Just s)
1385 Skip s -> go SPEC s acc
1386 Stop -> done acc >>= \r -> return $ Yield r Nothing
1388 step' _ _ Nothing = return Stop
1390{-# INLINE_NORMAL wordsBy #-}
1391wordsBy :: Monad m => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1392wordsBy predicate f (Stream step state) = Stream (stepOuter f) (Just state)
1396 {-# INLINE_LATE stepOuter #-}
1397 stepOuter (Fold fstep initial done) gst (Just st) = do
1398 res <- step (adaptState gst) st
1402 then return $ Skip (Just s)
1408 Skip s -> return $ Skip $ Just s
1414 res <- step (adaptState gst) stt
1418 then done acc >>= \r -> return $ Yield r (Just s)
1422 Skip s -> go SPEC s acc
1423 Stop -> done acc >>= \r -> return $ Yield r Nothing
1425 stepOuter _ _ Nothing = return Stop
1427-- String search algorithms:
1428-- http://www-igm.univ-mlv.fr/~lecroq/string/index.html
1431-- TODO can we unify the splitting operations using a splitting configuration
1432-- like in the split package.
1434data SplitStyle = Infix | Suffix | Prefix deriving (Eq, Show)
1436data SplitOptions = SplitOptions
1437 { style :: SplitStyle
1438 , withSep :: Bool -- ^ keep the separators in output
1439 -- , compact :: Bool -- ^ treat multiple consecutive separators as one
1440 -- , trimHead :: Bool -- ^ drop blank at head
1441 -- , trimTail :: Bool -- ^ drop blank at tail
1445data SplitOnState s a =
1450 | GO_KARP_RABIN s !(RB.Ring a) !(Ptr a)
1453{-# INLINE_NORMAL splitOn #-}
1455 :: forall m a b. (MonadIO m, Storable a, Enum a, Eq a)
1460splitOn patArr@Array{..} (Fold fstep initial done) (Stream step state) =
1461 Stream stepOuter GO_START
1465 patLen = A.length patArr
1466 maxIndex = patLen - 1
1467 elemBits = sizeOf (undefined :: a) * 8
1469 {-# INLINE_LATE stepOuter #-}
1470 stepOuter _ GO_START =
1472 then return $ Skip $ GO_EMPTY_PAT state
1475 r <- liftIO $ (A.unsafeIndexIO patArr 0)
1476 return $ Skip $ GO_SINGLE_PAT state r
1477 else if sizeOf (undefined :: a) * patLen
1478 <= sizeOf (undefined :: Word)
1479 then return $ Skip $ GO_SHORT_PAT state
1481 (rb, rhead) <- liftIO $ RB.new patLen
1482 return $ Skip $ GO_KARP_RABIN state rb rhead
1484 stepOuter gst (GO_SINGLE_PAT stt pat) = initial >>= go SPEC stt
1489 res <- step (adaptState gst) st
1495 return $ Yield r (GO_SINGLE_PAT s pat)
1496 else fstep acc x >>= go SPEC s
1497 Skip s -> go SPEC s acc
1498 Stop -> done acc >>= \r -> return $ Yield r GO_DONE
1500 stepOuter gst (GO_SHORT_PAT stt) = initial >>= go0 SPEC 0 (0 :: Word) stt
1505 mask = (1 `shiftL` (elemBits * patLen)) - 1
1507 addToWord wrd a = (wrd `shiftL` elemBits) .|. fromIntegral (fromEnum a)
1510 patWord = mask .&. A.foldl' addToWord 0 patArr
1512 go0 !_ !idx wrd st !acc = do
1513 res <- step (adaptState gst) st
1516 let wrd' = addToWord wrd x
1519 if wrd' .&. mask == patWord
1522 return $ Yield r (GO_SHORT_PAT s)
1523 else go1 SPEC wrd' s acc
1524 else go0 SPEC (idx + 1) wrd' s acc
1525 Skip s -> go0 SPEC idx wrd s acc
1528 then go2 wrd idx acc
1530 done acc' >>= \r -> return $ Yield r GO_DONE
1533 go1 !_ wrd st !acc = do
1534 res <- step (adaptState gst) st
1537 let wrd' = addToWord wrd x
1538 old = (mask .&. wrd) `shiftR` (elemBits * (patLen - 1))
1539 acc' <- fstep acc (toEnum $ fromIntegral old)
1540 if wrd' .&. mask == patWord
1541 then done acc' >>= \r -> return $ Yield r (GO_SHORT_PAT s)
1542 else go1 SPEC wrd' s acc'
1543 Skip s -> go1 SPEC wrd s acc
1545 acc' <- go2 wrd patLen acc
1546 done acc' >>= \r -> return $ Yield r GO_DONE
1548 go2 !wrd !n !acc | n > 0 = do
1549 let old = (mask .&. wrd) `shiftR` (elemBits * (n - 1))
1550 fstep acc (toEnum $ fromIntegral old) >>= go2 wrd (n - 1)
1551 go2 _ _ acc = return acc
1553 stepOuter gst (GO_KARP_RABIN stt rb rhead) = do
1554 initial >>= go0 SPEC 0 rhead stt
1558 k = 2891336453 :: Word32
1560 addCksum cksum a = cksum * k + fromIntegral (fromEnum a)
1561 deltaCksum cksum old new =
1562 addCksum cksum new - coeff * fromIntegral (fromEnum old)
1564 -- XXX shall we use a random starting hash or 1 instead of 0?
1565 patHash = A.foldl' addCksum 0 patArr
1568 go0 !_ !idx !rh st !acc = do
1569 res <- step (adaptState gst) st
1572 rh' <- liftIO $ RB.unsafeInsert rb rh x
1575 let fold = RB.unsafeFoldRing (RB.ringBound rb)
1576 let !ringHash = fold addCksum 0 rb
1577 if ringHash == patHash
1578 then go2 SPEC ringHash rh' s acc
1579 else go1 SPEC ringHash rh' s acc
1580 else go0 SPEC (idx + 1) rh' s acc
1581 Skip s -> go0 SPEC idx rh s acc
1583 !acc' <- if idx /= 0
1584 then RB.unsafeFoldRingM rh fstep acc rb
1586 done acc' >>= \r -> return $ Yield r GO_DONE
1588 -- XXX Theoretically this code can do 4 times faster if GHC generates
1589 -- optimal code. If we use just "(cksum' == patHash)" condition it goes
1590 -- 4x faster, as soon as we add the "RB.unsafeEqArray rb v" condition
1591 -- the generated code changes drastically and becomes 4x slower. Need
1592 -- to investigate what is going on with GHC.
1594 go1 !_ !cksum !rh st !acc = do
1595 res <- step (adaptState gst) st
1598 old <- liftIO $ peek rh
1599 let cksum' = deltaCksum cksum old x
1600 acc' <- fstep acc old
1602 if (cksum' == patHash)
1604 rh' <- liftIO (RB.unsafeInsert rb rh x)
1605 go2 SPEC cksum' rh' s acc'
1607 rh' <- liftIO (RB.unsafeInsert rb rh x)
1608 go1 SPEC cksum' rh' s acc'
1609 Skip s -> go1 SPEC cksum rh s acc
1611 acc' <- RB.unsafeFoldRingFullM rh fstep acc rb
1612 done acc' >>= \r -> return $ Yield r GO_DONE
1614 go2 !_ !cksum' !rh' s !acc' = do
1615 if RB.unsafeEqArray rb rh' patArr
1618 return $ Yield r (GO_KARP_RABIN s rb rhead)
1619 else go1 SPEC cksum' rh' s acc'
1621 stepOuter gst (GO_EMPTY_PAT st) = do
1622 res <- step (adaptState gst) st
1627 done acc' >>= \r -> return $ Yield r (GO_EMPTY_PAT s)
1628 Skip s -> return $ Skip (GO_EMPTY_PAT s)
1631 stepOuter _ GO_DONE = return Stop
1633{-# INLINE_NORMAL splitSuffixOn #-}
1635 :: forall m a b. (MonadIO m, Storable a, Enum a, Eq a)
1641splitSuffixOn withSep patArr@Array{..} (Fold fstep initial done)
1642 (Stream step state) =
1643 Stream stepOuter GO_START
1647 patLen = A.length patArr
1648 maxIndex = patLen - 1
1649 elemBits = sizeOf (undefined :: a) * 8
1651 {-# INLINE_LATE stepOuter #-}
1652 stepOuter _ GO_START =
1654 then return $ Skip $ GO_EMPTY_PAT state
1657 r <- liftIO $ (A.unsafeIndexIO patArr 0)
1658 return $ Skip $ GO_SINGLE_PAT state r
1659 else if sizeOf (undefined :: a) * patLen
1660 <= sizeOf (undefined :: Word)
1661 then return $ Skip $ GO_SHORT_PAT state
1663 (rb, rhead) <- liftIO $ RB.new patLen
1664 return $ Skip $ GO_KARP_RABIN state rb rhead
1666 stepOuter gst (GO_SINGLE_PAT stt pat) = do
1667 -- This first part is the only difference between splitOn and
1669 -- If the last element is a separator do not issue a blank segment.
1670 res <- step (adaptState gst) stt
1676 acc' <- if withSep then fstep acc x else return acc
1677 done acc' >>= \r -> return $ Yield r (GO_SINGLE_PAT s pat)
1678 else fstep acc x >>= go SPEC s
1679 Skip s -> return $ Skip $ (GO_SINGLE_PAT s pat)
1684 -- This is identical for splitOn and splitSuffixOn
1686 res <- step (adaptState gst) st
1691 acc' <- if withSep then fstep acc x else return acc
1693 return $ Yield r (GO_SINGLE_PAT s pat)
1694 else fstep acc x >>= go SPEC s
1695 Skip s -> go SPEC s acc
1696 Stop -> done acc >>= \r -> return $ Yield r GO_DONE
1698 stepOuter gst (GO_SHORT_PAT stt) = do
1700 -- Call "initial" only if the stream yields an element, otherwise we
1701 -- may call "initial" but never yield anything. initial may produce a
1702 -- side effect, therefore we will end up doing and discard a side
1707 res <- step (adaptState gst) stt
1711 let wrd' = addToWord wrd x
1712 acc' <- if withSep then fstep acc x else return acc
1715 if wrd' .&. mask == patWord
1716 then done acc' >>= \r -> return $ Yield r (GO_SHORT_PAT s)
1717 else go0 SPEC (idx + 1) wrd' s acc'
1718 else go0 SPEC (idx + 1) wrd' s acc'
1719 Skip s -> return $ Skip (GO_SHORT_PAT s)
1725 mask = (1 `shiftL` (elemBits * patLen)) - 1
1727 addToWord wrd a = (wrd `shiftL` elemBits) .|. fromIntegral (fromEnum a)
1730 patWord = mask .&. A.foldl' addToWord 0 patArr
1732 go0 !_ !idx wrd st !acc = do
1733 res <- step (adaptState gst) st
1736 let wrd' = addToWord wrd x
1737 acc' <- if withSep then fstep acc x else return acc
1740 if wrd' .&. mask == patWord
1743 return $ Yield r (GO_SHORT_PAT s)
1744 else go1 SPEC wrd' s acc'
1745 else go0 SPEC (idx + 1) wrd' s acc'
1746 Skip s -> go0 SPEC idx wrd s acc
1748 if (idx == maxIndex) && (wrd .&. mask == patWord)
1751 acc' <- if idx /= 0 && not withSep
1752 then go2 wrd idx acc
1754 done acc' >>= \r -> return $ Yield r GO_DONE
1757 go1 !_ wrd st !acc = do
1758 res <- step (adaptState gst) st
1761 let wrd' = addToWord wrd x
1762 old = (mask .&. wrd) `shiftR` (elemBits * (patLen - 1))
1765 else fstep acc (toEnum $ fromIntegral old)
1766 if wrd' .&. mask == patWord
1767 then done acc' >>= \r -> return $ Yield r (GO_SHORT_PAT s)
1768 else go1 SPEC wrd' s acc'
1769 Skip s -> go1 SPEC wrd s acc
1771 -- If the last sequence is a separator do not issue a blank
1773 if wrd .&. mask == patWord
1778 else go2 wrd patLen acc
1779 done acc' >>= \r -> return $ Yield r GO_DONE
1781 go2 !wrd !n !acc | n > 0 = do
1782 let old = (mask .&. wrd) `shiftR` (elemBits * (n - 1))
1783 fstep acc (toEnum $ fromIntegral old) >>= go2 wrd (n - 1)
1784 go2 _ _ acc = return acc
1786 stepOuter gst (GO_KARP_RABIN stt rb rhead) = do
1788 res <- step (adaptState gst) stt
1792 acc' <- if withSep then fstep acc x else return acc
1793 rh' <- liftIO (RB.unsafeInsert rb rhead x)
1796 let fold = RB.unsafeFoldRing (RB.ringBound rb)
1797 let !ringHash = fold addCksum 0 rb
1798 if ringHash == patHash
1799 then go2 SPEC ringHash rh' s acc'
1800 else go0 SPEC (idx + 1) rh' s acc'
1801 else go0 SPEC (idx + 1) rh' s acc'
1802 Skip s -> return $ Skip (GO_KARP_RABIN s rb rhead)
1807 k = 2891336453 :: Word32
1809 addCksum cksum a = cksum * k + fromIntegral (fromEnum a)
1810 deltaCksum cksum old new =
1811 addCksum cksum new - coeff * fromIntegral (fromEnum old)
1813 -- XXX shall we use a random starting hash or 1 instead of 0?
1814 patHash = A.foldl' addCksum 0 patArr
1817 go0 !_ !idx !rh st !acc = do
1818 res <- step (adaptState gst) st
1821 acc' <- if withSep then fstep acc x else return acc
1822 rh' <- liftIO (RB.unsafeInsert rb rh x)
1825 let fold = RB.unsafeFoldRing (RB.ringBound rb)
1826 let !ringHash = fold addCksum 0 rb
1827 if ringHash == patHash
1828 then go2 SPEC ringHash rh' s acc'
1829 else go1 SPEC ringHash rh' s acc'
1830 else go0 SPEC (idx + 1) rh' s acc'
1831 Skip s -> go0 SPEC idx rh s acc
1833 -- do not issue a blank segment when we end at pattern
1834 if (idx == maxIndex) && RB.unsafeEqArray rb rh patArr
1837 !acc' <- if idx /= 0 && not withSep
1838 then RB.unsafeFoldRingM rh fstep acc rb
1840 done acc' >>= \r -> return $ Yield r GO_DONE
1842 -- XXX Theoretically this code can do 4 times faster if GHC generates
1843 -- optimal code. If we use just "(cksum' == patHash)" condition it goes
1844 -- 4x faster, as soon as we add the "RB.unsafeEqArray rb v" condition
1845 -- the generated code changes drastically and becomes 4x slower. Need
1846 -- to investigate what is going on with GHC.
1848 go1 !_ !cksum !rh st !acc = do
1849 res <- step (adaptState gst) st
1852 old <- liftIO $ peek rh
1853 let cksum' = deltaCksum cksum old x
1858 if (cksum' == patHash)
1860 rh' <- liftIO (RB.unsafeInsert rb rh x)
1861 go2 SPEC cksum' rh' s acc'
1863 rh' <- liftIO (RB.unsafeInsert rb rh x)
1864 go1 SPEC cksum' rh' s acc'
1865 Skip s -> go1 SPEC cksum rh s acc
1867 if RB.unsafeEqArray rb rh patArr
1872 else RB.unsafeFoldRingFullM rh fstep acc rb
1873 done acc' >>= \r -> return $ Yield r GO_DONE
1875 go2 !_ !cksum' !rh' s !acc' = do
1876 if RB.unsafeEqArray rb rh' patArr
1879 return $ Yield r (GO_KARP_RABIN s rb rhead)
1880 else go1 SPEC cksum' rh' s acc'
1882 stepOuter gst (GO_EMPTY_PAT st) = do
1883 res <- step (adaptState gst) st
1888 done acc' >>= \r -> return $ Yield r (GO_EMPTY_PAT s)
1889 Skip s -> return $ Skip (GO_EMPTY_PAT s)
1892 stepOuter _ GO_DONE = return Stop
1894data SplitState s arr
1896 | SplitBuffering s arr
1897 | SplitSplitting s arr
1898 | SplitYielding arr (SplitState s arr)
1901-- XXX An alternative approach would be to use a partial fold (Fold m a b) to
1902-- split using a splitBy like combinator. The Fold would consume upto the
1903-- separator and return any leftover which can then be fed to the next fold.
1905-- We can revisit this once we have partial folds/parsers.
1907-- | Performs infix separator style splitting.
1908{-# INLINE_NORMAL splitInnerBy #-}
1911 => (f a -> m (f a, Maybe (f a))) -- splitter
1912 -> (f a -> f a -> m (f a)) -- joiner
1915splitInnerBy splitter joiner (Stream step1 state1) =
1916 (Stream step (SplitInitial state1))
1920 {-# INLINE_LATE step #-}
1921 step gst (SplitInitial st) = do
1925 (x1, mx2) <- splitter x
1926 return $ case mx2 of
1927 Nothing -> Skip (SplitBuffering s x1)
1928 Just x2 -> Skip (SplitYielding x1 (SplitSplitting s x2))
1929 Skip s -> return $ Skip (SplitInitial s)
1930 Stop -> return $ Stop
1932 step gst (SplitBuffering st buf) = do
1936 (x1, mx2) <- splitter x
1937 buf' <- joiner buf x1
1938 return $ case mx2 of
1939 Nothing -> Skip (SplitBuffering s buf')
1940 Just x2 -> Skip (SplitYielding buf' (SplitSplitting s x2))
1941 Skip s -> return $ Skip (SplitBuffering s buf)
1942 Stop -> return $ Skip (SplitYielding buf SplitFinishing)
1944 step _ (SplitSplitting st buf) = do
1945 (x1, mx2) <- splitter buf
1946 return $ case mx2 of
1947 Nothing -> Skip $ SplitBuffering st x1
1948 Just x2 -> Skip $ SplitYielding x1 (SplitSplitting st x2)
1950 step _ (SplitYielding x next) = return $ Yield x next
1951 step _ SplitFinishing = return $ Stop
1953-- | Performs infix separator style splitting.
1954{-# INLINE_NORMAL splitInnerBySuffix #-}
1956 :: (Monad m, Eq (f a), Monoid (f a))
1957 => (f a -> m (f a, Maybe (f a))) -- splitter
1958 -> (f a -> f a -> m (f a)) -- joiner
1961splitInnerBySuffix splitter joiner (Stream step1 state1) =
1962 (Stream step (SplitInitial state1))
1966 {-# INLINE_LATE step #-}
1967 step gst (SplitInitial st) = do
1971 (x1, mx2) <- splitter x
1972 return $ case mx2 of
1973 Nothing -> Skip (SplitBuffering s x1)
1974 Just x2 -> Skip (SplitYielding x1 (SplitSplitting s x2))
1975 Skip s -> return $ Skip (SplitInitial s)
1976 Stop -> return $ Stop
1978 step gst (SplitBuffering st buf) = do
1982 (x1, mx2) <- splitter x
1983 buf' <- joiner buf x1
1984 return $ case mx2 of
1985 Nothing -> Skip (SplitBuffering s buf')
1986 Just x2 -> Skip (SplitYielding buf' (SplitSplitting s x2))
1987 Skip s -> return $ Skip (SplitBuffering s buf)
1991 else Skip (SplitYielding buf SplitFinishing)
1993 step _ (SplitSplitting st buf) = do
1994 (x1, mx2) <- splitter buf
1995 return $ case mx2 of
1996 Nothing -> Skip $ SplitBuffering st x1
1997 Just x2 -> Skip $ SplitYielding x1 (SplitSplitting st x2)
1999 step _ (SplitYielding x next) = return $ Yield x next
2000 step _ SplitFinishing = return $ Stop
2002------------------------------------------------------------------------------
2004------------------------------------------------------------------------------
2006{-# INLINE_NORMAL isPrefixOf #-}
2007isPrefixOf :: (Eq a, Monad m) => Stream m a -> Stream m a -> m Bool
2008isPrefixOf (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
2010 go (sa, sb, Nothing) = do
2011 r <- stepa defState sa
2013 Yield x sa' -> go (sa', sb, Just x)
2014 Skip sa' -> go (sa', sb, Nothing)
2017 go (sa, sb, Just x) = do
2018 r <- stepb defState sb
2022 then go (sa, sb', Nothing)
2024 Skip sb' -> go (sa, sb', Just x)
2025 Stop -> return False
2027{-# INLINE_NORMAL isSubsequenceOf #-}
2028isSubsequenceOf :: (Eq a, Monad m) => Stream m a -> Stream m a -> m Bool
2029isSubsequenceOf (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
2031 go (sa, sb, Nothing) = do
2032 r <- stepa defState sa
2034 Yield x sa' -> go (sa', sb, Just x)
2035 Skip sa' -> go (sa', sb, Nothing)
2038 go (sa, sb, Just x) = do
2039 r <- stepb defState sb
2043 then go (sa, sb', Nothing)
2044 else go (sa, sb', Just x)
2045 Skip sb' -> go (sa, sb', Just x)
2046 Stop -> return False
2048{-# INLINE_NORMAL stripPrefix #-}
2051 => Stream m a -> Stream m a -> m (Maybe (Stream m a))
2052stripPrefix (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
2054 go (sa, sb, Nothing) = do
2055 r <- stepa defState sa
2057 Yield x sa' -> go (sa', sb, Just x)
2058 Skip sa' -> go (sa', sb, Nothing)
2059 Stop -> return $ Just (Stream stepb sb)
2061 go (sa, sb, Just x) = do
2062 r <- stepb defState sb
2066 then go (sa, sb', Nothing)
2068 Skip sb' -> go (sa, sb', Just x)
2069 Stop -> return Nothing
2071------------------------------------------------------------------------------
2073------------------------------------------------------------------------------
2075-- | Execute a monadic action for each element of the 'Stream'
2076{-# INLINE_NORMAL mapM_ #-}
2077mapM_ :: Monad m => (a -> m b) -> Stream m a -> m ()
2078mapM_ m = drain . mapM m
2080-------------------------------------------------------------------------------
2081-- Stream transformations using Unfolds
2082-------------------------------------------------------------------------------
2084-- Define a unique structure to use in inspection testing
2085data ConcatMapUState o i =
2087 | ConcatMapUInner o i
2089-- | @concatMapU unfold stream@ uses @unfold@ to map the input stream elements
2090-- to streams and then flattens the generated streams into a single output
2093-- This is like 'concatMap' but uses an unfold with an explicit state to
2094-- generate the stream instead of a 'Stream' type generator. This allows better
2095-- optimization via fusion. This can be many times more efficient than
2098{-# INLINE_NORMAL concatMapU #-}
2099concatMapU :: Monad m => Unfold m a b -> Stream m a -> Stream m b
2100concatMapU (Unfold istep inject) (Stream ostep ost) =
2101 Stream step (ConcatMapUOuter ost)
2103 {-# INLINE_LATE step #-}
2104 step gst (ConcatMapUOuter o) = do
2105 r <- ostep (adaptState gst) o
2109 i `seq` return (Skip (ConcatMapUInner o' i))
2110 Skip o' -> return $ Skip (ConcatMapUOuter o')
2111 Stop -> return $ Stop
2113 step _ (ConcatMapUInner o i) = do
2116 Yield x i' -> Yield x (ConcatMapUInner o i')
2117 Skip i' -> Skip (ConcatMapUInner o i')
2118 Stop -> Skip (ConcatMapUOuter o)
2120data ConcatUnfoldInterleaveState o i =
2121 ConcatUnfoldInterleaveOuter o [i]
2122 | ConcatUnfoldInterleaveInner o [i]
2123 | ConcatUnfoldInterleaveInnerL [i] [i]
2124 | ConcatUnfoldInterleaveInnerR [i] [i]
2126-- XXX use arrays to store state instead of lists.
2127-- XXX In general we can use different scheduling strategies e.g. how to
2128-- schedule the outer vs inner loop or assigning weights to different streams
2129-- or outer and inner loops.
2131-- After a yield, switch to the next stream. Do not switch streams on Skip.
2132-- Yield from outer stream switches to the inner stream.
2134-- There are two choices here, (1) exhaust the outer stream first and then
2135-- start yielding from the inner streams, this is much simpler to implement,
2136-- (2) yield at least one element from an inner stream before going back to
2137-- outer stream and opening the next stream from it.
2139-- Ideally, we need some scheduling bias to inner streams vs outer stream.
2140-- Maybe we can configure the behavior.
2142{-# INLINE_NORMAL concatUnfoldInterleave #-}
2143concatUnfoldInterleave :: Monad m => Unfold m a b -> Stream m a -> Stream m b
2144concatUnfoldInterleave (Unfold istep inject) (Stream ostep ost) =
2145 Stream step (ConcatUnfoldInterleaveOuter ost [])
2147 {-# INLINE_LATE step #-}
2148 step gst (ConcatUnfoldInterleaveOuter o ls) = do
2149 r <- ostep (adaptState gst) o
2153 i `seq` return (Skip (ConcatUnfoldInterleaveInner o' (i : ls)))
2154 Skip o' -> return $ Skip (ConcatUnfoldInterleaveOuter o' ls)
2155 Stop -> return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2157 step _ (ConcatUnfoldInterleaveInner _ []) = undefined
2158 step _ (ConcatUnfoldInterleaveInner o (st:ls)) = do
2161 Yield x s -> Yield x (ConcatUnfoldInterleaveOuter o (s:ls))
2162 Skip s -> Skip (ConcatUnfoldInterleaveInner o (s:ls))
2163 Stop -> Skip (ConcatUnfoldInterleaveOuter o ls)
2165 step _ (ConcatUnfoldInterleaveInnerL [] []) = return Stop
2166 step _ (ConcatUnfoldInterleaveInnerL [] rs) =
2167 return $ Skip (ConcatUnfoldInterleaveInnerR [] rs)
2169 step _ (ConcatUnfoldInterleaveInnerL (st:ls) rs) = do
2172 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerL ls (s:rs))
2173 Skip s -> Skip (ConcatUnfoldInterleaveInnerL (s:ls) rs)
2174 Stop -> Skip (ConcatUnfoldInterleaveInnerL ls rs)
2176 step _ (ConcatUnfoldInterleaveInnerR [] []) = return Stop
2177 step _ (ConcatUnfoldInterleaveInnerR ls []) =
2178 return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2180 step _ (ConcatUnfoldInterleaveInnerR ls (st:rs)) = do
2183 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerR (s:ls) rs)
2184 Skip s -> Skip (ConcatUnfoldInterleaveInnerR ls (s:rs))
2185 Stop -> Skip (ConcatUnfoldInterleaveInnerR ls rs)
2187-- XXX In general we can use different scheduling strategies e.g. how to
2188-- schedule the outer vs inner loop or assigning weights to different streams
2189-- or outer and inner loops.
2191-- This could be inefficient if the tasks are too small.
2193-- Compared to concatUnfoldInterleave this one switches streams on Skips.
2195{-# INLINE_NORMAL concatUnfoldRoundrobin #-}
2196concatUnfoldRoundrobin :: Monad m => Unfold m a b -> Stream m a -> Stream m b
2197concatUnfoldRoundrobin (Unfold istep inject) (Stream ostep ost) =
2198 Stream step (ConcatUnfoldInterleaveOuter ost [])
2200 {-# INLINE_LATE step #-}
2201 step gst (ConcatUnfoldInterleaveOuter o ls) = do
2202 r <- ostep (adaptState gst) o
2206 i `seq` return (Skip (ConcatUnfoldInterleaveInner o' (i : ls)))
2207 Skip o' -> return $ Skip (ConcatUnfoldInterleaveInner o' ls)
2208 Stop -> return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2210 step _ (ConcatUnfoldInterleaveInner o []) =
2211 return $ Skip (ConcatUnfoldInterleaveOuter o [])
2213 step _ (ConcatUnfoldInterleaveInner o (st:ls)) = do
2216 Yield x s -> Yield x (ConcatUnfoldInterleaveOuter o (s:ls))
2217 Skip s -> Skip (ConcatUnfoldInterleaveOuter o (s:ls))
2218 Stop -> Skip (ConcatUnfoldInterleaveOuter o ls)
2220 step _ (ConcatUnfoldInterleaveInnerL [] []) = return Stop
2221 step _ (ConcatUnfoldInterleaveInnerL [] rs) =
2222 return $ Skip (ConcatUnfoldInterleaveInnerR [] rs)
2224 step _ (ConcatUnfoldInterleaveInnerL (st:ls) rs) = do
2227 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerL ls (s:rs))
2228 Skip s -> Skip (ConcatUnfoldInterleaveInnerL ls (s:rs))
2229 Stop -> Skip (ConcatUnfoldInterleaveInnerL ls rs)
2231 step _ (ConcatUnfoldInterleaveInnerR [] []) = return Stop
2232 step _ (ConcatUnfoldInterleaveInnerR ls []) =
2233 return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2235 step _ (ConcatUnfoldInterleaveInnerR ls (st:rs)) = do
2238 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerR (s:ls) rs)
2239 Skip s -> Skip (ConcatUnfoldInterleaveInnerR (s:ls) rs)
2240 Stop -> Skip (ConcatUnfoldInterleaveInnerR ls rs)
2242data AppendState s1 s2 = AppendFirst s1 | AppendSecond s2
2244-- Note that this could be much faster compared to the CPS stream. However, as
2245-- the number of streams being composed increases this may become expensive.
2246-- Need to see where the breaking point is between the two.
2248{-# INLINE_NORMAL append #-}
2249append :: Monad m => Stream m a -> Stream m a -> Stream m a
2250append (Stream step1 state1) (Stream step2 state2) =
2251 Stream step (AppendFirst state1)
2255 {-# INLINE_LATE step #-}
2256 step gst (AppendFirst st) = do
2259 Yield a s -> Yield a (AppendFirst s)
2260 Skip s -> Skip (AppendFirst s)
2261 Stop -> Skip (AppendSecond state2)
2263 step gst (AppendSecond st) = do
2266 Yield a s -> Yield a (AppendSecond s)
2267 Skip s -> Skip (AppendSecond s)
2270data InterleaveState s1 s2 = InterleaveFirst s1 s2 | InterleaveSecond s1 s2
2271 | InterleaveSecondOnly s2 | InterleaveFirstOnly s1
2273{-# INLINE_NORMAL interleave #-}
2274interleave :: Monad m => Stream m a -> Stream m a -> Stream m a
2275interleave (Stream step1 state1) (Stream step2 state2) =
2276 Stream step (InterleaveFirst state1 state2)
2280 {-# INLINE_LATE step #-}
2281 step gst (InterleaveFirst st1 st2) = do
2284 Yield a s -> Yield a (InterleaveSecond s st2)
2285 Skip s -> Skip (InterleaveFirst s st2)
2286 Stop -> Skip (InterleaveSecondOnly st2)
2288 step gst (InterleaveSecond st1 st2) = do
2291 Yield a s -> Yield a (InterleaveFirst st1 s)
2292 Skip s -> Skip (InterleaveSecond st1 s)
2293 Stop -> Skip (InterleaveFirstOnly st1)
2295 step gst (InterleaveFirstOnly st1) = do
2298 Yield a s -> Yield a (InterleaveFirstOnly s)
2299 Skip s -> Skip (InterleaveFirstOnly s)
2302 step gst (InterleaveSecondOnly st2) = do
2305 Yield a s -> Yield a (InterleaveSecondOnly s)
2306 Skip s -> Skip (InterleaveSecondOnly s)
2309{-# INLINE_NORMAL interleaveMin #-}
2310interleaveMin :: Monad m => Stream m a -> Stream m a -> Stream m a
2311interleaveMin (Stream step1 state1) (Stream step2 state2) =
2312 Stream step (InterleaveFirst state1 state2)
2316 {-# INLINE_LATE step #-}
2317 step gst (InterleaveFirst st1 st2) = do
2320 Yield a s -> Yield a (InterleaveSecond s st2)
2321 Skip s -> Skip (InterleaveFirst s st2)
2324 step gst (InterleaveSecond st1 st2) = do
2327 Yield a s -> Yield a (InterleaveFirst st1 s)
2328 Skip s -> Skip (InterleaveSecond st1 s)
2331 step _ (InterleaveFirstOnly _) = undefined
2332 step _ (InterleaveSecondOnly _) = undefined
2334{-# INLINE_NORMAL interleaveSuffix #-}
2335interleaveSuffix :: Monad m => Stream m a -> Stream m a -> Stream m a
2336interleaveSuffix (Stream step1 state1) (Stream step2 state2) =
2337 Stream step (InterleaveFirst state1 state2)
2341 {-# INLINE_LATE step #-}
2342 step gst (InterleaveFirst st1 st2) = do
2345 Yield a s -> Yield a (InterleaveSecond s st2)
2346 Skip s -> Skip (InterleaveFirst s st2)
2349 step gst (InterleaveSecond st1 st2) = do
2352 Yield a s -> Yield a (InterleaveFirst st1 s)
2353 Skip s -> Skip (InterleaveSecond st1 s)
2354 Stop -> Skip (InterleaveFirstOnly st1)
2356 step gst (InterleaveFirstOnly st1) = do
2359 Yield a s -> Yield a (InterleaveFirstOnly s)
2360 Skip s -> Skip (InterleaveFirstOnly s)
2363 step _ (InterleaveSecondOnly _) = undefined
2365data InterleaveInfixState s1 s2 a
2366 = InterleaveInfixFirst s1 s2
2367 | InterleaveInfixSecondBuf s1 s2
2368 | InterleaveInfixSecondYield s1 s2 a
2369 | InterleaveInfixFirstYield s1 s2 a
2370 | InterleaveInfixFirstOnly s1
2372{-# INLINE_NORMAL interleaveInfix #-}
2373interleaveInfix :: Monad m => Stream m a -> Stream m a -> Stream m a
2374interleaveInfix (Stream step1 state1) (Stream step2 state2) =
2375 Stream step (InterleaveInfixFirst state1 state2)
2379 {-# INLINE_LATE step #-}
2380 step gst (InterleaveInfixFirst st1 st2) = do
2383 Yield a s -> Yield a (InterleaveInfixSecondBuf s st2)
2384 Skip s -> Skip (InterleaveInfixFirst s st2)
2387 step gst (InterleaveInfixSecondBuf st1 st2) = do
2390 Yield a s -> Skip (InterleaveInfixSecondYield st1 s a)
2391 Skip s -> Skip (InterleaveInfixSecondBuf st1 s)
2392 Stop -> Skip (InterleaveInfixFirstOnly st1)
2394 step gst (InterleaveInfixSecondYield st1 st2 x) = do
2397 Yield a s -> Yield x (InterleaveInfixFirstYield s st2 a)
2398 Skip s -> Skip (InterleaveInfixSecondYield s st2 x)
2401 step _ (InterleaveInfixFirstYield st1 st2 x) = do
2402 return $ Yield x (InterleaveInfixSecondBuf st1 st2)
2404 step gst (InterleaveInfixFirstOnly st1) = do
2407 Yield a s -> Yield a (InterleaveInfixFirstOnly s)
2408 Skip s -> Skip (InterleaveInfixFirstOnly s)
2411{-# INLINE_NORMAL roundRobin #-}
2412roundRobin :: Monad m => Stream m a -> Stream m a -> Stream m a
2413roundRobin (Stream step1 state1) (Stream step2 state2) =
2414 Stream step (InterleaveFirst state1 state2)
2418 {-# INLINE_LATE step #-}
2419 step gst (InterleaveFirst st1 st2) = do
2422 Yield a s -> Yield a (InterleaveSecond s st2)
2423 Skip s -> Skip (InterleaveSecond s st2)
2424 Stop -> Skip (InterleaveSecondOnly st2)
2426 step gst (InterleaveSecond st1 st2) = do
2429 Yield a s -> Yield a (InterleaveFirst st1 s)
2430 Skip s -> Skip (InterleaveFirst st1 s)
2431 Stop -> Skip (InterleaveFirstOnly st1)
2433 step gst (InterleaveSecondOnly st2) = do
2436 Yield a s -> Yield a (InterleaveSecondOnly s)
2437 Skip s -> Skip (InterleaveSecondOnly s)
2440 step gst (InterleaveFirstOnly st1) = do
2443 Yield a s -> Yield a (InterleaveFirstOnly s)
2444 Skip s -> Skip (InterleaveFirstOnly s)
2447data ICUState s1 s2 i1 i2 =
2452 | ICUFirstInner s1 s2 i1
2453 | ICUSecondInner s1 s2 i2
2454 | ICUFirstOnlyInner s1 i1
2455 | ICUSecondOnlyInner s2 i2
2457-- | Interleave streams (full streams, not the elements) unfolded from two
2458-- input streams and concat. Stop when the first stream stops. If the second
2459-- stream ends before the first one then first stream still keeps running alone
2460-- without any interleaving with the second stream.
2462-- [a1, a2, ... an] [b1, b2 ...]
2463-- => [streamA1, streamA2, ... streamAn] [streamB1, streamB2, ...]
2464-- => [streamA1, streamB1, streamA2...StreamAn, streamBn]
2465-- => [a11, a12, ...a1j, b11, b12, ...b1k, a21, a22, ...]
2467{-# INLINE_NORMAL gintercalateSuffix #-}
2470 => Unfold m a c -> Stream m a -> Unfold m b c -> Stream m b -> Stream m c
2472 (Unfold istep1 inject1) (Stream step1 state1)
2473 (Unfold istep2 inject2) (Stream step2 state2) =
2474 Stream step (ICUFirst state1 state2)
2478 {-# INLINE_LATE step #-}
2479 step gst (ICUFirst s1 s2) = do
2480 r <- step1 (adaptState gst) s1
2484 i `seq` return (Skip (ICUFirstInner s s2 i))
2485 Skip s -> return $ Skip (ICUFirst s s2)
2488 step gst (ICUFirstOnly s1) = do
2489 r <- step1 (adaptState gst) s1
2493 i `seq` return (Skip (ICUFirstOnlyInner s i))
2494 Skip s -> return $ Skip (ICUFirstOnly s)
2497 step _ (ICUFirstInner s1 s2 i1) = do
2500 Yield x i' -> Yield x (ICUFirstInner s1 s2 i')
2501 Skip i' -> Skip (ICUFirstInner s1 s2 i')
2502 Stop -> Skip (ICUSecond s1 s2)
2504 step _ (ICUFirstOnlyInner s1 i1) = do
2507 Yield x i' -> Yield x (ICUFirstOnlyInner s1 i')
2508 Skip i' -> Skip (ICUFirstOnlyInner s1 i')
2509 Stop -> Skip (ICUFirstOnly s1)
2511 step gst (ICUSecond s1 s2) = do
2512 r <- step2 (adaptState gst) s2
2516 i `seq` return (Skip (ICUSecondInner s1 s i))
2517 Skip s -> return $ Skip (ICUSecond s1 s)
2518 Stop -> return $ Skip (ICUFirstOnly s1)
2520 step _ (ICUSecondInner s1 s2 i2) = do
2523 Yield x i' -> Yield x (ICUSecondInner s1 s2 i')
2524 Skip i' -> Skip (ICUSecondInner s1 s2 i')
2525 Stop -> Skip (ICUFirst s1 s2)
2527 step _ (ICUSecondOnly _s2) = undefined
2528 step _ (ICUSecondOnlyInner _s2 _i2) = undefined
2530data InterposeSuffixState s1 i1 =
2531 InterposeSuffixFirst s1
2532 -- | InterposeSuffixFirstYield s1 i1
2533 | InterposeSuffixFirstInner s1 i1
2534 | InterposeSuffixSecond s1
2536-- Note that if an unfolded layer turns out to be nil we still emit the
2537-- separator effect. An alternate behavior could be to emit the separator
2538-- effect only if at least one element has been yielded by the unfolding.
2539-- However, that becomes a bit complicated, so we have chosen the former
2541{-# INLINE_NORMAL interposeSuffix #-}
2544 => m c -> Unfold m b c -> Stream m b -> Stream m c
2547 (Unfold istep1 inject1) (Stream step1 state1) =
2548 Stream step (InterposeSuffixFirst state1)
2552 {-# INLINE_LATE step #-}
2553 step gst (InterposeSuffixFirst s1) = do
2554 r <- step1 (adaptState gst) s1
2558 i `seq` return (Skip (InterposeSuffixFirstInner s i))
2559 -- i `seq` return (Skip (InterposeSuffixFirstYield s i))
2560 Skip s -> return $ Skip (InterposeSuffixFirst s)
2564 step _ (InterposeSuffixFirstYield s1 i1) = do
2567 Yield x i' -> Yield x (InterposeSuffixFirstInner s1 i')
2568 Skip i' -> Skip (InterposeSuffixFirstYield s1 i')
2569 Stop -> Skip (InterposeSuffixFirst s1)
2572 step _ (InterposeSuffixFirstInner s1 i1) = do
2575 Yield x i' -> Yield x (InterposeSuffixFirstInner s1 i')
2576 Skip i' -> Skip (InterposeSuffixFirstInner s1 i')
2577 Stop -> Skip (InterposeSuffixSecond s1)
2579 step _ (InterposeSuffixSecond s1) = do
2581 return $ Yield r (InterposeSuffixFirst s1)
2583data ICALState s1 s2 i1 i2 a =
2585 -- | ICALFirstYield s1 s2 i1
2586 | ICALFirstInner s1 s2 i1
2588 | ICALFirstOnlyInner s1 i1
2589 | ICALSecondInject s1 s2
2590 | ICALFirstInject s1 s2 i2
2591 -- | ICALFirstBuf s1 s2 i1 i2
2592 | ICALSecondInner s1 s2 i1 i2
2593 -- -- | ICALSecondInner s1 s2 i1 i2 a
2594 -- -- | ICALFirstResume s1 s2 i1 i2 a
2596-- | Interleave streams (full streams, not the elements) unfolded from two
2597-- input streams and concat. Stop when the first stream stops. If the second
2598-- stream ends before the first one then first stream still keeps running alone
2599-- without any interleaving with the second stream.
2601-- [a1, a2, ... an] [b1, b2 ...]
2602-- => [streamA1, streamA2, ... streamAn] [streamB1, streamB2, ...]
2603-- => [streamA1, streamB1, streamA2...StreamAn, streamBn]
2604-- => [a11, a12, ...a1j, b11, b12, ...b1k, a21, a22, ...]
2606{-# INLINE_NORMAL gintercalate #-}
2609 => Unfold m a c -> Stream m a -> Unfold m b c -> Stream m b -> Stream m c
2611 (Unfold istep1 inject1) (Stream step1 state1)
2612 (Unfold istep2 inject2) (Stream step2 state2) =
2613 Stream step (ICALFirst state1 state2)
2617 {-# INLINE_LATE step #-}
2618 step gst (ICALFirst s1 s2) = do
2619 r <- step1 (adaptState gst) s1
2623 i `seq` return (Skip (ICALFirstInner s s2 i))
2624 -- i `seq` return (Skip (ICALFirstYield s s2 i))
2625 Skip s -> return $ Skip (ICALFirst s s2)
2629 step _ (ICALFirstYield s1 s2 i1) = do
2632 Yield x i' -> Yield x (ICALFirstInner s1 s2 i')
2633 Skip i' -> Skip (ICALFirstYield s1 s2 i')
2634 Stop -> Skip (ICALFirst s1 s2)
2637 step _ (ICALFirstInner s1 s2 i1) = do
2640 Yield x i' -> Yield x (ICALFirstInner s1 s2 i')
2641 Skip i' -> Skip (ICALFirstInner s1 s2 i')
2642 Stop -> Skip (ICALSecondInject s1 s2)
2644 step gst (ICALFirstOnly s1) = do
2645 r <- step1 (adaptState gst) s1
2649 i `seq` return (Skip (ICALFirstOnlyInner s i))
2650 Skip s -> return $ Skip (ICALFirstOnly s)
2653 step _ (ICALFirstOnlyInner s1 i1) = do
2656 Yield x i' -> Yield x (ICALFirstOnlyInner s1 i')
2657 Skip i' -> Skip (ICALFirstOnlyInner s1 i')
2658 Stop -> Skip (ICALFirstOnly s1)
2660 -- We inject the second stream even before checking if the first stream
2661 -- would yield any more elements. There is no clear choice whether we
2662 -- should do this before or after that. Doing it after may make the state
2663 -- machine a bit simpler though.
2664 step gst (ICALSecondInject s1 s2) = do
2665 r <- step2 (adaptState gst) s2
2669 i `seq` return (Skip (ICALFirstInject s1 s i))
2670 Skip s -> return $ Skip (ICALSecondInject s1 s)
2671 Stop -> return $ Skip (ICALFirstOnly s1)
2673 step gst (ICALFirstInject s1 s2 i2) = do
2674 r <- step1 (adaptState gst) s1
2678 i `seq` return (Skip (ICALSecondInner s s2 i i2))
2679 -- i `seq` return (Skip (ICALFirstBuf s s2 i i2))
2680 Skip s -> return $ Skip (ICALFirstInject s s2 i2)
2684 step _ (ICALFirstBuf s1 s2 i1 i2) = do
2687 Yield x i' -> Skip (ICALSecondInner s1 s2 i' i2 x)
2688 Skip i' -> Skip (ICALFirstBuf s1 s2 i' i2)
2691 step _ (ICALSecondInner s1 s2 i1 i2 v) = do
2694 Yield x i' -> Yield x (ICALSecondInner s1 s2 i1 i' v)
2695 Skip i' -> Skip (ICALSecondInner s1 s2 i1 i' v)
2696 Stop -> Skip (ICALFirstResume s1 s2 i1 i2 v)
2699 step _ (ICALSecondInner s1 s2 i1 i2) = do
2702 Yield x i' -> Yield x (ICALSecondInner s1 s2 i1 i')
2703 Skip i' -> Skip (ICALSecondInner s1 s2 i1 i')
2704 Stop -> Skip (ICALFirstInner s1 s2 i1)
2705 -- Stop -> Skip (ICALFirstResume s1 s2 i1 i2)
2708 step _ (ICALFirstResume s1 s2 i1 i2 x) = do
2709 return $ Yield x (ICALFirstInner s1 s2 i1 i2)
2712data InterposeState s1 i1 a =
2714 -- | InterposeFirstYield s1 i1
2715 | InterposeFirstInner s1 i1
2716 | InterposeFirstInject s1
2717 -- | InterposeFirstBuf s1 i1
2718 | InterposeSecondYield s1 i1
2719 -- -- | InterposeSecondYield s1 i1 a
2720 -- -- | InterposeFirstResume s1 i1 a
2722-- Note that this only interposes the pure values, we may run many effects to
2723-- generate those values as some effects may not generate anything (Skip).
2724{-# INLINE_NORMAL interpose #-}
2725interpose :: Monad m => m c -> Unfold m b c -> Stream m b -> Stream m c
2728 (Unfold istep1 inject1) (Stream step1 state1) =
2729 Stream step (InterposeFirst state1)
2733 {-# INLINE_LATE step #-}
2734 step gst (InterposeFirst s1) = do
2735 r <- step1 (adaptState gst) s1
2739 i `seq` return (Skip (InterposeFirstInner s i))
2740 -- i `seq` return (Skip (InterposeFirstYield s i))
2741 Skip s -> return $ Skip (InterposeFirst s)
2745 step _ (InterposeFirstYield s1 i1) = do
2748 Yield x i' -> Yield x (InterposeFirstInner s1 i')
2749 Skip i' -> Skip (InterposeFirstYield s1 i')
2750 Stop -> Skip (InterposeFirst s1)
2753 step _ (InterposeFirstInner s1 i1) = do
2756 Yield x i' -> Yield x (InterposeFirstInner s1 i')
2757 Skip i' -> Skip (InterposeFirstInner s1 i')
2758 Stop -> Skip (InterposeFirstInject s1)
2760 step gst (InterposeFirstInject s1) = do
2761 r <- step1 (adaptState gst) s1
2765 -- i `seq` return (Skip (InterposeFirstBuf s i))
2766 i `seq` return (Skip (InterposeSecondYield s i))
2767 Skip s -> return $ Skip (InterposeFirstInject s)
2771 step _ (InterposeFirstBuf s1 i1) = do
2774 Yield x i' -> Skip (InterposeSecondYield s1 i' x)
2775 Skip i' -> Skip (InterposeFirstBuf s1 i')
2780 step _ (InterposeSecondYield s1 i1 v) = do
2782 return $ Yield r (InterposeFirstResume s1 i1 v)
2784 step _ (InterposeSecondYield s1 i1) = do
2786 return $ Yield r (InterposeFirstInner s1 i1)
2789 step _ (InterposeFirstResume s1 i1 v) = do
2790 return $ Yield v (InterposeFirstInner s1 i1)
2793------------------------------------------------------------------------------
2795------------------------------------------------------------------------------
2797data GbracketState s1 s2 v
2799 | GBracketNormal s1 v
2800 | GBracketException s2
2802-- | The most general bracketing and exception combinator. All other
2803-- combinators can be expressed in terms of this combinator. This can also be
2804-- used for cases which are not covered by the standard combinators.
2808{-# INLINE_NORMAL gbracket #-}
2812 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
2813 -> (c -> m d) -- ^ after, on normal stop
2814 -> (c -> e -> Stream m b) -- ^ on exception
2815 -> (c -> Stream m b) -- ^ stream generator
2817gbracket bef exc aft fexc fnormal =
2818 Stream step GBracketInit
2822 {-# INLINE_LATE step #-}
2823 step _ GBracketInit = do
2825 return $ Skip $ GBracketNormal (fnormal r) r
2827 step gst (GBracketNormal (UnStream step1 st) v) = do
2828 res <- exc $ step1 gst st
2830 Right r -> case r of
2832 return $ Yield x (GBracketNormal (Stream step1 s) v)
2833 Skip s -> return $ Skip (GBracketNormal (Stream step1 s) v)
2834 Stop -> aft v >> return Stop
2835 Left e -> return $ Skip (GBracketException (fexc v e))
2836 step gst (GBracketException (UnStream step1 st)) = do
2839 Yield x s -> return $ Yield x (GBracketException (Stream step1 s))
2840 Skip s -> return $ Skip (GBracketException (Stream step1 s))
2843-- | Run a side effect before the stream yields its first element.
2844{-# INLINE_NORMAL before #-}
2845before :: Monad m => m b -> Stream m a -> Stream m a
2846before action (Stream step state) = Stream step' Nothing
2850 {-# INLINE_LATE step' #-}
2851 step' _ Nothing = action >> return (Skip (Just state))
2853 step' gst (Just st) = do
2856 Yield x s -> return $ Yield x (Just s)
2857 Skip s -> return $ Skip (Just s)
2860-- | Run a side effect whenever the stream stops normally.
2861{-# INLINE_NORMAL after #-}
2862after :: Monad m => m b -> Stream m a -> Stream m a
2863after action (Stream step state) = Stream step' state
2867 {-# INLINE_LATE step' #-}
2871 Yield x s -> return $ Yield x s
2872 Skip s -> return $ Skip s
2873 Stop -> action >> return Stop
2875-- XXX These combinators are expensive due to the call to
2876-- onException/handle/try on each step. Therefore, when possible, they should
2877-- be called in an outer loop where we perform less iterations. For example, we
2878-- cannot call them on each iteration in a char stream, instead we can call
2879-- them when doing an IO on an array.
2881-- XXX For high performance error checks in busy streams we may need another
2882-- Error constructor in step.
2884-- | Run a side effect whenever the stream aborts due to an exception. The
2885-- exception is not caught, simply rethrown.
2886{-# INLINE_NORMAL onException #-}
2887onException :: MonadCatch m => m b -> Stream m a -> Stream m a
2888onException action str =
2889 gbracket (return ()) MC.try return
2890 (\_ (e :: MC.SomeException) -> nilM (action >> MC.throwM e))
2893{-# INLINE_NORMAL _onException #-}
2894_onException :: MonadCatch m => m b -> Stream m a -> Stream m a
2895_onException action (Stream step state) = Stream step' state
2899 {-# INLINE_LATE step' #-}
2901 res <- step gst st `MC.onException` action
2903 Yield x s -> return $ Yield x s
2904 Skip s -> return $ Skip s
2907-- XXX bracket is like concatMap, it generates a stream and then flattens it.
2908-- Like concatMap it has 10x worse performance compared to linear fused
2911-- | Run the first action before the stream starts and remember its output,
2912-- generate a stream using the output, run the second action providing the
2913-- remembered value as an argument whenever the stream ends normally or due to
2915{-# INLINE_NORMAL bracket #-}
2916bracket :: MonadCatch m => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
2917bracket bef aft bet =
2918 gbracket bef MC.try aft
2919 (\a (e :: SomeException) -> nilM (aft a >> MC.throwM e)) bet
2921data BracketState s v = BracketInit | BracketRun s v
2923{-# INLINE_NORMAL _bracket #-}
2924_bracket :: MonadCatch m => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
2925_bracket bef aft bet = Stream step' BracketInit
2929 {-# INLINE_LATE step' #-}
2930 step' _ BracketInit = bef >>= \x -> return (Skip (BracketRun (bet x) x))
2932 -- NOTE: It is important to use UnStream instead of the Stream pattern
2933 -- here, otherwise we get huge perf degradation, see note in concatMap.
2934 step' gst (BracketRun (UnStream step state) v) = do
2935 -- res <- step gst state `MC.onException` aft v
2936 res <- MC.try $ step gst state
2938 Left (e :: SomeException) -> aft v >> MC.throwM e >> return Stop
2939 Right r -> case r of
2940 Yield x s -> return $ Yield x (BracketRun (Stream step s) v)
2941 Skip s -> return $ Skip (BracketRun (Stream step s) v)
2942 Stop -> aft v >> return Stop
2944-- | Run a side effect whenever the stream stops normally or aborts due to an
2946{-# INLINE finally #-}
2947finally :: MonadCatch m => m b -> Stream m a -> Stream m a
2948-- finally action xs = after action $ onException action xs
2949finally action xs = bracket (return ()) (\_ -> action) (const xs)
2951-- | When evaluating a stream if an exception occurs, stream evaluation aborts
2952-- and the specified exception handler is run with the exception as argument.
2953{-# INLINE_NORMAL handle #-}
2954handle :: (MonadCatch m, Exception e)
2955 => (e -> Stream m a) -> Stream m a -> Stream m a
2957 gbracket (return ()) MC.try return (\_ e -> f e) (\_ -> str)
2959{-# INLINE_NORMAL _handle #-}
2960_handle :: (MonadCatch m, Exception e)
2961 => (e -> Stream m a) -> Stream m a -> Stream m a
2962_handle f (Stream step state) = Stream step' (Left state)
2966 {-# INLINE_LATE step' #-}
2967 step' gst (Left st) = do
2968 res <- MC.try $ step gst st
2970 Left e -> return $ Skip $ Right (f e)
2971 Right r -> case r of
2972 Yield x s -> return $ Yield x (Left s)
2973 Skip s -> return $ Skip (Left s)
2976 step' gst (Right (UnStream step1 st)) = do
2979 Yield x s -> return $ Yield x (Right (Stream step1 s))
2980 Skip s -> return $ Skip (Right (Stream step1 s))
2983-------------------------------------------------------------------------------
2984-- General transformation
2985-------------------------------------------------------------------------------
2987{-# INLINE_NORMAL transform #-}
2988transform :: Monad m => Pipe m a b -> Stream m a -> Stream m b
2989transform (Pipe pstep1 pstep2 pstate) (Stream step state) =
2990 Stream step' (Consume pstate, state)
2994 {-# INLINE_LATE step' #-}
2996 step' gst (Consume pst, st) = pst `seq` do
2997 r <- step (adaptState gst) st
3002 Pipe.Yield b pst' -> return $ Yield b (pst', s)
3003 Pipe.Continue pst' -> return $ Skip (pst', s)
3004 Skip s -> return $ Skip (Consume pst, s)
3007 step' _ (Produce pst, st) = pst `seq` do
3010 Pipe.Yield b pst' -> return $ Yield b (pst', st)
3011 Pipe.Continue pst' -> return $ Skip (pst', st)
3013------------------------------------------------------------------------------
3014-- Transformation by Folding (Scans)
3015------------------------------------------------------------------------------
3017------------------------------------------------------------------------------
3019------------------------------------------------------------------------------
3021-- XXX Is a prescan useful, discarding the last step does not sound useful? I
3022-- am not sure about the utility of this function, so this is implemented but
3023-- not exposed. We can expose it if someone provides good reasons why this is
3026-- XXX We have to execute the stream one step ahead to know that we are at the
3027-- last step. The vector implementation of prescan executes the last fold step
3028-- but does not yield the result. This means we have executed the effect but
3029-- discarded value. This does not sound right. In this implementation we are
3030-- not executing the last fold step.
3031{-# INLINE_NORMAL prescanlM' #-}
3032prescanlM' :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
3033prescanlM' f mz (Stream step state) = Stream step' (state, mz)
3035 {-# INLINE_LATE step' #-}
3036 step' gst (st, prev) = do
3037 r <- step (adaptState gst) st
3041 return $ Yield acc (s, f acc x)
3042 Skip s -> return $ Skip (s, prev)
3045{-# INLINE prescanl' #-}
3046prescanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
3047prescanl' f z = prescanlM' (\a b -> return (f a b)) (return z)
3049------------------------------------------------------------------------------
3050-- Monolithic postscans (postscan followed by a map)
3051------------------------------------------------------------------------------
3053-- The performance of a modular postscan followed by a map seems to be
3054-- equivalent to this monolithic scan followed by map therefore we may not need
3055-- this implementation. We just have it for performance comparison and in case
3056-- modular version does not perform well in some situation.
3058{-# INLINE_NORMAL postscanlMx' #-}
3059postscanlMx' :: Monad m
3060 => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> Stream m b
3061postscanlMx' fstep begin done (Stream step state) = do
3062 Stream step' (state, begin)
3064 {-# INLINE_LATE step' #-}
3065 step' gst (st, acc) = do
3066 r <- step (adaptState gst) st
3072 v `seq` y `seq` return (Yield v (s, return y))
3073 Skip s -> return $ Skip (s, acc)
3076{-# INLINE_NORMAL postscanlx' #-}
3077postscanlx' :: Monad m
3078 => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> Stream m b
3079postscanlx' fstep begin done s =
3080 postscanlMx' (\b a -> return (fstep b a)) (return begin) (return . done) s
3082-- XXX do we need consM strict to evaluate the begin value?
3083{-# INLINE scanlMx' #-}
3085 => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> Stream m b
3086scanlMx' fstep begin done s =
3087 (begin >>= \x -> x `seq` done x) `consM` postscanlMx' fstep begin done s
3089{-# INLINE scanlx' #-}
3091 => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> Stream m b
3092scanlx' fstep begin done s =
3093 scanlMx' (\b a -> return (fstep b a)) (return begin) (return . done) s
3095------------------------------------------------------------------------------
3097------------------------------------------------------------------------------
3099{-# INLINE_NORMAL postscanlM' #-}
3100postscanlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
3101postscanlM' fstep begin (Stream step state) =
3102 begin `seq` Stream step' (state, begin)
3104 {-# INLINE_LATE step' #-}
3105 step' gst (st, acc) = acc `seq` do
3106 r <- step (adaptState gst) st
3110 y `seq` return (Yield y (s, y))
3111 Skip s -> return $ Skip (s, acc)
3114{-# INLINE_NORMAL postscanl' #-}
3115postscanl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
3116postscanl' f = postscanlM' (\a b -> return (f a b))
3118{-# INLINE_NORMAL postscanlM #-}
3119postscanlM :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
3120postscanlM fstep begin (Stream step state) = Stream step' (state, begin)
3122 {-# INLINE_LATE step' #-}
3123 step' gst (st, acc) = do
3124 r <- step (adaptState gst) st
3128 return (Yield y (s, y))
3129 Skip s -> return $ Skip (s, acc)
3132{-# INLINE_NORMAL postscanl #-}
3133postscanl :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
3134postscanl f = postscanlM (\a b -> return (f a b))
3136{-# INLINE_NORMAL scanlM' #-}
3137scanlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
3138scanlM' fstep begin s = begin `seq` (begin `cons` postscanlM' fstep begin s)
3140{-# INLINE scanl' #-}
3141scanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
3142scanl' f = scanlM' (\a b -> return (f a b))
3144{-# INLINE_NORMAL scanlM #-}
3145scanlM :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
3146scanlM fstep begin s = begin `cons` postscanlM fstep begin s
3149scanl :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
3150scanl f = scanlM (\a b -> return (f a b))
3152{-# INLINE_NORMAL scanl1M #-}
3153scanl1M :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
3154scanl1M fstep (Stream step state) = Stream step' (state, Nothing)
3156 {-# INLINE_LATE step' #-}
3157 step' gst (st, Nothing) = do
3160 Yield x s -> return $ Yield x (s, Just x)
3161 Skip s -> return $ Skip (s, Nothing)
3164 step' gst (st, Just acc) = do
3169 return $ Yield z (s, Just z)
3170 Skip s -> return $ Skip (s, Just acc)
3173{-# INLINE scanl1 #-}
3174scanl1 :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
3175scanl1 f = scanl1M (\x y -> return (f x y))
3177{-# INLINE_NORMAL scanl1M' #-}
3178scanl1M' :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
3179scanl1M' fstep (Stream step state) = Stream step' (state, Nothing)
3181 {-# INLINE_LATE step' #-}
3182 step' gst (st, Nothing) = do
3185 Yield x s -> x `seq` return $ Yield x (s, Just x)
3186 Skip s -> return $ Skip (s, Nothing)
3189 step' gst (st, Just acc) = acc `seq` do
3194 z `seq` return $ Yield z (s, Just z)
3195 Skip s -> return $ Skip (s, Just acc)
3198{-# INLINE scanl1' #-}
3199scanl1' :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
3200scanl1' f = scanl1M' (\x y -> return (f x y))
3203tap :: Monad m => Fold m a b -> Stream m a -> Stream m a
3204tap (Fold fstep initial extract) (Stream step state) = Stream step' Nothing
3208 step' _ Nothing = do
3210 return $ Skip (Just (r, state))
3212 step' gst (Just (acc, st)) = acc `seq` do
3217 return $ Yield x (Just (acc', s))
3218 Skip s -> return $ Skip (Just (acc, s))
3223{-# INLINE_NORMAL tapRate #-}
3225 (MonadAsync m, MonadCatch m)
3230tapRate samplingRate action (Stream step state) = Stream step' Nothing
3232 {-# NOINLINE loop #-}
3233 loop countRef prev = do
3236 (do liftIO $ threadDelay (round $ samplingRate * 1000000)
3237 i <- liftIO $ readIORef countRef
3238 let !diff = i - prev
3241 (\(e :: AsyncException) -> do
3242 i <- liftIO $ readIORef countRef
3243 let !diff = i - prev
3245 throwM (MC.toException e))
3248 {-# INLINE_LATE step' #-}
3249 step' _ Nothing = do
3250 countRef <- liftIO $ newIORef 0
3251 tid <- fork $ loop countRef 0
3252 ref <- liftIO $ newIORef ()
3253 _ <- liftIO $ mkWeakIORef ref (killThread tid)
3254 return $ Skip (Just (countRef, tid, state, ref))
3256 step' gst (Just (countRef, tid, st, ref)) = do
3260 liftIO $ modifyIORef' countRef (+ 1)
3261 return $ Yield x (Just (countRef, tid, s, ref))
3262 Skip s -> return $ Skip (Just (countRef, tid, s, ref))
3264 liftIO $ killThread tid
3268-------------------------------------------------------------------------------
3270-------------------------------------------------------------------------------
3272{-# INLINE_NORMAL takeWhileM #-}
3273takeWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
3274takeWhileM f (Stream step state) = Stream step' state
3276 {-# INLINE_LATE step' #-}
3282 return $ if b then Yield x s else Stop
3283 Skip s -> return $ Skip s
3286{-# INLINE takeWhile #-}
3287takeWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
3288takeWhile f = takeWhileM (return . f)
3290{-# INLINE_NORMAL drop #-}
3291drop :: Monad m => Int -> Stream m a -> Stream m a
3292drop n (Stream step state) = Stream step' (state, Just n)
3294 {-# INLINE_LATE step' #-}
3295 step' gst (st, Just i)
3300 Yield _ s -> Skip (s, Just (i - 1))
3301 Skip s -> Skip (s, Just i)
3303 | otherwise = return $ Skip (st, Nothing)
3305 step' gst (st, Nothing) = do
3309 Yield x s -> Yield x (s, Nothing)
3310 Skip s -> Skip (s, Nothing)
3313data DropWhileState s a
3315 | DropWhileYield a s
3318{-# INLINE_NORMAL dropWhileM #-}
3319dropWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
3320dropWhileM f (Stream step state) = Stream step' (DropWhileDrop state)
3322 {-# INLINE_LATE step' #-}
3323 step' gst (DropWhileDrop st) = do
3329 then return $ Skip (DropWhileDrop s)
3330 else return $ Skip (DropWhileYield x s)
3331 Skip s -> return $ Skip (DropWhileDrop s)
3334 step' gst (DropWhileNext st) = do
3337 Yield x s -> return $ Skip (DropWhileYield x s)
3338 Skip s -> return $ Skip (DropWhileNext s)
3341 step' _ (DropWhileYield x st) = return $ Yield x (DropWhileNext st)
3343{-# INLINE dropWhile #-}
3344dropWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
3345dropWhile f = dropWhileM (return . f)
3347{-# INLINE_NORMAL filterM #-}
3348filterM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
3349filterM f (Stream step state) = Stream step' state
3351 {-# INLINE_LATE step' #-}
3360 Skip s -> return $ Skip s
3363{-# INLINE filter #-}
3364filter :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
3365filter f = filterM (return . f)
3367{-# INLINE_NORMAL uniq #-}
3368uniq :: (Eq a, Monad m) => Stream m a -> Stream m a
3369uniq (Stream step state) = Stream step' (Nothing, state)
3371 {-# INLINE_LATE step' #-}
3372 step' gst (Nothing, st) = do
3375 Yield x s -> return $ Yield x (Just x, s)
3376 Skip s -> return $ Skip (Nothing, s)
3378 step' gst (Just x, st) = do
3381 Yield y s | x == y -> return $ Skip (Just x, s)
3382 | otherwise -> return $ Yield y (Just y, s)
3383 Skip s -> return $ Skip (Just x, s)
3386------------------------------------------------------------------------------
3387-- Transformation by Mapping
3388------------------------------------------------------------------------------
3390{-# INLINE_NORMAL sequence #-}
3391sequence :: Monad m => Stream m (m a) -> Stream m a
3392sequence (Stream step state) = Stream step' state
3394 {-# INLINE_LATE step' #-}
3396 r <- step (adaptState gst) st
3398 Yield x s -> x >>= \a -> return (Yield a s)
3399 Skip s -> return $ Skip s
3402------------------------------------------------------------------------------
3404------------------------------------------------------------------------------
3406data LoopState x s = FirstYield s
3407 | InterspersingYield s
3410{-# INLINE_NORMAL intersperseM #-}
3411intersperseM :: Monad m => m a -> Stream m a -> Stream m a
3412intersperseM m (Stream step state) = Stream step' (FirstYield state)
3414 {-# INLINE_LATE step' #-}
3415 step' gst (FirstYield st) = do
3419 Yield x s -> Skip (YieldAndCarry x s)
3420 Skip s -> Skip (FirstYield s)
3423 step' gst (InterspersingYield st) = do
3428 return $ Yield a (YieldAndCarry x s)
3429 Skip s -> return $ Skip $ InterspersingYield s
3432 step' _ (YieldAndCarry x st) = return $ Yield x (InterspersingYield st)
3437 | SuffixYield a (SuffixState s a)
3439{-# INLINE_NORMAL intersperseSuffix #-}
3440intersperseSuffix :: forall m a. Monad m => m a -> Stream m a -> Stream m a
3441intersperseSuffix action (Stream step state) = Stream step' (SuffixElem state)
3443 {-# INLINE_LATE step' #-}
3444 step' gst (SuffixElem st) = do
3447 Yield x s -> Skip (SuffixYield x (SuffixSuffix s))
3448 Skip s -> Skip (SuffixElem s)
3451 step' _ (SuffixSuffix st) = do
3452 action >>= \r -> return $ Skip (SuffixYield r (SuffixElem st))
3454 step' _ (SuffixYield x next) = return $ Yield x next
3456data SuffixSpanState s a
3457 = SuffixSpanElem s Int
3458 | SuffixSpanSuffix s
3459 | SuffixSpanYield a (SuffixSpanState s a)
3463-- | intersperse after every n items
3464{-# INLINE_NORMAL intersperseSuffixBySpan #-}
3465intersperseSuffixBySpan :: forall m a. Monad m
3466 => Int -> m a -> Stream m a -> Stream m a
3467intersperseSuffixBySpan n action (Stream step state) =
3468 Stream step' (SuffixSpanElem state n)
3470 {-# INLINE_LATE step' #-}
3471 step' gst (SuffixSpanElem st i) | i > 0 = do
3474 Yield x s -> Skip (SuffixSpanYield x (SuffixSpanElem s (i - 1)))
3475 Skip s -> Skip (SuffixSpanElem s i)
3476 Stop -> if i == n then Stop else Skip SuffixSpanLast
3477 step' _ (SuffixSpanElem st _) = return $ Skip (SuffixSpanSuffix st)
3479 step' _ (SuffixSpanSuffix st) = do
3480 action >>= \r -> return $ Skip (SuffixSpanYield r (SuffixSpanElem st n))
3482 step' _ (SuffixSpanLast) = do
3483 action >>= \r -> return $ Skip (SuffixSpanYield r SuffixSpanStop)
3485 step' _ (SuffixSpanYield x next) = return $ Yield x next
3487 step' _ (SuffixSpanStop) = return Stop
3489{-# INLINE intersperse #-}
3490intersperse :: Monad m => a -> Stream m a -> Stream m a
3491intersperse a = intersperseM (return a)
3493{-# INLINE_NORMAL insertBy #-}
3494insertBy :: Monad m => (a -> a -> Ordering) -> a -> Stream m a -> Stream m a
3495insertBy cmp a (Stream step state) = Stream step' (state, False, Nothing)
3497 {-# INLINE_LATE step' #-}
3498 step' gst (st, False, _) = do
3501 Yield x s -> case cmp a x of
3502 GT -> return $ Yield x (s, False, Nothing)
3503 _ -> return $ Yield a (s, True, Just x)
3504 Skip s -> return $ Skip (s, False, Nothing)
3505 Stop -> return $ Yield a (st, True, Nothing)
3507 step' _ (_, True, Nothing) = return Stop
3509 step' gst (st, True, Just prev) = do
3512 Yield x s -> return $ Yield prev (s, True, Just x)
3513 Skip s -> return $ Skip (s, True, Just prev)
3514 Stop -> return $ Yield prev (st, True, Nothing)
3516------------------------------------------------------------------------------
3518------------------------------------------------------------------------------
3520{-# INLINE_NORMAL deleteBy #-}
3521deleteBy :: Monad m => (a -> a -> Bool) -> a -> Stream m a -> Stream m a
3522deleteBy eq x (Stream step state) = Stream step' (state, False)
3524 {-# INLINE_LATE step' #-}
3525 step' gst (st, False) = do
3528 Yield y s -> return $
3529 if eq x y then Skip (s, True) else Yield y (s, False)
3530 Skip s -> return $ Skip (s, False)
3533 step' gst (st, True) = do
3536 Yield y s -> return $ Yield y (s, True)
3537 Skip s -> return $ Skip (s, True)
3540------------------------------------------------------------------------------
3541-- Transformation by Map and Filter
3542------------------------------------------------------------------------------
3544-- XXX Will this always fuse properly?
3545{-# INLINE_NORMAL mapMaybe #-}
3546mapMaybe :: Monad m => (a -> Maybe b) -> Stream m a -> Stream m b
3547mapMaybe f = fmap fromJust . filter isJust . map f
3549{-# INLINE_NORMAL mapMaybeM #-}
3550mapMaybeM :: Monad m => (a -> m (Maybe b)) -> Stream m a -> Stream m b
3551mapMaybeM f = fmap fromJust . filter isJust . mapM f
3553------------------------------------------------------------------------------
3555------------------------------------------------------------------------------
3557{-# INLINE_NORMAL indexed #-}
3558indexed :: Monad m => Stream m a -> Stream m (Int, a)
3559indexed (Stream step state) = Stream step' (state, 0)
3561 {-# INLINE_LATE step' #-}
3562 step' gst (st, i) = i `seq` do
3563 r <- step (adaptState gst) st
3565 Yield x s -> return $ Yield (i, x) (s, i+1)
3566 Skip s -> return $ Skip (s, i)
3569{-# INLINE_NORMAL indexedR #-}
3570indexedR :: Monad m => Int -> Stream m a -> Stream m (Int, a)
3571indexedR m (Stream step state) = Stream step' (state, m)
3573 {-# INLINE_LATE step' #-}
3574 step' gst (st, i) = i `seq` do
3575 r <- step (adaptState gst) st
3577 Yield x s -> let i' = i - 1
3578 in return $ Yield (i, x) (s, i')
3579 Skip s -> return $ Skip (s, i)
3582{-# INLINE_NORMAL zipWithM #-}
3584 => (a -> b -> m c) -> Stream m a -> Stream m b -> Stream m c
3585zipWithM f (Stream stepa ta) (Stream stepb tb) = Stream step (ta, tb, Nothing)
3587 {-# INLINE_LATE step #-}
3588 step gst (sa, sb, Nothing) = do
3589 r <- stepa (adaptState gst) sa
3592 Yield x sa' -> Skip (sa', sb, Just x)
3593 Skip sa' -> Skip (sa', sb, Nothing)
3596 step gst (sa, sb, Just x) = do
3597 r <- stepb (adaptState gst) sb
3601 return $ Yield z (sa, sb', Nothing)
3602 Skip sb' -> return $ Skip (sa, sb', Just x)
3605#if __GLASGOW_HASKELL__ >= 801
3606{-# RULES "zipWithM xs xs"
3607 forall f xs. zipWithM @Identity f xs xs = mapM (\x -> f x x) xs #-}
3610{-# INLINE zipWith #-}
3611zipWith :: Monad m => (a -> b -> c) -> Stream m a -> Stream m b -> Stream m c
3612zipWith f = zipWithM (\a b -> return (f a b))
3614-- | Like 'zipWithM' but zips concurrently i.e. both the streams being zipped
3615-- are generated concurrently.
3617{-# INLINE_NORMAL zipAsyncWithM #-}
3618zipAsyncWithM :: MonadAsync m
3619 => (a -> b -> m c) -> Stream m a -> Stream m b -> Stream m c
3620zipAsyncWithM f m1 m2 = Stream step Nothing
3624 {-# INLINE_LATE step #-}
3625 step gst Nothing = do
3626 ma <- mkParallel (adaptState gst) m1
3627 mb <- mkParallel (adaptState gst) m2
3628 return $ Skip (Just $ zipWithM f ma mb)
3630 step gst (Just (UnStream step1 st)) = do
3633 Yield a s -> Yield a (Just $ Stream step1 s)
3634 Skip s -> Skip (Just $ Stream step1 s)
3637-- | Like 'zipWith' but zips concurrently i.e. both the streams being zipped
3638-- are generated concurrently.
3640{-# INLINE zipAsyncWith #-}
3641zipAsyncWith :: MonadAsync m
3642 => (a -> b -> c) -> Stream m a -> Stream m b -> Stream m c
3643zipAsyncWith f = zipAsyncWithM (\a b -> return (f a b))
3645------------------------------------------------------------------------------
3647------------------------------------------------------------------------------
3649{-# INLINE_NORMAL mergeByM #-}
3652 => (a -> a -> m Ordering) -> Stream m a -> Stream m a -> Stream m a
3653mergeByM cmp (Stream stepa ta) (Stream stepb tb) =
3654 Stream step (Just ta, Just tb, Nothing, Nothing)
3656 {-# INLINE_LATE step #-}
3658 -- one of the values is missing, and the corresponding stream is running
3659 step gst (Just sa, sb, Nothing, b) = do
3662 Yield a sa' -> Skip (Just sa', sb, Just a, b)
3663 Skip sa' -> Skip (Just sa', sb, Nothing, b)
3664 Stop -> Skip (Nothing, sb, Nothing, b)
3666 step gst (sa, Just sb, a, Nothing) = do
3669 Yield b sb' -> Skip (sa, Just sb', a, Just b)
3670 Skip sb' -> Skip (sa, Just sb', a, Nothing)
3671 Stop -> Skip (sa, Nothing, a, Nothing)
3673 -- both the values are available
3674 step _ (sa, sb, Just a, Just b) = do
3676 return $ case res of
3677 GT -> Yield b (sa, sb, Just a, Nothing)
3678 _ -> Yield a (sa, sb, Nothing, Just b)
3680 -- one of the values is missing, corresponding stream is done
3681 step _ (Nothing, sb, Nothing, Just b) =
3682 return $ Yield b (Nothing, sb, Nothing, Nothing)
3684 step _ (sa, Nothing, Just a, Nothing) =
3685 return $ Yield a (sa, Nothing, Nothing, Nothing)
3687 step _ (Nothing, Nothing, Nothing, Nothing) = return Stop
3689{-# INLINE mergeBy #-}
3692 => (a -> a -> Ordering) -> Stream m a -> Stream m a -> Stream m a
3693mergeBy cmp = mergeByM (\a b -> return $ cmp a b)
3695------------------------------------------------------------------------------
3696-- Transformation comprehensions
3697------------------------------------------------------------------------------
3699{-# INLINE_NORMAL the #-}
3700the :: (Eq a, Monad m) => Stream m a -> m (Maybe a)
3701the (Stream step state) = go state
3704 r <- step defState st
3706 Yield x s -> go' x s
3708 Stop -> return Nothing
3710 r <- step defState st
3712 Yield x s | x == n -> go' n s
3713 | otherwise -> return Nothing
3715 Stop -> return (Just n)
3717{-# INLINE runFold #-}
3718runFold :: (Monad m) => Fold m a b -> Stream m a -> m b
3719runFold (Fold step begin done) = foldlMx' step begin done
3721-------------------------------------------------------------------------------
3722-- Concurrent application and fold
3723-------------------------------------------------------------------------------
3725-- Using StreamD the worker stream producing code can fuse with the code to
3726-- queue output to the SVar giving some perf boost.
3728-- Note that StreamD can only be used in limited situations, specifically, we
3729-- cannot implement joinStreamVarPar using this.
3731-- XXX make sure that the SVar passed is a Parallel style SVar.
3733-- | Fold the supplied stream to the SVar asynchronously using Parallel
3734-- concurrency style.
3735-- {-# INLINE_NORMAL toSVarParallel #-}
3736{-# INLINE toSVarParallel #-}
3737toSVarParallel :: MonadAsync m
3738 => State t m a -> SVar t m a -> Stream m a -> m ()
3739toSVarParallel st sv xs =
3740 if svarInspectMode sv
3744 case getYieldLimit st of
3745 Nothing -> doFork (work Nothing)
3747 (handleChildException sv)
3748 Just _ -> doFork (workLim Nothing)
3750 (handleChildException sv)
3755 {-# NOINLINE work #-}
3756 work info = (runFold (FL.toParallelSVar sv info) xs)
3758 {-# NOINLINE workLim #-}
3759 workLim info = runFold (FL.toParallelSVarLimited sv info) xs
3761 {-# NOINLINE forkWithDiag #-}
3763 -- We do not use workerCount in case of ParallelVar but still there is
3764 -- no harm in maintaining it correctly.
3765 liftIO $ atomicModifyIORefCAS_ (workerCount sv) $ \n -> n + 1
3767 -- This allocation matters when significant number of workers are being
3768 -- sent. We allocate it only when needed. The overhead increases by 4x.
3770 case yieldRateInfo sv of
3771 Nothing -> return Nothing
3772 Just _ -> liftIO $ do
3773 cntRef <- newIORef 0
3774 t <- getTime Monotonic
3775 lat <- newIORef (0, t)
3776 return $ Just WorkerInfo
3777 { workerYieldMax = 0
3778 , workerYieldCount = cntRef
3779 , workerLatencyStart = lat
3782 case getYieldLimit st of
3783 Nothing -> doFork (work winfo)
3785 (handleChildException sv)
3786 Just _ -> doFork (workLim winfo)
3788 (handleChildException sv)
3791{-# INLINE_NORMAL mkParallel #-}
3792mkParallel :: MonadAsync m => State t m a -> Stream m a -> m (Stream m a)
3794 sv <- newParallelVar StopNone defState
3795 toSVarParallel st sv m
3796 return $ fromSVar sv
3798{-# INLINE_NORMAL applyParallel #-}
3799applyParallel :: MonadAsync m
3800 => (Stream m a -> Stream m b) -> Stream m a -> Stream m b
3801applyParallel f m = Stream step Nothing
3804 {-# INLINE_LATE step #-}
3805 step gst Nothing = do
3806 sv <- newParallelVar StopNone (adaptState gst)
3807 toSVarParallel (adaptState gst) sv m
3808 return $ Skip $ Just $ f (fromSVar sv)
3810 step gst (Just (UnStream step1 st)) = do
3813 Yield a s -> Yield a (Just $ Stream step1 s)
3814 Skip s -> Skip (Just $ Stream step1 s)
3817{-# INLINE foldParallel #-}
3818foldParallel :: (K.IsStream t, MonadAsync m) => (t m a -> m b) -> t m a -> m b
3819foldParallel f m = do
3820 sv <- newParallelVar StopNone defState
3821 toSVarParallel defState sv (toStreamD m)
3822 f $ fromStreamD $ fromSVar sv
3824-------------------------------------------------------------------------------
3826-------------------------------------------------------------------------------
3828-- | Create an SVar with a fold consumer that will fold any elements sent to it
3829-- using the supplied fold function.
3830{-# INLINE newFoldSVar #-}
3831newFoldSVar :: MonadAsync m => State t m a -> Fold m a b -> m (SVar t m a)
3832newFoldSVar stt f = do
3833 -- Buffer size for the SVar is derived from the current state
3834 sv <- newParallelVar StopAny (adaptState stt)
3835 -- Add the producer thread-id to the SVar.
3836 liftIO myThreadId >>= modifyThread sv
3837 void $ doFork (work sv) (svarMrun sv) (handleFoldException sv)
3842 {-# NOINLINE work #-}
3843 work sv = void $ runFold f $ fromProducer sv
3845data TapState sv st = TapInit | Tapping sv st | TapDone st
3847{-# INLINE_NORMAL tapAsync #-}
3848tapAsync :: MonadAsync m => Fold m a b -> Stream m a -> Stream m a
3849tapAsync f (Stream step1 state1) = Stream step TapInit
3853 -- In general, a Stop event would come equipped with the result
3854 -- of the fold. It is not used here but it would be useful in
3855 -- applicative and distribute.
3856 done <- fromConsumer svr
3857 when (not done) $ do
3858 liftIO $ withDiagMVar svr "teeToSVar: waiting to drain"
3859 $ takeMVar (outputDoorBellFromConsumer svr)
3863 liftIO $ sendStop svr Nothing
3864 -- drain/wait until a stop event arrives from the fold.
3867 {-# INLINE_LATE step #-}
3868 step gst TapInit = do
3869 sv <- newFoldSVar gst f
3870 return $ Skip (Tapping sv state1)
3872 step gst (Tapping sv st) = do
3876 done <- pushToFold sv a
3879 -- XXX we do not need to wait synchronously here
3881 return $ Yield a (TapDone s)
3882 else return $ Yield a (Tapping sv s)
3883 Skip s -> return $ Skip (Tapping sv s)
3888 step gst (TapDone st) = do
3891 Yield a s -> Yield a (TapDone s)
3892 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