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.Streams.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.Streams.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.Streams.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
149 -- ** Flattening nested streams
152 , ConcatMapUState (..)
154 , ConcatUnfoldInterleaveState (..)
155 , concatUnfoldInterleave
156 , concatUnfoldRoundrobin
159 , InterleaveState(..)
164 , roundRobin -- interleaveFair?/ParallelFair
197 -- | Transform a stream into another type.
214 -- ** By folding (scans)
257 , intersperseSuffixBySpan
283 -- * Transformation comprehensions
295 -- * UTF8 Encoding / Decoding transformations.
303 , resumeDecodeUtf8Either
305 , decodeUtf8ArraysLenient
307 -- * Concurrent Application
314import Control.Exception (Exception, SomeException, fromException)
315import Control.Monad (void, when)
316import Control.Monad.Catch (MonadCatch, throwM)
317import Control.Monad.IO.Class (MonadIO(..))
318import Control.Monad.Reader (ReaderT)
319import Control.Monad.State.Strict (StateT)
320import Control.Monad.Trans (MonadTrans(lift))
321import Data.Bits (shiftR, shiftL, (.|.), (.&.))
322import Data.Functor.Identity (Identity(..))
323import Data.IORef (newIORef, readIORef, mkWeakIORef, writeIORef)
324import Data.Maybe (fromJust, isJust, isNothing)
325import Data.Word (Word32)
326import Foreign.Ptr (Ptr)
327import Foreign.Storable (Storable(..))
328import GHC.Base (assert, unsafeChr, ord)
329import GHC.IO.Encoding.Failure (isSurrogate)
330import GHC.ForeignPtr (ForeignPtr (..))
331import GHC.Ptr (Ptr (..))
332import GHC.Types (SPEC(..))
333import GHC.Word (Word8(..))
334import System.IO.Unsafe (unsafePerformIO)
335import System.Mem (performMajorGC)
337 hiding (map, mapM, mapM_, repeat, foldr, last, take, filter,
338 takeWhile, drop, dropWhile, all, any, maximum, minimum, elem,
339 notElem, null, head, tail, zipWith, lookup, foldr1, sequence,
340 (!!), scanl, scanl1, concatMap, replicate, enumFromTo, concat,
343import qualified Control.Monad.Catch as MC
344import qualified Control.Monad.Reader as Reader
345import qualified Control.Monad.State.Strict as State
346import qualified Prelude
348import Streamly.Internal.Data.Atomics (atomicModifyIORefCAS_)
349import Streamly.Internal.Memory.Array.Types (Array(..))
350import Streamly.Internal.Data.Fold.Types (Fold(..))
351import Streamly.Internal.Data.Pipe.Types (Pipe(..), PipeState(..))
352import Streamly.Internal.Data.Time.Clock (Clock(Monotonic), getTime)
353import Streamly.Internal.Data.Unfold.Types (Unfold(..))
354import Streamly.Internal.Data.Strict (Tuple'(..))
356import Streamly.Internal.Data.Stream.StreamD.Type
357import Streamly.Internal.Data.SVar
359import qualified Streamly.Internal.Data.Pipe.Types as Pipe
360import qualified Streamly.Internal.Memory.Array.Types as A
361import qualified Streamly.Internal.Data.Fold as FL
362import qualified Streamly.Memory.Ring as RB
363import qualified Streamly.Streams.StreamK as K
365import Foreign.Ptr (plusPtr)
366import Foreign.ForeignPtr.Unsafe (unsafeForeignPtrToPtr)
367import Foreign.ForeignPtr (touchForeignPtr)
369------------------------------------------------------------------------------
371------------------------------------------------------------------------------
373-- | An empty 'Stream'.
374{-# INLINE_NORMAL nil #-}
375nil :: Monad m => Stream m a
376nil = Stream (\_ _ -> return Stop) ()
378-- | An empty 'Stream' with a side effect.
379{-# INLINE_NORMAL nilM #-}
380nilM :: Monad m => m b -> Stream m a
381nilM m = Stream (\_ _ -> m >> return Stop) ()
383{-# INLINE_NORMAL consM #-}
384consM :: Monad m => m a -> Stream m a -> Stream m a
385consM m (Stream step state) = Stream step1 Nothing
387 {-# INLINE_LATE step1 #-}
388 step1 _ Nothing = m >>= \x -> 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-- XXX implement in terms of consM?
398-- cons x = consM (return x)
400-- | Can fuse but has O(n^2) complexity.
401{-# INLINE_NORMAL cons #-}
402cons :: Monad m => a -> Stream m a -> Stream m a
403cons x (Stream step state) = Stream step1 Nothing
405 {-# INLINE_LATE step1 #-}
406 step1 _ Nothing = return $ Yield x (Just state)
407 step1 gst (Just st) = do
411 Yield a s -> Yield a (Just s)
412 Skip s -> Skip (Just s)
415-------------------------------------------------------------------------------
417-------------------------------------------------------------------------------
419-- Does not fuse, has the same performance as the StreamK version.
420{-# INLINE_NORMAL uncons #-}
421uncons :: Monad m => Stream m a -> m (Maybe (a, Stream m a))
422uncons (UnStream step state) = go state
425 r <- step defState st
427 Yield x s -> return $ Just (x, Stream step s)
429 Stop -> return Nothing
431------------------------------------------------------------------------------
432-- Generation by unfold
433------------------------------------------------------------------------------
435{-# INLINE_NORMAL unfoldrM #-}
436unfoldrM :: Monad m => (s -> m (Maybe (a, s))) -> s -> Stream m a
437unfoldrM next state = Stream step state
439 {-# INLINE_LATE step #-}
443 Just (x, s) -> Yield x s
446{-# INLINE_LATE unfoldr #-}
447unfoldr :: Monad m => (s -> Maybe (a, s)) -> s -> Stream m a
448unfoldr f = unfoldrM (return . f)
450-- | Convert an 'Unfold' into a 'Stream' by supplying it a seed.
452{-# INLINE_NORMAL unfold #-}
453unfold :: Monad m => Unfold m a b -> a -> Stream m b
454unfold (Unfold ustep inject) seed = Stream step Nothing
456 {-# INLINE_LATE step #-}
457 step _ Nothing = inject seed >>= return . Skip . Just
458 step _ (Just st) = do
461 Yield x s -> Yield x (Just s)
462 Skip s -> Skip (Just s)
465------------------------------------------------------------------------------
466-- Specialized Generation
467------------------------------------------------------------------------------
469repeatM :: Monad m => m a -> Stream m a
470repeatM x = Stream (\_ _ -> x >>= \r -> return $ Yield r ()) ()
472repeat :: Monad m => a -> Stream m a
473repeat x = Stream (\_ _ -> return $ Yield x ()) ()
475{-# INLINE_NORMAL replicateM #-}
476replicateM :: forall m a. Monad m => Int -> m a -> Stream m a
477replicateM n p = Stream step n
479 {-# INLINE_LATE step #-}
481 | i <= 0 = return Stop
484 return $ Yield x (i - 1)
486{-# INLINE_NORMAL replicate #-}
487replicate :: Monad m => Int -> a -> Stream m a
488replicate n x = replicateM n (return x)
490-- This would not work properly for floats, therefore we put an Integral
492-- | Can be used to enumerate unbounded integrals. This does not check for
493-- overflow or underflow for bounded integrals.
494{-# INLINE_NORMAL enumerateFromStepIntegral #-}
495enumerateFromStepIntegral :: (Integral a, Monad m) => a -> a -> Stream m a
496enumerateFromStepIntegral from stride =
497 from `seq` stride `seq` Stream step from
499 {-# INLINE_LATE step #-}
500 step _ !x = return $ Yield x $! (x + stride)
502-- We are assuming that "to" is constrained by the type to be within
504{-# INLINE enumerateFromToIntegral #-}
505enumerateFromToIntegral :: (Monad m, Integral a) => a -> a -> Stream m a
506enumerateFromToIntegral from to =
507 takeWhile (<= to) $ enumerateFromStepIntegral from 1
509{-# INLINE enumerateFromIntegral #-}
510enumerateFromIntegral :: (Monad m, Integral a, Bounded a) => a -> Stream m a
511enumerateFromIntegral from = enumerateFromToIntegral from maxBound
513data EnumState a = EnumInit | EnumYield a a a | EnumStop
515{-# INLINE_NORMAL enumerateFromThenToIntegralUp #-}
516enumerateFromThenToIntegralUp
517 :: (Monad m, Integral a)
518 => a -> a -> a -> Stream m a
519enumerateFromThenToIntegralUp from next to = Stream step EnumInit
521 {-# INLINE_LATE step #-}
527 else Yield from EnumStop
528 else -- from <= next <= to
529 let stride = next - from
530 in Skip $ EnumYield from stride (to - stride)
532 step _ (EnumYield x stride toMinus) =
535 then Yield x EnumStop
536 else Yield x $ EnumYield (x + stride) stride toMinus
538 step _ EnumStop = return Stop
540{-# INLINE_NORMAL enumerateFromThenToIntegralDn #-}
541enumerateFromThenToIntegralDn
542 :: (Monad m, Integral a)
543 => a -> a -> a -> Stream m a
544enumerateFromThenToIntegralDn from next to = Stream step EnumInit
546 {-# INLINE_LATE step #-}
548 return $ if to > next
551 else Yield from EnumStop
552 else -- from >= next >= to
553 let stride = next - from
554 in Skip $ EnumYield from stride (to - stride)
556 step _ (EnumYield x stride toMinus) =
559 then Yield x EnumStop
560 else Yield x $ EnumYield (x + stride) stride toMinus
562 step _ EnumStop = return Stop
564{-# INLINE_NORMAL enumerateFromThenToIntegral #-}
565enumerateFromThenToIntegral
566 :: (Monad m, Integral a)
567 => a -> a -> a -> Stream m a
568enumerateFromThenToIntegral from next to
569 | next >= from = enumerateFromThenToIntegralUp from next to
570 | otherwise = enumerateFromThenToIntegralDn from next to
572{-# INLINE_NORMAL enumerateFromThenIntegral #-}
573enumerateFromThenIntegral
574 :: (Monad m, Integral a, Bounded a)
575 => a -> a -> Stream m a
576enumerateFromThenIntegral from next =
578 then enumerateFromThenToIntegralUp from next maxBound
579 else enumerateFromThenToIntegralDn from next minBound
581-- For floating point numbers if the increment is less than the precision then
582-- it just gets lost. Therefore we cannot always increment it correctly by just
584-- 9007199254740992 + 1 + 1 :: Double => 9.007199254740992e15
585-- 9007199254740992 + 2 :: Double => 9.007199254740994e15
587-- Instead we accumulate the increment counter and compute the increment
588-- every time before adding it to the starting number.
590-- This works for Integrals as well as floating point numbers, but
591-- enumerateFromStepIntegral is faster for integrals.
592{-# INLINE_NORMAL enumerateFromStepNum #-}
593enumerateFromStepNum :: (Monad m, Num a) => a -> a -> Stream m a
594enumerateFromStepNum from stride = Stream step 0
596 {-# INLINE_LATE step #-}
597 step _ !i = return $ (Yield $! (from + i * stride)) $! (i + 1)
599{-# INLINE_NORMAL numFrom #-}
600numFrom :: (Monad m, Num a) => a -> Stream m a
601numFrom from = enumerateFromStepNum from 1
603{-# INLINE_NORMAL numFromThen #-}
604numFromThen :: (Monad m, Num a) => a -> a -> Stream m a
605numFromThen from next = enumerateFromStepNum from (next - from)
607-- We cannot write a general function for Num. The only way to write code
608-- portable between the two is to use a 'Real' constraint and convert between
609-- Fractional and Integral using fromRational which is horribly slow.
610{-# INLINE_NORMAL enumerateFromToFractional #-}
611enumerateFromToFractional
612 :: (Monad m, Fractional a, Ord a)
613 => a -> a -> Stream m a
614enumerateFromToFractional from to =
615 takeWhile (<= to + 1 / 2) $ enumerateFromStepNum from 1
617{-# INLINE_NORMAL enumerateFromThenToFractional #-}
618enumerateFromThenToFractional
619 :: (Monad m, Fractional a, Ord a)
620 => a -> a -> a -> Stream m a
621enumerateFromThenToFractional from next to =
622 takeWhile predicate $ numFromThen from next
624 mid = (next - from) / 2
625 predicate | next >= from = (<= to + mid)
626 | otherwise = (>= to + mid)
628-------------------------------------------------------------------------------
629-- Generation by Conversion
630-------------------------------------------------------------------------------
632{-# INLINE_NORMAL fromIndicesM #-}
633fromIndicesM :: Monad m => (Int -> m a) -> Stream m a
634fromIndicesM gen = Stream step 0
636 {-# INLINE_LATE step #-}
639 return $ Yield x (i + 1)
641{-# INLINE fromIndices #-}
642fromIndices :: Monad m => (Int -> a) -> Stream m a
643fromIndices gen = fromIndicesM (return . gen)
645{-# INLINE_NORMAL generateM #-}
646generateM :: Monad m => Int -> (Int -> m a) -> Stream m a
647generateM n gen = n `seq` Stream step 0
649 {-# INLINE_LATE step #-}
650 step _ i | i < n = do
652 return $ Yield x (i + 1)
653 | otherwise = return Stop
655{-# INLINE generate #-}
656generate :: Monad m => Int -> (Int -> a) -> Stream m a
657generate n gen = generateM n (return . gen)
659-- XXX we need the MonadAsync constraint because of a rewrite rule.
660-- | Convert a list of monadic actions to a 'Stream'
661{-# INLINE_LATE fromListM #-}
662fromListM :: MonadAsync m => [m a] -> Stream m a
663fromListM = Stream step
665 {-# INLINE_LATE step #-}
666 step _ (m:ms) = m >>= \x -> return $ Yield x ms
667 step _ [] = return Stop
669{-# INLINE toStreamD #-}
670toStreamD :: (K.IsStream t, Monad m) => t m a -> Stream m a
671toStreamD = fromStreamK . K.toStream
673-------------------------------------------------------------------------------
674-- Generation from SVar
675-------------------------------------------------------------------------------
677data FromSVarState t m a =
679 | FromSVarRead (SVar t m a)
680 | FromSVarLoop (SVar t m a) [ChildEvent a]
681 | FromSVarDone (SVar t m a)
683{-# INLINE_NORMAL fromSVar #-}
684fromSVar :: (MonadAsync m) => SVar t m a -> Stream m a
685fromSVar svar = Stream step FromSVarInit
688 {-# INLINE_LATE step #-}
689 step _ FromSVarInit = do
690 ref <- liftIO $ newIORef ()
691 _ <- liftIO $ mkWeakIORef ref hook
692 -- when this copy of svar gets garbage collected "ref" will get
693 -- garbage collected and our GC hook will be called.
694 let sv = svar{svarRef = Just ref}
695 return $ Skip (FromSVarRead sv)
699 {-# NOINLINE hook #-}
701 when (svarInspectMode svar) $ do
702 r <- liftIO $ readIORef (svarStopTime (svarStats svar))
704 printSVar svar "SVar Garbage Collected"
706 -- If there are any SVars referenced by this SVar a GC will prompt
707 -- them to be cleaned up quickly.
708 when (svarInspectMode svar) performMajorGC
710 step _ (FromSVarRead sv) = do
711 list <- readOutputQ sv
712 -- Reversing the output is important to guarantee that we process the
713 -- outputs in the same order as they were generated by the constituent
715 return $ Skip $ FromSVarLoop sv (Prelude.reverse list)
717 step _ (FromSVarLoop sv []) = do
718 done <- postProcess sv
719 return $ Skip $ if done
720 then (FromSVarDone sv)
721 else (FromSVarRead sv)
723 step _ (FromSVarLoop sv (ev : es)) = do
725 ChildYield a -> return $ Yield a (FromSVarLoop sv es)
726 ChildStop tid e -> do
730 stop <- shouldStop tid
733 liftIO (cleanupSVar sv)
734 return $ Skip (FromSVarDone sv)
735 else return $ Skip (FromSVarLoop sv es)
737 case fromException ex of
739 return $ Skip (FromSVarLoop sv es)
740 Nothing -> liftIO (cleanupSVar sv) >> throwM ex
744 case svarStopStyle sv of
745 StopNone -> return False
746 StopAny -> return True
748 sid <- liftIO $ readIORef (svarStopBy sv)
749 return $ if tid == sid then True else False
751 step _ (FromSVarDone sv) = do
752 when (svarInspectMode sv) $ do
753 t <- liftIO $ getTime Monotonic
754 liftIO $ writeIORef (svarStopTime (svarStats sv)) (Just t)
755 liftIO $ printSVar sv "SVar Done"
758-------------------------------------------------------------------------------
759-- Hoisting the inner monad
760-------------------------------------------------------------------------------
762{-# INLINE_NORMAL hoist #-}
763hoist :: Monad n => (forall x. m x -> n x) -> Stream m a -> Stream n a
764hoist f (Stream step state) = (Stream step' state)
766 {-# INLINE_LATE step' #-}
768 r <- f $ step (adaptState gst) st
770 Yield x s -> Yield x s
774{-# INLINE generally #-}
775generally :: Monad m => Stream Identity a -> Stream m a
776generally = hoist (return . runIdentity)
778{-# INLINE_NORMAL liftInner #-}
779liftInner :: (Monad m, MonadTrans t, Monad (t m))
780 => Stream m a -> Stream (t m) a
781liftInner (Stream step state) = Stream step' state
783 {-# INLINE_LATE step' #-}
785 r <- lift $ step (adaptState gst) st
787 Yield x s -> Yield x s
791{-# INLINE_NORMAL runReaderT #-}
792runReaderT :: Monad m => s -> Stream (ReaderT s m) a -> Stream m a
793runReaderT sval (Stream step state) = Stream step' state
795 {-# INLINE_LATE step' #-}
797 r <- Reader.runReaderT (step (adaptState gst) st) sval
799 Yield x s -> Yield x s
803{-# INLINE_NORMAL evalStateT #-}
804evalStateT :: Monad m => s -> Stream (StateT s m) a -> Stream m a
805evalStateT sval (Stream step state) = Stream step' (state, sval)
807 {-# INLINE_LATE step' #-}
808 step' gst (st, sv) = do
809 (r, sv') <- State.runStateT (step (adaptState gst) st) sv
811 Yield x s -> Yield x (s, sv')
812 Skip s -> Skip (s, sv')
815{-# INLINE_NORMAL runStateT #-}
816runStateT :: Monad m => s -> Stream (StateT s m) a -> Stream m (s, a)
817runStateT sval (Stream step state) = Stream step' (state, sval)
819 {-# INLINE_LATE step' #-}
820 step' gst (st, sv) = do
821 (r, sv') <- State.runStateT (step (adaptState gst) st) sv
823 Yield x s -> Yield (sv', x) (s, sv')
824 Skip s -> Skip (s, sv')
827------------------------------------------------------------------------------
828-- Elimination by Folds
829------------------------------------------------------------------------------
831------------------------------------------------------------------------------
833------------------------------------------------------------------------------
835{-# INLINE_NORMAL foldr1 #-}
836foldr1 :: Monad m => (a -> a -> a) -> Stream m a -> m (Maybe a)
840 Nothing -> return Nothing
841 Just (h, t) -> fmap Just (foldr f h t)
843------------------------------------------------------------------------------
845------------------------------------------------------------------------------
847{-# INLINE_NORMAL foldlT #-}
848foldlT :: (Monad m, Monad (s m), MonadTrans s)
849 => (s m b -> a -> s m b) -> s m b -> Stream m a -> s m b
850foldlT fstep begin (Stream step state) = go SPEC begin state
853 r <- lift $ step defState st
855 Yield x s -> go SPEC (fstep acc x) s
856 Skip s -> go SPEC acc s
859-- Note, this is going to have horrible performance, because of the nature of
860-- the stream type (i.e. direct stream vs CPS). Its only for reference, it is
861-- likely be practically unusable.
862{-# INLINE_NORMAL foldlS #-}
864 => (Stream m b -> a -> Stream m b) -> Stream m b -> Stream m a -> Stream m b
865foldlS fstep begin (Stream step state) = Stream step' (Left (state, begin))
867 step' gst (Left (st, acc)) = do
868 r <- step (adaptState gst) st
870 Yield x s -> Skip (Left (s, fstep acc x))
871 Skip s -> Skip (Left (s, acc))
872 Stop -> Skip (Right acc)
874 step' gst (Right (Stream stp stt)) = do
875 r <- stp (adaptState gst) stt
877 Yield x s -> Yield x (Right (Stream stp s))
878 Skip s -> Skip (Right (Stream stp s))
881------------------------------------------------------------------------------
883------------------------------------------------------------------------------
885-- | Run a streaming composition, discard the results.
886{-# INLINE_LATE drain #-}
887drain :: Monad m => Stream m a -> m ()
888-- drain = foldrM (\_ xs -> xs) (return ())
889drain (Stream step state) = go SPEC state
892 r <- step defState st
894 Yield _ s -> go SPEC s
898{-# INLINE_NORMAL null #-}
899null :: Monad m => Stream m a -> m Bool
900null m = foldrM (\_ _ -> return False) (return True) m
902{-# INLINE_NORMAL head #-}
903head :: Monad m => Stream m a -> m (Maybe a)
904head m = foldrM (\x _ -> return (Just x)) (return Nothing) m
906{-# INLINE_NORMAL headElse #-}
907headElse :: Monad m => a -> Stream m a -> m a
908headElse a m = foldrM (\x _ -> return x) (return a) m
910-- Does not fuse, has the same performance as the StreamK version.
911{-# INLINE_NORMAL tail #-}
912tail :: Monad m => Stream m a -> m (Maybe (Stream m a))
913tail (UnStream step state) = go state
916 r <- step defState st
918 Yield _ s -> return (Just $ Stream step s)
920 Stop -> return Nothing
922-- XXX will it fuse? need custom impl?
923{-# INLINE_NORMAL last #-}
924last :: Monad m => Stream m a -> m (Maybe a)
925last = foldl' (\_ y -> Just y) Nothing
927{-# INLINE_NORMAL elem #-}
928elem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
929-- elem e m = foldrM (\x xs -> if x == e then return True else xs) (return False) m
930elem e (Stream step state) = go state
933 r <- step defState st
936 | x == e -> return True
941{-# INLINE_NORMAL notElem #-}
942notElem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
943notElem e s = fmap not (elem e s)
945{-# INLINE_NORMAL all #-}
946all :: Monad m => (a -> Bool) -> Stream m a -> m Bool
947-- all p m = foldrM (\x xs -> if p x then xs else return False) (return True) m
948all p (Stream step state) = go state
951 r <- step defState st
955 | otherwise -> return False
959{-# INLINE_NORMAL any #-}
960any :: Monad m => (a -> Bool) -> Stream m a -> m Bool
961-- any p m = foldrM (\x xs -> if p x then return True else xs) (return False) m
962any p (Stream step state) = go state
965 r <- step defState st
973{-# INLINE_NORMAL maximum #-}
974maximum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
975maximum (Stream step state) = go Nothing state
978 r <- step defState st
980 Yield x s -> go (Just x) s
981 Skip s -> go Nothing s
982 Stop -> return Nothing
983 go (Just acc) st = do
984 r <- step defState st
987 | acc <= x -> go (Just x) s
988 | otherwise -> go (Just acc) s
989 Skip s -> go (Just acc) s
990 Stop -> return (Just acc)
992{-# INLINE_NORMAL maximumBy #-}
993maximumBy :: Monad m => (a -> a -> Ordering) -> Stream m a -> m (Maybe a)
994maximumBy cmp (Stream step state) = go Nothing state
997 r <- step defState st
999 Yield x s -> go (Just x) s
1000 Skip s -> go Nothing s
1001 Stop -> return Nothing
1002 go (Just acc) st = do
1003 r <- step defState st
1005 Yield x s -> case cmp acc x of
1006 GT -> go (Just acc) s
1008 Skip s -> go (Just acc) s
1009 Stop -> return (Just acc)
1011{-# INLINE_NORMAL minimum #-}
1012minimum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
1013minimum (Stream step state) = go Nothing state
1016 r <- step defState st
1018 Yield x s -> go (Just x) s
1019 Skip s -> go Nothing s
1020 Stop -> return Nothing
1021 go (Just acc) st = do
1022 r <- step defState st
1025 | acc <= x -> go (Just acc) s
1026 | otherwise -> go (Just x) s
1027 Skip s -> go (Just acc) s
1028 Stop -> return (Just acc)
1030{-# INLINE_NORMAL minimumBy #-}
1031minimumBy :: Monad m => (a -> a -> Ordering) -> Stream m a -> m (Maybe a)
1032minimumBy cmp (Stream step state) = go Nothing state
1035 r <- step defState st
1037 Yield x s -> go (Just x) s
1038 Skip s -> go Nothing s
1039 Stop -> return Nothing
1040 go (Just acc) st = do
1041 r <- step defState st
1043 Yield x s -> case cmp acc x of
1045 _ -> go (Just acc) s
1046 Skip s -> go (Just acc) s
1047 Stop -> return (Just acc)
1049{-# INLINE_NORMAL (!!) #-}
1050(!!) :: (Monad m) => Stream m a -> Int -> m (Maybe a)
1051(Stream step state) !! i = go i state
1054 r <- step defState st
1056 Yield x s | n < 0 -> return Nothing
1057 | n == 0 -> return $ Just x
1058 | otherwise -> go (n - 1) s
1060 Stop -> return Nothing
1062{-# INLINE_NORMAL lookup #-}
1063lookup :: (Monad m, Eq a) => a -> Stream m (a, b) -> m (Maybe b)
1064lookup e m = foldrM (\(a, b) xs -> if e == a then return (Just b) else xs)
1067{-# INLINE_NORMAL findM #-}
1068findM :: Monad m => (a -> m Bool) -> Stream m a -> m (Maybe a)
1069findM p m = foldrM (\x xs -> p x >>= \r -> if r then return (Just x) else xs)
1073find :: Monad m => (a -> Bool) -> Stream m a -> m (Maybe a)
1074find p = findM (return . p)
1076{-# INLINE_NORMAL findIndices #-}
1077findIndices :: Monad m => (a -> Bool) -> Stream m a -> Stream m Int
1078findIndices p (Stream step state) = Stream step' (state, 0)
1080 {-# INLINE_LATE step' #-}
1081 step' gst (st, i) = i `seq` do
1082 r <- step (adaptState gst) st
1084 Yield x s -> if p x then Yield i (s, i+1) else Skip (s, i+1)
1085 Skip s -> Skip (s, i)
1088{-# INLINE toListRev #-}
1089toListRev :: Monad m => Stream m a -> m [a]
1090toListRev = foldl' (flip (:)) []
1092-- We can implement reverse as:
1094-- > reverse = foldlS (flip cons) nil
1096-- However, this implementation is unusable because of the horrible performance
1097-- of cons. So we just convert it to a list first and then stream from the
1100-- XXX Maybe we can use an Array instead of a list here?
1101{-# INLINE_NORMAL reverse #-}
1102reverse :: Monad m => Stream m a -> Stream m a
1103reverse m = Stream step Nothing
1105 {-# INLINE_LATE step #-}
1108 return $ Skip (Just xs)
1109 step _ (Just (x:xs)) = return $ Yield x (Just xs)
1110 step _ (Just []) = return Stop
1112-- Much faster reverse for Storables
1113{-# INLINE_NORMAL reverse' #-}
1114reverse' :: forall m a. (MonadIO m, Storable a) => Stream m a -> Stream m a
1116-- This commented implementation copies the whole stream into one single array
1117-- and then streams from that array, this is 3-4x faster than the chunked code
1118-- that follows. Though this could be problematic due to unbounded large
1119-- allocations. We need to figure out why the chunked code is slower and if we
1120-- can optimize the chunked code to work as fast as this one. It may be a
1122import Foreign.ForeignPtr (touchForeignPtr)
1123import Foreign.ForeignPtr.Unsafe (unsafeForeignPtrToPtr)
1124import Foreign.Ptr (Ptr, plusPtr)
1125reverse' m = Stream step Nothing
1127 {-# INLINE_LATE step #-}
1129 arr <- A.fromStreamD m
1130 let p = aEnd arr `plusPtr` negate (sizeOf (undefined :: a))
1131 return $ Skip $ Just (aStart arr, p)
1133 step _ (Just (start, p)) | p < unsafeForeignPtrToPtr start = return Stop
1135 step _ (Just (start, p)) = do
1136 let !x = A.unsafeInlineIO $ do
1138 touchForeignPtr start
1140 next = p `plusPtr` negate (sizeOf (undefined :: a))
1141 return $ Yield x (Just (start, next))
1148 $ A.fromStreamDArraysOf A.defaultChunkSize m
1151------------------------------------------------------------------------------
1152-- Grouping/Splitting
1153------------------------------------------------------------------------------
1155{-# INLINE_NORMAL splitSuffixBy' #-}
1156splitSuffixBy' :: Monad m
1157 => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1158splitSuffixBy' predicate f (Stream step state) =
1159 Stream (stepOuter f) (Just state)
1163 {-# INLINE_LATE stepOuter #-}
1164 stepOuter (Fold fstep initial done) gst (Just st) = do
1165 res <- step (adaptState gst) st
1171 then done acc' >>= \val -> return $ Yield val (Just s)
1174 Skip s -> return $ Skip $ Just s
1180 res <- step (adaptState gst) stt
1185 then done acc' >>= \val -> return $ Yield val (Just s)
1187 Skip s -> go SPEC s acc
1188 Stop -> done acc >>= \val -> return $ Yield val Nothing
1190 stepOuter _ _ Nothing = return Stop
1192{-# INLINE_NORMAL groupsBy #-}
1198groupsBy cmp f (Stream step state) = Stream (stepOuter f) (Just state, Nothing)
1202 {-# INLINE_LATE stepOuter #-}
1203 stepOuter (Fold fstep initial done) gst (Just st, Nothing) = do
1204 res <- step (adaptState gst) st
1211 Skip s -> return $ Skip $ (Just s, Nothing)
1216 go !_ prev stt !acc = do
1217 res <- step (adaptState gst) stt
1224 else done acc >>= \r -> return $ Yield r (Just s, Just x)
1225 Skip s -> go SPEC prev s acc
1226 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1228 stepOuter (Fold fstep initial done) gst (Just st, Just prev) = do
1230 acc' <- fstep acc prev
1235 -- XXX code duplicated from the previous equation
1237 res <- step (adaptState gst) stt
1244 else done acc >>= \r -> return $ Yield r (Just s, Just x)
1245 Skip s -> go SPEC s acc
1246 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1248 stepOuter _ _ (Nothing,_) = return Stop
1250{-# INLINE_NORMAL groupsRollingBy #-}
1251groupsRollingBy :: Monad m
1256groupsRollingBy cmp f (Stream step state) =
1257 Stream (stepOuter f) (Just state, Nothing)
1260 {-# INLINE_LATE stepOuter #-}
1261 stepOuter (Fold fstep initial done) gst (Just st, Nothing) = do
1262 res <- step (adaptState gst) st
1269 Skip s -> return $ Skip $ (Just s, Nothing)
1273 go !_ prev stt !acc = do
1274 res <- step (adaptState gst) stt
1282 done acc >>= \r -> return $ Yield r (Just s, Just x)
1283 Skip s -> go SPEC prev s acc
1284 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1286 stepOuter (Fold fstep initial done) gst (Just st, Just prev') = do
1288 acc' <- fstep acc prev'
1289 go SPEC prev' st acc'
1292 go !_ prevv stt !acc = do
1293 res <- step (adaptState gst) stt
1300 else done acc >>= \r -> return $ Yield r (Just s, Just x)
1301 Skip s -> go SPEC prevv s acc
1302 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1304 stepOuter _ _ (Nothing, _) = return Stop
1306{-# INLINE_NORMAL splitBy #-}
1307splitBy :: Monad m => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1308splitBy predicate f (Stream step state) = Stream (step' f) (Just state)
1312 {-# INLINE_LATE step' #-}
1313 step' (Fold fstep initial done) gst (Just st) = initial >>= go SPEC st
1318 res <- step (adaptState gst) stt
1322 then done acc >>= \r -> return $ Yield r (Just s)
1326 Skip s -> go SPEC s acc
1327 Stop -> done acc >>= \r -> return $ Yield r Nothing
1329 step' _ _ Nothing = return Stop
1331-- XXX requires -funfolding-use-threshold=150 in lines-unlines benchmark
1332{-# INLINE_NORMAL splitSuffixBy #-}
1333splitSuffixBy :: Monad m
1334 => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1335splitSuffixBy predicate f (Stream step state) = Stream (step' f) (Just state)
1339 {-# INLINE_LATE step' #-}
1340 step' (Fold fstep initial done) gst (Just st) = do
1341 res <- step (adaptState gst) st
1346 then done acc >>= \val -> return $ Yield val (Just s)
1351 Skip s -> return $ Skip $ Just s
1357 res <- step (adaptState gst) stt
1361 then done acc >>= \r -> return $ Yield r (Just s)
1365 Skip s -> go SPEC s acc
1366 Stop -> done acc >>= \r -> return $ Yield r Nothing
1368 step' _ _ Nothing = return Stop
1370{-# INLINE_NORMAL wordsBy #-}
1371wordsBy :: Monad m => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1372wordsBy predicate f (Stream step state) = Stream (stepOuter f) (Just state)
1376 {-# INLINE_LATE stepOuter #-}
1377 stepOuter (Fold fstep initial done) gst (Just st) = do
1378 res <- step (adaptState gst) st
1382 then return $ Skip (Just s)
1388 Skip s -> return $ Skip $ Just s
1394 res <- step (adaptState gst) stt
1398 then done acc >>= \r -> return $ Yield r (Just s)
1402 Skip s -> go SPEC s acc
1403 Stop -> done acc >>= \r -> return $ Yield r Nothing
1405 stepOuter _ _ Nothing = return Stop
1407-- String search algorithms:
1408-- http://www-igm.univ-mlv.fr/~lecroq/string/index.html
1411-- TODO can we unify the splitting operations using a splitting configuration
1412-- like in the split package.
1414data SplitStyle = Infix | Suffix | Prefix deriving (Eq, Show)
1416data SplitOptions = SplitOptions
1417 { style :: SplitStyle
1418 , withSep :: Bool -- ^ keep the separators in output
1419 -- , compact :: Bool -- ^ treat multiple consecutive separators as one
1420 -- , trimHead :: Bool -- ^ drop blank at head
1421 -- , trimTail :: Bool -- ^ drop blank at tail
1425data SplitOnState s a =
1430 | GO_KARP_RABIN s !(RB.Ring a) !(Ptr a)
1433{-# INLINE_NORMAL splitOn #-}
1435 :: forall m a b. (MonadIO m, Storable a, Enum a, Eq a)
1440splitOn patArr@Array{..} (Fold fstep initial done) (Stream step state) =
1441 Stream stepOuter GO_START
1445 patLen = A.length patArr
1446 maxIndex = patLen - 1
1447 elemBits = sizeOf (undefined :: a) * 8
1449 {-# INLINE_LATE stepOuter #-}
1450 stepOuter _ GO_START =
1452 then return $ Skip $ GO_EMPTY_PAT state
1455 r <- liftIO $ (A.unsafeIndexIO patArr 0)
1456 return $ Skip $ GO_SINGLE_PAT state r
1457 else if sizeOf (undefined :: a) * patLen
1458 <= sizeOf (undefined :: Word)
1459 then return $ Skip $ GO_SHORT_PAT state
1461 (rb, rhead) <- liftIO $ RB.new patLen
1462 return $ Skip $ GO_KARP_RABIN state rb rhead
1464 stepOuter gst (GO_SINGLE_PAT stt pat) = initial >>= go SPEC stt
1469 res <- step (adaptState gst) st
1475 return $ Yield r (GO_SINGLE_PAT s pat)
1476 else fstep acc x >>= go SPEC s
1477 Skip s -> go SPEC s acc
1478 Stop -> done acc >>= \r -> return $ Yield r GO_DONE
1480 stepOuter gst (GO_SHORT_PAT stt) = initial >>= go0 SPEC 0 (0 :: Word) stt
1485 mask = (1 `shiftL` (elemBits * patLen)) - 1
1487 addToWord wrd a = (wrd `shiftL` elemBits) .|. fromIntegral (fromEnum a)
1490 patWord = mask .&. A.foldl' addToWord 0 patArr
1492 go0 !_ !idx wrd st !acc = do
1493 res <- step (adaptState gst) st
1496 let wrd' = addToWord wrd x
1499 if wrd' .&. mask == patWord
1502 return $ Yield r (GO_SHORT_PAT s)
1503 else go1 SPEC wrd' s acc
1504 else go0 SPEC (idx + 1) wrd' s acc
1505 Skip s -> go0 SPEC idx wrd s acc
1508 then go2 wrd idx acc
1510 done acc' >>= \r -> return $ Yield r GO_DONE
1513 go1 !_ wrd st !acc = do
1514 res <- step (adaptState gst) st
1517 let wrd' = addToWord wrd x
1518 old = (mask .&. wrd) `shiftR` (elemBits * (patLen - 1))
1519 acc' <- fstep acc (toEnum $ fromIntegral old)
1520 if wrd' .&. mask == patWord
1521 then done acc' >>= \r -> return $ Yield r (GO_SHORT_PAT s)
1522 else go1 SPEC wrd' s acc'
1523 Skip s -> go1 SPEC wrd s acc
1525 acc' <- go2 wrd patLen acc
1526 done acc' >>= \r -> return $ Yield r GO_DONE
1528 go2 !wrd !n !acc | n > 0 = do
1529 let old = (mask .&. wrd) `shiftR` (elemBits * (n - 1))
1530 fstep acc (toEnum $ fromIntegral old) >>= go2 wrd (n - 1)
1531 go2 _ _ acc = return acc
1533 stepOuter gst (GO_KARP_RABIN stt rb rhead) = do
1534 initial >>= go0 SPEC 0 rhead stt
1538 k = 2891336453 :: Word32
1540 addCksum cksum a = cksum * k + fromIntegral (fromEnum a)
1541 deltaCksum cksum old new =
1542 addCksum cksum new - coeff * fromIntegral (fromEnum old)
1544 -- XXX shall we use a random starting hash or 1 instead of 0?
1545 patHash = A.foldl' addCksum 0 patArr
1548 go0 !_ !idx !rh st !acc = do
1549 res <- step (adaptState gst) st
1552 rh' <- liftIO $ RB.unsafeInsert rb rh x
1555 let fold = RB.unsafeFoldRing (RB.ringBound rb)
1556 let !ringHash = fold addCksum 0 rb
1557 if ringHash == patHash
1558 then go2 SPEC ringHash rh' s acc
1559 else go1 SPEC ringHash rh' s acc
1560 else go0 SPEC (idx + 1) rh' s acc
1561 Skip s -> go0 SPEC idx rh s acc
1563 !acc' <- if idx /= 0
1564 then RB.unsafeFoldRingM rh fstep acc rb
1566 done acc' >>= \r -> return $ Yield r GO_DONE
1568 -- XXX Theoretically this code can do 4 times faster if GHC generates
1569 -- optimal code. If we use just "(cksum' == patHash)" condition it goes
1570 -- 4x faster, as soon as we add the "RB.unsafeEqArray rb v" condition
1571 -- the generated code changes drastically and becomes 4x slower. Need
1572 -- to investigate what is going on with GHC.
1574 go1 !_ !cksum !rh st !acc = do
1575 res <- step (adaptState gst) st
1578 old <- liftIO $ peek rh
1579 let cksum' = deltaCksum cksum old x
1580 acc' <- fstep acc old
1582 if (cksum' == patHash)
1584 rh' <- liftIO (RB.unsafeInsert rb rh x)
1585 go2 SPEC cksum' rh' s acc'
1587 rh' <- liftIO (RB.unsafeInsert rb rh x)
1588 go1 SPEC cksum' rh' s acc'
1589 Skip s -> go1 SPEC cksum rh s acc
1591 acc' <- RB.unsafeFoldRingFullM rh fstep acc rb
1592 done acc' >>= \r -> return $ Yield r GO_DONE
1594 go2 !_ !cksum' !rh' s !acc' = do
1595 if RB.unsafeEqArray rb rh' patArr
1598 return $ Yield r (GO_KARP_RABIN s rb rhead)
1599 else go1 SPEC cksum' rh' s acc'
1601 stepOuter gst (GO_EMPTY_PAT st) = do
1602 res <- step (adaptState gst) st
1607 done acc' >>= \r -> return $ Yield r (GO_EMPTY_PAT s)
1608 Skip s -> return $ Skip (GO_EMPTY_PAT s)
1611 stepOuter _ GO_DONE = return Stop
1613{-# INLINE_NORMAL splitSuffixOn #-}
1615 :: forall m a b. (MonadIO m, Storable a, Enum a, Eq a)
1621splitSuffixOn withSep patArr@Array{..} (Fold fstep initial done)
1622 (Stream step state) =
1623 Stream stepOuter GO_START
1627 patLen = A.length patArr
1628 maxIndex = patLen - 1
1629 elemBits = sizeOf (undefined :: a) * 8
1631 {-# INLINE_LATE stepOuter #-}
1632 stepOuter _ GO_START =
1634 then return $ Skip $ GO_EMPTY_PAT state
1637 r <- liftIO $ (A.unsafeIndexIO patArr 0)
1638 return $ Skip $ GO_SINGLE_PAT state r
1639 else if sizeOf (undefined :: a) * patLen
1640 <= sizeOf (undefined :: Word)
1641 then return $ Skip $ GO_SHORT_PAT state
1643 (rb, rhead) <- liftIO $ RB.new patLen
1644 return $ Skip $ GO_KARP_RABIN state rb rhead
1646 stepOuter gst (GO_SINGLE_PAT stt pat) = do
1647 -- This first part is the only difference between splitOn and
1649 -- If the last element is a separator do not issue a blank segment.
1650 res <- step (adaptState gst) stt
1656 acc' <- if withSep then fstep acc x else return acc
1657 done acc' >>= \r -> return $ Yield r (GO_SINGLE_PAT s pat)
1658 else fstep acc x >>= go SPEC s
1659 Skip s -> return $ Skip $ (GO_SINGLE_PAT s pat)
1664 -- This is identical for splitOn and splitSuffixOn
1666 res <- step (adaptState gst) st
1671 acc' <- if withSep then fstep acc x else return acc
1673 return $ Yield r (GO_SINGLE_PAT s pat)
1674 else fstep acc x >>= go SPEC s
1675 Skip s -> go SPEC s acc
1676 Stop -> done acc >>= \r -> return $ Yield r GO_DONE
1678 stepOuter gst (GO_SHORT_PAT stt) = do
1680 -- Call "initial" only if the stream yields an element, otherwise we
1681 -- may call "initial" but never yield anything. initial may produce a
1682 -- side effect, therefore we will end up doing and discard a side
1687 res <- step (adaptState gst) stt
1691 let wrd' = addToWord wrd x
1692 acc' <- if withSep then fstep acc x else return acc
1695 if wrd' .&. mask == patWord
1696 then done acc' >>= \r -> return $ Yield r (GO_SHORT_PAT s)
1697 else go0 SPEC (idx + 1) wrd' s acc'
1698 else go0 SPEC (idx + 1) wrd' s acc'
1699 Skip s -> return $ Skip (GO_SHORT_PAT s)
1705 mask = (1 `shiftL` (elemBits * patLen)) - 1
1707 addToWord wrd a = (wrd `shiftL` elemBits) .|. fromIntegral (fromEnum a)
1710 patWord = mask .&. A.foldl' addToWord 0 patArr
1712 go0 !_ !idx wrd st !acc = do
1713 res <- step (adaptState gst) st
1716 let wrd' = addToWord wrd x
1717 acc' <- if withSep then fstep acc x else return acc
1720 if wrd' .&. mask == patWord
1723 return $ Yield r (GO_SHORT_PAT s)
1724 else go1 SPEC wrd' s acc'
1725 else go0 SPEC (idx + 1) wrd' s acc'
1726 Skip s -> go0 SPEC idx wrd s acc
1728 if (idx == maxIndex) && (wrd .&. mask == patWord)
1731 acc' <- if idx /= 0 && not withSep
1732 then go2 wrd idx acc
1734 done acc' >>= \r -> return $ Yield r GO_DONE
1737 go1 !_ wrd st !acc = do
1738 res <- step (adaptState gst) st
1741 let wrd' = addToWord wrd x
1742 old = (mask .&. wrd) `shiftR` (elemBits * (patLen - 1))
1745 else fstep acc (toEnum $ fromIntegral old)
1746 if wrd' .&. mask == patWord
1747 then done acc' >>= \r -> return $ Yield r (GO_SHORT_PAT s)
1748 else go1 SPEC wrd' s acc'
1749 Skip s -> go1 SPEC wrd s acc
1751 -- If the last sequence is a separator do not issue a blank
1753 if wrd .&. mask == patWord
1758 else go2 wrd patLen acc
1759 done acc' >>= \r -> return $ Yield r GO_DONE
1761 go2 !wrd !n !acc | n > 0 = do
1762 let old = (mask .&. wrd) `shiftR` (elemBits * (n - 1))
1763 fstep acc (toEnum $ fromIntegral old) >>= go2 wrd (n - 1)
1764 go2 _ _ acc = return acc
1766 stepOuter gst (GO_KARP_RABIN stt rb rhead) = do
1768 res <- step (adaptState gst) stt
1772 acc' <- if withSep then fstep acc x else return acc
1773 rh' <- liftIO (RB.unsafeInsert rb rhead x)
1776 let fold = RB.unsafeFoldRing (RB.ringBound rb)
1777 let !ringHash = fold addCksum 0 rb
1778 if ringHash == patHash
1779 then go2 SPEC ringHash rh' s acc'
1780 else go0 SPEC (idx + 1) rh' s acc'
1781 else go0 SPEC (idx + 1) rh' s acc'
1782 Skip s -> return $ Skip (GO_KARP_RABIN s rb rhead)
1787 k = 2891336453 :: Word32
1789 addCksum cksum a = cksum * k + fromIntegral (fromEnum a)
1790 deltaCksum cksum old new =
1791 addCksum cksum new - coeff * fromIntegral (fromEnum old)
1793 -- XXX shall we use a random starting hash or 1 instead of 0?
1794 patHash = A.foldl' addCksum 0 patArr
1797 go0 !_ !idx !rh st !acc = do
1798 res <- step (adaptState gst) st
1801 acc' <- if withSep then fstep acc x else return acc
1802 rh' <- liftIO (RB.unsafeInsert rb rh x)
1805 let fold = RB.unsafeFoldRing (RB.ringBound rb)
1806 let !ringHash = fold addCksum 0 rb
1807 if ringHash == patHash
1808 then go2 SPEC ringHash rh' s acc'
1809 else go1 SPEC ringHash rh' s acc'
1810 else go0 SPEC (idx + 1) rh' s acc'
1811 Skip s -> go0 SPEC idx rh s acc
1813 -- do not issue a blank segment when we end at pattern
1814 if (idx == maxIndex) && RB.unsafeEqArray rb rh patArr
1817 !acc' <- if idx /= 0 && not withSep
1818 then RB.unsafeFoldRingM rh fstep acc rb
1820 done acc' >>= \r -> return $ Yield r GO_DONE
1822 -- XXX Theoretically this code can do 4 times faster if GHC generates
1823 -- optimal code. If we use just "(cksum' == patHash)" condition it goes
1824 -- 4x faster, as soon as we add the "RB.unsafeEqArray rb v" condition
1825 -- the generated code changes drastically and becomes 4x slower. Need
1826 -- to investigate what is going on with GHC.
1828 go1 !_ !cksum !rh st !acc = do
1829 res <- step (adaptState gst) st
1832 old <- liftIO $ peek rh
1833 let cksum' = deltaCksum cksum old x
1838 if (cksum' == patHash)
1840 rh' <- liftIO (RB.unsafeInsert rb rh x)
1841 go2 SPEC cksum' rh' s acc'
1843 rh' <- liftIO (RB.unsafeInsert rb rh x)
1844 go1 SPEC cksum' rh' s acc'
1845 Skip s -> go1 SPEC cksum rh s acc
1847 if RB.unsafeEqArray rb rh patArr
1852 else RB.unsafeFoldRingFullM rh fstep acc rb
1853 done acc' >>= \r -> return $ Yield r GO_DONE
1855 go2 !_ !cksum' !rh' s !acc' = do
1856 if RB.unsafeEqArray rb rh' patArr
1859 return $ Yield r (GO_KARP_RABIN s rb rhead)
1860 else go1 SPEC cksum' rh' s acc'
1862 stepOuter gst (GO_EMPTY_PAT st) = do
1863 res <- step (adaptState gst) st
1868 done acc' >>= \r -> return $ Yield r (GO_EMPTY_PAT s)
1869 Skip s -> return $ Skip (GO_EMPTY_PAT s)
1872 stepOuter _ GO_DONE = return Stop
1874data SplitState s arr
1876 | SplitBuffering s arr
1877 | SplitSplitting s arr
1878 | SplitYielding arr (SplitState s arr)
1881-- XXX An alternative approach would be to use a partial fold (Fold m a b) to
1882-- split using a splitBy like combinator. The Fold would consume upto the
1883-- separator and return any leftover which can then be fed to the next fold.
1885-- We can revisit this once we have partial folds/parsers.
1887-- | Performs infix separator style splitting.
1888{-# INLINE_NORMAL splitInnerBy #-}
1891 => (f a -> m (f a, Maybe (f a))) -- splitter
1892 -> (f a -> f a -> m (f a)) -- joiner
1895splitInnerBy splitter joiner (Stream step1 state1) =
1896 (Stream step (SplitInitial state1))
1900 {-# INLINE_LATE step #-}
1901 step gst (SplitInitial st) = do
1905 (x1, mx2) <- splitter x
1906 return $ case mx2 of
1907 Nothing -> Skip (SplitBuffering s x1)
1908 Just x2 -> Skip (SplitYielding x1 (SplitSplitting s x2))
1909 Skip s -> return $ Skip (SplitInitial s)
1910 Stop -> return $ Stop
1912 step gst (SplitBuffering st buf) = do
1916 (x1, mx2) <- splitter x
1917 buf' <- joiner buf x1
1918 return $ case mx2 of
1919 Nothing -> Skip (SplitBuffering s buf')
1920 Just x2 -> Skip (SplitYielding buf' (SplitSplitting s x2))
1921 Skip s -> return $ Skip (SplitBuffering s buf)
1922 Stop -> return $ Skip (SplitYielding buf SplitFinishing)
1924 step _ (SplitSplitting st buf) = do
1925 (x1, mx2) <- splitter buf
1926 return $ case mx2 of
1927 Nothing -> Skip $ SplitBuffering st x1
1928 Just x2 -> Skip $ SplitYielding x1 (SplitSplitting st x2)
1930 step _ (SplitYielding x next) = return $ Yield x next
1931 step _ SplitFinishing = return $ Stop
1933-- | Performs infix separator style splitting.
1934{-# INLINE_NORMAL splitInnerBySuffix #-}
1936 :: (Monad m, Eq (f a), Monoid (f a))
1937 => (f a -> m (f a, Maybe (f a))) -- splitter
1938 -> (f a -> f a -> m (f a)) -- joiner
1941splitInnerBySuffix splitter joiner (Stream step1 state1) =
1942 (Stream step (SplitInitial state1))
1946 {-# INLINE_LATE step #-}
1947 step gst (SplitInitial st) = do
1951 (x1, mx2) <- splitter x
1952 return $ case mx2 of
1953 Nothing -> Skip (SplitBuffering s x1)
1954 Just x2 -> Skip (SplitYielding x1 (SplitSplitting s x2))
1955 Skip s -> return $ Skip (SplitInitial s)
1956 Stop -> return $ Stop
1958 step gst (SplitBuffering st buf) = do
1962 (x1, mx2) <- splitter x
1963 buf' <- joiner buf x1
1964 return $ case mx2 of
1965 Nothing -> Skip (SplitBuffering s buf')
1966 Just x2 -> Skip (SplitYielding buf' (SplitSplitting s x2))
1967 Skip s -> return $ Skip (SplitBuffering s buf)
1971 else Skip (SplitYielding buf SplitFinishing)
1973 step _ (SplitSplitting st buf) = do
1974 (x1, mx2) <- splitter buf
1975 return $ case mx2 of
1976 Nothing -> Skip $ SplitBuffering st x1
1977 Just x2 -> Skip $ SplitYielding x1 (SplitSplitting st x2)
1979 step _ (SplitYielding x next) = return $ Yield x next
1980 step _ SplitFinishing = return $ Stop
1982------------------------------------------------------------------------------
1984------------------------------------------------------------------------------
1986{-# INLINE_NORMAL isPrefixOf #-}
1987isPrefixOf :: (Eq a, Monad m) => Stream m a -> Stream m a -> m Bool
1988isPrefixOf (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
1990 go (sa, sb, Nothing) = do
1991 r <- stepa defState sa
1993 Yield x sa' -> go (sa', sb, Just x)
1994 Skip sa' -> go (sa', sb, Nothing)
1997 go (sa, sb, Just x) = do
1998 r <- stepb defState sb
2002 then go (sa, sb', Nothing)
2004 Skip sb' -> go (sa, sb', Just x)
2005 Stop -> return False
2007{-# INLINE_NORMAL isSubsequenceOf #-}
2008isSubsequenceOf :: (Eq a, Monad m) => Stream m a -> Stream m a -> m Bool
2009isSubsequenceOf (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
2011 go (sa, sb, Nothing) = do
2012 r <- stepa defState sa
2014 Yield x sa' -> go (sa', sb, Just x)
2015 Skip sa' -> go (sa', sb, Nothing)
2018 go (sa, sb, Just x) = do
2019 r <- stepb defState sb
2023 then go (sa, sb', Nothing)
2024 else go (sa, sb', Just x)
2025 Skip sb' -> go (sa, sb', Just x)
2026 Stop -> return False
2028{-# INLINE_NORMAL stripPrefix #-}
2031 => Stream m a -> Stream m a -> m (Maybe (Stream m a))
2032stripPrefix (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
2034 go (sa, sb, Nothing) = do
2035 r <- stepa defState sa
2037 Yield x sa' -> go (sa', sb, Just x)
2038 Skip sa' -> go (sa', sb, Nothing)
2039 Stop -> return $ Just (Stream stepb sb)
2041 go (sa, sb, Just x) = do
2042 r <- stepb defState sb
2046 then go (sa, sb', Nothing)
2048 Skip sb' -> go (sa, sb', Just x)
2049 Stop -> return Nothing
2051------------------------------------------------------------------------------
2053------------------------------------------------------------------------------
2055-- | Execute a monadic action for each element of the 'Stream'
2056{-# INLINE_NORMAL mapM_ #-}
2057mapM_ :: Monad m => (a -> m b) -> Stream m a -> m ()
2058mapM_ m = drain . mapM m
2060-------------------------------------------------------------------------------
2061-- Stream transformations using Unfolds
2062-------------------------------------------------------------------------------
2064-- Define a unique structure to use in inspection testing
2065data ConcatMapUState o i =
2067 | ConcatMapUInner o i
2069-- | @concatMapU unfold stream@ uses @unfold@ to map the input stream elements
2070-- to streams and then flattens the generated streams into a single output
2073-- This is like 'concatMap' but uses an unfold with an explicit state to
2074-- generate the stream instead of a 'Stream' type generator. This allows better
2075-- optimization via fusion. This can be many times more efficient than
2078{-# INLINE_NORMAL concatMapU #-}
2079concatMapU :: Monad m => Unfold m a b -> Stream m a -> Stream m b
2080concatMapU (Unfold istep inject) (Stream ostep ost) =
2081 Stream step (ConcatMapUOuter ost)
2083 {-# INLINE_LATE step #-}
2084 step gst (ConcatMapUOuter o) = do
2085 r <- ostep (adaptState gst) o
2089 i `seq` return (Skip (ConcatMapUInner o' i))
2090 Skip o' -> return $ Skip (ConcatMapUOuter o')
2091 Stop -> return $ Stop
2093 step _ (ConcatMapUInner o i) = do
2096 Yield x i' -> Yield x (ConcatMapUInner o i')
2097 Skip i' -> Skip (ConcatMapUInner o i')
2098 Stop -> Skip (ConcatMapUOuter o)
2100data ConcatUnfoldInterleaveState o i =
2101 ConcatUnfoldInterleaveOuter o [i]
2102 | ConcatUnfoldInterleaveInner o [i]
2103 | ConcatUnfoldInterleaveInnerL [i] [i]
2104 | ConcatUnfoldInterleaveInnerR [i] [i]
2106-- XXX use arrays to store state instead of lists.
2107-- XXX In general we can use different scheduling strategies e.g. how to
2108-- schedule the outer vs inner loop or assigning weights to different streams
2109-- or outer and inner loops.
2111-- After a yield, switch to the next stream. Do not switch streams on Skip.
2112-- Yield from outer stream switches to the inner stream.
2114-- There are two choices here, (1) exhaust the outer stream first and then
2115-- start yielding from the inner streams, this is much simpler to implement,
2116-- (2) yield at least one element from an inner stream before going back to
2117-- outer stream and opening the next stream from it.
2119-- Ideally, we need some scheduling bias to inner streams vs outer stream.
2120-- Maybe we can configure the behavior.
2122{-# INLINE_NORMAL concatUnfoldInterleave #-}
2123concatUnfoldInterleave :: Monad m => Unfold m a b -> Stream m a -> Stream m b
2124concatUnfoldInterleave (Unfold istep inject) (Stream ostep ost) =
2125 Stream step (ConcatUnfoldInterleaveOuter ost [])
2127 {-# INLINE_LATE step #-}
2128 step gst (ConcatUnfoldInterleaveOuter o ls) = do
2129 r <- ostep (adaptState gst) o
2133 i `seq` return (Skip (ConcatUnfoldInterleaveInner o' (i : ls)))
2134 Skip o' -> return $ Skip (ConcatUnfoldInterleaveOuter o' ls)
2135 Stop -> return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2137 step _ (ConcatUnfoldInterleaveInner _ []) = undefined
2138 step _ (ConcatUnfoldInterleaveInner o (st:ls)) = do
2141 Yield x s -> Yield x (ConcatUnfoldInterleaveOuter o (s:ls))
2142 Skip s -> Skip (ConcatUnfoldInterleaveInner o (s:ls))
2143 Stop -> Skip (ConcatUnfoldInterleaveOuter o ls)
2145 step _ (ConcatUnfoldInterleaveInnerL [] []) = return Stop
2146 step _ (ConcatUnfoldInterleaveInnerL [] rs) =
2147 return $ Skip (ConcatUnfoldInterleaveInnerR [] rs)
2149 step _ (ConcatUnfoldInterleaveInnerL (st:ls) rs) = do
2152 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerL ls (s:rs))
2153 Skip s -> Skip (ConcatUnfoldInterleaveInnerL (s:ls) rs)
2154 Stop -> Skip (ConcatUnfoldInterleaveInnerL ls rs)
2156 step _ (ConcatUnfoldInterleaveInnerR [] []) = return Stop
2157 step _ (ConcatUnfoldInterleaveInnerR ls []) =
2158 return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2160 step _ (ConcatUnfoldInterleaveInnerR ls (st:rs)) = do
2163 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerR (s:ls) rs)
2164 Skip s -> Skip (ConcatUnfoldInterleaveInnerR ls (s:rs))
2165 Stop -> Skip (ConcatUnfoldInterleaveInnerR ls rs)
2167-- XXX In general we can use different scheduling strategies e.g. how to
2168-- schedule the outer vs inner loop or assigning weights to different streams
2169-- or outer and inner loops.
2171-- This could be inefficient if the tasks are too small.
2173-- Compared to concatUnfoldInterleave this one switches streams on Skips.
2175{-# INLINE_NORMAL concatUnfoldRoundrobin #-}
2176concatUnfoldRoundrobin :: Monad m => Unfold m a b -> Stream m a -> Stream m b
2177concatUnfoldRoundrobin (Unfold istep inject) (Stream ostep ost) =
2178 Stream step (ConcatUnfoldInterleaveOuter ost [])
2180 {-# INLINE_LATE step #-}
2181 step gst (ConcatUnfoldInterleaveOuter o ls) = do
2182 r <- ostep (adaptState gst) o
2186 i `seq` return (Skip (ConcatUnfoldInterleaveInner o' (i : ls)))
2187 Skip o' -> return $ Skip (ConcatUnfoldInterleaveInner o' ls)
2188 Stop -> return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2190 step _ (ConcatUnfoldInterleaveInner o []) =
2191 return $ Skip (ConcatUnfoldInterleaveOuter o [])
2193 step _ (ConcatUnfoldInterleaveInner o (st:ls)) = do
2196 Yield x s -> Yield x (ConcatUnfoldInterleaveOuter o (s:ls))
2197 Skip s -> Skip (ConcatUnfoldInterleaveOuter o (s:ls))
2198 Stop -> Skip (ConcatUnfoldInterleaveOuter o ls)
2200 step _ (ConcatUnfoldInterleaveInnerL [] []) = return Stop
2201 step _ (ConcatUnfoldInterleaveInnerL [] rs) =
2202 return $ Skip (ConcatUnfoldInterleaveInnerR [] rs)
2204 step _ (ConcatUnfoldInterleaveInnerL (st:ls) rs) = do
2207 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerL ls (s:rs))
2208 Skip s -> Skip (ConcatUnfoldInterleaveInnerL ls (s:rs))
2209 Stop -> Skip (ConcatUnfoldInterleaveInnerL ls rs)
2211 step _ (ConcatUnfoldInterleaveInnerR [] []) = return Stop
2212 step _ (ConcatUnfoldInterleaveInnerR ls []) =
2213 return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
2215 step _ (ConcatUnfoldInterleaveInnerR ls (st:rs)) = do
2218 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerR (s:ls) rs)
2219 Skip s -> Skip (ConcatUnfoldInterleaveInnerR (s:ls) rs)
2220 Stop -> Skip (ConcatUnfoldInterleaveInnerR ls rs)
2222data AppendState s1 s2 = AppendFirst s1 | AppendSecond s2
2224-- Note that this could be much faster compared to the CPS stream. However, as
2225-- the number of streams being composed increases this may become expensive.
2226-- Need to see where the breaking point is between the two.
2228{-# INLINE_NORMAL append #-}
2229append :: Monad m => Stream m a -> Stream m a -> Stream m a
2230append (Stream step1 state1) (Stream step2 state2) =
2231 Stream step (AppendFirst state1)
2235 {-# INLINE_LATE step #-}
2236 step gst (AppendFirst st) = do
2239 Yield a s -> Yield a (AppendFirst s)
2240 Skip s -> Skip (AppendFirst s)
2241 Stop -> Skip (AppendSecond state2)
2243 step gst (AppendSecond st) = do
2246 Yield a s -> Yield a (AppendSecond s)
2247 Skip s -> Skip (AppendSecond s)
2250data InterleaveState s1 s2 = InterleaveFirst s1 s2 | InterleaveSecond s1 s2
2251 | InterleaveSecondOnly s2 | InterleaveFirstOnly s1
2253{-# INLINE_NORMAL interleave #-}
2254interleave :: Monad m => Stream m a -> Stream m a -> Stream m a
2255interleave (Stream step1 state1) (Stream step2 state2) =
2256 Stream step (InterleaveFirst state1 state2)
2260 {-# INLINE_LATE step #-}
2261 step gst (InterleaveFirst st1 st2) = do
2264 Yield a s -> Yield a (InterleaveSecond s st2)
2265 Skip s -> Skip (InterleaveFirst s st2)
2266 Stop -> Skip (InterleaveSecondOnly st2)
2268 step gst (InterleaveSecond st1 st2) = do
2271 Yield a s -> Yield a (InterleaveFirst st1 s)
2272 Skip s -> Skip (InterleaveSecond st1 s)
2273 Stop -> Skip (InterleaveFirstOnly st1)
2275 step gst (InterleaveFirstOnly st1) = do
2278 Yield a s -> Yield a (InterleaveFirstOnly s)
2279 Skip s -> Skip (InterleaveFirstOnly s)
2282 step gst (InterleaveSecondOnly st2) = do
2285 Yield a s -> Yield a (InterleaveSecondOnly s)
2286 Skip s -> Skip (InterleaveSecondOnly s)
2289{-# INLINE_NORMAL interleaveMin #-}
2290interleaveMin :: Monad m => Stream m a -> Stream m a -> Stream m a
2291interleaveMin (Stream step1 state1) (Stream step2 state2) =
2292 Stream step (InterleaveFirst state1 state2)
2296 {-# INLINE_LATE step #-}
2297 step gst (InterleaveFirst st1 st2) = do
2300 Yield a s -> Yield a (InterleaveSecond s st2)
2301 Skip s -> Skip (InterleaveFirst s st2)
2304 step gst (InterleaveSecond st1 st2) = do
2307 Yield a s -> Yield a (InterleaveFirst st1 s)
2308 Skip s -> Skip (InterleaveSecond st1 s)
2311 step _ (InterleaveFirstOnly _) = undefined
2312 step _ (InterleaveSecondOnly _) = undefined
2314{-# INLINE_NORMAL interleaveSuffix #-}
2315interleaveSuffix :: Monad m => Stream m a -> Stream m a -> Stream m a
2316interleaveSuffix (Stream step1 state1) (Stream step2 state2) =
2317 Stream step (InterleaveFirst state1 state2)
2321 {-# INLINE_LATE step #-}
2322 step gst (InterleaveFirst st1 st2) = do
2325 Yield a s -> Yield a (InterleaveSecond s st2)
2326 Skip s -> Skip (InterleaveFirst s st2)
2329 step gst (InterleaveSecond st1 st2) = do
2332 Yield a s -> Yield a (InterleaveFirst st1 s)
2333 Skip s -> Skip (InterleaveSecond st1 s)
2334 Stop -> Skip (InterleaveFirstOnly st1)
2336 step gst (InterleaveFirstOnly st1) = do
2339 Yield a s -> Yield a (InterleaveFirstOnly s)
2340 Skip s -> Skip (InterleaveFirstOnly s)
2343 step _ (InterleaveSecondOnly _) = undefined
2345data InterleaveInfixState s1 s2 a
2346 = InterleaveInfixFirst s1 s2
2347 | InterleaveInfixSecondBuf s1 s2
2348 | InterleaveInfixSecondYield s1 s2 a
2349 | InterleaveInfixFirstYield s1 s2 a
2350 | InterleaveInfixFirstOnly s1
2352{-# INLINE_NORMAL interleaveInfix #-}
2353interleaveInfix :: Monad m => Stream m a -> Stream m a -> Stream m a
2354interleaveInfix (Stream step1 state1) (Stream step2 state2) =
2355 Stream step (InterleaveInfixFirst state1 state2)
2359 {-# INLINE_LATE step #-}
2360 step gst (InterleaveInfixFirst st1 st2) = do
2363 Yield a s -> Yield a (InterleaveInfixSecondBuf s st2)
2364 Skip s -> Skip (InterleaveInfixFirst s st2)
2367 step gst (InterleaveInfixSecondBuf st1 st2) = do
2370 Yield a s -> Skip (InterleaveInfixSecondYield st1 s a)
2371 Skip s -> Skip (InterleaveInfixSecondBuf st1 s)
2372 Stop -> Skip (InterleaveInfixFirstOnly st1)
2374 step gst (InterleaveInfixSecondYield st1 st2 x) = do
2377 Yield a s -> Yield x (InterleaveInfixFirstYield s st2 a)
2378 Skip s -> Skip (InterleaveInfixSecondYield s st2 x)
2381 step _ (InterleaveInfixFirstYield st1 st2 x) = do
2382 return $ Yield x (InterleaveInfixSecondBuf st1 st2)
2384 step gst (InterleaveInfixFirstOnly st1) = do
2387 Yield a s -> Yield a (InterleaveInfixFirstOnly s)
2388 Skip s -> Skip (InterleaveInfixFirstOnly s)
2391{-# INLINE_NORMAL roundRobin #-}
2392roundRobin :: Monad m => Stream m a -> Stream m a -> Stream m a
2393roundRobin (Stream step1 state1) (Stream step2 state2) =
2394 Stream step (InterleaveFirst state1 state2)
2398 {-# INLINE_LATE step #-}
2399 step gst (InterleaveFirst st1 st2) = do
2402 Yield a s -> Yield a (InterleaveSecond s st2)
2403 Skip s -> Skip (InterleaveSecond s st2)
2404 Stop -> Skip (InterleaveSecondOnly st2)
2406 step gst (InterleaveSecond st1 st2) = do
2409 Yield a s -> Yield a (InterleaveFirst st1 s)
2410 Skip s -> Skip (InterleaveFirst st1 s)
2411 Stop -> Skip (InterleaveFirstOnly st1)
2413 step gst (InterleaveSecondOnly st2) = do
2416 Yield a s -> Yield a (InterleaveSecondOnly s)
2417 Skip s -> Skip (InterleaveSecondOnly s)
2420 step gst (InterleaveFirstOnly st1) = do
2423 Yield a s -> Yield a (InterleaveFirstOnly s)
2424 Skip s -> Skip (InterleaveFirstOnly s)
2427data ICUState s1 s2 i1 i2 =
2432 | ICUFirstInner s1 s2 i1
2433 | ICUSecondInner s1 s2 i2
2434 | ICUFirstOnlyInner s1 i1
2435 | ICUSecondOnlyInner s2 i2
2437-- | Interleave streams (full streams, not the elements) unfolded from two
2438-- input streams and concat. Stop when the first stream stops. If the second
2439-- stream ends before the first one then first stream still keeps running alone
2440-- without any interleaving with the second stream.
2442-- [a1, a2, ... an] [b1, b2 ...]
2443-- => [streamA1, streamA2, ... streamAn] [streamB1, streamB2, ...]
2444-- => [streamA1, streamB1, streamA2...StreamAn, streamBn]
2445-- => [a11, a12, ...a1j, b11, b12, ...b1k, a21, a22, ...]
2447{-# INLINE_NORMAL gintercalateSuffix #-}
2450 => Unfold m a c -> Stream m a -> Unfold m b c -> Stream m b -> Stream m c
2452 (Unfold istep1 inject1) (Stream step1 state1)
2453 (Unfold istep2 inject2) (Stream step2 state2) =
2454 Stream step (ICUFirst state1 state2)
2458 {-# INLINE_LATE step #-}
2459 step gst (ICUFirst s1 s2) = do
2460 r <- step1 (adaptState gst) s1
2464 i `seq` return (Skip (ICUFirstInner s s2 i))
2465 Skip s -> return $ Skip (ICUFirst s s2)
2468 step gst (ICUFirstOnly s1) = do
2469 r <- step1 (adaptState gst) s1
2473 i `seq` return (Skip (ICUFirstOnlyInner s i))
2474 Skip s -> return $ Skip (ICUFirstOnly s)
2477 step _ (ICUFirstInner s1 s2 i1) = do
2480 Yield x i' -> Yield x (ICUFirstInner s1 s2 i')
2481 Skip i' -> Skip (ICUFirstInner s1 s2 i')
2482 Stop -> Skip (ICUSecond s1 s2)
2484 step _ (ICUFirstOnlyInner s1 i1) = do
2487 Yield x i' -> Yield x (ICUFirstOnlyInner s1 i')
2488 Skip i' -> Skip (ICUFirstOnlyInner s1 i')
2489 Stop -> Skip (ICUFirstOnly s1)
2491 step gst (ICUSecond s1 s2) = do
2492 r <- step2 (adaptState gst) s2
2496 i `seq` return (Skip (ICUSecondInner s1 s i))
2497 Skip s -> return $ Skip (ICUSecond s1 s)
2498 Stop -> return $ Skip (ICUFirstOnly s1)
2500 step _ (ICUSecondInner s1 s2 i2) = do
2503 Yield x i' -> Yield x (ICUSecondInner s1 s2 i')
2504 Skip i' -> Skip (ICUSecondInner s1 s2 i')
2505 Stop -> Skip (ICUFirst s1 s2)
2507 step _ (ICUSecondOnly _s2) = undefined
2508 step _ (ICUSecondOnlyInner _s2 _i2) = undefined
2510data InterposeSuffixState s1 i1 =
2511 InterposeSuffixFirst s1
2512 -- | InterposeSuffixFirstYield s1 i1
2513 | InterposeSuffixFirstInner s1 i1
2514 | InterposeSuffixSecond s1
2516-- Note that if an unfolded layer turns out to be nil we still emit the
2517-- separator effect. An alternate behavior could be to emit the separator
2518-- effect only if at least one element has been yielded by the unfolding.
2519-- However, that becomes a bit complicated, so we have chosen the former
2521{-# INLINE_NORMAL interposeSuffix #-}
2524 => m c -> Unfold m b c -> Stream m b -> Stream m c
2527 (Unfold istep1 inject1) (Stream step1 state1) =
2528 Stream step (InterposeSuffixFirst state1)
2532 {-# INLINE_LATE step #-}
2533 step gst (InterposeSuffixFirst s1) = do
2534 r <- step1 (adaptState gst) s1
2538 i `seq` return (Skip (InterposeSuffixFirstInner s i))
2539 -- i `seq` return (Skip (InterposeSuffixFirstYield s i))
2540 Skip s -> return $ Skip (InterposeSuffixFirst s)
2544 step _ (InterposeSuffixFirstYield s1 i1) = do
2547 Yield x i' -> Yield x (InterposeSuffixFirstInner s1 i')
2548 Skip i' -> Skip (InterposeSuffixFirstYield s1 i')
2549 Stop -> Skip (InterposeSuffixFirst s1)
2552 step _ (InterposeSuffixFirstInner s1 i1) = do
2555 Yield x i' -> Yield x (InterposeSuffixFirstInner s1 i')
2556 Skip i' -> Skip (InterposeSuffixFirstInner s1 i')
2557 Stop -> Skip (InterposeSuffixSecond s1)
2559 step _ (InterposeSuffixSecond s1) = do
2561 return $ Yield r (InterposeSuffixFirst s1)
2563data ICALState s1 s2 i1 i2 a =
2565 -- | ICALFirstYield s1 s2 i1
2566 | ICALFirstInner s1 s2 i1
2568 | ICALFirstOnlyInner s1 i1
2569 | ICALSecondInject s1 s2
2570 | ICALFirstInject s1 s2 i2
2571 -- | ICALFirstBuf s1 s2 i1 i2
2572 | ICALSecondInner s1 s2 i1 i2
2573 -- -- | ICALSecondInner s1 s2 i1 i2 a
2574 -- -- | ICALFirstResume s1 s2 i1 i2 a
2576-- | Interleave streams (full streams, not the elements) unfolded from two
2577-- input streams and concat. Stop when the first stream stops. If the second
2578-- stream ends before the first one then first stream still keeps running alone
2579-- without any interleaving with the second stream.
2581-- [a1, a2, ... an] [b1, b2 ...]
2582-- => [streamA1, streamA2, ... streamAn] [streamB1, streamB2, ...]
2583-- => [streamA1, streamB1, streamA2...StreamAn, streamBn]
2584-- => [a11, a12, ...a1j, b11, b12, ...b1k, a21, a22, ...]
2586{-# INLINE_NORMAL gintercalate #-}
2589 => Unfold m a c -> Stream m a -> Unfold m b c -> Stream m b -> Stream m c
2591 (Unfold istep1 inject1) (Stream step1 state1)
2592 (Unfold istep2 inject2) (Stream step2 state2) =
2593 Stream step (ICALFirst state1 state2)
2597 {-# INLINE_LATE step #-}
2598 step gst (ICALFirst s1 s2) = do
2599 r <- step1 (adaptState gst) s1
2603 i `seq` return (Skip (ICALFirstInner s s2 i))
2604 -- i `seq` return (Skip (ICALFirstYield s s2 i))
2605 Skip s -> return $ Skip (ICALFirst s s2)
2609 step _ (ICALFirstYield s1 s2 i1) = do
2612 Yield x i' -> Yield x (ICALFirstInner s1 s2 i')
2613 Skip i' -> Skip (ICALFirstYield s1 s2 i')
2614 Stop -> Skip (ICALFirst s1 s2)
2617 step _ (ICALFirstInner s1 s2 i1) = do
2620 Yield x i' -> Yield x (ICALFirstInner s1 s2 i')
2621 Skip i' -> Skip (ICALFirstInner s1 s2 i')
2622 Stop -> Skip (ICALSecondInject s1 s2)
2624 step gst (ICALFirstOnly s1) = do
2625 r <- step1 (adaptState gst) s1
2629 i `seq` return (Skip (ICALFirstOnlyInner s i))
2630 Skip s -> return $ Skip (ICALFirstOnly s)
2633 step _ (ICALFirstOnlyInner s1 i1) = do
2636 Yield x i' -> Yield x (ICALFirstOnlyInner s1 i')
2637 Skip i' -> Skip (ICALFirstOnlyInner s1 i')
2638 Stop -> Skip (ICALFirstOnly s1)
2640 -- We inject the second stream even before checking if the first stream
2641 -- would yield any more elements. There is no clear choice whether we
2642 -- should do this before or after that. Doing it after may make the state
2643 -- machine a bit simpler though.
2644 step gst (ICALSecondInject s1 s2) = do
2645 r <- step2 (adaptState gst) s2
2649 i `seq` return (Skip (ICALFirstInject s1 s i))
2650 Skip s -> return $ Skip (ICALSecondInject s1 s)
2651 Stop -> return $ Skip (ICALFirstOnly s1)
2653 step gst (ICALFirstInject s1 s2 i2) = do
2654 r <- step1 (adaptState gst) s1
2658 i `seq` return (Skip (ICALSecondInner s s2 i i2))
2659 -- i `seq` return (Skip (ICALFirstBuf s s2 i i2))
2660 Skip s -> return $ Skip (ICALFirstInject s s2 i2)
2664 step _ (ICALFirstBuf s1 s2 i1 i2) = do
2667 Yield x i' -> Skip (ICALSecondInner s1 s2 i' i2 x)
2668 Skip i' -> Skip (ICALFirstBuf s1 s2 i' i2)
2671 step _ (ICALSecondInner s1 s2 i1 i2 v) = do
2674 Yield x i' -> Yield x (ICALSecondInner s1 s2 i1 i' v)
2675 Skip i' -> Skip (ICALSecondInner s1 s2 i1 i' v)
2676 Stop -> Skip (ICALFirstResume s1 s2 i1 i2 v)
2679 step _ (ICALSecondInner s1 s2 i1 i2) = do
2682 Yield x i' -> Yield x (ICALSecondInner s1 s2 i1 i')
2683 Skip i' -> Skip (ICALSecondInner s1 s2 i1 i')
2684 Stop -> Skip (ICALFirstInner s1 s2 i1)
2685 -- Stop -> Skip (ICALFirstResume s1 s2 i1 i2)
2688 step _ (ICALFirstResume s1 s2 i1 i2 x) = do
2689 return $ Yield x (ICALFirstInner s1 s2 i1 i2)
2692data InterposeState s1 i1 a =
2694 -- | InterposeFirstYield s1 i1
2695 | InterposeFirstInner s1 i1
2696 | InterposeFirstInject s1
2697 -- | InterposeFirstBuf s1 i1
2698 | InterposeSecondYield s1 i1
2699 -- -- | InterposeSecondYield s1 i1 a
2700 -- -- | InterposeFirstResume s1 i1 a
2702-- Note that this only interposes the pure values, we may run many effects to
2703-- generate those values as some effects may not generate anything (Skip).
2704{-# INLINE_NORMAL interpose #-}
2705interpose :: Monad m => m c -> Unfold m b c -> Stream m b -> Stream m c
2708 (Unfold istep1 inject1) (Stream step1 state1) =
2709 Stream step (InterposeFirst state1)
2713 {-# INLINE_LATE step #-}
2714 step gst (InterposeFirst s1) = do
2715 r <- step1 (adaptState gst) s1
2719 i `seq` return (Skip (InterposeFirstInner s i))
2720 -- i `seq` return (Skip (InterposeFirstYield s i))
2721 Skip s -> return $ Skip (InterposeFirst s)
2725 step _ (InterposeFirstYield s1 i1) = do
2728 Yield x i' -> Yield x (InterposeFirstInner s1 i')
2729 Skip i' -> Skip (InterposeFirstYield s1 i')
2730 Stop -> Skip (InterposeFirst s1)
2733 step _ (InterposeFirstInner s1 i1) = do
2736 Yield x i' -> Yield x (InterposeFirstInner s1 i')
2737 Skip i' -> Skip (InterposeFirstInner s1 i')
2738 Stop -> Skip (InterposeFirstInject s1)
2740 step gst (InterposeFirstInject s1) = do
2741 r <- step1 (adaptState gst) s1
2745 -- i `seq` return (Skip (InterposeFirstBuf s i))
2746 i `seq` return (Skip (InterposeSecondYield s i))
2747 Skip s -> return $ Skip (InterposeFirstInject s)
2751 step _ (InterposeFirstBuf s1 i1) = do
2754 Yield x i' -> Skip (InterposeSecondYield s1 i' x)
2755 Skip i' -> Skip (InterposeFirstBuf s1 i')
2760 step _ (InterposeSecondYield s1 i1 v) = do
2762 return $ Yield r (InterposeFirstResume s1 i1 v)
2764 step _ (InterposeSecondYield s1 i1) = do
2766 return $ Yield r (InterposeFirstInner s1 i1)
2769 step _ (InterposeFirstResume s1 i1 v) = do
2770 return $ Yield v (InterposeFirstInner s1 i1)
2773------------------------------------------------------------------------------
2775------------------------------------------------------------------------------
2777data GbracketState s1 s2 v
2779 | GBracketNormal s1 v
2780 | GBracketException s2
2782-- | The most general bracketing and exception combinator. All other
2783-- combinators can be expressed in terms of this combinator. This can also be
2784-- used for cases which are not covered by the standard combinators.
2788{-# INLINE_NORMAL gbracket #-}
2792 -> (forall s. m s -> m (Either e s)) -- ^ try (exception handling)
2793 -> (c -> m d) -- ^ after, on normal stop
2794 -> (c -> e -> Stream m b) -- ^ on exception
2795 -> (c -> Stream m b) -- ^ stream generator
2797gbracket bef exc aft fexc fnormal =
2798 Stream step GBracketInit
2802 {-# INLINE_LATE step #-}
2803 step _ GBracketInit = do
2805 return $ Skip $ GBracketNormal (fnormal r) r
2807 step gst (GBracketNormal (UnStream step1 st) v) = do
2808 res <- exc $ step1 gst st
2810 Right r -> case r of
2812 return $ Yield x (GBracketNormal (Stream step1 s) v)
2813 Skip s -> return $ Skip (GBracketNormal (Stream step1 s) v)
2814 Stop -> aft v >> return Stop
2815 Left e -> return $ Skip (GBracketException (fexc v e))
2816 step gst (GBracketException (UnStream step1 st)) = do
2819 Yield x s -> return $ Yield x (GBracketException (Stream step1 s))
2820 Skip s -> return $ Skip (GBracketException (Stream step1 s))
2823-- | Run a side effect before the stream yields its first element.
2824{-# INLINE_NORMAL before #-}
2825before :: Monad m => m b -> Stream m a -> Stream m a
2826before action (Stream step state) = Stream step' Nothing
2830 {-# INLINE_LATE step' #-}
2831 step' _ Nothing = action >> return (Skip (Just state))
2833 step' gst (Just st) = do
2836 Yield x s -> return $ Yield x (Just s)
2837 Skip s -> return $ Skip (Just s)
2840-- | Run a side effect whenever the stream stops normally.
2841{-# INLINE_NORMAL after #-}
2842after :: Monad m => m b -> Stream m a -> Stream m a
2843after action (Stream step state) = Stream step' state
2847 {-# INLINE_LATE step' #-}
2851 Yield x s -> return $ Yield x s
2852 Skip s -> return $ Skip s
2853 Stop -> action >> return Stop
2855-- XXX These combinators are expensive due to the call to
2856-- onException/handle/try on each step. Therefore, when possible, they should
2857-- be called in an outer loop where we perform less iterations. For example, we
2858-- cannot call them on each iteration in a char stream, instead we can call
2859-- them when doing an IO on an array.
2861-- XXX For high performance error checks in busy streams we may need another
2862-- Error constructor in step.
2864-- | Run a side effect whenever the stream aborts due to an exception. The
2865-- exception is not caught, simply rethrown.
2866{-# INLINE_NORMAL onException #-}
2867onException :: MonadCatch m => m b -> Stream m a -> Stream m a
2868onException action str =
2869 gbracket (return ()) MC.try return
2870 (\_ (e :: MC.SomeException) -> nilM (action >> MC.throwM e))
2873{-# INLINE_NORMAL _onException #-}
2874_onException :: MonadCatch m => m b -> Stream m a -> Stream m a
2875_onException action (Stream step state) = Stream step' state
2879 {-# INLINE_LATE step' #-}
2881 res <- step gst st `MC.onException` action
2883 Yield x s -> return $ Yield x s
2884 Skip s -> return $ Skip s
2887-- XXX bracket is like concatMap, it generates a stream and then flattens it.
2888-- Like concatMap it has 10x worse performance compared to linear fused
2891-- | Run the first action before the stream starts and remember its output,
2892-- generate a stream using the output, run the second action providing the
2893-- remembered value as an argument whenever the stream ends normally or due to
2895{-# INLINE_NORMAL bracket #-}
2896bracket :: MonadCatch m => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
2897bracket bef aft bet =
2898 gbracket bef MC.try aft
2899 (\a (e :: SomeException) -> nilM (aft a >> MC.throwM e)) bet
2901data BracketState s v = BracketInit | BracketRun s v
2903{-# INLINE_NORMAL _bracket #-}
2904_bracket :: MonadCatch m => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
2905_bracket bef aft bet = Stream step' BracketInit
2909 {-# INLINE_LATE step' #-}
2910 step' _ BracketInit = bef >>= \x -> return (Skip (BracketRun (bet x) x))
2912 -- NOTE: It is important to use UnStream instead of the Stream pattern
2913 -- here, otherwise we get huge perf degradation, see note in concatMap.
2914 step' gst (BracketRun (UnStream step state) v) = do
2915 -- res <- step gst state `MC.onException` aft v
2916 res <- MC.try $ step gst state
2918 Left (e :: SomeException) -> aft v >> MC.throwM e >> return Stop
2919 Right r -> case r of
2920 Yield x s -> return $ Yield x (BracketRun (Stream step s) v)
2921 Skip s -> return $ Skip (BracketRun (Stream step s) v)
2922 Stop -> aft v >> return Stop
2924-- | Run a side effect whenever the stream stops normally or aborts due to an
2926{-# INLINE finally #-}
2927finally :: MonadCatch m => m b -> Stream m a -> Stream m a
2928-- finally action xs = after action $ onException action xs
2929finally action xs = bracket (return ()) (\_ -> action) (const xs)
2931-- | When evaluating a stream if an exception occurs, stream evaluation aborts
2932-- and the specified exception handler is run with the exception as argument.
2933{-# INLINE_NORMAL handle #-}
2934handle :: (MonadCatch m, Exception e)
2935 => (e -> Stream m a) -> Stream m a -> Stream m a
2937 gbracket (return ()) MC.try return (\_ e -> f e) (\_ -> str)
2939{-# INLINE_NORMAL _handle #-}
2940_handle :: (MonadCatch m, Exception e)
2941 => (e -> Stream m a) -> Stream m a -> Stream m a
2942_handle f (Stream step state) = Stream step' (Left state)
2946 {-# INLINE_LATE step' #-}
2947 step' gst (Left st) = do
2948 res <- MC.try $ step gst st
2950 Left e -> return $ Skip $ Right (f e)
2951 Right r -> case r of
2952 Yield x s -> return $ Yield x (Left s)
2953 Skip s -> return $ Skip (Left s)
2956 step' gst (Right (UnStream step1 st)) = do
2959 Yield x s -> return $ Yield x (Right (Stream step1 s))
2960 Skip s -> return $ Skip (Right (Stream step1 s))
2963-------------------------------------------------------------------------------
2964-- General transformation
2965-------------------------------------------------------------------------------
2967{-# INLINE_NORMAL transform #-}
2968transform :: Monad m => Pipe m a b -> Stream m a -> Stream m b
2969transform (Pipe pstep1 pstep2 pstate) (Stream step state) =
2970 Stream step' (Consume pstate, state)
2974 {-# INLINE_LATE step' #-}
2976 step' gst (Consume pst, st) = pst `seq` do
2977 r <- step (adaptState gst) st
2982 Pipe.Yield b pst' -> return $ Yield b (pst', s)
2983 Pipe.Continue pst' -> return $ Skip (pst', s)
2984 Skip s -> return $ Skip (Consume pst, s)
2987 step' _ (Produce pst, st) = pst `seq` do
2990 Pipe.Yield b pst' -> return $ Yield b (pst', st)
2991 Pipe.Continue pst' -> return $ Skip (pst', st)
2993------------------------------------------------------------------------------
2994-- Transformation by Folding (Scans)
2995------------------------------------------------------------------------------
2997------------------------------------------------------------------------------
2999------------------------------------------------------------------------------
3001-- XXX Is a prescan useful, discarding the last step does not sound useful? I
3002-- am not sure about the utility of this function, so this is implemented but
3003-- not exposed. We can expose it if someone provides good reasons why this is
3006-- XXX We have to execute the stream one step ahead to know that we are at the
3007-- last step. The vector implementation of prescan executes the last fold step
3008-- but does not yield the result. This means we have executed the effect but
3009-- discarded value. This does not sound right. In this implementation we are
3010-- not executing the last fold step.
3011{-# INLINE_NORMAL prescanlM' #-}
3012prescanlM' :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
3013prescanlM' f mz (Stream step state) = Stream step' (state, mz)
3015 {-# INLINE_LATE step' #-}
3016 step' gst (st, prev) = do
3017 r <- step (adaptState gst) st
3021 return $ Yield acc (s, f acc x)
3022 Skip s -> return $ Skip (s, prev)
3025{-# INLINE prescanl' #-}
3026prescanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
3027prescanl' f z = prescanlM' (\a b -> return (f a b)) (return z)
3029------------------------------------------------------------------------------
3030-- Monolithic postscans (postscan followed by a map)
3031------------------------------------------------------------------------------
3033-- The performance of a modular postscan followed by a map seems to be
3034-- equivalent to this monolithic scan followed by map therefore we may not need
3035-- this implementation. We just have it for performance comparison and in case
3036-- modular version does not perform well in some situation.
3038{-# INLINE_NORMAL postscanlMx' #-}
3039postscanlMx' :: Monad m
3040 => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> Stream m b
3041postscanlMx' fstep begin done (Stream step state) = do
3042 Stream step' (state, begin)
3044 {-# INLINE_LATE step' #-}
3045 step' gst (st, acc) = do
3046 r <- step (adaptState gst) st
3052 v `seq` y `seq` return (Yield v (s, return y))
3053 Skip s -> return $ Skip (s, acc)
3056{-# INLINE_NORMAL postscanlx' #-}
3057postscanlx' :: Monad m
3058 => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> Stream m b
3059postscanlx' fstep begin done s =
3060 postscanlMx' (\b a -> return (fstep b a)) (return begin) (return . done) s
3062-- XXX do we need consM strict to evaluate the begin value?
3063{-# INLINE scanlMx' #-}
3065 => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> Stream m b
3066scanlMx' fstep begin done s =
3067 (begin >>= \x -> x `seq` done x) `consM` postscanlMx' fstep begin done s
3069{-# INLINE scanlx' #-}
3071 => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> Stream m b
3072scanlx' fstep begin done s =
3073 scanlMx' (\b a -> return (fstep b a)) (return begin) (return . done) s
3075------------------------------------------------------------------------------
3077------------------------------------------------------------------------------
3079{-# INLINE_NORMAL postscanlM' #-}
3080postscanlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
3081postscanlM' fstep begin (Stream step state) =
3082 begin `seq` Stream step' (state, begin)
3084 {-# INLINE_LATE step' #-}
3085 step' gst (st, acc) = acc `seq` do
3086 r <- step (adaptState gst) st
3090 y `seq` return (Yield y (s, y))
3091 Skip s -> return $ Skip (s, acc)
3094{-# INLINE_NORMAL postscanl' #-}
3095postscanl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
3096postscanl' f = postscanlM' (\a b -> return (f a b))
3098{-# INLINE_NORMAL postscanlM #-}
3099postscanlM :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
3100postscanlM fstep begin (Stream step state) = Stream step' (state, begin)
3102 {-# INLINE_LATE step' #-}
3103 step' gst (st, acc) = do
3104 r <- step (adaptState gst) st
3108 return (Yield y (s, y))
3109 Skip s -> return $ Skip (s, acc)
3112{-# INLINE_NORMAL postscanl #-}
3113postscanl :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
3114postscanl f = postscanlM (\a b -> return (f a b))
3116{-# INLINE_NORMAL scanlM' #-}
3117scanlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
3118scanlM' fstep begin s = begin `seq` (begin `cons` postscanlM' fstep begin s)
3120{-# INLINE scanl' #-}
3121scanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
3122scanl' f = scanlM' (\a b -> return (f a b))
3124{-# INLINE_NORMAL scanlM #-}
3125scanlM :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
3126scanlM fstep begin s = begin `cons` postscanlM fstep begin s
3129scanl :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
3130scanl f = scanlM (\a b -> return (f a b))
3132{-# INLINE_NORMAL scanl1M #-}
3133scanl1M :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
3134scanl1M fstep (Stream step state) = Stream step' (state, Nothing)
3136 {-# INLINE_LATE step' #-}
3137 step' gst (st, Nothing) = do
3140 Yield x s -> return $ Yield x (s, Just x)
3141 Skip s -> return $ Skip (s, Nothing)
3144 step' gst (st, Just acc) = do
3149 return $ Yield z (s, Just z)
3150 Skip s -> return $ Skip (s, Just acc)
3153{-# INLINE scanl1 #-}
3154scanl1 :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
3155scanl1 f = scanl1M (\x y -> return (f x y))
3157{-# INLINE_NORMAL scanl1M' #-}
3158scanl1M' :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
3159scanl1M' fstep (Stream step state) = Stream step' (state, Nothing)
3161 {-# INLINE_LATE step' #-}
3162 step' gst (st, Nothing) = do
3165 Yield x s -> x `seq` return $ Yield x (s, Just x)
3166 Skip s -> return $ Skip (s, Nothing)
3169 step' gst (st, Just acc) = acc `seq` do
3174 z `seq` return $ Yield z (s, Just z)
3175 Skip s -> return $ Skip (s, Just acc)
3178{-# INLINE scanl1' #-}
3179scanl1' :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
3180scanl1' f = scanl1M' (\x y -> return (f x y))
3183tap :: Monad m => Fold m a b -> Stream m a -> Stream m a
3184tap (Fold fstep initial extract) (Stream step state) = Stream step' Nothing
3188 step' _ Nothing = do
3190 return $ Skip (Just (r, state))
3192 step' gst (Just (acc, st)) = do
3197 return $ Yield x (Just (acc', s))
3198 Skip s -> return $ Skip (Just (acc, s))
3203-------------------------------------------------------------------------------
3205-------------------------------------------------------------------------------
3207{-# INLINE_NORMAL takeWhileM #-}
3208takeWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
3209takeWhileM f (Stream step state) = Stream step' state
3211 {-# INLINE_LATE step' #-}
3217 return $ if b then Yield x s else Stop
3218 Skip s -> return $ Skip s
3221{-# INLINE takeWhile #-}
3222takeWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
3223takeWhile f = takeWhileM (return . f)
3225{-# INLINE_NORMAL drop #-}
3226drop :: Monad m => Int -> Stream m a -> Stream m a
3227drop n (Stream step state) = Stream step' (state, Just n)
3229 {-# INLINE_LATE step' #-}
3230 step' gst (st, Just i)
3235 Yield _ s -> Skip (s, Just (i - 1))
3236 Skip s -> Skip (s, Just i)
3238 | otherwise = return $ Skip (st, Nothing)
3240 step' gst (st, Nothing) = do
3244 Yield x s -> Yield x (s, Nothing)
3245 Skip s -> Skip (s, Nothing)
3248data DropWhileState s a
3250 | DropWhileYield a s
3253{-# INLINE_NORMAL dropWhileM #-}
3254dropWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
3255dropWhileM f (Stream step state) = Stream step' (DropWhileDrop state)
3257 {-# INLINE_LATE step' #-}
3258 step' gst (DropWhileDrop st) = do
3264 then return $ Skip (DropWhileDrop s)
3265 else return $ Skip (DropWhileYield x s)
3266 Skip s -> return $ Skip (DropWhileDrop s)
3269 step' gst (DropWhileNext st) = do
3272 Yield x s -> return $ Skip (DropWhileYield x s)
3273 Skip s -> return $ Skip (DropWhileNext s)
3276 step' _ (DropWhileYield x st) = return $ Yield x (DropWhileNext st)
3278{-# INLINE dropWhile #-}
3279dropWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
3280dropWhile f = dropWhileM (return . f)
3282{-# INLINE_NORMAL filterM #-}
3283filterM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
3284filterM f (Stream step state) = Stream step' state
3286 {-# INLINE_LATE step' #-}
3295 Skip s -> return $ Skip s
3298{-# INLINE filter #-}
3299filter :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
3300filter f = filterM (return . f)
3302{-# INLINE_NORMAL uniq #-}
3303uniq :: (Eq a, Monad m) => Stream m a -> Stream m a
3304uniq (Stream step state) = Stream step' (Nothing, state)
3306 {-# INLINE_LATE step' #-}
3307 step' gst (Nothing, st) = do
3310 Yield x s -> return $ Yield x (Just x, s)
3311 Skip s -> return $ Skip (Nothing, s)
3313 step' gst (Just x, st) = do
3316 Yield y s | x == y -> return $ Skip (Just x, s)
3317 | otherwise -> return $ Yield y (Just y, s)
3318 Skip s -> return $ Skip (Just x, s)
3321------------------------------------------------------------------------------
3322-- Transformation by Mapping
3323------------------------------------------------------------------------------
3325{-# INLINE_NORMAL sequence #-}
3326sequence :: Monad m => Stream m (m a) -> Stream m a
3327sequence (Stream step state) = Stream step' state
3329 {-# INLINE_LATE step' #-}
3331 r <- step (adaptState gst) st
3333 Yield x s -> x >>= \a -> return (Yield a s)
3334 Skip s -> return $ Skip s
3337------------------------------------------------------------------------------
3339------------------------------------------------------------------------------
3341data LoopState x s = FirstYield s
3342 | InterspersingYield s
3345{-# INLINE_NORMAL intersperseM #-}
3346intersperseM :: Monad m => m a -> Stream m a -> Stream m a
3347intersperseM m (Stream step state) = Stream step' (FirstYield state)
3349 {-# INLINE_LATE step' #-}
3350 step' gst (FirstYield st) = do
3354 Yield x s -> Skip (YieldAndCarry x s)
3355 Skip s -> Skip (FirstYield s)
3358 step' gst (InterspersingYield st) = do
3363 return $ Yield a (YieldAndCarry x s)
3364 Skip s -> return $ Skip $ InterspersingYield s
3367 step' _ (YieldAndCarry x st) = return $ Yield x (InterspersingYield st)
3372 | SuffixYield a (SuffixState s a)
3374{-# INLINE_NORMAL intersperseSuffix #-}
3375intersperseSuffix :: forall m a. Monad m => m a -> Stream m a -> Stream m a
3376intersperseSuffix action (Stream step state) = Stream step' (SuffixElem state)
3378 {-# INLINE_LATE step' #-}
3379 step' gst (SuffixElem st) = do
3382 Yield x s -> Skip (SuffixYield x (SuffixSuffix s))
3383 Skip s -> Skip (SuffixElem s)
3386 step' _ (SuffixSuffix st) = do
3387 action >>= \r -> return $ Skip (SuffixYield r (SuffixElem st))
3389 step' _ (SuffixYield x next) = return $ Yield x next
3391data SuffixSpanState s a
3392 = SuffixSpanElem s Int
3393 | SuffixSpanSuffix s
3394 | SuffixSpanYield a (SuffixSpanState s a)
3398-- | intersperse after every n items
3399{-# INLINE_NORMAL intersperseSuffixBySpan #-}
3400intersperseSuffixBySpan :: forall m a. Monad m
3401 => Int -> m a -> Stream m a -> Stream m a
3402intersperseSuffixBySpan n action (Stream step state) =
3403 Stream step' (SuffixSpanElem state n)
3405 {-# INLINE_LATE step' #-}
3406 step' gst (SuffixSpanElem st i) | i > 0 = do
3409 Yield x s -> Skip (SuffixSpanYield x (SuffixSpanElem s (i - 1)))
3410 Skip s -> Skip (SuffixSpanElem s i)
3411 Stop -> if i == n then Stop else Skip SuffixSpanLast
3412 step' _ (SuffixSpanElem st _) = return $ Skip (SuffixSpanSuffix st)
3414 step' _ (SuffixSpanSuffix st) = do
3415 action >>= \r -> return $ Skip (SuffixSpanYield r (SuffixSpanElem st n))
3417 step' _ (SuffixSpanLast) = do
3418 action >>= \r -> return $ Skip (SuffixSpanYield r SuffixSpanStop)
3420 step' _ (SuffixSpanYield x next) = return $ Yield x next
3422 step' _ (SuffixSpanStop) = return Stop
3424{-# INLINE intersperse #-}
3425intersperse :: Monad m => a -> Stream m a -> Stream m a
3426intersperse a = intersperseM (return a)
3428{-# INLINE_NORMAL insertBy #-}
3429insertBy :: Monad m => (a -> a -> Ordering) -> a -> Stream m a -> Stream m a
3430insertBy cmp a (Stream step state) = Stream step' (state, False, Nothing)
3432 {-# INLINE_LATE step' #-}
3433 step' gst (st, False, _) = do
3436 Yield x s -> case cmp a x of
3437 GT -> return $ Yield x (s, False, Nothing)
3438 _ -> return $ Yield a (s, True, Just x)
3439 Skip s -> return $ Skip (s, False, Nothing)
3440 Stop -> return $ Yield a (st, True, Nothing)
3442 step' _ (_, True, Nothing) = return Stop
3444 step' gst (st, True, Just prev) = do
3447 Yield x s -> return $ Yield prev (s, True, Just x)
3448 Skip s -> return $ Skip (s, True, Just prev)
3449 Stop -> return $ Yield prev (st, True, Nothing)
3451------------------------------------------------------------------------------
3453------------------------------------------------------------------------------
3455{-# INLINE_NORMAL deleteBy #-}
3456deleteBy :: Monad m => (a -> a -> Bool) -> a -> Stream m a -> Stream m a
3457deleteBy eq x (Stream step state) = Stream step' (state, False)
3459 {-# INLINE_LATE step' #-}
3460 step' gst (st, False) = do
3463 Yield y s -> return $
3464 if eq x y then Skip (s, True) else Yield y (s, False)
3465 Skip s -> return $ Skip (s, False)
3468 step' gst (st, True) = do
3471 Yield y s -> return $ Yield y (s, True)
3472 Skip s -> return $ Skip (s, True)
3475------------------------------------------------------------------------------
3476-- Transformation by Map and Filter
3477------------------------------------------------------------------------------
3479-- XXX Will this always fuse properly?
3480{-# INLINE_NORMAL mapMaybe #-}
3481mapMaybe :: Monad m => (a -> Maybe b) -> Stream m a -> Stream m b
3482mapMaybe f = fmap fromJust . filter isJust . map f
3484{-# INLINE_NORMAL mapMaybeM #-}
3485mapMaybeM :: Monad m => (a -> m (Maybe b)) -> Stream m a -> Stream m b
3486mapMaybeM f = fmap fromJust . filter isJust . mapM f
3488------------------------------------------------------------------------------
3490------------------------------------------------------------------------------
3492{-# INLINE_NORMAL indexed #-}
3493indexed :: Monad m => Stream m a -> Stream m (Int, a)
3494indexed (Stream step state) = Stream step' (state, 0)
3496 {-# INLINE_LATE step' #-}
3497 step' gst (st, i) = i `seq` do
3498 r <- step (adaptState gst) st
3500 Yield x s -> return $ Yield (i, x) (s, i+1)
3501 Skip s -> return $ Skip (s, i)
3504{-# INLINE_NORMAL indexedR #-}
3505indexedR :: Monad m => Int -> Stream m a -> Stream m (Int, a)
3506indexedR m (Stream step state) = Stream step' (state, m)
3508 {-# INLINE_LATE step' #-}
3509 step' gst (st, i) = i `seq` do
3510 r <- step (adaptState gst) st
3512 Yield x s -> let i' = i - 1
3513 in return $ Yield (i, x) (s, i')
3514 Skip s -> return $ Skip (s, i)
3517{-# INLINE_NORMAL zipWithM #-}
3519 => (a -> b -> m c) -> Stream m a -> Stream m b -> Stream m c
3520zipWithM f (Stream stepa ta) (Stream stepb tb) = Stream step (ta, tb, Nothing)
3522 {-# INLINE_LATE step #-}
3523 step gst (sa, sb, Nothing) = do
3524 r <- stepa (adaptState gst) sa
3527 Yield x sa' -> Skip (sa', sb, Just x)
3528 Skip sa' -> Skip (sa', sb, Nothing)
3531 step gst (sa, sb, Just x) = do
3532 r <- stepb (adaptState gst) sb
3536 return $ Yield z (sa, sb', Nothing)
3537 Skip sb' -> return $ Skip (sa, sb', Just x)
3540#if __GLASGOW_HASKELL__ >= 801
3541{-# RULES "zipWithM xs xs"
3542 forall f xs. zipWithM @Identity f xs xs = mapM (\x -> f x x) xs #-}
3545{-# INLINE zipWith #-}
3546zipWith :: Monad m => (a -> b -> c) -> Stream m a -> Stream m b -> Stream m c
3547zipWith f = zipWithM (\a b -> return (f a b))
3549-- | Like 'zipWithM' but zips concurrently i.e. both the streams being zipped
3550-- are generated concurrently.
3552{-# INLINE_NORMAL zipAsyncWithM #-}
3553zipAsyncWithM :: MonadAsync m
3554 => (a -> b -> m c) -> Stream m a -> Stream m b -> Stream m c
3555zipAsyncWithM f m1 m2 = Stream step Nothing
3559 {-# INLINE_LATE step #-}
3560 step gst Nothing = do
3561 ma <- mkParallel (adaptState gst) m1
3562 mb <- mkParallel (adaptState gst) m2
3563 return $ Skip (Just $ zipWithM f ma mb)
3565 step gst (Just (UnStream step1 st)) = do
3568 Yield a s -> Yield a (Just $ Stream step1 s)
3569 Skip s -> Skip (Just $ Stream step1 s)
3572-- | Like 'zipWith' but zips concurrently i.e. both the streams being zipped
3573-- are generated concurrently.
3575{-# INLINE zipAsyncWith #-}
3576zipAsyncWith :: MonadAsync m
3577 => (a -> b -> c) -> Stream m a -> Stream m b -> Stream m c
3578zipAsyncWith f = zipAsyncWithM (\a b -> return (f a b))
3580------------------------------------------------------------------------------
3582------------------------------------------------------------------------------
3584{-# INLINE_NORMAL mergeByM #-}
3587 => (a -> a -> m Ordering) -> Stream m a -> Stream m a -> Stream m a
3588mergeByM cmp (Stream stepa ta) (Stream stepb tb) =
3589 Stream step (Just ta, Just tb, Nothing, Nothing)
3591 {-# INLINE_LATE step #-}
3593 -- one of the values is missing, and the corresponding stream is running
3594 step gst (Just sa, sb, Nothing, b) = do
3597 Yield a sa' -> Skip (Just sa', sb, Just a, b)
3598 Skip sa' -> Skip (Just sa', sb, Nothing, b)
3599 Stop -> Skip (Nothing, sb, Nothing, b)
3601 step gst (sa, Just sb, a, Nothing) = do
3604 Yield b sb' -> Skip (sa, Just sb', a, Just b)
3605 Skip sb' -> Skip (sa, Just sb', a, Nothing)
3606 Stop -> Skip (sa, Nothing, a, Nothing)
3608 -- both the values are available
3609 step _ (sa, sb, Just a, Just b) = do
3611 return $ case res of
3612 GT -> Yield b (sa, sb, Just a, Nothing)
3613 _ -> Yield a (sa, sb, Nothing, Just b)
3615 -- one of the values is missing, corresponding stream is done
3616 step _ (Nothing, sb, Nothing, Just b) =
3617 return $ Yield b (Nothing, sb, Nothing, Nothing)
3619 step _ (sa, Nothing, Just a, Nothing) =
3620 return $ Yield a (sa, Nothing, Nothing, Nothing)
3622 step _ (Nothing, Nothing, Nothing, Nothing) = return Stop
3624{-# INLINE mergeBy #-}
3627 => (a -> a -> Ordering) -> Stream m a -> Stream m a -> Stream m a
3628mergeBy cmp = mergeByM (\a b -> return $ cmp a b)
3630------------------------------------------------------------------------------
3631-- Transformation comprehensions
3632------------------------------------------------------------------------------
3634{-# INLINE_NORMAL the #-}
3635the :: (Eq a, Monad m) => Stream m a -> m (Maybe a)
3636the (Stream step state) = go state
3639 r <- step defState st
3641 Yield x s -> go' x s
3643 Stop -> return Nothing
3645 r <- step defState st
3647 Yield x s | x == n -> go' n s
3648 | otherwise -> return Nothing
3650 Stop -> return (Just n)
3652--------------------------------------------------------------------------------
3653-- UTF8 Encoding / Decoding
3654--------------------------------------------------------------------------------
3656-- UTF-8 primitives, Lifted from GHC.IO.Encoding.UTF8.
3659ord2 :: Char -> WList
3660ord2 c = assert (n >= 0x80 && n <= 0x07ff) (WCons x1 (WCons x2 WNil))
3663 x1 = fromIntegral $ (n `shiftR` 6) + 0xC0
3664 x2 = fromIntegral $ (n .&. 0x3F) + 0x80
3667ord3 :: Char -> WList
3668ord3 c = assert (n >= 0x0800 && n <= 0xffff) (WCons x1 (WCons x2 (WCons x3 WNil)))
3671 x1 = fromIntegral $ (n `shiftR` 12) + 0xE0
3672 x2 = fromIntegral $ ((n `shiftR` 6) .&. 0x3F) + 0x80
3673 x3 = fromIntegral $ (n .&. 0x3F) + 0x80
3676ord4 :: Char -> WList
3677ord4 c = assert (n >= 0x10000) (WCons x1 (WCons x2 (WCons x3 (WCons x4 WNil))))
3680 x1 = fromIntegral $ (n `shiftR` 18) + 0xF0
3681 x2 = fromIntegral $ ((n `shiftR` 12) .&. 0x3F) + 0x80
3682 x3 = fromIntegral $ ((n `shiftR` 6) .&. 0x3F) + 0x80
3683 x4 = fromIntegral $ (n .&. 0x3F) + 0x80
3685data CodingFailureMode
3686 = TransliterateCodingFailure
3687 | ErrorOnCodingFailure
3690{-# INLINE replacementChar #-}
3691replacementChar :: Char
3692replacementChar = '\xFFFD'
3694-- Int helps in cheaper conversion from Int to Char
3696type DecodeState = Word8
3698-- See http://bjoern.hoehrmann.de/utf-8/decoder/dfa/ for details.
3700{-# INLINE runFold #-}
3701runFold :: (Monad m) => Fold m a b -> Stream m a -> m b
3702runFold (Fold step begin done) = foldlMx' step begin done
3704-- XXX Use names decodeSuccess = 0, decodeFailure = 12
3706decodeTable :: [Word8]
3708 -- The first part of the table maps bytes to character classes that
3709 -- to reduce the size of the transition table and create bitmasks.
3710 0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0, 0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,
3711 0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0, 0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,
3712 0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0, 0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,
3713 0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0, 0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,
3714 1,1,1,1,1,1,1,1,1,1,1,1,1,1,1,1, 9,9,9,9,9,9,9,9,9,9,9,9,9,9,9,9,
3715 7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7, 7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,7,
3716 8,8,2,2,2,2,2,2,2,2,2,2,2,2,2,2, 2,2,2,2,2,2,2,2,2,2,2,2,2,2,2,2,
3717 10,3,3,3,3,3,3,3,3,3,3,3,3,4,3,3, 11,6,6,6,5,8,8,8,8,8,8,8,8,8,8,8,
3719 -- The second part is a transition table that maps a combination
3720 -- of a state of the automaton and a character class to a state.
3721 0,12,24,36,60,96,84,12,12,12,48,72, 12,12,12,12,12,12,12,12,12,12,12,12,
3722 12, 0,12,12,12,12,12, 0,12, 0,12,12, 12,24,12,12,12,12,12,24,12,24,12,12,
3723 12,12,12,12,12,12,12,24,12,12,12,12, 12,24,12,12,12,12,12,12,12,24,12,12,
3724 12,12,12,12,12,12,12,36,12,36,12,12, 12,36,12,12,12,12,12,36,12,36,12,12,
3725 12,36,12,12,12,12,12,12,12,12,12,12
3728utf8d :: A.Array Word8
3731 -- Aligning to cacheline makes a barely noticeable difference
3732 -- XXX currently alignment is not implemented for unmanaged allocation
3733 $ runFold (A.writeNAlignedUnmanaged 64 (length decodeTable))
3734 (fromList decodeTable)
3736-- | Return element at the specified index without checking the bounds.
3737-- and without touching the foreign ptr.
3738{-# INLINE_NORMAL unsafePeekElemOff #-}
3739unsafePeekElemOff :: forall a. Storable a => Ptr a -> Int -> a
3740unsafePeekElemOff p i = let !x = A.unsafeInlineIO $ peekElemOff p i in x
3742-- decode is split into two separate cases to avoid branching instructions.
3743-- From the higher level flow we already know which case we are in so we can
3744-- call the appropriate decode function.
3746-- When the state is 0
3747{-# INLINE decode0 #-}
3748decode0 :: Ptr Word8 -> Word8 -> Tuple' DecodeState CodePoint
3750 let !t = table `unsafePeekElemOff` fromIntegral byte
3751 !codep' = (0xff `shiftR` (fromIntegral t)) .&. fromIntegral byte
3752 !state' = table `unsafePeekElemOff` (256 + fromIntegral t)
3753 in assert ((byte > 0x7f || error showByte)
3754 && (state' /= 0 || error (showByte ++ showTable)))
3755 (Tuple' state' codep')
3760 let !(Ptr addr) = table
3761 end = table `plusPtr` 364
3762 in A.Array (ForeignPtr addr undefined) end end :: A.Array Word8
3763 showByte = "Streamly: decode0: byte: " ++ show byte
3764 showTable = " table: " ++ show utf8table
3766-- When the state is not 0
3767{-# INLINE decode1 #-}
3773 -> Tuple' DecodeState CodePoint
3774decode1 table state codep byte =
3775 -- Remember codep is Int type!
3776 -- Can it be unsafe to convert the resulting Int to Char?
3777 let !t = table `unsafePeekElemOff` fromIntegral byte
3778 !codep' = (fromIntegral byte .&. 0x3f) .|. (codep `shiftL` 6)
3779 !state' = table `unsafePeekElemOff`
3780 (256 + fromIntegral state + fromIntegral t)
3781 in assert (codep' <= 0x10FFFF
3782 || error (showByte ++ showState state codep))
3783 (Tuple' state' codep')
3787 let !(Ptr addr) = table
3788 end = table `plusPtr` 364
3789 in A.Array (ForeignPtr addr undefined) end end :: A.Array Word8
3790 showByte = "Streamly: decode1: byte: " ++ show byte
3792 " state: " ++ show st ++
3793 " codepoint: " ++ show cp ++
3794 " table: " ++ show utf8table
3796-- We can divide the errors in three general categories:
3797-- * A non-starter was encountered in a begin state
3798-- * A starter was encountered without completing a codepoint
3799-- * The last codepoint was not complete (input underflow)
3801data DecodeError = DecodeError !DecodeState !CodePoint deriving Show
3804 = FreshPointDecodeInit s
3805 | FreshPointDecodeInit1 s Word8
3806 | FreshPointDecodeFirst s Word8
3807 | FreshPointDecoding s !DecodeState !CodePoint
3808 | YieldAndContinue a (FreshPoint s a)
3811-- XXX Add proper error messages
3812-- XXX Implement this in terms of decodeUtf8Either
3813{-# INLINE_NORMAL decodeUtf8With #-}
3814decodeUtf8With :: Monad m => CodingFailureMode -> Stream m Word8 -> Stream m Char
3815decodeUtf8With cfm (Stream step state) =
3816 let Array p _ _ = utf8d
3817 !ptr = (unsafeForeignPtrToPtr p)
3818 in Stream (step' ptr) (FreshPointDecodeInit state)
3820 {-# INLINE transliterateOrError #-}
3821 transliterateOrError e s =
3823 ErrorOnCodingFailure -> error e
3824 TransliterateCodingFailure -> YieldAndContinue replacementChar s
3825 {-# INLINE inputUnderflow #-}
3828 ErrorOnCodingFailure ->
3829 error "Streamly.Streams.StreamD.decodeUtf8With: Input Underflow"
3830 TransliterateCodingFailure -> YieldAndContinue replacementChar Done
3831 {-# INLINE_LATE step' #-}
3832 step' _ gst (FreshPointDecodeInit st) = do
3833 r <- step (adaptState gst) st
3835 Yield x s -> Skip (FreshPointDecodeInit1 s x)
3836 Skip s -> Skip (FreshPointDecodeInit s)
3839 step' _ _ (FreshPointDecodeInit1 st x) = do
3840 -- Note: It is important to use a ">" instead of a "<=" test
3841 -- here for GHC to generate code layout for default branch
3842 -- prediction for the common case. This is fragile and might
3843 -- change with the compiler versions, we need a more reliable
3844 -- "likely" primitive to control branch predication.
3847 return $ Skip $ YieldAndContinue
3848 (unsafeChr (fromIntegral x))
3849 (FreshPointDecodeInit st)
3850 -- Using a separate state here generates a jump to a
3851 -- separate code block in the core which seems to perform
3852 -- slightly better for the non-ascii case.
3853 True -> return $ Skip $ FreshPointDecodeFirst st x
3855 -- XXX should we merge it with FreshPointDecodeInit1?
3856 step' table _ (FreshPointDecodeFirst st x) = do
3857 let (Tuple' sv cp) = decode0 table x
3862 transliterateOrError
3863 "Streamly.Streams.StreamD.decodeUtf8With: Invalid UTF8 codepoint encountered"
3864 (FreshPointDecodeInit st)
3865 0 -> error "unreachable state"
3866 _ -> Skip (FreshPointDecoding st sv cp)
3868 -- We recover by trying the new byte x a starter of a new codepoint.
3869 -- XXX need to use the same recovery in array decoding routine as well
3870 step' table gst (FreshPointDecoding st statePtr codepointPtr) = do
3871 r <- step (adaptState gst) st
3874 let (Tuple' sv cp) = decode1 table statePtr codepointPtr x
3877 0 -> Skip $ YieldAndContinue (unsafeChr cp)
3878 (FreshPointDecodeInit s)
3881 transliterateOrError
3882 "Streamly.Streams.StreamD.decodeUtf8With: Invalid UTF8 codepoint encountered"
3883 (FreshPointDecodeInit1 s x)
3884 _ -> Skip (FreshPointDecoding s sv cp)
3885 Skip s -> return $ Skip (FreshPointDecoding s statePtr codepointPtr)
3886 Stop -> return $ Skip inputUnderflow
3888 step' _ _ (YieldAndContinue c s) = return $ Yield c s
3889 step' _ _ Done = return Stop
3891{-# INLINE decodeUtf8 #-}
3892decodeUtf8 :: Monad m => Stream m Word8 -> Stream m Char
3893decodeUtf8 = decodeUtf8With ErrorOnCodingFailure
3895{-# INLINE decodeUtf8Lenient #-}
3896decodeUtf8Lenient :: Monad m => Stream m Word8 -> Stream m Char
3897decodeUtf8Lenient = decodeUtf8With TransliterateCodingFailure
3899{-# INLINE_NORMAL resumeDecodeUtf8Either #-}
3900resumeDecodeUtf8Either
3905 -> Stream m (Either DecodeError Char)
3906resumeDecodeUtf8Either dst codep (Stream step state) =
3907 let Array p _ _ = utf8d
3908 !ptr = (unsafeForeignPtrToPtr p)
3911 then FreshPointDecodeInit state
3912 else FreshPointDecoding state dst codep
3913 in Stream (step' ptr) stt
3915 {-# INLINE_LATE step' #-}
3916 step' _ gst (FreshPointDecodeInit st) = do
3917 r <- step (adaptState gst) st
3919 Yield x s -> Skip (FreshPointDecodeInit1 s x)
3920 Skip s -> Skip (FreshPointDecodeInit s)
3923 step' _ _ (FreshPointDecodeInit1 st x) = do
3924 -- Note: It is important to use a ">" instead of a "<=" test
3925 -- here for GHC to generate code layout for default branch
3926 -- prediction for the common case. This is fragile and might
3927 -- change with the compiler versions, we need a more reliable
3928 -- "likely" primitive to control branch predication.
3931 return $ Skip $ YieldAndContinue
3932 (Right $ unsafeChr (fromIntegral x))
3933 (FreshPointDecodeInit st)
3934 -- Using a separate state here generates a jump to a
3935 -- separate code block in the core which seems to perform
3936 -- slightly better for the non-ascii case.
3937 True -> return $ Skip $ FreshPointDecodeFirst st x
3939 -- XXX should we merge it with FreshPointDecodeInit1?
3940 step' table _ (FreshPointDecodeFirst st x) = do
3941 let (Tuple' sv cp) = decode0 table x
3945 Skip $ YieldAndContinue (Left $ DecodeError 0 (fromIntegral x))
3946 (FreshPointDecodeInit st)
3947 0 -> error "unreachable state"
3948 _ -> Skip (FreshPointDecoding st sv cp)
3950 -- We recover by trying the new byte x a starter of a new codepoint.
3951 -- XXX need to use the same recovery in array decoding routine as well
3952 step' table gst (FreshPointDecoding st statePtr codepointPtr) = do
3953 r <- step (adaptState gst) st
3956 let (Tuple' sv cp) = decode1 table statePtr codepointPtr x
3959 0 -> Skip $ YieldAndContinue (Right $ unsafeChr cp)
3960 (FreshPointDecodeInit s)
3962 Skip $ YieldAndContinue (Left $ DecodeError statePtr codepointPtr)
3963 (FreshPointDecodeInit1 s x)
3964 _ -> Skip (FreshPointDecoding s sv cp)
3965 Skip s -> return $ Skip (FreshPointDecoding s statePtr codepointPtr)
3966 Stop -> return $ Skip $ YieldAndContinue (Left $ DecodeError statePtr codepointPtr) Done
3968 step' _ _ (YieldAndContinue c s) = return $ Yield c s
3969 step' _ _ Done = return Stop
3971{-# INLINE_NORMAL decodeUtf8Either #-}
3972decodeUtf8Either :: Monad m
3973 => Stream m Word8 -> Stream m (Either DecodeError Char)
3974decodeUtf8Either = resumeDecodeUtf8Either 0 0
3976data FlattenState s a
3977 = OuterLoop s !(Maybe (DecodeState, CodePoint))
3978 | InnerLoopDecodeInit s (ForeignPtr a) !(Ptr a) !(Ptr a)
3979 | InnerLoopDecodeFirst s (ForeignPtr a) !(Ptr a) !(Ptr a) Word8
3980 | InnerLoopDecoding s (ForeignPtr a) !(Ptr a) !(Ptr a)
3981 !DecodeState !CodePoint
3982 | YAndC !Char (FlattenState s a) -- These constructors can be
3983 -- encoded in the FreshPoint
3984 -- type, I prefer to keep these
3985 -- flat even though that means
3986 -- coming up with new names
3989-- The normal decodeUtf8 above should fuse with flattenArrays
3990-- to create this exact code but it doesn't for some reason, as of now this
3991-- remains the fastest way I could figure out to decodeUtf8.
3993-- XXX Add Proper error messages
3994{-# INLINE_NORMAL decodeUtf8ArraysWith #-}
3995decodeUtf8ArraysWith ::
3997 => CodingFailureMode
3998 -> Stream m (A.Array Word8)
4000decodeUtf8ArraysWith cfm (Stream step state) =
4001 let Array p _ _ = utf8d
4002 !ptr = (unsafeForeignPtrToPtr p)
4003 in Stream (step' ptr) (OuterLoop state Nothing)
4005 {-# INLINE transliterateOrError #-}
4006 transliterateOrError e s =
4008 ErrorOnCodingFailure -> error e
4009 TransliterateCodingFailure -> YAndC replacementChar s
4010 {-# INLINE inputUnderflow #-}
4013 ErrorOnCodingFailure ->
4015 "Streamly.Streams.StreamD.decodeUtf8ArraysWith: Input Underflow"
4016 TransliterateCodingFailure -> YAndC replacementChar D
4017 {-# INLINE_LATE step' #-}
4018 step' _ gst (OuterLoop st Nothing) = do
4019 r <- step (adaptState gst) st
4022 Yield A.Array {..} s ->
4023 let p = unsafeForeignPtrToPtr aStart
4024 in Skip (InnerLoopDecodeInit s aStart p aEnd)
4025 Skip s -> Skip (OuterLoop s Nothing)
4027 step' _ gst (OuterLoop st dst@(Just (ds, cp))) = do
4028 r <- step (adaptState gst) st
4031 Yield A.Array {..} s ->
4032 let p = unsafeForeignPtrToPtr aStart
4033 in Skip (InnerLoopDecoding s aStart p aEnd ds cp)
4034 Skip s -> Skip (OuterLoop s dst)
4035 Stop -> Skip inputUnderflow
4036 step' _ _ (InnerLoopDecodeInit st startf p end)
4038 liftIO $ touchForeignPtr startf
4039 return $ Skip $ OuterLoop st Nothing
4040 step' _ _ (InnerLoopDecodeInit st startf p end) = do
4041 x <- liftIO $ peek p
4042 -- Note: It is important to use a ">" instead of a "<=" test here for
4043 -- GHC to generate code layout for default branch prediction for the
4044 -- common case. This is fragile and might change with the compiler
4045 -- versions, we need a more reliable "likely" primitive to control
4046 -- branch predication.
4049 return $ Skip $ YAndC
4050 (unsafeChr (fromIntegral x))
4051 (InnerLoopDecodeInit st startf (p `plusPtr` 1) end)
4052 -- Using a separate state here generates a jump to a separate code
4053 -- block in the core which seems to perform slightly better for the
4055 True -> return $ Skip $ InnerLoopDecodeFirst st startf p end x
4057 step' table _ (InnerLoopDecodeFirst st startf p end x) = do
4058 let (Tuple' sv cp) = decode0 table x
4063 transliterateOrError
4064 "Streamly.Streams.StreamD.decodeUtf8ArraysWith: Invalid UTF8 codepoint encountered"
4065 (InnerLoopDecodeInit st startf (p `plusPtr` 1) end)
4066 0 -> error "unreachable state"
4067 _ -> Skip (InnerLoopDecoding st startf (p `plusPtr` 1) end sv cp)
4068 step' _ _ (InnerLoopDecoding st startf p end sv cp)
4070 liftIO $ touchForeignPtr startf
4071 return $ Skip $ OuterLoop st (Just (sv, cp))
4072 step' table _ (InnerLoopDecoding st startf p end statePtr codepointPtr) = do
4073 x <- liftIO $ peek p
4074 let (Tuple' sv cp) = decode1 table statePtr codepointPtr x
4081 (InnerLoopDecodeInit st startf (p `plusPtr` 1) end)
4084 transliterateOrError
4085 "Streamly.Streams.StreamD.decodeUtf8ArraysWith: Invalid UTF8 codepoint encountered"
4086 (InnerLoopDecodeInit st startf (p `plusPtr` 1) end)
4087 _ -> Skip (InnerLoopDecoding st startf (p `plusPtr` 1) end sv cp)
4088 step' _ _ (YAndC c s) = return $ Yield c s
4089 step' _ _ D = return Stop
4091{-# INLINE decodeUtf8Arrays #-}
4094 => Stream m (A.Array Word8)
4096decodeUtf8Arrays = decodeUtf8ArraysWith ErrorOnCodingFailure
4098{-# INLINE decodeUtf8ArraysLenient #-}
4099decodeUtf8ArraysLenient ::
4101 => Stream m (A.Array Word8)
4103decodeUtf8ArraysLenient = decodeUtf8ArraysWith TransliterateCodingFailure
4105data WList = WCons !Word8 !WList | WNil
4107data EncodeState s = EncodeState s !WList
4109-- More yield points improve performance, but I am not sure if they can cause
4110-- too much code bloat or some trouble with fusion. So keeping only two yield
4111-- points for now, one for the ascii chars (fast path) and one for all other
4112-- paths (slow path).
4113{-# INLINE_NORMAL encodeUtf8 #-}
4114encodeUtf8 :: Monad m => Stream m Char -> Stream m Word8
4115encodeUtf8 (Stream step state) = Stream step' (EncodeState state WNil)
4117 {-# INLINE_LATE step' #-}
4118 step' gst (EncodeState st WNil) = do
4119 r <- step (adaptState gst) st
4126 Yield (fromIntegral x) (EncodeState s WNil)
4127 | x <= 0x7FF -> Skip (EncodeState s (ord2 c))
4131 "Streamly.Streams.StreamD.encodeUtf8: Encountered a surrogate"
4132 else Skip (EncodeState s (ord3 c))
4133 | otherwise -> Skip (EncodeState s (ord4 c))
4134 Skip s -> Skip (EncodeState s WNil)
4136 step' _ (EncodeState s (WCons x xs)) = return $ Yield x (EncodeState s xs)
4138-------------------------------------------------------------------------------
4139-- Concurrent application and fold
4140-------------------------------------------------------------------------------
4142-- Using StreamD the worker stream producing code can fuse with the code to
4143-- queue output to the SVar giving some perf boost.
4145-- Note that StreamD can only be used in limited situations, specifically, we
4146-- cannot implement joinStreamVarPar using this.
4148-- XXX make sure that the SVar passed is a Parallel style SVar.
4150-- | Fold the supplied stream to the SVar asynchronously using Parallel
4151-- concurrency style.
4152-- {-# INLINE_NORMAL toSVarParallel #-}
4153{-# INLINE toSVarParallel #-}
4154toSVarParallel :: MonadAsync m
4155 => State t m a -> SVar t m a -> Stream m a -> m ()
4156toSVarParallel st sv xs =
4157 if svarInspectMode sv
4161 case getYieldLimit st of
4162 Nothing -> doFork (work Nothing)
4164 (handleChildException sv)
4165 Just _ -> doFork (workLim Nothing)
4167 (handleChildException sv)
4172 {-# NOINLINE work #-}
4173 work info = (runFold (FL.toParallelSVar sv info) xs)
4175 {-# NOINLINE workLim #-}
4176 workLim info = runFold (FL.toParallelSVarLimited sv info) xs
4178 {-# NOINLINE forkWithDiag #-}
4180 -- We do not use workerCount in case of ParallelVar but still there is
4181 -- no harm in maintaining it correctly.
4182 liftIO $ atomicModifyIORefCAS_ (workerCount sv) $ \n -> n + 1
4184 -- This allocation matters when significant number of workers are being
4185 -- sent. We allocate it only when needed. The overhead increases by 4x.
4187 case yieldRateInfo sv of
4188 Nothing -> return Nothing
4189 Just _ -> liftIO $ do
4190 cntRef <- newIORef 0
4191 t <- getTime Monotonic
4192 lat <- newIORef (0, t)
4193 return $ Just WorkerInfo
4194 { workerYieldMax = 0
4195 , workerYieldCount = cntRef
4196 , workerLatencyStart = lat
4199 case getYieldLimit st of
4200 Nothing -> doFork (work winfo)
4202 (handleChildException sv)
4203 Just _ -> doFork (workLim winfo)
4205 (handleChildException sv)
4208{-# INLINE_NORMAL mkParallel #-}
4209mkParallel :: MonadAsync m => State t m a -> Stream m a -> m (Stream m a)
4211 sv <- newParallelVar StopNone defState
4212 toSVarParallel st sv m
4213 return $ fromSVar sv
4215{-# INLINE_NORMAL applyParallel #-}
4216applyParallel :: MonadAsync m
4217 => (Stream m a -> Stream m b) -> Stream m a -> Stream m b
4218applyParallel f m = Stream step Nothing
4221 {-# INLINE_LATE step #-}
4222 step gst Nothing = do
4223 sv <- newParallelVar StopNone (adaptState gst)
4224 toSVarParallel (adaptState gst) sv m
4225 return $ Skip $ Just $ f (fromSVar sv)
4227 step gst (Just (UnStream step1 st)) = do
4230 Yield a s -> Yield a (Just $ Stream step1 s)
4231 Skip s -> Skip (Just $ Stream step1 s)
4234{-# INLINE foldParallel #-}
4235foldParallel :: (K.IsStream t, MonadAsync m) => (t m a -> m b) -> t m a -> m b
4236foldParallel f m = do
4237 sv <- newParallelVar StopNone defState
4238 toSVarParallel defState sv (toStreamD m)
4239 f $ fromStreamD $ fromSVar sv
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