-- | -- Module : Streamly.Internal.Data.Stream.StreamD.Nesting -- Copyright : (c) 2018 Composewell Technologies -- (c) Roman Leshchinskiy 2008-2010 -- License : BSD-3-Clause -- Maintainer : streamly@composewell.com -- Stability : experimental -- Portability : GHC -- -- This module contains transformations involving multiple streams, unfolds or -- folds. There are two types of transformations generational or eliminational. -- Generational transformations are like the "Generate" module but they -- generate a stream by combining streams instead of elements. Eliminational -- transformations are like the "Eliminate" module but they transform a stream -- by eliminating parts of the stream instead of eliminating the whole stream. -- -- These combinators involve transformation, generation, elimination so can be -- classified under any of those. -- -- Ultimately these operations should be supported by Unfolds, Pipes and Folds, -- and this module may become redundant. -- The zipWithM combinator in this module has been adapted from the vector -- package (c) Roman Leshchinskiy. -- module Streamly.Internal.Data.Stream.StreamD.Nesting ( -- * Generate -- | Combining streams to generate streams. -- ** Combine Two Streams -- | Functions ending in the shape: -- -- @t m a -> t m a -> t m a@. -- *** Appending -- | Append a stream after another. A special case of concatMap or -- unfoldMany. AppendState(..) , append -- *** Interleaving -- | Interleave elements from two streams alternately. A special case of -- unfoldManyInterleave. , InterleaveState(..) , interleave , interleaveMin , interleaveSuffix , interleaveInfix -- *** Scheduling -- | Execute streams alternately irrespective of whether they generate -- elements or not. Note 'interleave' would execute a stream until it -- yields an element. A special case of unfoldManyRoundRobin. , roundRobin -- interleaveFair?/ParallelFair -- *** Zipping -- | Zip corresponding elements of two streams. , zipWith , zipWithM -- *** Merging -- | Interleave elements from two streams based on a condition. , mergeBy , mergeByM -- ** Combine N Streams -- | Functions generally ending in these shapes: -- -- @ -- concat: f (t m a) -> t m a -- concatMap: (a -> t m b) -> t m a -> t m b -- unfoldMany: Unfold m a b -> t m a -> t m b -- @ -- *** ConcatMap -- | Generate streams by mapping a stream generator on each element of an -- input stream, append the resulting streams and flatten. , concatMap , concatMapM -- *** ConcatUnfold -- | Generate streams by using an unfold on each element of an input -- stream, append the resulting streams and flatten. A special case of -- gintercalate. , unfoldMany , ConcatUnfoldInterleaveState (..) , unfoldManyInterleave , unfoldManyRoundRobin -- *** Interpose -- | Like unfoldMany but intersperses an effect between the streams. A -- special case of gintercalate. , interpose , interposeSuffix -- *** Intercalate -- | Like unfoldMany but intersperses streams from another source between -- the streams from the first source. , gintercalate , gintercalateSuffix -- * Eliminate -- | Folding and Parsing chunks of streams to eliminate nested streams. -- Functions generally ending in these shapes: -- -- @ -- f (Fold m a b) -> t m a -> t m b -- f (Parser m a b) -> t m a -> t m b -- @ -- ** Folding -- | Apply folds on a stream. , foldMany , foldIterateM -- ** Parsing -- | Parsing is opposite to flattening. 'parseMany' is dual to concatMap or -- unfoldMany. concatMap generates a stream from single values in a -- stream and flattens, parseMany does the opposite of flattening by -- splitting the stream and then folds each such split to single value in -- the output stream. , parseMany , parseIterate -- ** Grouping -- | Group segments of a stream and fold. Special case of parsing. , chunksOf , groupsOf2 , groupsBy , groupsRollingBy -- ** Splitting -- | A special case of parsing. , wordsBy , splitOnSeq , splitOnSuffixSeq -- * Transform (Nested Containers) -- | Opposite to compact in ArrayStream , splitInnerBy , splitInnerBySuffix ) where #include "inline.hs" import Control.Exception (assert) import Control.Monad.Catch (MonadThrow, throwM) import Control.Monad.IO.Class (MonadIO(..)) import Data.Bits (shiftR, shiftL, (.|.), (.&.)) #if __GLASGOW_HASKELL__ >= 801 import Data.Functor.Identity ( Identity ) #endif import Data.Word (Word32) import Foreign.Storable (Storable(..)) import Fusion.Plugin.Types (Fuse(..)) import GHC.Types (SPEC(..)) import Streamly.Internal.Data.Array.Foreign.Type (Array(..)) import Streamly.Internal.Data.Fold.Type (Fold(..)) import Streamly.Internal.Data.Parser (ParseError(..)) import Streamly.Internal.Data.SVar.Type (adaptState) import Streamly.Internal.Data.Unfold.Type (Unfold(..)) import qualified Streamly.Internal.Data.Array.Foreign.Type as A import qualified Streamly.Internal.Data.Fold.Type as FL import qualified Streamly.Internal.Data.Parser as PR import qualified Streamly.Internal.Data.Parser.ParserD as PRD import qualified Streamly.Internal.Ring.Foreign as RB import Streamly.Internal.Data.Stream.StreamD.Type import Prelude hiding (concatMap, mapM, zipWith) ------------------------------------------------------------------------------ -- Appending ------------------------------------------------------------------------------ data AppendState s1 s2 = AppendFirst s1 | AppendSecond s2 -- Note that this could be much faster compared to the CPS stream. However, as -- the number of streams being composed increases this may become expensive. -- Need to see where the breaking point is between the two. -- {-# INLINE_NORMAL append #-} append :: Monad m => Stream m a -> Stream m a -> Stream m a append (Stream step1 state1) (Stream step2 state2) = Stream step (AppendFirst state1) where {-# INLINE_LATE step #-} step gst (AppendFirst st) = do r <- step1 gst st return $ case r of Yield a s -> Yield a (AppendFirst s) Skip s -> Skip (AppendFirst s) Stop -> Skip (AppendSecond state2) step gst (AppendSecond st) = do r <- step2 gst st return $ case r of Yield a s -> Yield a (AppendSecond s) Skip s -> Skip (AppendSecond s) Stop -> Stop ------------------------------------------------------------------------------ -- Interleaving ------------------------------------------------------------------------------ data InterleaveState s1 s2 = InterleaveFirst s1 s2 | InterleaveSecond s1 s2 | InterleaveSecondOnly s2 | InterleaveFirstOnly s1 {-# INLINE_NORMAL interleave #-} interleave :: Monad m => Stream m a -> Stream m a -> Stream m a interleave (Stream step1 state1) (Stream step2 state2) = Stream step (InterleaveFirst state1 state2) where {-# INLINE_LATE step #-} step gst (InterleaveFirst st1 st2) = do r <- step1 gst st1 return $ case r of Yield a s -> Yield a (InterleaveSecond s st2) Skip s -> Skip (InterleaveFirst s st2) Stop -> Skip (InterleaveSecondOnly st2) step gst (InterleaveSecond st1 st2) = do r <- step2 gst st2 return $ case r of Yield a s -> Yield a (InterleaveFirst st1 s) Skip s -> Skip (InterleaveSecond st1 s) Stop -> Skip (InterleaveFirstOnly st1) step gst (InterleaveFirstOnly st1) = do r <- step1 gst st1 return $ case r of Yield a s -> Yield a (InterleaveFirstOnly s) Skip s -> Skip (InterleaveFirstOnly s) Stop -> Stop step gst (InterleaveSecondOnly st2) = do r <- step2 gst st2 return $ case r of Yield a s -> Yield a (InterleaveSecondOnly s) Skip s -> Skip (InterleaveSecondOnly s) Stop -> Stop {-# INLINE_NORMAL interleaveMin #-} interleaveMin :: Monad m => Stream m a -> Stream m a -> Stream m a interleaveMin (Stream step1 state1) (Stream step2 state2) = Stream step (InterleaveFirst state1 state2) where {-# INLINE_LATE step #-} step gst (InterleaveFirst st1 st2) = do r <- step1 gst st1 return $ case r of Yield a s -> Yield a (InterleaveSecond s st2) Skip s -> Skip (InterleaveFirst s st2) Stop -> Stop step gst (InterleaveSecond st1 st2) = do r <- step2 gst st2 return $ case r of Yield a s -> Yield a (InterleaveFirst st1 s) Skip s -> Skip (InterleaveSecond st1 s) Stop -> Stop step _ (InterleaveFirstOnly _) = undefined step _ (InterleaveSecondOnly _) = undefined {-# INLINE_NORMAL interleaveSuffix #-} interleaveSuffix :: Monad m => Stream m a -> Stream m a -> Stream m a interleaveSuffix (Stream step1 state1) (Stream step2 state2) = Stream step (InterleaveFirst state1 state2) where {-# INLINE_LATE step #-} step gst (InterleaveFirst st1 st2) = do r <- step1 gst st1 return $ case r of Yield a s -> Yield a (InterleaveSecond s st2) Skip s -> Skip (InterleaveFirst s st2) Stop -> Stop step gst (InterleaveSecond st1 st2) = do r <- step2 gst st2 return $ case r of Yield a s -> Yield a (InterleaveFirst st1 s) Skip s -> Skip (InterleaveSecond st1 s) Stop -> Skip (InterleaveFirstOnly st1) step gst (InterleaveFirstOnly st1) = do r <- step1 gst st1 return $ case r of Yield a s -> Yield a (InterleaveFirstOnly s) Skip s -> Skip (InterleaveFirstOnly s) Stop -> Stop step _ (InterleaveSecondOnly _) = undefined data InterleaveInfixState s1 s2 a = InterleaveInfixFirst s1 s2 | InterleaveInfixSecondBuf s1 s2 | InterleaveInfixSecondYield s1 s2 a | InterleaveInfixFirstYield s1 s2 a | InterleaveInfixFirstOnly s1 {-# INLINE_NORMAL interleaveInfix #-} interleaveInfix :: Monad m => Stream m a -> Stream m a -> Stream m a interleaveInfix (Stream step1 state1) (Stream step2 state2) = Stream step (InterleaveInfixFirst state1 state2) where {-# INLINE_LATE step #-} step gst (InterleaveInfixFirst st1 st2) = do r <- step1 gst st1 return $ case r of Yield a s -> Yield a (InterleaveInfixSecondBuf s st2) Skip s -> Skip (InterleaveInfixFirst s st2) Stop -> Stop step gst (InterleaveInfixSecondBuf st1 st2) = do r <- step2 gst st2 return $ case r of Yield a s -> Skip (InterleaveInfixSecondYield st1 s a) Skip s -> Skip (InterleaveInfixSecondBuf st1 s) Stop -> Skip (InterleaveInfixFirstOnly st1) step gst (InterleaveInfixSecondYield st1 st2 x) = do r <- step1 gst st1 return $ case r of Yield a s -> Yield x (InterleaveInfixFirstYield s st2 a) Skip s -> Skip (InterleaveInfixSecondYield s st2 x) Stop -> Stop step _ (InterleaveInfixFirstYield st1 st2 x) = do return $ Yield x (InterleaveInfixSecondBuf st1 st2) step gst (InterleaveInfixFirstOnly st1) = do r <- step1 gst st1 return $ case r of Yield a s -> Yield a (InterleaveInfixFirstOnly s) Skip s -> Skip (InterleaveInfixFirstOnly s) Stop -> Stop ------------------------------------------------------------------------------ -- Scheduling ------------------------------------------------------------------------------ {-# INLINE_NORMAL roundRobin #-} roundRobin :: Monad m => Stream m a -> Stream m a -> Stream m a roundRobin (Stream step1 state1) (Stream step2 state2) = Stream step (InterleaveFirst state1 state2) where {-# INLINE_LATE step #-} step gst (InterleaveFirst st1 st2) = do r <- step1 gst st1 return $ case r of Yield a s -> Yield a (InterleaveSecond s st2) Skip s -> Skip (InterleaveSecond s st2) Stop -> Skip (InterleaveSecondOnly st2) step gst (InterleaveSecond st1 st2) = do r <- step2 gst st2 return $ case r of Yield a s -> Yield a (InterleaveFirst st1 s) Skip s -> Skip (InterleaveFirst st1 s) Stop -> Skip (InterleaveFirstOnly st1) step gst (InterleaveSecondOnly st2) = do r <- step2 gst st2 return $ case r of Yield a s -> Yield a (InterleaveSecondOnly s) Skip s -> Skip (InterleaveSecondOnly s) Stop -> Stop step gst (InterleaveFirstOnly st1) = do r <- step1 gst st1 return $ case r of Yield a s -> Yield a (InterleaveFirstOnly s) Skip s -> Skip (InterleaveFirstOnly s) Stop -> Stop ------------------------------------------------------------------------------ -- Zipping ------------------------------------------------------------------------------ {-# INLINE_NORMAL zipWithM #-} zipWithM :: Monad m => (a -> b -> m c) -> Stream m a -> Stream m b -> Stream m c zipWithM f (Stream stepa ta) (Stream stepb tb) = Stream step (ta, tb, Nothing) where {-# INLINE_LATE step #-} step gst (sa, sb, Nothing) = do r <- stepa (adaptState gst) sa return $ case r of Yield x sa' -> Skip (sa', sb, Just x) Skip sa' -> Skip (sa', sb, Nothing) Stop -> Stop step gst (sa, sb, Just x) = do r <- stepb (adaptState gst) sb case r of Yield y sb' -> do z <- f x y return $ Yield z (sa, sb', Nothing) Skip sb' -> return $ Skip (sa, sb', Just x) Stop -> return Stop #if __GLASGOW_HASKELL__ >= 801 {-# RULES "zipWithM xs xs" forall f xs. zipWithM @Identity f xs xs = mapM (\x -> f x x) xs #-} #endif {-# INLINE zipWith #-} zipWith :: Monad m => (a -> b -> c) -> Stream m a -> Stream m b -> Stream m c zipWith f = zipWithM (\a b -> return (f a b)) ------------------------------------------------------------------------------ -- Merging ------------------------------------------------------------------------------ {-# INLINE_NORMAL mergeByM #-} mergeByM :: (Monad m) => (a -> a -> m Ordering) -> Stream m a -> Stream m a -> Stream m a mergeByM cmp (Stream stepa ta) (Stream stepb tb) = Stream step (Just ta, Just tb, Nothing, Nothing) where {-# INLINE_LATE step #-} -- one of the values is missing, and the corresponding stream is running step gst (Just sa, sb, Nothing, b) = do r <- stepa gst sa return $ case r of Yield a sa' -> Skip (Just sa', sb, Just a, b) Skip sa' -> Skip (Just sa', sb, Nothing, b) Stop -> Skip (Nothing, sb, Nothing, b) step gst (sa, Just sb, a, Nothing) = do r <- stepb gst sb return $ case r of Yield b sb' -> Skip (sa, Just sb', a, Just b) Skip sb' -> Skip (sa, Just sb', a, Nothing) Stop -> Skip (sa, Nothing, a, Nothing) -- both the values are available step _ (sa, sb, Just a, Just b) = do res <- cmp a b return $ case res of GT -> Yield b (sa, sb, Just a, Nothing) _ -> Yield a (sa, sb, Nothing, Just b) -- one of the values is missing, corresponding stream is done step _ (Nothing, sb, Nothing, Just b) = return $ Yield b (Nothing, sb, Nothing, Nothing) step _ (sa, Nothing, Just a, Nothing) = return $ Yield a (sa, Nothing, Nothing, Nothing) step _ (Nothing, Nothing, Nothing, Nothing) = return Stop {-# INLINE mergeBy #-} mergeBy :: (Monad m) => (a -> a -> Ordering) -> Stream m a -> Stream m a -> Stream m a mergeBy cmp = mergeByM (\a b -> return $ cmp a b) ------------------------------------------------------------------------------ -- Combine N Streams - unfoldMany ------------------------------------------------------------------------------ data ConcatUnfoldInterleaveState o i = ConcatUnfoldInterleaveOuter o [i] | ConcatUnfoldInterleaveInner o [i] | ConcatUnfoldInterleaveInnerL [i] [i] | ConcatUnfoldInterleaveInnerR [i] [i] -- XXX use arrays to store state instead of lists. -- XXX In general we can use different scheduling strategies e.g. how to -- schedule the outer vs inner loop or assigning weights to different streams -- or outer and inner loops. -- After a yield, switch to the next stream. Do not switch streams on Skip. -- Yield from outer stream switches to the inner stream. -- -- There are two choices here, (1) exhaust the outer stream first and then -- start yielding from the inner streams, this is much simpler to implement, -- (2) yield at least one element from an inner stream before going back to -- outer stream and opening the next stream from it. -- -- Ideally, we need some scheduling bias to inner streams vs outer stream. -- Maybe we can configure the behavior. -- {-# INLINE_NORMAL unfoldManyInterleave #-} unfoldManyInterleave :: Monad m => Unfold m a b -> Stream m a -> Stream m b unfoldManyInterleave (Unfold istep inject) (Stream ostep ost) = Stream step (ConcatUnfoldInterleaveOuter ost []) where {-# INLINE_LATE step #-} step gst (ConcatUnfoldInterleaveOuter o ls) = do r <- ostep (adaptState gst) o case r of Yield a o' -> do i <- inject a i `seq` return (Skip (ConcatUnfoldInterleaveInner o' (i : ls))) Skip o' -> return $ Skip (ConcatUnfoldInterleaveOuter o' ls) Stop -> return $ Skip (ConcatUnfoldInterleaveInnerL ls []) step _ (ConcatUnfoldInterleaveInner _ []) = undefined step _ (ConcatUnfoldInterleaveInner o (st:ls)) = do r <- istep st return $ case r of Yield x s -> Yield x (ConcatUnfoldInterleaveOuter o (s:ls)) Skip s -> Skip (ConcatUnfoldInterleaveInner o (s:ls)) Stop -> Skip (ConcatUnfoldInterleaveOuter o ls) step _ (ConcatUnfoldInterleaveInnerL [] []) = return Stop step _ (ConcatUnfoldInterleaveInnerL [] rs) = return $ Skip (ConcatUnfoldInterleaveInnerR [] rs) step _ (ConcatUnfoldInterleaveInnerL (st:ls) rs) = do r <- istep st return $ case r of Yield x s -> Yield x (ConcatUnfoldInterleaveInnerL ls (s:rs)) Skip s -> Skip (ConcatUnfoldInterleaveInnerL (s:ls) rs) Stop -> Skip (ConcatUnfoldInterleaveInnerL ls rs) step _ (ConcatUnfoldInterleaveInnerR [] []) = return Stop step _ (ConcatUnfoldInterleaveInnerR ls []) = return $ Skip (ConcatUnfoldInterleaveInnerL ls []) step _ (ConcatUnfoldInterleaveInnerR ls (st:rs)) = do r <- istep st return $ case r of Yield x s -> Yield x (ConcatUnfoldInterleaveInnerR (s:ls) rs) Skip s -> Skip (ConcatUnfoldInterleaveInnerR ls (s:rs)) Stop -> Skip (ConcatUnfoldInterleaveInnerR ls rs) -- XXX In general we can use different scheduling strategies e.g. how to -- schedule the outer vs inner loop or assigning weights to different streams -- or outer and inner loops. -- -- This could be inefficient if the tasks are too small. -- -- Compared to unfoldManyInterleave this one switches streams on Skips. -- {-# INLINE_NORMAL unfoldManyRoundRobin #-} unfoldManyRoundRobin :: Monad m => Unfold m a b -> Stream m a -> Stream m b unfoldManyRoundRobin (Unfold istep inject) (Stream ostep ost) = Stream step (ConcatUnfoldInterleaveOuter ost []) where {-# INLINE_LATE step #-} step gst (ConcatUnfoldInterleaveOuter o ls) = do r <- ostep (adaptState gst) o case r of Yield a o' -> do i <- inject a i `seq` return (Skip (ConcatUnfoldInterleaveInner o' (i : ls))) Skip o' -> return $ Skip (ConcatUnfoldInterleaveInner o' ls) Stop -> return $ Skip (ConcatUnfoldInterleaveInnerL ls []) step _ (ConcatUnfoldInterleaveInner o []) = return $ Skip (ConcatUnfoldInterleaveOuter o []) step _ (ConcatUnfoldInterleaveInner o (st:ls)) = do r <- istep st return $ case r of Yield x s -> Yield x (ConcatUnfoldInterleaveOuter o (s:ls)) Skip s -> Skip (ConcatUnfoldInterleaveOuter o (s:ls)) Stop -> Skip (ConcatUnfoldInterleaveOuter o ls) step _ (ConcatUnfoldInterleaveInnerL [] []) = return Stop step _ (ConcatUnfoldInterleaveInnerL [] rs) = return $ Skip (ConcatUnfoldInterleaveInnerR [] rs) step _ (ConcatUnfoldInterleaveInnerL (st:ls) rs) = do r <- istep st return $ case r of Yield x s -> Yield x (ConcatUnfoldInterleaveInnerL ls (s:rs)) Skip s -> Skip (ConcatUnfoldInterleaveInnerL ls (s:rs)) Stop -> Skip (ConcatUnfoldInterleaveInnerL ls rs) step _ (ConcatUnfoldInterleaveInnerR [] []) = return Stop step _ (ConcatUnfoldInterleaveInnerR ls []) = return $ Skip (ConcatUnfoldInterleaveInnerL ls []) step _ (ConcatUnfoldInterleaveInnerR ls (st:rs)) = do r <- istep st return $ case r of Yield x s -> Yield x (ConcatUnfoldInterleaveInnerR (s:ls) rs) Skip s -> Skip (ConcatUnfoldInterleaveInnerR (s:ls) rs) Stop -> Skip (ConcatUnfoldInterleaveInnerR ls rs) ------------------------------------------------------------------------------ -- Combine N Streams - interpose ------------------------------------------------------------------------------ {-# ANN type InterposeSuffixState Fuse #-} data InterposeSuffixState s1 i1 = InterposeSuffixFirst s1 -- | InterposeSuffixFirstYield s1 i1 | InterposeSuffixFirstInner s1 i1 | InterposeSuffixSecond s1 -- Note that if an unfolded layer turns out to be nil we still emit the -- separator effect. An alternate behavior could be to emit the separator -- effect only if at least one element has been yielded by the unfolding. -- However, that becomes a bit complicated, so we have chosen the former -- behvaior for now. {-# INLINE_NORMAL interposeSuffix #-} interposeSuffix :: Monad m => m c -> Unfold m b c -> Stream m b -> Stream m c interposeSuffix action (Unfold istep1 inject1) (Stream step1 state1) = Stream step (InterposeSuffixFirst state1) where {-# INLINE_LATE step #-} step gst (InterposeSuffixFirst s1) = do r <- step1 (adaptState gst) s1 case r of Yield a s -> do i <- inject1 a i `seq` return (Skip (InterposeSuffixFirstInner s i)) -- i `seq` return (Skip (InterposeSuffixFirstYield s i)) Skip s -> return $ Skip (InterposeSuffixFirst s) Stop -> return Stop {- step _ (InterposeSuffixFirstYield s1 i1) = do r <- istep1 i1 return $ case r of Yield x i' -> Yield x (InterposeSuffixFirstInner s1 i') Skip i' -> Skip (InterposeSuffixFirstYield s1 i') Stop -> Skip (InterposeSuffixFirst s1) -} step _ (InterposeSuffixFirstInner s1 i1) = do r <- istep1 i1 return $ case r of Yield x i' -> Yield x (InterposeSuffixFirstInner s1 i') Skip i' -> Skip (InterposeSuffixFirstInner s1 i') Stop -> Skip (InterposeSuffixSecond s1) step _ (InterposeSuffixSecond s1) = do r <- action return $ Yield r (InterposeSuffixFirst s1) {-# ANN type InterposeState Fuse #-} data InterposeState s1 i1 a = InterposeFirst s1 -- | InterposeFirstYield s1 i1 | InterposeFirstInner s1 i1 | InterposeFirstInject s1 -- | InterposeFirstBuf s1 i1 | InterposeSecondYield s1 i1 -- -- | InterposeSecondYield s1 i1 a -- -- | InterposeFirstResume s1 i1 a -- Note that this only interposes the pure values, we may run many effects to -- generate those values as some effects may not generate anything (Skip). {-# INLINE_NORMAL interpose #-} interpose :: Monad m => m c -> Unfold m b c -> Stream m b -> Stream m c interpose action (Unfold istep1 inject1) (Stream step1 state1) = Stream step (InterposeFirst state1) where {-# INLINE_LATE step #-} step gst (InterposeFirst s1) = do r <- step1 (adaptState gst) s1 case r of Yield a s -> do i <- inject1 a i `seq` return (Skip (InterposeFirstInner s i)) -- i `seq` return (Skip (InterposeFirstYield s i)) Skip s -> return $ Skip (InterposeFirst s) Stop -> return Stop {- step _ (InterposeFirstYield s1 i1) = do r <- istep1 i1 return $ case r of Yield x i' -> Yield x (InterposeFirstInner s1 i') Skip i' -> Skip (InterposeFirstYield s1 i') Stop -> Skip (InterposeFirst s1) -} step _ (InterposeFirstInner s1 i1) = do r <- istep1 i1 return $ case r of Yield x i' -> Yield x (InterposeFirstInner s1 i') Skip i' -> Skip (InterposeFirstInner s1 i') Stop -> Skip (InterposeFirstInject s1) step gst (InterposeFirstInject s1) = do r <- step1 (adaptState gst) s1 case r of Yield a s -> do i <- inject1 a -- i `seq` return (Skip (InterposeFirstBuf s i)) i `seq` return (Skip (InterposeSecondYield s i)) Skip s -> return $ Skip (InterposeFirstInject s) Stop -> return Stop {- step _ (InterposeFirstBuf s1 i1) = do r <- istep1 i1 return $ case r of Yield x i' -> Skip (InterposeSecondYield s1 i' x) Skip i' -> Skip (InterposeFirstBuf s1 i') Stop -> Stop -} {- step _ (InterposeSecondYield s1 i1 v) = do r <- action return $ Yield r (InterposeFirstResume s1 i1 v) -} step _ (InterposeSecondYield s1 i1) = do r <- action return $ Yield r (InterposeFirstInner s1 i1) {- step _ (InterposeFirstResume s1 i1 v) = do return $ Yield v (InterposeFirstInner s1 i1) -} ------------------------------------------------------------------------------ -- Combine N Streams - intercalate ------------------------------------------------------------------------------ data ICUState s1 s2 i1 i2 = ICUFirst s1 s2 | ICUSecond s1 s2 | ICUSecondOnly s2 | ICUFirstOnly s1 | ICUFirstInner s1 s2 i1 | ICUSecondInner s1 s2 i2 | ICUFirstOnlyInner s1 i1 | ICUSecondOnlyInner s2 i2 -- | Interleave streams (full streams, not the elements) unfolded from two -- input streams and concat. Stop when the first stream stops. If the second -- stream ends before the first one then first stream still keeps running alone -- without any interleaving with the second stream. -- -- [a1, a2, ... an] [b1, b2 ...] -- => [streamA1, streamA2, ... streamAn] [streamB1, streamB2, ...] -- => [streamA1, streamB1, streamA2...StreamAn, streamBn] -- => [a11, a12, ...a1j, b11, b12, ...b1k, a21, a22, ...] -- {-# INLINE_NORMAL gintercalateSuffix #-} gintercalateSuffix :: Monad m => Unfold m a c -> Stream m a -> Unfold m b c -> Stream m b -> Stream m c gintercalateSuffix (Unfold istep1 inject1) (Stream step1 state1) (Unfold istep2 inject2) (Stream step2 state2) = Stream step (ICUFirst state1 state2) where {-# INLINE_LATE step #-} step gst (ICUFirst s1 s2) = do r <- step1 (adaptState gst) s1 case r of Yield a s -> do i <- inject1 a i `seq` return (Skip (ICUFirstInner s s2 i)) Skip s -> return $ Skip (ICUFirst s s2) Stop -> return Stop step gst (ICUFirstOnly s1) = do r <- step1 (adaptState gst) s1 case r of Yield a s -> do i <- inject1 a i `seq` return (Skip (ICUFirstOnlyInner s i)) Skip s -> return $ Skip (ICUFirstOnly s) Stop -> return Stop step _ (ICUFirstInner s1 s2 i1) = do r <- istep1 i1 return $ case r of Yield x i' -> Yield x (ICUFirstInner s1 s2 i') Skip i' -> Skip (ICUFirstInner s1 s2 i') Stop -> Skip (ICUSecond s1 s2) step _ (ICUFirstOnlyInner s1 i1) = do r <- istep1 i1 return $ case r of Yield x i' -> Yield x (ICUFirstOnlyInner s1 i') Skip i' -> Skip (ICUFirstOnlyInner s1 i') Stop -> Skip (ICUFirstOnly s1) step gst (ICUSecond s1 s2) = do r <- step2 (adaptState gst) s2 case r of Yield a s -> do i <- inject2 a i `seq` return (Skip (ICUSecondInner s1 s i)) Skip s -> return $ Skip (ICUSecond s1 s) Stop -> return $ Skip (ICUFirstOnly s1) step _ (ICUSecondInner s1 s2 i2) = do r <- istep2 i2 return $ case r of Yield x i' -> Yield x (ICUSecondInner s1 s2 i') Skip i' -> Skip (ICUSecondInner s1 s2 i') Stop -> Skip (ICUFirst s1 s2) step _ (ICUSecondOnly _s2) = undefined step _ (ICUSecondOnlyInner _s2 _i2) = undefined data ICALState s1 s2 i1 i2 a = ICALFirst s1 s2 -- | ICALFirstYield s1 s2 i1 | ICALFirstInner s1 s2 i1 | ICALFirstOnly s1 | ICALFirstOnlyInner s1 i1 | ICALSecondInject s1 s2 | ICALFirstInject s1 s2 i2 -- | ICALFirstBuf s1 s2 i1 i2 | ICALSecondInner s1 s2 i1 i2 -- -- | ICALSecondInner s1 s2 i1 i2 a -- -- | ICALFirstResume s1 s2 i1 i2 a -- | Interleave streams (full streams, not the elements) unfolded from two -- input streams and concat. Stop when the first stream stops. If the second -- stream ends before the first one then first stream still keeps running alone -- without any interleaving with the second stream. -- -- [a1, a2, ... an] [b1, b2 ...] -- => [streamA1, streamA2, ... streamAn] [streamB1, streamB2, ...] -- => [streamA1, streamB1, streamA2...StreamAn, streamBn] -- => [a11, a12, ...a1j, b11, b12, ...b1k, a21, a22, ...] -- {-# INLINE_NORMAL gintercalate #-} gintercalate :: Monad m => Unfold m a c -> Stream m a -> Unfold m b c -> Stream m b -> Stream m c gintercalate (Unfold istep1 inject1) (Stream step1 state1) (Unfold istep2 inject2) (Stream step2 state2) = Stream step (ICALFirst state1 state2) where {-# INLINE_LATE step #-} step gst (ICALFirst s1 s2) = do r <- step1 (adaptState gst) s1 case r of Yield a s -> do i <- inject1 a i `seq` return (Skip (ICALFirstInner s s2 i)) -- i `seq` return (Skip (ICALFirstYield s s2 i)) Skip s -> return $ Skip (ICALFirst s s2) Stop -> return Stop {- step _ (ICALFirstYield s1 s2 i1) = do r <- istep1 i1 return $ case r of Yield x i' -> Yield x (ICALFirstInner s1 s2 i') Skip i' -> Skip (ICALFirstYield s1 s2 i') Stop -> Skip (ICALFirst s1 s2) -} step _ (ICALFirstInner s1 s2 i1) = do r <- istep1 i1 return $ case r of Yield x i' -> Yield x (ICALFirstInner s1 s2 i') Skip i' -> Skip (ICALFirstInner s1 s2 i') Stop -> Skip (ICALSecondInject s1 s2) step gst (ICALFirstOnly s1) = do r <- step1 (adaptState gst) s1 case r of Yield a s -> do i <- inject1 a i `seq` return (Skip (ICALFirstOnlyInner s i)) Skip s -> return $ Skip (ICALFirstOnly s) Stop -> return Stop step _ (ICALFirstOnlyInner s1 i1) = do r <- istep1 i1 return $ case r of Yield x i' -> Yield x (ICALFirstOnlyInner s1 i') Skip i' -> Skip (ICALFirstOnlyInner s1 i') Stop -> Skip (ICALFirstOnly s1) -- We inject the second stream even before checking if the first stream -- would yield any more elements. There is no clear choice whether we -- should do this before or after that. Doing it after may make the state -- machine a bit simpler though. step gst (ICALSecondInject s1 s2) = do r <- step2 (adaptState gst) s2 case r of Yield a s -> do i <- inject2 a i `seq` return (Skip (ICALFirstInject s1 s i)) Skip s -> return $ Skip (ICALSecondInject s1 s) Stop -> return $ Skip (ICALFirstOnly s1) step gst (ICALFirstInject s1 s2 i2) = do r <- step1 (adaptState gst) s1 case r of Yield a s -> do i <- inject1 a i `seq` return (Skip (ICALSecondInner s s2 i i2)) -- i `seq` return (Skip (ICALFirstBuf s s2 i i2)) Skip s -> return $ Skip (ICALFirstInject s s2 i2) Stop -> return Stop {- step _ (ICALFirstBuf s1 s2 i1 i2) = do r <- istep1 i1 return $ case r of Yield x i' -> Skip (ICALSecondInner s1 s2 i' i2 x) Skip i' -> Skip (ICALFirstBuf s1 s2 i' i2) Stop -> Stop step _ (ICALSecondInner s1 s2 i1 i2 v) = do r <- istep2 i2 return $ case r of Yield x i' -> Yield x (ICALSecondInner s1 s2 i1 i' v) Skip i' -> Skip (ICALSecondInner s1 s2 i1 i' v) Stop -> Skip (ICALFirstResume s1 s2 i1 i2 v) -} step _ (ICALSecondInner s1 s2 i1 i2) = do r <- istep2 i2 return $ case r of Yield x i' -> Yield x (ICALSecondInner s1 s2 i1 i') Skip i' -> Skip (ICALSecondInner s1 s2 i1 i') Stop -> Skip (ICALFirstInner s1 s2 i1) -- Stop -> Skip (ICALFirstResume s1 s2 i1 i2) {- step _ (ICALFirstResume s1 s2 i1 i2 x) = do return $ Yield x (ICALFirstInner s1 s2 i1 i2) -} ------------------------------------------------------------------------------ -- Folding ------------------------------------------------------------------------------ {-# ANN type FIterState Fuse #-} data FIterState s f m a b = FIterInit s f | forall fs. FIterStream s (fs -> a -> m (FL.Step fs b)) fs (fs -> m b) | FIterYield b (FIterState s f m a b) | FIterStop {-# INLINE_NORMAL foldIterateM #-} foldIterateM :: Monad m => (b -> m (FL.Fold m a b)) -> b -> Stream m a -> Stream m b foldIterateM func seed0 (Stream step state) = Stream stepOuter (FIterInit state seed0) where {-# INLINE iterStep #-} iterStep from st fstep extract = do res <- from return $ Skip $ case res of FL.Partial fs -> FIterStream st fstep fs extract FL.Done fb -> FIterYield fb $ FIterInit st fb {-# INLINE_LATE stepOuter #-} stepOuter _ (FIterInit st seed) = do (FL.Fold fstep initial extract) <- func seed iterStep initial st fstep extract stepOuter gst (FIterStream st fstep fs extract) = do r <- step (adaptState gst) st case r of Yield x s -> do iterStep (fstep fs x) s fstep extract Skip s -> return $ Skip $ FIterStream s fstep fs extract Stop -> do b <- extract fs return $ Skip $ FIterYield b FIterStop stepOuter _ (FIterYield a next) = return $ Yield a next stepOuter _ FIterStop = return Stop ------------------------------------------------------------------------------ -- Parsing ------------------------------------------------------------------------------ {-# ANN type ParseChunksState Fuse #-} data ParseChunksState x inpBuf st pst = ParseChunksInit inpBuf st | ParseChunksInitLeftOver inpBuf | ParseChunksStream st inpBuf !pst | ParseChunksBuf inpBuf st inpBuf !pst | ParseChunksYield x (ParseChunksState x inpBuf st pst) {-# INLINE_NORMAL parseMany #-} parseMany :: MonadThrow m => PRD.Parser m a b -> Stream m a -> Stream m b parseMany (PRD.Parser pstep initial extract) (Stream step state) = Stream stepOuter (ParseChunksInit [] state) where {-# INLINE_LATE stepOuter #-} -- Buffer is empty, get the first element from the stream, initialize the -- fold and then go to stream processing loop. stepOuter gst (ParseChunksInit [] st) = do r <- step (adaptState gst) st case r of Yield x s -> do res <- initial case res of PRD.IPartial ps -> return $ Skip $ ParseChunksBuf [x] s [] ps PRD.IDone pb -> let next = ParseChunksInit [x] s in return $ Skip $ ParseChunksYield pb next PRD.IError err -> throwM $ ParseError err Skip s -> return $ Skip $ ParseChunksInit [] s Stop -> return Stop -- Buffer is not empty, go to buffered processing loop stepOuter _ (ParseChunksInit src st) = do res <- initial case res of PRD.IPartial ps -> return $ Skip $ ParseChunksBuf src st [] ps PRD.IDone pb -> let next = ParseChunksInit src st in return $ Skip $ ParseChunksYield pb next PRD.IError err -> throwM $ ParseError err -- XXX we just discard any leftover input at the end stepOuter _ (ParseChunksInitLeftOver _) = return Stop -- Buffer is empty, process elements from the stream stepOuter gst (ParseChunksStream st buf pst) = do r <- step (adaptState gst) st case r of Yield x s -> do pRes <- pstep pst x case pRes of PR.Partial 0 pst1 -> return $ Skip $ ParseChunksStream s [] pst1 PR.Partial n pst1 -> do assert (n <= length (x:buf)) (return ()) let src0 = Prelude.take n (x:buf) src = Prelude.reverse src0 return $ Skip $ ParseChunksBuf src s [] pst1 PR.Continue 0 pst1 -> return $ Skip $ ParseChunksStream s (x:buf) pst1 PR.Continue n pst1 -> do assert (n <= length (x:buf)) (return ()) let (src0, buf1) = splitAt n (x:buf) src = Prelude.reverse src0 return $ Skip $ ParseChunksBuf src s buf1 pst1 PR.Done 0 b -> do return $ Skip $ ParseChunksYield b (ParseChunksInit [] s) PR.Done n b -> do assert (n <= length (x:buf)) (return ()) let src = Prelude.reverse (Prelude.take n (x:buf)) return $ Skip $ ParseChunksYield b (ParseChunksInit src s) PR.Error err -> throwM $ ParseError err Skip s -> return $ Skip $ ParseChunksStream s buf pst Stop -> do b <- extract pst let src = Prelude.reverse buf return $ Skip $ ParseChunksYield b (ParseChunksInitLeftOver src) -- go back to stream processing mode stepOuter _ (ParseChunksBuf [] s buf pst) = return $ Skip $ ParseChunksStream s buf pst -- buffered processing loop stepOuter _ (ParseChunksBuf (x:xs) s buf pst) = do pRes <- pstep pst x case pRes of PR.Partial 0 pst1 -> return $ Skip $ ParseChunksBuf xs s [] pst1 PR.Partial n pst1 -> do assert (n <= length (x:buf)) (return ()) let src0 = Prelude.take n (x:buf) src = Prelude.reverse src0 ++ xs return $ Skip $ ParseChunksBuf src s [] pst1 PR.Continue 0 pst1 -> return $ Skip $ ParseChunksBuf xs s (x:buf) pst1 PR.Continue n pst1 -> do assert (n <= length (x:buf)) (return ()) let (src0, buf1) = splitAt n (x:buf) src = Prelude.reverse src0 ++ xs return $ Skip $ ParseChunksBuf src s buf1 pst1 PR.Done 0 b -> return $ Skip $ ParseChunksYield b (ParseChunksInit xs s) PR.Done n b -> do assert (n <= length (x:buf)) (return ()) let src = Prelude.reverse (Prelude.take n (x:buf)) ++ xs return $ Skip $ ParseChunksYield b (ParseChunksInit src s) PR.Error err -> throwM $ ParseError err stepOuter _ (ParseChunksYield a next) = return $ Yield a next {-# ANN type ConcatParseState Fuse #-} data ConcatParseState b inpBuf st p m a = ConcatParseInit inpBuf st p | ConcatParseInitLeftOver inpBuf | forall s. ConcatParseStream st inpBuf (s -> a -> m (PRD.Step s b)) s (s -> m b) | forall s. ConcatParseBuf inpBuf st inpBuf (s -> a -> m (PRD.Step s b)) s (s -> m b) | ConcatParseYield b (ConcatParseState b inpBuf st p m a) {-# INLINE_NORMAL parseIterate #-} parseIterate :: MonadThrow m => (b -> PRD.Parser m a b) -> b -> Stream m a -> Stream m b parseIterate func seed (Stream step state) = Stream stepOuter (ConcatParseInit [] state (func seed)) where {-# INLINE_LATE stepOuter #-} -- Buffer is empty, go to stream processing loop stepOuter _ (ConcatParseInit [] st (PRD.Parser pstep initial extract)) = do res <- initial case res of PRD.IPartial ps -> return $ Skip $ ConcatParseStream st [] pstep ps extract PRD.IDone pb -> let next = ConcatParseInit [] st (func pb) in return $ Skip $ ConcatParseYield pb next PRD.IError err -> throwM $ ParseError err -- Buffer is not empty, go to buffered processing loop stepOuter _ (ConcatParseInit src st (PRD.Parser pstep initial extract)) = do res <- initial case res of PRD.IPartial ps -> return $ Skip $ ConcatParseBuf src st [] pstep ps extract PRD.IDone pb -> let next = ConcatParseInit src st (func pb) in return $ Skip $ ConcatParseYield pb next PRD.IError err -> throwM $ ParseError err -- XXX we just discard any leftover input at the end stepOuter _ (ConcatParseInitLeftOver _) = return Stop -- Buffer is empty process elements from the stream stepOuter gst (ConcatParseStream st buf pstep pst extract) = do r <- step (adaptState gst) st case r of Yield x s -> do pRes <- pstep pst x case pRes of PR.Partial 0 pst1 -> return $ Skip $ ConcatParseStream s [] pstep pst1 extract PR.Partial n pst1 -> do assert (n <= length (x:buf)) (return ()) let src0 = Prelude.take n (x:buf) src = Prelude.reverse src0 return $ Skip $ ConcatParseBuf src s [] pstep pst1 extract -- PR.Continue 0 pst1 -> -- return $ Skip $ ConcatParseStream s (x:buf) pst1 PR.Continue n pst1 -> do assert (n <= length (x:buf)) (return ()) let (src0, buf1) = splitAt n (x:buf) src = Prelude.reverse src0 return $ Skip $ ConcatParseBuf src s buf1 pstep pst1 extract -- XXX Specialize for Stop 0 common case? PR.Done n b -> do assert (n <= length (x:buf)) (return ()) let src = Prelude.reverse (Prelude.take n (x:buf)) return $ Skip $ ConcatParseYield b (ConcatParseInit src s (func b)) PR.Error err -> throwM $ ParseError err Skip s -> return $ Skip $ ConcatParseStream s buf pstep pst extract Stop -> do b <- extract pst let src = Prelude.reverse buf return $ Skip $ ConcatParseYield b (ConcatParseInitLeftOver src) -- go back to stream processing mode stepOuter _ (ConcatParseBuf [] s buf pstep ps extract) = return $ Skip $ ConcatParseStream s buf pstep ps extract -- buffered processing loop stepOuter _ (ConcatParseBuf (x:xs) s buf pstep pst extract) = do pRes <- pstep pst x case pRes of PR.Partial 0 pst1 -> return $ Skip $ ConcatParseBuf xs s [] pstep pst1 extract PR.Partial n pst1 -> do assert (n <= length (x:buf)) (return ()) let src0 = Prelude.take n (x:buf) src = Prelude.reverse src0 ++ xs return $ Skip $ ConcatParseBuf src s [] pstep pst1 extract -- PR.Continue 0 pst1 -> return $ Skip $ ConcatParseBuf xs s (x:buf) pst1 PR.Continue n pst1 -> do assert (n <= length (x:buf)) (return ()) let (src0, buf1) = splitAt n (x:buf) src = Prelude.reverse src0 ++ xs return $ Skip $ ConcatParseBuf src s buf1 pstep pst1 extract -- XXX Specialize for Stop 0 common case? PR.Done n b -> do assert (n <= length (x:buf)) (return ()) let src = Prelude.reverse (Prelude.take n (x:buf)) ++ xs return $ Skip $ ConcatParseYield b (ConcatParseInit src s (func b)) PR.Error err -> throwM $ ParseError err stepOuter _ (ConcatParseYield a next) = return $ Yield a next ------------------------------------------------------------------------------ -- Grouping ------------------------------------------------------------------------------ data GroupByState st fs a b = GroupingInit st | GroupingDo st !fs | GroupingInitWith st !a | GroupingDoWith st !fs !a | GroupingYield !b (GroupByState st fs a b) | GroupingDone {-# INLINE_NORMAL groupsBy #-} groupsBy :: Monad m => (a -> a -> Bool) -> Fold m a b -> Stream m a -> Stream m b {- groupsBy eq fld = parseMany (PRD.groupBy eq fld) -} groupsBy cmp (Fold fstep initial done) (Stream step state) = Stream stepOuter (GroupingInit state) where {-# INLINE_LATE stepOuter #-} stepOuter _ (GroupingInit st) = do -- XXX Note that if the stream stops without yielding a single element -- in the group we discard the "initial" effect. res <- initial return $ case res of FL.Partial s -> Skip $ GroupingDo st s FL.Done b -> Yield b $ GroupingInit st stepOuter gst (GroupingDo st fs) = do res <- step (adaptState gst) st case res of Yield x s -> do r <- fstep fs x case r of FL.Partial fs1 -> go SPEC x s fs1 FL.Done b -> return $ Yield b (GroupingInit s) Skip s -> return $ Skip $ GroupingDo s fs Stop -> return Stop where go !_ prev stt !acc = do res <- step (adaptState gst) stt case res of Yield x s -> do if cmp x prev then do r <- fstep acc x case r of FL.Partial fs1 -> go SPEC prev s fs1 FL.Done b -> return $ Yield b (GroupingInit s) else do r <- done acc return $ Yield r (GroupingInitWith s x) Skip s -> go SPEC prev s acc Stop -> done acc >>= \r -> return $ Yield r GroupingDone stepOuter _ (GroupingInitWith st x) = do res <- initial return $ case res of FL.Partial s -> Skip $ GroupingDoWith st s x FL.Done b -> Yield b $ GroupingInitWith st x stepOuter gst (GroupingDoWith st fs prev) = do res <- fstep fs prev case res of FL.Partial fs1 -> go SPEC st fs1 FL.Done b -> return $ Yield b (GroupingInit st) where -- XXX code duplicated from the previous equation go !_ stt !acc = do res <- step (adaptState gst) stt case res of Yield x s -> do if cmp x prev then do r <- fstep acc x case r of FL.Partial fs1 -> go SPEC s fs1 FL.Done b -> return $ Yield b (GroupingInit s) else do r <- done acc return $ Yield r (GroupingInitWith s x) Skip s -> go SPEC s acc Stop -> done acc >>= \r -> return $ Yield r GroupingDone stepOuter _ (GroupingYield _ _) = error "groupsBy: Unreachable" stepOuter _ GroupingDone = return Stop {-# INLINE_NORMAL groupsRollingBy #-} groupsRollingBy :: Monad m => (a -> a -> Bool) -> Fold m a b -> Stream m a -> Stream m b {- groupsRollingBy eq fld = parseMany (PRD.groupByRolling eq fld) -} groupsRollingBy cmp (Fold fstep initial done) (Stream step state) = Stream stepOuter (GroupingInit state) where {-# INLINE_LATE stepOuter #-} stepOuter _ (GroupingInit st) = do -- XXX Note that if the stream stops without yielding a single element -- in the group we discard the "initial" effect. res <- initial return $ case res of FL.Partial fs -> Skip $ GroupingDo st fs FL.Done fb -> Yield fb $ GroupingInit st stepOuter gst (GroupingDo st fs) = do res <- step (adaptState gst) st case res of Yield x s -> do r <- fstep fs x case r of FL.Partial fs1 -> go SPEC x s fs1 FL.Done fb -> return $ Yield fb (GroupingInit s) Skip s -> return $ Skip $ GroupingDo s fs Stop -> return Stop where go !_ prev stt !acc = do res <- step (adaptState gst) stt case res of Yield x s -> do if cmp prev x then do r <- fstep acc x case r of FL.Partial fs1 -> go SPEC x s fs1 FL.Done b -> return $ Yield b (GroupingInit s) else do r <- done acc return $ Yield r (GroupingInitWith s x) Skip s -> go SPEC prev s acc Stop -> done acc >>= \r -> return $ Yield r GroupingDone stepOuter _ (GroupingInitWith st x) = do res <- initial return $ case res of FL.Partial s -> Skip $ GroupingDoWith st s x FL.Done b -> Yield b $ GroupingInitWith st x stepOuter gst (GroupingDoWith st fs previous) = do res <- fstep fs previous case res of FL.Partial s -> go SPEC previous st s FL.Done b -> return $ Yield b (GroupingInit st) where -- XXX GHC: groupsBy has one less parameter in this go loop and it -- fuses. However, groupsRollingBy does not fuse, removing the prev -- parameter makes it fuse. Something needs to be fixed in GHC. The -- workaround for this is noted in the comments below. go !_ prev !stt !acc = do res <- step (adaptState gst) stt case res of Yield x s -> do if cmp prev x then do r <- fstep acc x case r of FL.Partial fs1 -> go SPEC x s fs1 FL.Done b -> return $ Yield b (GroupingInit st) else do {- r <- done acc return $ Yield r (GroupingInitWith s x) -} -- The code above does not let groupBy fuse. We use the -- alternative code below instead. Instead of jumping -- to GroupingInitWith state, we unroll the code of -- GroupingInitWith state here to help GHC with stream -- fusion. result <- initial r <- done acc return $ Yield r $ case result of FL.Partial fsi -> GroupingDoWith s fsi x FL.Done b -> GroupingYield b (GroupingInit s) Skip s -> go SPEC prev s acc Stop -> done acc >>= \r -> return $ Yield r GroupingDone stepOuter _ (GroupingYield r next) = return $ Yield r next stepOuter _ GroupingDone = return Stop ------------------------------------------------------------------------------ -- Splitting - by a predicate ------------------------------------------------------------------------------ data WordsByState st fs b = WordsByInit st | WordsByDo st !fs | WordsByDone | WordsByYield !b (WordsByState st fs b) {-# INLINE_NORMAL wordsBy #-} wordsBy :: Monad m => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b wordsBy predicate (Fold fstep initial done) (Stream step state) = Stream stepOuter (WordsByInit state) where {-# INLINE_LATE stepOuter #-} stepOuter _ (WordsByInit st) = do res <- initial return $ case res of FL.Partial s -> Skip $ WordsByDo st s FL.Done b -> Yield b (WordsByInit st) stepOuter gst (WordsByDo st fs) = do res <- step (adaptState gst) st case res of Yield x s -> do if predicate x then do resi <- initial return $ case resi of FL.Partial fs1 -> Skip $ WordsByDo s fs1 FL.Done b -> Yield b (WordsByInit s) else do r <- fstep fs x case r of FL.Partial fs1 -> go SPEC s fs1 FL.Done b -> return $ Yield b (WordsByInit s) Skip s -> return $ Skip $ WordsByDo s fs Stop -> return Stop where go !_ stt !acc = do res <- step (adaptState gst) stt case res of Yield x s -> do if predicate x then do {- r <- done acc return $ Yield r (WordsByInit s) -} -- The above code does not fuse well. Need to check why -- GHC is not able to simplify it well. Using the code -- below, instead of jumping through the WordsByInit -- state always, we directly go to WordsByDo state in -- the common case of Partial. resi <- initial r <- done acc return $ Yield r $ case resi of FL.Partial fs1 -> WordsByDo s fs1 FL.Done b -> WordsByYield b (WordsByInit s) else do r <- fstep acc x case r of FL.Partial fs1 -> go SPEC s fs1 FL.Done b -> return $ Yield b (WordsByInit s) Skip s -> go SPEC s acc Stop -> done acc >>= \r -> return $ Yield r WordsByDone stepOuter _ WordsByDone = return Stop stepOuter _ (WordsByYield b next) = return $ Yield b next ------------------------------------------------------------------------------ -- Splitting on a sequence ------------------------------------------------------------------------------ -- String search algorithms: -- http://www-igm.univ-mlv.fr/~lecroq/string/index.html {- -- TODO can we unify the splitting operations using a splitting configuration -- like in the split package. -- data SplitStyle = Infix | Suffix | Prefix deriving (Eq, Show) data SplitOptions = SplitOptions { style :: SplitStyle , withSep :: Bool -- ^ keep the separators in output -- , compact :: Bool -- ^ treat multiple consecutive separators as one -- , trimHead :: Bool -- ^ drop blank at head -- , trimTail :: Bool -- ^ drop blank at tail } -} -- XXX using "fs" as the last arg in Constructors may simplify the code a bit, -- because we can use the constructor directly without having to create "jump" -- functions. {-# ANN type SplitOnSeqState Fuse #-} data SplitOnSeqState rb rh ck w fs s b x = SplitOnSeqInit | SplitOnSeqYield b (SplitOnSeqState rb rh ck w fs s b x) | SplitOnSeqDone | SplitOnSeqEmpty !fs s | SplitOnSeqSingle !fs s x | SplitOnSeqWordInit !fs s | SplitOnSeqWordLoop !w s !fs | SplitOnSeqWordDone Int !fs !w | SplitOnSeqKRInit Int !fs s rb !rh | SplitOnSeqKRLoop fs s rb !rh !ck | SplitOnSeqKRCheck fs s rb !rh | SplitOnSeqKRDone Int !fs rb !rh | SplitOnSeqReinit (fs -> SplitOnSeqState rb rh ck w fs s b x) {-# INLINE_NORMAL splitOnSeq #-} splitOnSeq :: forall m a b. (MonadIO m, Storable a, Enum a, Eq a) => Array a -> Fold m a b -> Stream m a -> Stream m b splitOnSeq patArr (Fold fstep initial done) (Stream step state) = Stream stepOuter SplitOnSeqInit where patLen = A.length patArr maxIndex = patLen - 1 elemBits = sizeOf (undefined :: a) * 8 -- For word pattern case wordMask :: Word wordMask = (1 `shiftL` (elemBits * patLen)) - 1 elemMask :: Word elemMask = (1 `shiftL` elemBits) - 1 wordPat :: Word wordPat = wordMask .&. A.foldl' addToWord 0 patArr addToWord wd a = (wd `shiftL` elemBits) .|. fromIntegral (fromEnum a) -- For Rabin-Karp search k = 2891336453 :: Word32 coeff = k ^ patLen addCksum cksum a = cksum * k + fromIntegral (fromEnum a) deltaCksum cksum old new = addCksum cksum new - coeff * fromIntegral (fromEnum old) -- XXX shall we use a random starting hash or 1 instead of 0? patHash = A.foldl' addCksum 0 patArr skip = return . Skip nextAfterInit nextGen stepRes = case stepRes of FL.Partial s -> nextGen s FL.Done b -> SplitOnSeqYield b (SplitOnSeqReinit nextGen) {-# INLINE yieldProceed #-} yieldProceed nextGen fs = initial >>= skip . SplitOnSeqYield fs . nextAfterInit nextGen {-# INLINE_LATE stepOuter #-} stepOuter _ SplitOnSeqInit = do res <- initial case res of FL.Partial acc -> if patLen == 0 then return $ Skip $ SplitOnSeqEmpty acc state else if patLen == 1 then do pat <- liftIO $ A.unsafeIndexIO patArr 0 return $ Skip $ SplitOnSeqSingle acc state pat else if sizeOf (undefined :: a) * patLen <= sizeOf (undefined :: Word) then return $ Skip $ SplitOnSeqWordInit acc state else do (rb, rhead) <- liftIO $ RB.new patLen skip $ SplitOnSeqKRInit 0 acc state rb rhead FL.Done b -> skip $ SplitOnSeqYield b SplitOnSeqInit stepOuter _ (SplitOnSeqYield x next) = return $ Yield x next --------------------------- -- Checkpoint --------------------------- stepOuter _ (SplitOnSeqReinit nextGen) = initial >>= skip . nextAfterInit nextGen --------------------------- -- Empty pattern --------------------------- stepOuter gst (SplitOnSeqEmpty acc st) = do res <- step (adaptState gst) st case res of Yield x s -> do r <- fstep acc x b1 <- case r of FL.Partial acc1 -> done acc1 FL.Done b -> return b let jump c = SplitOnSeqEmpty c s in yieldProceed jump b1 Skip s -> skip (SplitOnSeqEmpty acc s) Stop -> return Stop ----------------- -- Done ----------------- stepOuter _ SplitOnSeqDone = return Stop ----------------- -- Single Pattern ----------------- stepOuter gst (SplitOnSeqSingle fs st pat) = do res <- step (adaptState gst) st case res of Yield x s -> do let jump c = SplitOnSeqSingle c s pat if pat == x then done fs >>= yieldProceed jump else do r <- fstep fs x case r of FL.Partial fs1 -> skip $ jump fs1 FL.Done b -> yieldProceed jump b Skip s -> return $ Skip $ SplitOnSeqSingle fs s pat Stop -> do r <- done fs return $ Skip $ SplitOnSeqYield r SplitOnSeqDone --------------------------- -- Short Pattern - Shift Or --------------------------- stepOuter _ (SplitOnSeqWordDone 0 fs _) = do r <- done fs skip $ SplitOnSeqYield r SplitOnSeqDone stepOuter _ (SplitOnSeqWordDone n fs wrd) = do let old = elemMask .&. (wrd `shiftR` (elemBits * (n - 1))) r <- fstep fs (toEnum $ fromIntegral old) case r of FL.Partial fs1 -> skip $ SplitOnSeqWordDone (n - 1) fs1 wrd FL.Done b -> do let jump c = SplitOnSeqWordDone (n - 1) c wrd yieldProceed jump b stepOuter gst (SplitOnSeqWordInit fs st0) = go SPEC 0 0 st0 where {-# INLINE go #-} go !_ !idx !wrd !st = do res <- step (adaptState gst) st case res of Yield x s -> do let wrd1 = addToWord wrd x if idx == maxIndex then do if wrd1 .&. wordMask == wordPat then do let jump c = SplitOnSeqWordInit c s done fs >>= yieldProceed jump else skip $ SplitOnSeqWordLoop wrd1 s fs else go SPEC (idx + 1) wrd1 s Skip s -> go SPEC idx wrd s Stop -> do if idx /= 0 then skip $ SplitOnSeqWordDone idx fs wrd else do r <- done fs skip $ SplitOnSeqYield r SplitOnSeqDone stepOuter gst (SplitOnSeqWordLoop wrd0 st0 fs0) = go SPEC wrd0 st0 fs0 where {-# INLINE go #-} go !_ !wrd !st !fs = do res <- step (adaptState gst) st case res of Yield x s -> do let jump c = SplitOnSeqWordInit c s wrd1 = addToWord wrd x old = (wordMask .&. wrd) `shiftR` (elemBits * (patLen - 1)) r <- fstep fs (toEnum $ fromIntegral old) case r of FL.Partial fs1 -> do if wrd1 .&. wordMask == wordPat then done fs1 >>= yieldProceed jump else go SPEC wrd1 s fs1 FL.Done b -> yieldProceed jump b Skip s -> go SPEC wrd s fs Stop -> skip $ SplitOnSeqWordDone patLen fs wrd ------------------------------- -- General Pattern - Karp Rabin ------------------------------- stepOuter gst (SplitOnSeqKRInit idx fs st rb rh) = do res <- step (adaptState gst) st case res of Yield x s -> do rh1 <- liftIO $ RB.unsafeInsert rb rh x if idx == maxIndex then do let fld = RB.unsafeFoldRing (RB.ringBound rb) let !ringHash = fld addCksum 0 rb if ringHash == patHash then skip $ SplitOnSeqKRCheck fs s rb rh1 else skip $ SplitOnSeqKRLoop fs s rb rh1 ringHash else skip $ SplitOnSeqKRInit (idx + 1) fs s rb rh1 Skip s -> skip $ SplitOnSeqKRInit idx fs s rb rh Stop -> do skip $ SplitOnSeqKRDone idx fs rb (RB.startOf rb) -- XXX The recursive "go" is more efficient than the state based recursion -- code commented out below. Perhaps its more efficient because of -- factoring out "rb" outside the loop. -- stepOuter gst (SplitOnSeqKRLoop fs0 st0 rb rh0 cksum0) = go SPEC fs0 st0 rh0 cksum0 where go !_ !fs !st !rh !cksum = do res <- step (adaptState gst) st case res of Yield x s -> do old <- liftIO $ peek rh let cksum1 = deltaCksum cksum old x r <- fstep fs old case r of FL.Partial fs1 -> do rh1 <- liftIO (RB.unsafeInsert rb rh x) if cksum1 == patHash then skip $ SplitOnSeqKRCheck fs1 s rb rh1 else go SPEC fs1 s rh1 cksum1 FL.Done b -> do let rst = RB.startOf rb jump c = SplitOnSeqKRInit 0 c s rb rst yieldProceed jump b Skip s -> go SPEC fs s rh cksum Stop -> skip $ SplitOnSeqKRDone patLen fs rb rh -- XXX The following code is 5 times slower compared to the recursive loop -- based code above. Need to investigate why. One possibility is that the -- go loop above does not thread around the ring buffer (rb). This code may -- be causing the state to bloat and getting allocated on each iteration. -- We can check the cmm/asm code to confirm. If so a good GHC solution to -- such problem is needed. One way to avoid this could be to use unboxed -- mutable state? {- stepOuter gst (SplitOnSeqKRLoop fs st rb rh cksum) = do res <- step (adaptState gst) st case res of Yield x s -> do old <- liftIO $ peek rh let cksum1 = deltaCksum cksum old x fs1 <- fstep fs old if (cksum1 == patHash) then do r <- done fs1 skip $ SplitOnSeqYield r $ SplitOnSeqKRInit 0 s rb rh else do rh1 <- liftIO (RB.unsafeInsert rb rh x) skip $ SplitOnSeqKRLoop fs1 s rb rh1 cksum1 Skip s -> skip $ SplitOnSeqKRLoop fs s rb rh cksum Stop -> skip $ SplitOnSeqKRDone patLen fs rb rh -} stepOuter _ (SplitOnSeqKRCheck fs st rb rh) = do if RB.unsafeEqArray rb rh patArr then do r <- done fs let rst = RB.startOf rb jump c = SplitOnSeqKRInit 0 c st rb rst yieldProceed jump r else skip $ SplitOnSeqKRLoop fs st rb rh patHash stepOuter _ (SplitOnSeqKRDone 0 fs _ _) = do r <- done fs skip $ SplitOnSeqYield r SplitOnSeqDone stepOuter _ (SplitOnSeqKRDone n fs rb rh) = do old <- liftIO $ peek rh let rh1 = RB.advance rb rh r <- fstep fs old case r of FL.Partial fs1 -> skip $ SplitOnSeqKRDone (n - 1) fs1 rb rh1 FL.Done b -> do let jump c = SplitOnSeqKRDone (n - 1) c rb rh1 yieldProceed jump b {-# ANN type SplitOnSuffixSeqState Fuse #-} data SplitOnSuffixSeqState rb rh ck w fs s b x = SplitOnSuffixSeqInit | SplitOnSuffixSeqYield b (SplitOnSuffixSeqState rb rh ck w fs s b x) | SplitOnSuffixSeqDone | SplitOnSuffixSeqEmpty !fs s | SplitOnSuffixSeqSingleInit !fs s x | SplitOnSuffixSeqSingle !fs s x | SplitOnSuffixSeqWordInit !fs s | SplitOnSuffixSeqWordLoop !w s !fs | SplitOnSuffixSeqWordDone Int !fs !w | SplitOnSuffixSeqKRInit Int !fs s rb !rh | SplitOnSuffixSeqKRInit1 !fs s rb !rh | SplitOnSuffixSeqKRLoop fs s rb !rh !ck | SplitOnSuffixSeqKRCheck fs s rb !rh | SplitOnSuffixSeqKRDone Int !fs rb !rh | SplitOnSuffixSeqReinit (fs -> SplitOnSuffixSeqState rb rh ck w fs s b x) {-# INLINE_NORMAL splitOnSuffixSeq #-} splitOnSuffixSeq :: forall m a b. (MonadIO m, Storable a, Enum a, Eq a) => Bool -> Array a -> Fold m a b -> Stream m a -> Stream m b splitOnSuffixSeq withSep patArr (Fold fstep initial done) (Stream step state) = Stream stepOuter SplitOnSuffixSeqInit where patLen = A.length patArr maxIndex = patLen - 1 elemBits = sizeOf (undefined :: a) * 8 -- For word pattern case wordMask :: Word wordMask = (1 `shiftL` (elemBits * patLen)) - 1 elemMask :: Word elemMask = (1 `shiftL` elemBits) - 1 wordPat :: Word wordPat = wordMask .&. A.foldl' addToWord 0 patArr addToWord wd a = (wd `shiftL` elemBits) .|. fromIntegral (fromEnum a) nextAfterInit nextGen stepRes = case stepRes of FL.Partial s -> nextGen s FL.Done b -> SplitOnSuffixSeqYield b (SplitOnSuffixSeqReinit nextGen) {-# INLINE yieldProceed #-} yieldProceed nextGen fs = initial >>= skip . SplitOnSuffixSeqYield fs . nextAfterInit nextGen -- For single element pattern case {-# INLINE processYieldSingle #-} processYieldSingle pat x s fs = do let jump c = SplitOnSuffixSeqSingleInit c s pat if pat == x then do r <- if withSep then fstep fs x else return $ FL.Partial fs b1 <- case r of FL.Partial fs1 -> done fs1 FL.Done b -> return b yieldProceed jump b1 else do r <- fstep fs x case r of FL.Partial fs1 -> skip $ SplitOnSuffixSeqSingle fs1 s pat FL.Done b -> yieldProceed jump b -- For Rabin-Karp search k = 2891336453 :: Word32 coeff = k ^ patLen addCksum cksum a = cksum * k + fromIntegral (fromEnum a) deltaCksum cksum old new = addCksum cksum new - coeff * fromIntegral (fromEnum old) -- XXX shall we use a random starting hash or 1 instead of 0? patHash = A.foldl' addCksum 0 patArr skip = return . Skip {-# INLINE_LATE stepOuter #-} stepOuter _ SplitOnSuffixSeqInit = do res <- initial case res of FL.Partial fs -> if patLen == 0 then skip $ SplitOnSuffixSeqEmpty fs state else if patLen == 1 then do pat <- liftIO $ A.unsafeIndexIO patArr 0 skip $ SplitOnSuffixSeqSingleInit fs state pat else if sizeOf (undefined :: a) * patLen <= sizeOf (undefined :: Word) then skip $ SplitOnSuffixSeqWordInit fs state else do (rb, rhead) <- liftIO $ RB.new patLen skip $ SplitOnSuffixSeqKRInit 0 fs state rb rhead FL.Done fb -> skip $ SplitOnSuffixSeqYield fb SplitOnSuffixSeqInit stepOuter _ (SplitOnSuffixSeqYield x next) = return $ Yield x next --------------------------- -- Reinit --------------------------- stepOuter _ (SplitOnSuffixSeqReinit nextGen) = initial >>= skip . nextAfterInit nextGen --------------------------- -- Empty pattern --------------------------- stepOuter gst (SplitOnSuffixSeqEmpty acc st) = do res <- step (adaptState gst) st case res of Yield x s -> do let jump c = SplitOnSuffixSeqEmpty c s r <- fstep acc x b1 <- case r of FL.Partial fs -> done fs FL.Done b -> return b yieldProceed jump b1 Skip s -> skip (SplitOnSuffixSeqEmpty acc s) Stop -> return Stop ----------------- -- Done ----------------- stepOuter _ SplitOnSuffixSeqDone = return Stop ----------------- -- Single Pattern ----------------- stepOuter gst (SplitOnSuffixSeqSingleInit fs st pat) = do res <- step (adaptState gst) st case res of Yield x s -> processYieldSingle pat x s fs Skip s -> skip $ SplitOnSuffixSeqSingleInit fs s pat Stop -> return Stop stepOuter gst (SplitOnSuffixSeqSingle fs st pat) = do res <- step (adaptState gst) st case res of Yield x s -> processYieldSingle pat x s fs Skip s -> skip $ SplitOnSuffixSeqSingle fs s pat Stop -> do r <- done fs skip $ SplitOnSuffixSeqYield r SplitOnSuffixSeqDone --------------------------- -- Short Pattern - Shift Or --------------------------- stepOuter _ (SplitOnSuffixSeqWordDone 0 fs _) = do r <- done fs skip $ SplitOnSuffixSeqYield r SplitOnSuffixSeqDone stepOuter _ (SplitOnSuffixSeqWordDone n fs wrd) = do let old = elemMask .&. (wrd `shiftR` (elemBits * (n - 1))) r <- fstep fs (toEnum $ fromIntegral old) case r of FL.Partial fs1 -> skip $ SplitOnSuffixSeqWordDone (n - 1) fs1 wrd FL.Done b -> do let jump c = SplitOnSuffixSeqWordDone (n - 1) c wrd yieldProceed jump b stepOuter gst (SplitOnSuffixSeqWordInit fs0 st0) = do res <- step (adaptState gst) st0 case res of Yield x s -> do let wrd = addToWord 0 x r <- if withSep then fstep fs0 x else return $ FL.Partial fs0 case r of FL.Partial fs1 -> go SPEC 1 wrd s fs1 FL.Done b -> do let jump c = SplitOnSuffixSeqWordInit c s yieldProceed jump b Skip s -> skip (SplitOnSuffixSeqWordInit fs0 s) Stop -> return Stop where {-# INLINE go #-} go !_ !idx !wrd !st !fs = do res <- step (adaptState gst) st case res of Yield x s -> do let jump c = SplitOnSuffixSeqWordInit c s let wrd1 = addToWord wrd x r <- if withSep then fstep fs x else return $ FL.Partial fs case r of FL.Partial fs1 -> if idx /= maxIndex then go SPEC (idx + 1) wrd1 s fs1 else if wrd1 .&. wordMask /= wordPat then skip $ SplitOnSuffixSeqWordLoop wrd1 s fs1 else do done fs >>= yieldProceed jump FL.Done b -> yieldProceed jump b Skip s -> go SPEC idx wrd s fs Stop -> skip $ SplitOnSuffixSeqWordDone idx fs wrd stepOuter gst (SplitOnSuffixSeqWordLoop wrd0 st0 fs0) = go SPEC wrd0 st0 fs0 where {-# INLINE go #-} go !_ !wrd !st !fs = do res <- step (adaptState gst) st case res of Yield x s -> do let jump c = SplitOnSuffixSeqWordInit c s wrd1 = addToWord wrd x old = (wordMask .&. wrd) `shiftR` (elemBits * (patLen - 1)) r <- if withSep then fstep fs x else fstep fs (toEnum $ fromIntegral old) case r of FL.Partial fs1 -> if wrd1 .&. wordMask == wordPat then done fs1 >>= yieldProceed jump else go SPEC wrd1 s fs1 FL.Done b -> yieldProceed jump b Skip s -> go SPEC wrd s fs Stop -> if wrd .&. wordMask == wordPat then return Stop else if withSep then do r <- done fs skip $ SplitOnSuffixSeqYield r SplitOnSuffixSeqDone else skip $ SplitOnSuffixSeqWordDone patLen fs wrd ------------------------------- -- General Pattern - Karp Rabin ------------------------------- stepOuter gst (SplitOnSuffixSeqKRInit idx0 fs st0 rb rh0) = do res <- step (adaptState gst) st0 case res of Yield x s -> do rh1 <- liftIO $ RB.unsafeInsert rb rh0 x r <- if withSep then fstep fs x else return $ FL.Partial fs case r of FL.Partial fs1 -> skip $ SplitOnSuffixSeqKRInit1 fs1 s rb rh1 FL.Done b -> do let rst = RB.startOf rb jump c = SplitOnSuffixSeqKRInit 0 c s rb rst yieldProceed jump b Skip s -> skip $ SplitOnSuffixSeqKRInit idx0 fs s rb rh0 Stop -> return Stop stepOuter gst (SplitOnSuffixSeqKRInit1 fs0 st0 rb rh0) = do go SPEC 1 rh0 st0 fs0 where go !_ !idx !rh st !fs = do res <- step (adaptState gst) st case res of Yield x s -> do rh1 <- liftIO (RB.unsafeInsert rb rh x) r <- if withSep then fstep fs x else return $ FL.Partial fs case r of FL.Partial fs1 -> if idx /= maxIndex then go SPEC (idx + 1) rh1 s fs1 else skip $ let fld = RB.unsafeFoldRing (RB.ringBound rb) !ringHash = fld addCksum 0 rb in if ringHash == patHash then SplitOnSuffixSeqKRCheck fs1 s rb rh1 else SplitOnSuffixSeqKRLoop fs1 s rb rh1 ringHash FL.Done b -> do let rst = RB.startOf rb jump c = SplitOnSuffixSeqKRInit 0 c s rb rst yieldProceed jump b Skip s -> go SPEC idx rh s fs Stop -> do -- do not issue a blank segment when we end at pattern if (idx == maxIndex) && RB.unsafeEqArray rb rh patArr then return Stop else if withSep then do r <- done fs skip $ SplitOnSuffixSeqYield r SplitOnSuffixSeqDone else skip $ SplitOnSuffixSeqKRDone idx fs rb (RB.startOf rb) stepOuter gst (SplitOnSuffixSeqKRLoop fs0 st0 rb rh0 cksum0) = go SPEC fs0 st0 rh0 cksum0 where go !_ !fs !st !rh !cksum = do res <- step (adaptState gst) st case res of Yield x s -> do old <- liftIO $ peek rh rh1 <- liftIO (RB.unsafeInsert rb rh x) let cksum1 = deltaCksum cksum old x r <- if withSep then fstep fs x else fstep fs old case r of FL.Partial fs1 -> if (cksum1 /= patHash) then go SPEC fs1 s rh1 cksum1 else skip $ SplitOnSuffixSeqKRCheck fs1 s rb rh1 FL.Done b -> do let rst = RB.startOf rb jump c = SplitOnSuffixSeqKRInit 0 c s rb rst yieldProceed jump b Skip s -> go SPEC fs s rh cksum Stop -> if RB.unsafeEqArray rb rh patArr then return Stop else if withSep then do r <- done fs skip $ SplitOnSuffixSeqYield r SplitOnSuffixSeqDone else skip $ SplitOnSuffixSeqKRDone patLen fs rb rh stepOuter _ (SplitOnSuffixSeqKRCheck fs st rb rh) = do if RB.unsafeEqArray rb rh patArr then do r <- done fs let rst = RB.startOf rb jump c = SplitOnSuffixSeqKRInit 0 c st rb rst yieldProceed jump r else skip $ SplitOnSuffixSeqKRLoop fs st rb rh patHash stepOuter _ (SplitOnSuffixSeqKRDone 0 fs _ _) = do r <- done fs skip $ SplitOnSuffixSeqYield r SplitOnSuffixSeqDone stepOuter _ (SplitOnSuffixSeqKRDone n fs rb rh) = do old <- liftIO $ peek rh let rh1 = RB.advance rb rh r <- fstep fs old case r of FL.Partial fs1 -> skip $ SplitOnSuffixSeqKRDone (n - 1) fs1 rb rh1 FL.Done b -> do let jump c = SplitOnSuffixSeqKRDone (n - 1) c rb rh1 yieldProceed jump b ------------------------------------------------------------------------------ -- Nested Container Transformation ------------------------------------------------------------------------------ {-# ANN type SplitState Fuse #-} data SplitState s arr = SplitInitial s | SplitBuffering s arr | SplitSplitting s arr | SplitYielding arr (SplitState s arr) | SplitFinishing -- XXX An alternative approach would be to use a partial fold (Fold m a b) to -- split using a splitBy like combinator. The Fold would consume upto the -- separator and return any leftover which can then be fed to the next fold. -- -- We can revisit this once we have partial folds/parsers. -- -- | Performs infix separator style splitting. {-# INLINE_NORMAL splitInnerBy #-} splitInnerBy :: Monad m => (f a -> m (f a, Maybe (f a))) -- splitter -> (f a -> f a -> m (f a)) -- joiner -> Stream m (f a) -> Stream m (f a) splitInnerBy splitter joiner (Stream step1 state1) = (Stream step (SplitInitial state1)) where {-# INLINE_LATE step #-} step gst (SplitInitial st) = do r <- step1 gst st case r of Yield x s -> do (x1, mx2) <- splitter x return $ case mx2 of Nothing -> Skip (SplitBuffering s x1) Just x2 -> Skip (SplitYielding x1 (SplitSplitting s x2)) Skip s -> return $ Skip (SplitInitial s) Stop -> return $ Stop step gst (SplitBuffering st buf) = do r <- step1 gst st case r of Yield x s -> do (x1, mx2) <- splitter x buf' <- joiner buf x1 return $ case mx2 of Nothing -> Skip (SplitBuffering s buf') Just x2 -> Skip (SplitYielding buf' (SplitSplitting s x2)) Skip s -> return $ Skip (SplitBuffering s buf) Stop -> return $ Skip (SplitYielding buf SplitFinishing) step _ (SplitSplitting st buf) = do (x1, mx2) <- splitter buf return $ case mx2 of Nothing -> Skip $ SplitBuffering st x1 Just x2 -> Skip $ SplitYielding x1 (SplitSplitting st x2) step _ (SplitYielding x next) = return $ Yield x next step _ SplitFinishing = return $ Stop -- | Performs infix separator style splitting. {-# INLINE_NORMAL splitInnerBySuffix #-} splitInnerBySuffix :: (Monad m, Eq (f a), Monoid (f a)) => (f a -> m (f a, Maybe (f a))) -- splitter -> (f a -> f a -> m (f a)) -- joiner -> Stream m (f a) -> Stream m (f a) splitInnerBySuffix splitter joiner (Stream step1 state1) = (Stream step (SplitInitial state1)) where {-# INLINE_LATE step #-} step gst (SplitInitial st) = do r <- step1 gst st case r of Yield x s -> do (x1, mx2) <- splitter x return $ case mx2 of Nothing -> Skip (SplitBuffering s x1) Just x2 -> Skip (SplitYielding x1 (SplitSplitting s x2)) Skip s -> return $ Skip (SplitInitial s) Stop -> return $ Stop step gst (SplitBuffering st buf) = do r <- step1 gst st case r of Yield x s -> do (x1, mx2) <- splitter x buf' <- joiner buf x1 return $ case mx2 of Nothing -> Skip (SplitBuffering s buf') Just x2 -> Skip (SplitYielding buf' (SplitSplitting s x2)) Skip s -> return $ Skip (SplitBuffering s buf) Stop -> return $ if buf == mempty then Stop else Skip (SplitYielding buf SplitFinishing) step _ (SplitSplitting st buf) = do (x1, mx2) <- splitter buf return $ case mx2 of Nothing -> Skip $ SplitBuffering st x1 Just x2 -> Skip $ SplitYielding x1 (SplitSplitting st x2) step _ (SplitYielding x next) = return $ Yield x next step _ SplitFinishing = return Stop