git.iko.soy/streamly
on da36a01 help sign in
streamly/core/src/Streamly/Internal/Data/Stream/StreamD/Nesting.hs
2587 lines 97 kB history raw
1-- |
2-- Module : Streamly.Internal.Data.Stream.StreamD.Nesting
3-- Copyright : (c) 2018 Composewell Technologies
4-- (c) Roman Leshchinskiy 2008-2010
5-- License : BSD-3-Clause
6-- Maintainer : streamly@composewell.com
7-- Stability : experimental
8-- Portability : GHC
9--
10-- This module contains transformations involving multiple streams, unfolds or
11-- folds. There are two types of transformations generational or eliminational.
12-- Generational transformations are like the "Generate" module but they
13-- generate a stream by combining streams instead of elements. Eliminational
14-- transformations are like the "Eliminate" module but they transform a stream
15-- by eliminating parts of the stream instead of eliminating the whole stream.
16--
17-- These combinators involve transformation, generation, elimination so can be
18-- classified under any of those.
19--
20-- Ultimately these operations should be supported by Unfolds, Pipes and Folds,
21-- and this module may become redundant.
23-- The zipWithM combinator in this module has been adapted from the vector
24-- package (c) Roman Leshchinskiy.
25--
26module Streamly.Internal.Data.Stream.StreamD.Nesting
27 (
28 -- * Generate
29 -- | Combining streams to generate streams.
31 -- ** Combine Two Streams
32 -- | Functions ending in the shape:
33 --
34 -- @t m a -> t m a -> t m a@.
36 -- *** Appending
37 -- | Append a stream after another. A special case of concatMap or
38 -- unfoldMany.
39 AppendState(..)
40 , append
42 -- *** Interleaving
43 -- | Interleave elements from two streams alternately. A special case of
44 -- unfoldManyInterleave.
45 , InterleaveState(..)
46 , interleave
47 , interleaveMin
48 , interleaveSuffix
49 , interleaveInfix
51 -- *** Scheduling
52 -- | Execute streams alternately irrespective of whether they generate
53 -- elements or not. Note 'interleave' would execute a stream until it
54 -- yields an element. A special case of unfoldManyRoundRobin.
55 , roundRobin -- interleaveFair?/ParallelFair
57 -- *** Zipping
58 -- | Zip corresponding elements of two streams.
59 , zipWith
60 , zipWithM
62 -- *** Merging
63 -- | Interleave elements from two streams based on a condition.
64 , mergeBy
65 , mergeByM
67 -- ** Combine N Streams
68 -- | Functions generally ending in these shapes:
69 --
70 -- @
71 -- concat: f (t m a) -> t m a
72 -- concatMap: (a -> t m b) -> t m a -> t m b
73 -- unfoldMany: Unfold m a b -> t m a -> t m b
74 -- @
76 -- *** ConcatMap
77 -- | Generate streams by mapping a stream generator on each element of an
78 -- input stream, append the resulting streams and flatten.
79 , concatMap
80 , concatMapM
82 -- *** ConcatUnfold
83 -- | Generate streams by using an unfold on each element of an input
84 -- stream, append the resulting streams and flatten. A special case of
85 -- gintercalate.
86 , unfoldMany
87 , ConcatUnfoldInterleaveState (..)
88 , unfoldManyInterleave
89 , unfoldManyRoundRobin
91 -- *** Interpose
92 -- | Like unfoldMany but intersperses an effect between the streams. A
93 -- special case of gintercalate.
94 , interpose
95 , interposeSuffix
97 -- *** Intercalate
98 -- | Like unfoldMany but intersperses streams from another source between
99 -- the streams from the first source.
100 , gintercalate
101 , gintercalateSuffix
103 -- * Eliminate
104 -- | Folding and Parsing chunks of streams to eliminate nested streams.
105 -- Functions generally ending in these shapes:
106 --
107 -- @
108 -- f (Fold m a b) -> t m a -> t m b
109 -- f (Parser m a b) -> t m a -> t m b
110 -- @
112 -- ** Folding
113 -- | Apply folds on a stream.
114 , foldMany
115 , refoldMany
116 , foldIterateM
117 , refoldIterateM
119 -- ** Parsing
120 -- | Parsing is opposite to flattening. 'parseMany' is dual to concatMap or
121 -- unfoldMany. concatMap generates a stream from single values in a
122 -- stream and flattens, parseMany does the opposite of flattening by
123 -- splitting the stream and then folds each such split to single value in
124 -- the output stream.
125 , parseMany
126 , parseIterate
128 -- ** Grouping
129 -- | Group segments of a stream and fold. Special case of parsing.
130 , chunksOf
131 , groupsBy
132 , groupsRollingBy
134 -- ** Splitting
135 -- | A special case of parsing.
136 , wordsBy
137 , splitOnSeq
138 , splitOnSuffixSeq
139 , sliceOnSuffix
141 -- * Transform (Nested Containers)
142 -- | Opposite to compact in ArrayStream
143 , splitInnerBy
144 , splitInnerBySuffix
145 , intersectBySorted
146 )
147where
149#include "inline.hs"
150#include "ArrayMacros.h"
152import Control.Exception (assert)
153import Control.Monad.Catch (MonadThrow, throwM)
154import Control.Monad.IO.Class (MonadIO(..))
155import Data.Bits (shiftR, shiftL, (.|.), (.&.))
156import Data.Functor.Identity ( Identity )
157import Data.Word (Word32)
158import Foreign.Storable (peek, sizeOf)
159import Fusion.Plugin.Types (Fuse(..))
160import GHC.Types (SPEC(..))
162import Streamly.Internal.Data.Array.Unboxed.Type (Array(..))
163import Streamly.Internal.Data.Fold.Step (Step(..))
164import Streamly.Internal.Data.Fold.Type (Fold(..))
165import Streamly.Internal.Data.Parser (ParseError(..))
166import Streamly.Internal.Data.Refold.Type (Refold(..))
167import Streamly.Internal.Data.SVar.Type (adaptState)
168import Streamly.Internal.Data.Tuple.Strict (Tuple'(..))
169import Streamly.Internal.Data.Unboxed (Unboxed)
170import Streamly.Internal.Data.Unfold.Type (Unfold(..))
172import qualified Streamly.Internal.Data.Array.Unboxed.Type as A
173import qualified Streamly.Internal.Data.Fold as FL
174import qualified Streamly.Internal.Data.Parser as PR
175import qualified Streamly.Internal.Data.Parser.ParserD as PRD
176import qualified Streamly.Internal.Data.Ring.Unboxed as RB
178import Streamly.Internal.Data.Stream.StreamD.Type
180import Prelude hiding (concatMap, mapM, zipWith)
182------------------------------------------------------------------------------
183-- Appending
184------------------------------------------------------------------------------
186data AppendState s1 s2 = AppendFirst s1 | AppendSecond s2
188-- Note that this could be much faster compared to the CPS stream. However, as
189-- the number of streams being composed increases this may become expensive.
190-- Need to see where the breaking point is between the two.
191--
192{-# INLINE_NORMAL append #-}
193append :: Monad m => Stream m a -> Stream m a -> Stream m a
194append (Stream step1 state1) (Stream step2 state2) =
195 Stream step (AppendFirst state1)
197 where
199 {-# INLINE_LATE step #-}
200 step gst (AppendFirst st) = do
201 r <- step1 gst st
202 return $ case r of
203 Yield a s -> Yield a (AppendFirst s)
204 Skip s -> Skip (AppendFirst s)
205 Stop -> Skip (AppendSecond state2)
207 step gst (AppendSecond st) = do
208 r <- step2 gst st
209 return $ case r of
210 Yield a s -> Yield a (AppendSecond s)
211 Skip s -> Skip (AppendSecond s)
212 Stop -> Stop
214------------------------------------------------------------------------------
215-- Interleaving
216------------------------------------------------------------------------------
218data InterleaveState s1 s2 = InterleaveFirst s1 s2 | InterleaveSecond s1 s2
219 | InterleaveSecondOnly s2 | InterleaveFirstOnly s1
221{-# INLINE_NORMAL interleave #-}
222interleave :: Monad m => Stream m a -> Stream m a -> Stream m a
223interleave (Stream step1 state1) (Stream step2 state2) =
224 Stream step (InterleaveFirst state1 state2)
226 where
228 {-# INLINE_LATE step #-}
229 step gst (InterleaveFirst st1 st2) = do
230 r <- step1 gst st1
231 return $ case r of
232 Yield a s -> Yield a (InterleaveSecond s st2)
233 Skip s -> Skip (InterleaveFirst s st2)
234 Stop -> Skip (InterleaveSecondOnly st2)
236 step gst (InterleaveSecond st1 st2) = do
237 r <- step2 gst st2
238 return $ case r of
239 Yield a s -> Yield a (InterleaveFirst st1 s)
240 Skip s -> Skip (InterleaveSecond st1 s)
241 Stop -> Skip (InterleaveFirstOnly st1)
243 step gst (InterleaveFirstOnly st1) = do
244 r <- step1 gst st1
245 return $ case r of
246 Yield a s -> Yield a (InterleaveFirstOnly s)
247 Skip s -> Skip (InterleaveFirstOnly s)
248 Stop -> Stop
250 step gst (InterleaveSecondOnly st2) = do
251 r <- step2 gst st2
252 return $ case r of
253 Yield a s -> Yield a (InterleaveSecondOnly s)
254 Skip s -> Skip (InterleaveSecondOnly s)
255 Stop -> Stop
257{-# INLINE_NORMAL interleaveMin #-}
258interleaveMin :: Monad m => Stream m a -> Stream m a -> Stream m a
259interleaveMin (Stream step1 state1) (Stream step2 state2) =
260 Stream step (InterleaveFirst state1 state2)
262 where
264 {-# INLINE_LATE step #-}
265 step gst (InterleaveFirst st1 st2) = do
266 r <- step1 gst st1
267 return $ case r of
268 Yield a s -> Yield a (InterleaveSecond s st2)
269 Skip s -> Skip (InterleaveFirst s st2)
270 Stop -> Stop
272 step gst (InterleaveSecond st1 st2) = do
273 r <- step2 gst st2
274 return $ case r of
275 Yield a s -> Yield a (InterleaveFirst st1 s)
276 Skip s -> Skip (InterleaveSecond st1 s)
277 Stop -> Stop
279 step _ (InterleaveFirstOnly _) = undefined
280 step _ (InterleaveSecondOnly _) = undefined
282{-# INLINE_NORMAL interleaveSuffix #-}
283interleaveSuffix :: Monad m => Stream m a -> Stream m a -> Stream m a
284interleaveSuffix (Stream step1 state1) (Stream step2 state2) =
285 Stream step (InterleaveFirst state1 state2)
287 where
289 {-# INLINE_LATE step #-}
290 step gst (InterleaveFirst st1 st2) = do
291 r <- step1 gst st1
292 return $ case r of
293 Yield a s -> Yield a (InterleaveSecond s st2)
294 Skip s -> Skip (InterleaveFirst s st2)
295 Stop -> Stop
297 step gst (InterleaveSecond st1 st2) = do
298 r <- step2 gst st2
299 return $ case r of
300 Yield a s -> Yield a (InterleaveFirst st1 s)
301 Skip s -> Skip (InterleaveSecond st1 s)
302 Stop -> Skip (InterleaveFirstOnly st1)
304 step gst (InterleaveFirstOnly st1) = do
305 r <- step1 gst st1
306 return $ case r of
307 Yield a s -> Yield a (InterleaveFirstOnly s)
308 Skip s -> Skip (InterleaveFirstOnly s)
309 Stop -> Stop
311 step _ (InterleaveSecondOnly _) = undefined
313data InterleaveInfixState s1 s2 a
314 = InterleaveInfixFirst s1 s2
315 | InterleaveInfixSecondBuf s1 s2
316 | InterleaveInfixSecondYield s1 s2 a
317 | InterleaveInfixFirstYield s1 s2 a
318 | InterleaveInfixFirstOnly s1
320{-# INLINE_NORMAL interleaveInfix #-}
321interleaveInfix :: Monad m => Stream m a -> Stream m a -> Stream m a
322interleaveInfix (Stream step1 state1) (Stream step2 state2) =
323 Stream step (InterleaveInfixFirst state1 state2)
325 where
327 {-# INLINE_LATE step #-}
328 step gst (InterleaveInfixFirst st1 st2) = do
329 r <- step1 gst st1
330 return $ case r of
331 Yield a s -> Yield a (InterleaveInfixSecondBuf s st2)
332 Skip s -> Skip (InterleaveInfixFirst s st2)
333 Stop -> Stop
335 step gst (InterleaveInfixSecondBuf st1 st2) = do
336 r <- step2 gst st2
337 return $ case r of
338 Yield a s -> Skip (InterleaveInfixSecondYield st1 s a)
339 Skip s -> Skip (InterleaveInfixSecondBuf st1 s)
340 Stop -> Skip (InterleaveInfixFirstOnly st1)
342 step gst (InterleaveInfixSecondYield st1 st2 x) = do
343 r <- step1 gst st1
344 return $ case r of
345 Yield a s -> Yield x (InterleaveInfixFirstYield s st2 a)
346 Skip s -> Skip (InterleaveInfixSecondYield s st2 x)
347 Stop -> Stop
349 step _ (InterleaveInfixFirstYield st1 st2 x) = do
350 return $ Yield x (InterleaveInfixSecondBuf st1 st2)
352 step gst (InterleaveInfixFirstOnly st1) = do
353 r <- step1 gst st1
354 return $ case r of
355 Yield a s -> Yield a (InterleaveInfixFirstOnly s)
356 Skip s -> Skip (InterleaveInfixFirstOnly s)
357 Stop -> Stop
359------------------------------------------------------------------------------
360-- Scheduling
361------------------------------------------------------------------------------
363{-# INLINE_NORMAL roundRobin #-}
364roundRobin :: Monad m => Stream m a -> Stream m a -> Stream m a
365roundRobin (Stream step1 state1) (Stream step2 state2) =
366 Stream step (InterleaveFirst state1 state2)
368 where
370 {-# INLINE_LATE step #-}
371 step gst (InterleaveFirst st1 st2) = do
372 r <- step1 gst st1
373 return $ case r of
374 Yield a s -> Yield a (InterleaveSecond s st2)
375 Skip s -> Skip (InterleaveSecond s st2)
376 Stop -> Skip (InterleaveSecondOnly st2)
378 step gst (InterleaveSecond st1 st2) = do
379 r <- step2 gst st2
380 return $ case r of
381 Yield a s -> Yield a (InterleaveFirst st1 s)
382 Skip s -> Skip (InterleaveFirst st1 s)
383 Stop -> Skip (InterleaveFirstOnly st1)
385 step gst (InterleaveSecondOnly st2) = do
386 r <- step2 gst st2
387 return $ case r of
388 Yield a s -> Yield a (InterleaveSecondOnly s)
389 Skip s -> Skip (InterleaveSecondOnly s)
390 Stop -> Stop
392 step gst (InterleaveFirstOnly st1) = do
393 r <- step1 gst st1
394 return $ case r of
395 Yield a s -> Yield a (InterleaveFirstOnly s)
396 Skip s -> Skip (InterleaveFirstOnly s)
397 Stop -> Stop
399------------------------------------------------------------------------------
400-- Zipping
401------------------------------------------------------------------------------
403{-# INLINE_NORMAL zipWithM #-}
404zipWithM :: Monad m
405 => (a -> b -> m c) -> Stream m a -> Stream m b -> Stream m c
406zipWithM f (Stream stepa ta) (Stream stepb tb) = Stream step (ta, tb, Nothing)
407 where
408 {-# INLINE_LATE step #-}
409 step gst (sa, sb, Nothing) = do
410 r <- stepa (adaptState gst) sa
411 return $
412 case r of
413 Yield x sa' -> Skip (sa', sb, Just x)
414 Skip sa' -> Skip (sa', sb, Nothing)
415 Stop -> Stop
417 step gst (sa, sb, Just x) = do
418 r <- stepb (adaptState gst) sb
419 case r of
420 Yield y sb' -> do
421 z <- f x y
422 return $ Yield z (sa, sb', Nothing)
423 Skip sb' -> return $ Skip (sa, sb', Just x)
424 Stop -> return Stop
426{-# RULES "zipWithM xs xs"
427 forall f xs. zipWithM @Identity f xs xs = mapM (\x -> f x x) xs #-}
429{-# INLINE zipWith #-}
430zipWith :: Monad m => (a -> b -> c) -> Stream m a -> Stream m b -> Stream m c
431zipWith f = zipWithM (\a b -> return (f a b))
433------------------------------------------------------------------------------
434-- Merging
435------------------------------------------------------------------------------
437{-# INLINE_NORMAL mergeByM #-}
438mergeByM
439 :: (Monad m)
440 => (a -> a -> m Ordering) -> Stream m a -> Stream m a -> Stream m a
441mergeByM cmp (Stream stepa ta) (Stream stepb tb) =
442 Stream step (Just ta, Just tb, Nothing, Nothing)
443 where
444 {-# INLINE_LATE step #-}
446 -- one of the values is missing, and the corresponding stream is running
447 step gst (Just sa, sb, Nothing, b) = do
448 r <- stepa gst sa
449 return $ case r of
450 Yield a sa' -> Skip (Just sa', sb, Just a, b)
451 Skip sa' -> Skip (Just sa', sb, Nothing, b)
452 Stop -> Skip (Nothing, sb, Nothing, b)
454 step gst (sa, Just sb, a, Nothing) = do
455 r <- stepb gst sb
456 return $ case r of
457 Yield b sb' -> Skip (sa, Just sb', a, Just b)
458 Skip sb' -> Skip (sa, Just sb', a, Nothing)
459 Stop -> Skip (sa, Nothing, a, Nothing)
461 -- both the values are available
462 step _ (sa, sb, Just a, Just b) = do
463 res <- cmp a b
464 return $ case res of
465 GT -> Yield b (sa, sb, Just a, Nothing)
466 _ -> Yield a (sa, sb, Nothing, Just b)
468 -- one of the values is missing, corresponding stream is done
469 step _ (Nothing, sb, Nothing, Just b) =
470 return $ Yield b (Nothing, sb, Nothing, Nothing)
472 step _ (sa, Nothing, Just a, Nothing) =
473 return $ Yield a (sa, Nothing, Nothing, Nothing)
475 step _ (Nothing, Nothing, Nothing, Nothing) = return Stop
477{-# INLINE mergeBy #-}
478mergeBy
479 :: (Monad m)
480 => (a -> a -> Ordering) -> Stream m a -> Stream m a -> Stream m a
481mergeBy cmp = mergeByM (\a b -> return $ cmp a b)
483-------------------------------------------------------------------------------
484-- Intersection of sorted streams
485-------------------------------------------------------------------------------
487-- Assuming the streams are sorted in ascending order
488{-# INLINE_NORMAL intersectBySorted #-}
489intersectBySorted :: Monad m
490 => (a -> a -> Ordering) -> Stream m a -> Stream m a -> Stream m a
491intersectBySorted cmp (Stream stepa ta) (Stream stepb tb) =
492 Stream step
493 ( ta -- left stream state
494 , tb -- right stream state
495 , Nothing -- left value
496 , Nothing -- right value
497 )
499 where
501 {-# INLINE_LATE step #-}
502 -- step 1, fetch the first value
503 step gst (sa, sb, Nothing, b) = do
504 r <- stepa gst sa
505 return $ case r of
506 Yield a sa' -> Skip (sa', sb, Just a, b) -- step 2/3
507 Skip sa' -> Skip (sa', sb, Nothing, b)
508 Stop -> Stop
510 -- step 2, fetch the second value
511 step gst (sa, sb, a@(Just _), Nothing) = do
512 r <- stepb gst sb
513 return $ case r of
514 Yield b sb' -> Skip (sa, sb', a, Just b) -- step 3
515 Skip sb' -> Skip (sa, sb', a, Nothing)
516 Stop -> Stop
518 -- step 3, compare the two values
519 step _ (sa, sb, Just a, Just b) = do
520 let res = cmp a b
521 return $ case res of
522 GT -> Skip (sa, sb, Just a, Nothing) -- step 2
523 LT -> Skip (sa, sb, Nothing, Just b) -- step 1
524 EQ -> Yield a (sa, sb, Nothing, Just b) -- step 1
526------------------------------------------------------------------------------
527-- Combine N Streams - unfoldMany
528------------------------------------------------------------------------------
530data ConcatUnfoldInterleaveState o i =
531 ConcatUnfoldInterleaveOuter o [i]
532 | ConcatUnfoldInterleaveInner o [i]
533 | ConcatUnfoldInterleaveInnerL [i] [i]
534 | ConcatUnfoldInterleaveInnerR [i] [i]
536-- XXX use arrays to store state instead of lists?
537--
538-- XXX In general we can use different scheduling strategies e.g. how to
539-- schedule the outer vs inner loop or assigning weights to different streams
540-- or outer and inner loops.
542-- After a yield, switch to the next stream. Do not switch streams on Skip.
543-- Yield from outer stream switches to the inner stream.
544--
545-- There are two choices here, (1) exhaust the outer stream first and then
546-- start yielding from the inner streams, this is much simpler to implement,
547-- (2) yield at least one element from an inner stream before going back to
548-- outer stream and opening the next stream from it.
549--
550-- Ideally, we need some scheduling bias to inner streams vs outer stream.
551-- Maybe we can configure the behavior.
552--
553-- XXX Instead of using "concatPairsWith wSerial" we can implement an N-way
554-- interleaving CPS combinator which behaves like unfoldManyInterleave. Instead
555-- of pairing up the streams we just need to go yielding one element from each
556-- stream and storing the remaining streams and then keep doing rounds through
557-- those in a round robin fashion. This would be much like wAsync.
558--
559-- See 'Streamly.Internal.Data.Stream.unfoldInterleave' documentation for more
560-- details.
561--
562{-# INLINE_NORMAL unfoldManyInterleave #-}
563unfoldManyInterleave :: Monad m => Unfold m a b -> Stream m a -> Stream m b
564unfoldManyInterleave (Unfold istep inject) (Stream ostep ost) =
565 Stream step (ConcatUnfoldInterleaveOuter ost [])
567 where
569 {-# INLINE_LATE step #-}
570 step gst (ConcatUnfoldInterleaveOuter o ls) = do
571 r <- ostep (adaptState gst) o
572 case r of
573 Yield a o' -> do
574 i <- inject a
575 i `seq` return (Skip (ConcatUnfoldInterleaveInner o' (i : ls)))
576 Skip o' -> return $ Skip (ConcatUnfoldInterleaveOuter o' ls)
577 Stop -> return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
579 step _ (ConcatUnfoldInterleaveInner _ []) = undefined
580 step _ (ConcatUnfoldInterleaveInner o (st:ls)) = do
581 r <- istep st
582 return $ case r of
583 Yield x s -> Yield x (ConcatUnfoldInterleaveOuter o (s:ls))
584 Skip s -> Skip (ConcatUnfoldInterleaveInner o (s:ls))
585 Stop -> Skip (ConcatUnfoldInterleaveOuter o ls)
587 step _ (ConcatUnfoldInterleaveInnerL [] []) = return Stop
588 step _ (ConcatUnfoldInterleaveInnerL [] rs) =
589 return $ Skip (ConcatUnfoldInterleaveInnerR [] rs)
591 step _ (ConcatUnfoldInterleaveInnerL (st:ls) rs) = do
592 r <- istep st
593 return $ case r of
594 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerL ls (s:rs))
595 Skip s -> Skip (ConcatUnfoldInterleaveInnerL (s:ls) rs)
596 Stop -> Skip (ConcatUnfoldInterleaveInnerL ls rs)
598 step _ (ConcatUnfoldInterleaveInnerR [] []) = return Stop
599 step _ (ConcatUnfoldInterleaveInnerR ls []) =
600 return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
602 step _ (ConcatUnfoldInterleaveInnerR ls (st:rs)) = do
603 r <- istep st
604 return $ case r of
605 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerR (s:ls) rs)
606 Skip s -> Skip (ConcatUnfoldInterleaveInnerR ls (s:rs))
607 Stop -> Skip (ConcatUnfoldInterleaveInnerR ls rs)
609-- XXX In general we can use different scheduling strategies e.g. how to
610-- schedule the outer vs inner loop or assigning weights to different streams
611-- or outer and inner loops.
612--
613-- This could be inefficient if the tasks are too small.
614--
615-- Compared to unfoldManyInterleave this one switches streams on Skips.
616--
617{-# INLINE_NORMAL unfoldManyRoundRobin #-}
618unfoldManyRoundRobin :: Monad m => Unfold m a b -> Stream m a -> Stream m b
619unfoldManyRoundRobin (Unfold istep inject) (Stream ostep ost) =
620 Stream step (ConcatUnfoldInterleaveOuter ost [])
621 where
622 {-# INLINE_LATE step #-}
623 step gst (ConcatUnfoldInterleaveOuter o ls) = do
624 r <- ostep (adaptState gst) o
625 case r of
626 Yield a o' -> do
627 i <- inject a
628 i `seq` return (Skip (ConcatUnfoldInterleaveInner o' (i : ls)))
629 Skip o' -> return $ Skip (ConcatUnfoldInterleaveInner o' ls)
630 Stop -> return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
632 step _ (ConcatUnfoldInterleaveInner o []) =
633 return $ Skip (ConcatUnfoldInterleaveOuter o [])
635 step _ (ConcatUnfoldInterleaveInner o (st:ls)) = do
636 r <- istep st
637 return $ case r of
638 Yield x s -> Yield x (ConcatUnfoldInterleaveOuter o (s:ls))
639 Skip s -> Skip (ConcatUnfoldInterleaveOuter o (s:ls))
640 Stop -> Skip (ConcatUnfoldInterleaveOuter o ls)
642 step _ (ConcatUnfoldInterleaveInnerL [] []) = return Stop
643 step _ (ConcatUnfoldInterleaveInnerL [] rs) =
644 return $ Skip (ConcatUnfoldInterleaveInnerR [] rs)
646 step _ (ConcatUnfoldInterleaveInnerL (st:ls) rs) = do
647 r <- istep st
648 return $ case r of
649 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerL ls (s:rs))
650 Skip s -> Skip (ConcatUnfoldInterleaveInnerL ls (s:rs))
651 Stop -> Skip (ConcatUnfoldInterleaveInnerL ls rs)
653 step _ (ConcatUnfoldInterleaveInnerR [] []) = return Stop
654 step _ (ConcatUnfoldInterleaveInnerR ls []) =
655 return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
657 step _ (ConcatUnfoldInterleaveInnerR ls (st:rs)) = do
658 r <- istep st
659 return $ case r of
660 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerR (s:ls) rs)
661 Skip s -> Skip (ConcatUnfoldInterleaveInnerR (s:ls) rs)
662 Stop -> Skip (ConcatUnfoldInterleaveInnerR ls rs)
664------------------------------------------------------------------------------
665-- Combine N Streams - interpose
666------------------------------------------------------------------------------
668{-# ANN type InterposeSuffixState Fuse #-}
669data InterposeSuffixState s1 i1 =
670 InterposeSuffixFirst s1
671 -- | InterposeSuffixFirstYield s1 i1
672 | InterposeSuffixFirstInner s1 i1
673 | InterposeSuffixSecond s1
675-- Note that if an unfolded layer turns out to be nil we still emit the
676-- separator effect. An alternate behavior could be to emit the separator
677-- effect only if at least one element has been yielded by the unfolding.
678-- However, that becomes a bit complicated, so we have chosen the former
679-- behvaior for now.
680{-# INLINE_NORMAL interposeSuffix #-}
681interposeSuffix
682 :: Monad m
683 => m c -> Unfold m b c -> Stream m b -> Stream m c
684interposeSuffix
685 action
686 (Unfold istep1 inject1) (Stream step1 state1) =
687 Stream step (InterposeSuffixFirst state1)
689 where
691 {-# INLINE_LATE step #-}
692 step gst (InterposeSuffixFirst s1) = do
693 r <- step1 (adaptState gst) s1
694 case r of
695 Yield a s -> do
696 i <- inject1 a
697 i `seq` return (Skip (InterposeSuffixFirstInner s i))
698 -- i `seq` return (Skip (InterposeSuffixFirstYield s i))
699 Skip s -> return $ Skip (InterposeSuffixFirst s)
700 Stop -> return Stop
702 {-
703 step _ (InterposeSuffixFirstYield s1 i1) = do
704 r <- istep1 i1
705 return $ case r of
706 Yield x i' -> Yield x (InterposeSuffixFirstInner s1 i')
707 Skip i' -> Skip (InterposeSuffixFirstYield s1 i')
708 Stop -> Skip (InterposeSuffixFirst s1)
709 -}
711 step _ (InterposeSuffixFirstInner s1 i1) = do
712 r <- istep1 i1
713 return $ case r of
714 Yield x i' -> Yield x (InterposeSuffixFirstInner s1 i')
715 Skip i' -> Skip (InterposeSuffixFirstInner s1 i')
716 Stop -> Skip (InterposeSuffixSecond s1)
718 step _ (InterposeSuffixSecond s1) = do
719 r <- action
720 return $ Yield r (InterposeSuffixFirst s1)
722{-# ANN type InterposeState Fuse #-}
723data InterposeState s1 i1 a =
724 InterposeFirst s1
725 -- | InterposeFirstYield s1 i1
726 | InterposeFirstInner s1 i1
727 | InterposeFirstInject s1
728 -- | InterposeFirstBuf s1 i1
729 | InterposeSecondYield s1 i1
730 -- -- | InterposeSecondYield s1 i1 a
731 -- -- | InterposeFirstResume s1 i1 a
733-- Note that this only interposes the pure values, we may run many effects to
734-- generate those values as some effects may not generate anything (Skip).
735{-# INLINE_NORMAL interpose #-}
736interpose :: Monad m => m c -> Unfold m b c -> Stream m b -> Stream m c
737interpose
738 action
739 (Unfold istep1 inject1) (Stream step1 state1) =
740 Stream step (InterposeFirst state1)
742 where
744 {-# INLINE_LATE step #-}
745 step gst (InterposeFirst s1) = do
746 r <- step1 (adaptState gst) s1
747 case r of
748 Yield a s -> do
749 i <- inject1 a
750 i `seq` return (Skip (InterposeFirstInner s i))
751 -- i `seq` return (Skip (InterposeFirstYield s i))
752 Skip s -> return $ Skip (InterposeFirst s)
753 Stop -> return Stop
755 {-
756 step _ (InterposeFirstYield s1 i1) = do
757 r <- istep1 i1
758 return $ case r of
759 Yield x i' -> Yield x (InterposeFirstInner s1 i')
760 Skip i' -> Skip (InterposeFirstYield s1 i')
761 Stop -> Skip (InterposeFirst s1)
762 -}
764 step _ (InterposeFirstInner s1 i1) = do
765 r <- istep1 i1
766 return $ case r of
767 Yield x i' -> Yield x (InterposeFirstInner s1 i')
768 Skip i' -> Skip (InterposeFirstInner s1 i')
769 Stop -> Skip (InterposeFirstInject s1)
771 step gst (InterposeFirstInject s1) = do
772 r <- step1 (adaptState gst) s1
773 case r of
774 Yield a s -> do
775 i <- inject1 a
776 -- i `seq` return (Skip (InterposeFirstBuf s i))
777 i `seq` return (Skip (InterposeSecondYield s i))
778 Skip s -> return $ Skip (InterposeFirstInject s)
779 Stop -> return Stop
781 {-
782 step _ (InterposeFirstBuf s1 i1) = do
783 r <- istep1 i1
784 return $ case r of
785 Yield x i' -> Skip (InterposeSecondYield s1 i' x)
786 Skip i' -> Skip (InterposeFirstBuf s1 i')
787 Stop -> Stop
788 -}
790 {-
791 step _ (InterposeSecondYield s1 i1 v) = do
792 r <- action
793 return $ Yield r (InterposeFirstResume s1 i1 v)
794 -}
795 step _ (InterposeSecondYield s1 i1) = do
796 r <- action
797 return $ Yield r (InterposeFirstInner s1 i1)
799 {-
800 step _ (InterposeFirstResume s1 i1 v) = do
801 return $ Yield v (InterposeFirstInner s1 i1)
802 -}
804------------------------------------------------------------------------------
805-- Combine N Streams - intercalate
806------------------------------------------------------------------------------
808data ICUState s1 s2 i1 i2 =
809 ICUFirst s1 s2
810 | ICUSecond s1 s2
811 | ICUSecondOnly s2
812 | ICUFirstOnly s1
813 | ICUFirstInner s1 s2 i1
814 | ICUSecondInner s1 s2 i2
815 | ICUFirstOnlyInner s1 i1
816 | ICUSecondOnlyInner s2 i2
818-- | Interleave streams (full streams, not the elements) unfolded from two
819-- input streams and concat. Stop when the first stream stops. If the second
820-- stream ends before the first one then first stream still keeps running alone
821-- without any interleaving with the second stream.
822--
823-- [a1, a2, ... an] [b1, b2 ...]
824-- => [streamA1, streamA2, ... streamAn] [streamB1, streamB2, ...]
825-- => [streamA1, streamB1, streamA2...StreamAn, streamBn]
826-- => [a11, a12, ...a1j, b11, b12, ...b1k, a21, a22, ...]
827--
828{-# INLINE_NORMAL gintercalateSuffix #-}
829gintercalateSuffix
830 :: Monad m
831 => Unfold m a c -> Stream m a -> Unfold m b c -> Stream m b -> Stream m c
832gintercalateSuffix
833 (Unfold istep1 inject1) (Stream step1 state1)
834 (Unfold istep2 inject2) (Stream step2 state2) =
835 Stream step (ICUFirst state1 state2)
837 where
839 {-# INLINE_LATE step #-}
840 step gst (ICUFirst s1 s2) = do
841 r <- step1 (adaptState gst) s1
842 case r of
843 Yield a s -> do
844 i <- inject1 a
845 i `seq` return (Skip (ICUFirstInner s s2 i))
846 Skip s -> return $ Skip (ICUFirst s s2)
847 Stop -> return Stop
849 step gst (ICUFirstOnly s1) = do
850 r <- step1 (adaptState gst) s1
851 case r of
852 Yield a s -> do
853 i <- inject1 a
854 i `seq` return (Skip (ICUFirstOnlyInner s i))
855 Skip s -> return $ Skip (ICUFirstOnly s)
856 Stop -> return Stop
858 step _ (ICUFirstInner s1 s2 i1) = do
859 r <- istep1 i1
860 return $ case r of
861 Yield x i' -> Yield x (ICUFirstInner s1 s2 i')
862 Skip i' -> Skip (ICUFirstInner s1 s2 i')
863 Stop -> Skip (ICUSecond s1 s2)
865 step _ (ICUFirstOnlyInner s1 i1) = do
866 r <- istep1 i1
867 return $ case r of
868 Yield x i' -> Yield x (ICUFirstOnlyInner s1 i')
869 Skip i' -> Skip (ICUFirstOnlyInner s1 i')
870 Stop -> Skip (ICUFirstOnly s1)
872 step gst (ICUSecond s1 s2) = do
873 r <- step2 (adaptState gst) s2
874 case r of
875 Yield a s -> do
876 i <- inject2 a
877 i `seq` return (Skip (ICUSecondInner s1 s i))
878 Skip s -> return $ Skip (ICUSecond s1 s)
879 Stop -> return $ Skip (ICUFirstOnly s1)
881 step _ (ICUSecondInner s1 s2 i2) = do
882 r <- istep2 i2
883 return $ case r of
884 Yield x i' -> Yield x (ICUSecondInner s1 s2 i')
885 Skip i' -> Skip (ICUSecondInner s1 s2 i')
886 Stop -> Skip (ICUFirst s1 s2)
888 step _ (ICUSecondOnly _s2) = undefined
889 step _ (ICUSecondOnlyInner _s2 _i2) = undefined
891data ICALState s1 s2 i1 i2 a =
892 ICALFirst s1 s2
893 -- | ICALFirstYield s1 s2 i1
894 | ICALFirstInner s1 s2 i1
895 | ICALFirstOnly s1
896 | ICALFirstOnlyInner s1 i1
897 | ICALSecondInject s1 s2
898 | ICALFirstInject s1 s2 i2
899 -- | ICALFirstBuf s1 s2 i1 i2
900 | ICALSecondInner s1 s2 i1 i2
901 -- -- | ICALSecondInner s1 s2 i1 i2 a
902 -- -- | ICALFirstResume s1 s2 i1 i2 a
904-- | Interleave streams (full streams, not the elements) unfolded from two
905-- input streams and concat. Stop when the first stream stops. If the second
906-- stream ends before the first one then first stream still keeps running alone
907-- without any interleaving with the second stream.
908--
909-- [a1, a2, ... an] [b1, b2 ...]
910-- => [streamA1, streamA2, ... streamAn] [streamB1, streamB2, ...]
911-- => [streamA1, streamB1, streamA2...StreamAn, streamBn]
912-- => [a11, a12, ...a1j, b11, b12, ...b1k, a21, a22, ...]
913--
914{-# INLINE_NORMAL gintercalate #-}
915gintercalate
916 :: Monad m
917 => Unfold m a c -> Stream m a -> Unfold m b c -> Stream m b -> Stream m c
918gintercalate
919 (Unfold istep1 inject1) (Stream step1 state1)
920 (Unfold istep2 inject2) (Stream step2 state2) =
921 Stream step (ICALFirst state1 state2)
923 where
925 {-# INLINE_LATE step #-}
926 step gst (ICALFirst s1 s2) = do
927 r <- step1 (adaptState gst) s1
928 case r of
929 Yield a s -> do
930 i <- inject1 a
931 i `seq` return (Skip (ICALFirstInner s s2 i))
932 -- i `seq` return (Skip (ICALFirstYield s s2 i))
933 Skip s -> return $ Skip (ICALFirst s s2)
934 Stop -> return Stop
936 {-
937 step _ (ICALFirstYield s1 s2 i1) = do
938 r <- istep1 i1
939 return $ case r of
940 Yield x i' -> Yield x (ICALFirstInner s1 s2 i')
941 Skip i' -> Skip (ICALFirstYield s1 s2 i')
942 Stop -> Skip (ICALFirst s1 s2)
943 -}
945 step _ (ICALFirstInner s1 s2 i1) = do
946 r <- istep1 i1
947 return $ case r of
948 Yield x i' -> Yield x (ICALFirstInner s1 s2 i')
949 Skip i' -> Skip (ICALFirstInner s1 s2 i')
950 Stop -> Skip (ICALSecondInject s1 s2)
952 step gst (ICALFirstOnly s1) = do
953 r <- step1 (adaptState gst) s1
954 case r of
955 Yield a s -> do
956 i <- inject1 a
957 i `seq` return (Skip (ICALFirstOnlyInner s i))
958 Skip s -> return $ Skip (ICALFirstOnly s)
959 Stop -> return Stop
961 step _ (ICALFirstOnlyInner s1 i1) = do
962 r <- istep1 i1
963 return $ case r of
964 Yield x i' -> Yield x (ICALFirstOnlyInner s1 i')
965 Skip i' -> Skip (ICALFirstOnlyInner s1 i')
966 Stop -> Skip (ICALFirstOnly s1)
968 -- We inject the second stream even before checking if the first stream
969 -- would yield any more elements. There is no clear choice whether we
970 -- should do this before or after that. Doing it after may make the state
971 -- machine a bit simpler though.
972 step gst (ICALSecondInject s1 s2) = do
973 r <- step2 (adaptState gst) s2
974 case r of
975 Yield a s -> do
976 i <- inject2 a
977 i `seq` return (Skip (ICALFirstInject s1 s i))
978 Skip s -> return $ Skip (ICALSecondInject s1 s)
979 Stop -> return $ Skip (ICALFirstOnly s1)
981 step gst (ICALFirstInject s1 s2 i2) = do
982 r <- step1 (adaptState gst) s1
983 case r of
984 Yield a s -> do
985 i <- inject1 a
986 i `seq` return (Skip (ICALSecondInner s s2 i i2))
987 -- i `seq` return (Skip (ICALFirstBuf s s2 i i2))
988 Skip s -> return $ Skip (ICALFirstInject s s2 i2)
989 Stop -> return Stop
991 {-
992 step _ (ICALFirstBuf s1 s2 i1 i2) = do
993 r <- istep1 i1
994 return $ case r of
995 Yield x i' -> Skip (ICALSecondInner s1 s2 i' i2 x)
996 Skip i' -> Skip (ICALFirstBuf s1 s2 i' i2)
997 Stop -> Stop
999 step _ (ICALSecondInner s1 s2 i1 i2 v) = do
1000 r <- istep2 i2
1001 return $ case r of
1002 Yield x i' -> Yield x (ICALSecondInner s1 s2 i1 i' v)
1003 Skip i' -> Skip (ICALSecondInner s1 s2 i1 i' v)
1004 Stop -> Skip (ICALFirstResume s1 s2 i1 i2 v)
1005 -}
1007 step _ (ICALSecondInner s1 s2 i1 i2) = do
1008 r <- istep2 i2
1009 return $ case r of
1010 Yield x i' -> Yield x (ICALSecondInner s1 s2 i1 i')
1011 Skip i' -> Skip (ICALSecondInner s1 s2 i1 i')
1012 Stop -> Skip (ICALFirstInner s1 s2 i1)
1013 -- Stop -> Skip (ICALFirstResume s1 s2 i1 i2)
1015 {-
1016 step _ (ICALFirstResume s1 s2 i1 i2 x) = do
1017 return $ Yield x (ICALFirstInner s1 s2 i1 i2)
1018 -}
1020------------------------------------------------------------------------------
1021-- Folding
1022------------------------------------------------------------------------------
1024{-# ANN type FIterState Fuse #-}
1025data FIterState s f m a b
1026 = FIterInit s f
1027 | forall fs. FIterStream s (fs -> a -> m (FL.Step fs b)) fs (fs -> m b)
1028 | FIterYield b (FIterState s f m a b)
1029 | FIterStop
1031{-# INLINE_NORMAL foldIterateM #-}
1032foldIterateM ::
1033 Monad m => (b -> m (FL.Fold m a b)) -> m b -> Stream m a -> Stream m b
1034foldIterateM func seed0 (Stream step state) =
1035 Stream stepOuter (FIterInit state seed0)
1037 where
1039 {-# INLINE iterStep #-}
1040 iterStep from st fstep extract = do
1041 res <- from
1042 return
1043 $ Skip
1044 $ case res of
1045 FL.Partial fs -> FIterStream st fstep fs extract
1046 FL.Done fb -> FIterYield fb $ FIterInit st (return fb)
1048 {-# INLINE_LATE stepOuter #-}
1049 stepOuter _ (FIterInit st seed) = do
1050 (FL.Fold fstep initial extract) <- seed >>= func
1051 iterStep initial st fstep extract
1052 stepOuter gst (FIterStream st fstep fs extract) = do
1053 r <- step (adaptState gst) st
1054 case r of
1055 Yield x s -> do
1056 iterStep (fstep fs x) s fstep extract
1057 Skip s -> return $ Skip $ FIterStream s fstep fs extract
1058 Stop -> do
1059 b <- extract fs
1060 return $ Skip $ FIterYield b FIterStop
1061 stepOuter _ (FIterYield a next) = return $ Yield a next
1062 stepOuter _ FIterStop = return Stop
1064{-# ANN type CIterState Fuse #-}
1065data CIterState s f fs b
1066 = CIterInit s f
1067 | CIterConsume s fs
1068 | CIterYield b (CIterState s f fs b)
1069 | CIterStop
1071-- | Like 'foldIterateM' but using the 'Refold' type instead. This could be
1072-- much more efficient due to stream fusion.
1074-- /Internal/
1075{-# INLINE_NORMAL refoldIterateM #-}
1076refoldIterateM ::
1077 Monad m => Refold m b a b -> m b -> Stream m a -> Stream m b
1078refoldIterateM (Refold fstep finject fextract) initial (Stream step state) =
1079 Stream stepOuter (CIterInit state initial)
1081 where
1083 {-# INLINE iterStep #-}
1084 iterStep st action = do
1085 res <- action
1086 return
1087 $ Skip
1088 $ case res of
1089 FL.Partial fs -> CIterConsume st fs
1090 FL.Done fb -> CIterYield fb $ CIterInit st (return fb)
1092 {-# INLINE_LATE stepOuter #-}
1093 stepOuter _ (CIterInit st action) = do
1094 iterStep st (action >>= finject)
1095 stepOuter gst (CIterConsume st fs) = do
1096 r <- step (adaptState gst) st
1097 case r of
1098 Yield x s -> iterStep s (fstep fs x)
1099 Skip s -> return $ Skip $ CIterConsume s fs
1100 Stop -> do
1101 b <- fextract fs
1102 return $ Skip $ CIterYield b CIterStop
1103 stepOuter _ (CIterYield a next) = return $ Yield a next
1104 stepOuter _ CIterStop = return Stop
1106-- "n" elements at the end are dropped by the fold.
1107{-# INLINE sliceBy #-}
1108sliceBy :: Monad m => Fold m a Int -> Int -> Refold m (Int, Int) a (Int, Int)
1109sliceBy (Fold step1 initial1 extract1) n = Refold step inject extract
1111 where
1113 inject (i, len) = do
1114 r <- initial1
1115 return $ case r of
1116 Partial s -> Partial $ Tuple' (i + len + n) s
1117 Done l -> Done (i, l)
1119 step (Tuple' i s) x = do
1120 r <- step1 s x
1121 return $ case r of
1122 Partial s1 -> Partial $ Tuple' i s1
1123 Done len -> Done (i, len)
1125 extract (Tuple' i s) = (i,) <$> extract1 s
1127{-# INLINE sliceOnSuffix #-}
1128sliceOnSuffix :: Monad m => (a -> Bool) -> Stream m a -> Stream m (Int, Int)
1129sliceOnSuffix predicate =
1130 -- Scan the stream with the given refold
1131 refoldIterateM
1132 (sliceBy (FL.takeEndBy_ predicate FL.length) 1)
1133 (return (-1, 0))
1135------------------------------------------------------------------------------
1136-- Parsing
1137------------------------------------------------------------------------------
1139{-# ANN type ParseChunksState Fuse #-}
1140data ParseChunksState x inpBuf st pst =
1141 ParseChunksInit inpBuf st
1142 | ParseChunksInitBuf inpBuf
1143 | ParseChunksInitLeftOver inpBuf
1144 | ParseChunksStream st inpBuf !pst
1145 | ParseChunksStop inpBuf !pst
1146 | ParseChunksBuf inpBuf st inpBuf !pst
1147 | ParseChunksExtract inpBuf inpBuf !pst
1148 | ParseChunksYield x (ParseChunksState x inpBuf st pst)
1150-- XXX This is in fact parseMany1 (a la foldMany1). Do we need a parseMany as
1151-- well?
1152{-# INLINE_NORMAL parseMany #-}
1153parseMany
1154 :: MonadThrow m
1155 => PRD.Parser m a b
1156 -> Stream m a
1157 -> Stream m b
1158parseMany (PRD.Parser pstep initial extract) (Stream step state) =
1159 Stream stepOuter (ParseChunksInit [] state)
1161 where
1163 {-# INLINE_LATE stepOuter #-}
1164 -- Buffer is empty, get the first element from the stream, initialize the
1165 -- fold and then go to stream processing loop.
1166 stepOuter gst (ParseChunksInit [] st) = do
1167 r <- step (adaptState gst) st
1168 case r of
1169 Yield x s -> do
1170 res <- initial
1171 case res of
1172 PRD.IPartial ps ->
1173 return $ Skip $ ParseChunksBuf [x] s [] ps
1174 PRD.IDone pb ->
1175 let next = ParseChunksInit [x] s
1176 in return $ Skip $ ParseChunksYield pb next
1177 PRD.IError err -> throwM $ ParseError err
1178 Skip s -> return $ Skip $ ParseChunksInit [] s
1179 Stop -> return Stop
1181 -- Buffer is not empty, go to buffered processing loop
1182 stepOuter _ (ParseChunksInit src st) = do
1183 res <- initial
1184 case res of
1185 PRD.IPartial ps ->
1186 return $ Skip $ ParseChunksBuf src st [] ps
1187 PRD.IDone pb ->
1188 let next = ParseChunksInit src st
1189 in return $ Skip $ ParseChunksYield pb next
1190 PRD.IError err -> throwM $ ParseError err
1192 -- This is simplified ParseChunksInit
1193 stepOuter _ (ParseChunksInitBuf src) = do
1194 res <- initial
1195 case res of
1196 PRD.IPartial ps ->
1197 return $ Skip $ ParseChunksExtract src [] ps
1198 PRD.IDone pb ->
1199 let next = ParseChunksInitBuf src
1200 in return $ Skip $ ParseChunksYield pb next
1201 PRD.IError err -> throwM $ ParseError err
1203 -- XXX we just discard any leftover input at the end
1204 stepOuter _ (ParseChunksInitLeftOver _) = return Stop
1206 -- Buffer is empty, process elements from the stream
1207 stepOuter gst (ParseChunksStream st buf pst) = do
1208 r <- step (adaptState gst) st
1209 case r of
1210 Yield x s -> do
1211 pRes <- pstep pst x
1212 case pRes of
1213 PR.Partial 0 pst1 ->
1214 return $ Skip $ ParseChunksStream s [] pst1
1215 PR.Partial n pst1 -> do
1216 assert (n <= length (x:buf)) (return ())
1217 let src0 = Prelude.take n (x:buf)
1218 src = Prelude.reverse src0
1219 return $ Skip $ ParseChunksBuf src s [] pst1
1220 PR.Continue 0 pst1 ->
1221 return $ Skip $ ParseChunksStream s (x:buf) pst1
1222 PR.Continue n pst1 -> do
1223 assert (n <= length (x:buf)) (return ())
1224 let (src0, buf1) = splitAt n (x:buf)
1225 src = Prelude.reverse src0
1226 return $ Skip $ ParseChunksBuf src s buf1 pst1
1227 PR.Done 0 b -> do
1228 return $ Skip $
1229 ParseChunksYield b (ParseChunksInit [] s)
1230 PR.Done n b -> do
1231 assert (n <= length (x:buf)) (return ())
1232 let src = Prelude.reverse (Prelude.take n (x:buf))
1233 return $ Skip $
1234 ParseChunksYield b (ParseChunksInit src s)
1235 PR.Error err -> throwM $ ParseError err
1236 Skip s -> return $ Skip $ ParseChunksStream s buf pst
1237 Stop -> return $ Skip $ ParseChunksStop buf pst
1239 -- go back to stream processing mode
1240 stepOuter _ (ParseChunksBuf [] s buf pst) =
1241 return $ Skip $ ParseChunksStream s buf pst
1243 -- buffered processing loop
1244 stepOuter _ (ParseChunksBuf (x:xs) s buf pst) = do
1245 pRes <- pstep pst x
1246 case pRes of
1247 PR.Partial 0 pst1 ->
1248 return $ Skip $ ParseChunksBuf xs s [] pst1
1249 PR.Partial n pst1 -> do
1250 assert (n <= length (x:buf)) (return ())
1251 let src0 = Prelude.take n (x:buf)
1252 src = Prelude.reverse src0 ++ xs
1253 return $ Skip $ ParseChunksBuf src s [] pst1
1254 PR.Continue 0 pst1 ->
1255 return $ Skip $ ParseChunksBuf xs s (x:buf) pst1
1256 PR.Continue n pst1 -> do
1257 assert (n <= length (x:buf)) (return ())
1258 let (src0, buf1) = splitAt n (x:buf)
1259 src = Prelude.reverse src0 ++ xs
1260 return $ Skip $ ParseChunksBuf src s buf1 pst1
1261 PR.Done 0 b ->
1262 return $ Skip $ ParseChunksYield b (ParseChunksInit xs s)
1263 PR.Done n b -> do
1264 assert (n <= length (x:buf)) (return ())
1265 let src = Prelude.reverse (Prelude.take n (x:buf)) ++ xs
1266 return $ Skip $ ParseChunksYield b (ParseChunksInit src s)
1267 PR.Error err -> throwM $ ParseError err
1269 -- This is simplified ParseChunksBuf
1270 stepOuter _ (ParseChunksExtract [] buf pst) =
1271 return $ Skip $ ParseChunksStop buf pst
1273 stepOuter _ (ParseChunksExtract (x:xs) buf pst) = do
1274 pRes <- pstep pst x
1275 case pRes of
1276 PR.Partial 0 pst1 ->
1277 return $ Skip $ ParseChunksExtract xs [] pst1
1278 PR.Partial n pst1 -> do
1279 assert (n <= length (x:buf)) (return ())
1280 let src0 = Prelude.take n (x:buf)
1281 src = Prelude.reverse src0 ++ xs
1282 return $ Skip $ ParseChunksExtract src [] pst1
1283 PR.Continue 0 pst1 ->
1284 return $ Skip $ ParseChunksExtract xs (x:buf) pst1
1285 PR.Continue n pst1 -> do
1286 assert (n <= length (x:buf)) (return ())
1287 let (src0, buf1) = splitAt n (x:buf)
1288 src = Prelude.reverse src0 ++ xs
1289 return $ Skip $ ParseChunksExtract src buf1 pst1
1290 PR.Done 0 b ->
1291 return $ Skip $ ParseChunksYield b (ParseChunksInitBuf xs)
1292 PR.Done n b -> do
1293 assert (n <= length (x:buf)) (return ())
1294 let src = Prelude.reverse (Prelude.take n (x:buf)) ++ xs
1295 return $ Skip $ ParseChunksYield b (ParseChunksInitBuf src)
1296 PR.Error err -> throwM $ ParseError err
1298 -- This is simplified ParseChunksExtract
1299 stepOuter _ (ParseChunksStop buf pst) = do
1300 pRes <- extract pst
1301 case pRes of
1302 PR.Partial _ _ -> error "Bug: parseMany: Partial in extract"
1303 PR.Continue 0 _ ->
1304 error "parseMany: extract, Continue 0 creates infinite loop"
1305 PR.Continue n pst1 -> do
1306 assert (n <= length buf) (return ())
1307 let (src0, buf1) = splitAt n buf
1308 src = Prelude.reverse src0
1309 return $ Skip $ ParseChunksExtract src buf1 pst1
1310 PR.Done 0 b -> do
1311 return $ Skip $
1312 ParseChunksYield b (ParseChunksInitLeftOver [])
1313 PR.Done n b -> do
1314 assert (n <= length buf) (return ())
1315 let src = Prelude.reverse (Prelude.take n buf)
1316 return $ Skip $
1317 ParseChunksYield b (ParseChunksInitBuf src)
1318 PR.Error err -> throwM $ ParseError err
1320 stepOuter _ (ParseChunksYield a next) = return $ Yield a next
1322{-# ANN type ConcatParseState Fuse #-}
1323data ConcatParseState b inpBuf st p m a =
1324 ConcatParseInit inpBuf st p
1325 | ConcatParseInitBuf inpBuf p
1326 | ConcatParseInitLeftOver inpBuf
1327 | forall s. ConcatParseStop
1328 inpBuf (s -> a -> m (PRD.Step s b)) s (s -> m (PRD.Step s b))
1329 | forall s. ConcatParseStream
1330 st inpBuf (s -> a -> m (PRD.Step s b)) s (s -> m (PRD.Step s b))
1331 | forall s. ConcatParseBuf
1332 inpBuf st inpBuf (s -> a -> m (PRD.Step s b)) s (s -> m (PRD.Step s b))
1333 | forall s. ConcatParseExtract
1334 inpBuf inpBuf (s -> a -> m (PRD.Step s b)) s (s -> m (PRD.Step s b))
1335 | ConcatParseYield b (ConcatParseState b inpBuf st p m a)
1337-- XXX Review the changes
1338{-# INLINE_NORMAL parseIterate #-}
1339parseIterate
1340 :: MonadThrow m
1341 => (b -> PRD.Parser m a b)
1342 -> b
1343 -> Stream m a
1344 -> Stream m b
1345parseIterate func seed (Stream step state) =
1346 Stream stepOuter (ConcatParseInit [] state (func seed))
1348 where
1350 {-# INLINE_LATE stepOuter #-}
1351 -- Buffer is empty, go to stream processing loop
1352 stepOuter _ (ConcatParseInit [] st (PRD.Parser pstep initial extract)) = do
1353 res <- initial
1354 case res of
1355 PRD.IPartial ps ->
1356 return $ Skip $ ConcatParseStream st [] pstep ps extract
1357 PRD.IDone pb ->
1358 let next = ConcatParseInit [] st (func pb)
1359 in return $ Skip $ ConcatParseYield pb next
1360 PRD.IError err -> throwM $ ParseError err
1362 -- Buffer is not empty, go to buffered processing loop
1363 stepOuter _ (ConcatParseInit src st
1364 (PRD.Parser pstep initial extract)) = do
1365 res <- initial
1366 case res of
1367 PRD.IPartial ps ->
1368 return $ Skip $ ConcatParseBuf src st [] pstep ps extract
1369 PRD.IDone pb ->
1370 let next = ConcatParseInit src st (func pb)
1371 in return $ Skip $ ConcatParseYield pb next
1372 PRD.IError err -> throwM $ ParseError err
1374 -- This is simplified ConcatParseInit
1375 stepOuter _ (ConcatParseInitBuf src
1376 (PRD.Parser pstep initial extract)) = do
1377 res <- initial
1378 case res of
1379 PRD.IPartial ps ->
1380 return $ Skip $ ConcatParseExtract src [] pstep ps extract
1381 PRD.IDone pb ->
1382 let next = ConcatParseInitBuf src (func pb)
1383 in return $ Skip $ ConcatParseYield pb next
1384 PRD.IError err -> throwM $ ParseError err
1386 -- XXX we just discard any leftover input at the end
1387 stepOuter _ (ConcatParseInitLeftOver _) = return Stop
1389 -- Buffer is empty process elements from the stream
1390 stepOuter gst (ConcatParseStream st buf pstep pst extract) = do
1391 r <- step (adaptState gst) st
1392 case r of
1393 Yield x s -> do
1394 pRes <- pstep pst x
1395 case pRes of
1396 PR.Partial 0 pst1 ->
1397 return $ Skip $ ConcatParseStream s [] pstep pst1 extract
1398 PR.Partial n pst1 -> do
1399 assert (n <= length (x:buf)) (return ())
1400 let src0 = Prelude.take n (x:buf)
1401 src = Prelude.reverse src0
1402 return $ Skip $ ConcatParseBuf src s [] pstep pst1 extract
1403 -- PR.Continue 0 pst1 ->
1404 -- return $ Skip $ ConcatParseStream s (x:buf) pst1
1405 PR.Continue n pst1 -> do
1406 assert (n <= length (x:buf)) (return ())
1407 let (src0, buf1) = splitAt n (x:buf)
1408 src = Prelude.reverse src0
1409 return $ Skip $ ConcatParseBuf src s buf1 pstep pst1 extract
1410 -- XXX Specialize for Stop 0 common case?
1411 PR.Done n b -> do
1412 assert (n <= length (x:buf)) (return ())
1413 let src = Prelude.reverse (Prelude.take n (x:buf))
1414 return $ Skip $
1415 ConcatParseYield b (ConcatParseInit src s (func b))
1416 PR.Error err -> throwM $ ParseError err
1417 Skip s -> return $ Skip $ ConcatParseStream s buf pstep pst extract
1418 Stop -> return $ Skip $ ConcatParseStop buf pstep pst extract
1420 -- go back to stream processing mode
1421 stepOuter _ (ConcatParseBuf [] s buf pstep ps extract) =
1422 return $ Skip $ ConcatParseStream s buf pstep ps extract
1424 -- buffered processing loop
1425 stepOuter _ (ConcatParseBuf (x:xs) s buf pstep pst extract) = do
1426 pRes <- pstep pst x
1427 case pRes of
1428 PR.Partial 0 pst1 ->
1429 return $ Skip $ ConcatParseBuf xs s [] pstep pst1 extract
1430 PR.Partial n pst1 -> do
1431 assert (n <= length (x:buf)) (return ())
1432 let src0 = Prelude.take n (x:buf)
1433 src = Prelude.reverse src0 ++ xs
1434 return $ Skip $ ConcatParseBuf src s [] pstep pst1 extract
1435 -- PR.Continue 0 pst1 -> return $ Skip $ ConcatParseBuf xs s (x:buf) pst1
1436 PR.Continue n pst1 -> do
1437 assert (n <= length (x:buf)) (return ())
1438 let (src0, buf1) = splitAt n (x:buf)
1439 src = Prelude.reverse src0 ++ xs
1440 return $ Skip $ ConcatParseBuf src s buf1 pstep pst1 extract
1441 -- XXX Specialize for Stop 0 common case?
1442 PR.Done n b -> do
1443 assert (n <= length (x:buf)) (return ())
1444 let src = Prelude.reverse (Prelude.take n (x:buf)) ++ xs
1445 return $ Skip $ ConcatParseYield b
1446 (ConcatParseInit src s (func b))
1447 PR.Error err -> throwM $ ParseError err
1449 -- This is simplified ConcatParseBuf
1450 stepOuter _ (ConcatParseExtract [] buf pstep pst extract) =
1451 return $ Skip $ ConcatParseStop buf pstep pst extract
1453 stepOuter _ (ConcatParseExtract (x:xs) buf pstep pst extract) = do
1454 pRes <- pstep pst x
1455 case pRes of
1456 PR.Partial 0 pst1 ->
1457 return $ Skip $ ConcatParseExtract xs [] pstep pst1 extract
1458 PR.Partial n pst1 -> do
1459 assert (n <= length (x:buf)) (return ())
1460 let src0 = Prelude.take n (x:buf)
1461 src = Prelude.reverse src0 ++ xs
1462 return $ Skip $ ConcatParseExtract src [] pstep pst1 extract
1463 PR.Continue 0 pst1 ->
1464 return $ Skip $ ConcatParseExtract xs (x:buf) pstep pst1 extract
1465 PR.Continue n pst1 -> do
1466 assert (n <= length (x:buf)) (return ())
1467 let (src0, buf1) = splitAt n (x:buf)
1468 src = Prelude.reverse src0 ++ xs
1469 return $ Skip $ ConcatParseExtract src buf1 pstep pst1 extract
1470 PR.Done 0 b ->
1471 return $ Skip $ ConcatParseYield b (ConcatParseInitBuf xs (func b))
1472 PR.Done n b -> do
1473 assert (n <= length (x:buf)) (return ())
1474 let src = Prelude.reverse (Prelude.take n (x:buf)) ++ xs
1475 return $ Skip $ ConcatParseYield b (ConcatParseInitBuf src (func b))
1476 PR.Error err -> throwM $ ParseError err
1478 -- This is simplified ConcatParseExtract
1479 stepOuter _ (ConcatParseStop buf pstep pst extract) = do
1480 pRes <- extract pst
1481 case pRes of
1482 PR.Partial _ _ -> error "Bug: parseIterate: Partial in extract"
1483 PR.Continue 0 _ ->
1484 error "parseIterate: extract, Continue 0 creates infinite loop"
1485 PR.Continue n pst1 -> do
1486 assert (n <= length buf) (return ())
1487 let (src0, buf1) = splitAt n buf
1488 src = Prelude.reverse src0
1489 return $ Skip $ ConcatParseExtract src buf1 pstep pst1 extract
1490 PR.Done 0 b -> do
1491 return $ Skip $
1492 ConcatParseYield b (ConcatParseInitLeftOver [])
1493 PR.Done n b -> do
1494 assert (n <= length buf) (return ())
1495 let src = Prelude.reverse (Prelude.take n buf)
1496 return $ Skip $
1497 ConcatParseYield b (ConcatParseInitBuf src (func b))
1498 PR.Error err -> throwM $ ParseError err
1500 stepOuter _ (ConcatParseYield a next) = return $ Yield a next
1502------------------------------------------------------------------------------
1503-- Grouping
1504------------------------------------------------------------------------------
1506data GroupByState st fs a b
1507 = GroupingInit st
1508 | GroupingDo st !fs
1509 | GroupingInitWith st !a
1510 | GroupingDoWith st !fs !a
1511 | GroupingYield !b (GroupByState st fs a b)
1512 | GroupingDone
1514{-# INLINE_NORMAL groupsBy #-}
1515groupsBy :: Monad m
1516 => (a -> a -> Bool)
1517 -> Fold m a b
1518 -> Stream m a
1519 -> Stream m b
1521groupsBy eq fld = parseMany (PRD.groupBy eq fld)
1523groupsBy cmp (Fold fstep initial done) (Stream step state) =
1524 Stream stepOuter (GroupingInit state)
1526 where
1528 {-# INLINE_LATE stepOuter #-}
1529 stepOuter _ (GroupingInit st) = do
1530 -- XXX Note that if the stream stops without yielding a single element
1531 -- in the group we discard the "initial" effect.
1532 res <- initial
1533 return
1534 $ case res of
1535 FL.Partial s -> Skip $ GroupingDo st s
1536 FL.Done b -> Yield b $ GroupingInit st
1537 stepOuter gst (GroupingDo st fs) = do
1538 res <- step (adaptState gst) st
1539 case res of
1540 Yield x s -> do
1541 r <- fstep fs x
1542 case r of
1543 FL.Partial fs1 -> go SPEC x s fs1
1544 FL.Done b -> return $ Yield b (GroupingInit s)
1545 Skip s -> return $ Skip $ GroupingDo s fs
1546 Stop -> return Stop
1548 where
1550 go !_ prev stt !acc = do
1551 res <- step (adaptState gst) stt
1552 case res of
1553 Yield x s -> do
1554 if cmp x prev
1555 then do
1556 r <- fstep acc x
1557 case r of
1558 FL.Partial fs1 -> go SPEC prev s fs1
1559 FL.Done b -> return $ Yield b (GroupingInit s)
1560 else do
1561 r <- done acc
1562 return $ Yield r (GroupingInitWith s x)
1563 Skip s -> go SPEC prev s acc
1564 Stop -> done acc >>= \r -> return $ Yield r GroupingDone
1565 stepOuter _ (GroupingInitWith st x) = do
1566 res <- initial
1567 return
1568 $ case res of
1569 FL.Partial s -> Skip $ GroupingDoWith st s x
1570 FL.Done b -> Yield b $ GroupingInitWith st x
1571 stepOuter gst (GroupingDoWith st fs prev) = do
1572 res <- fstep fs prev
1573 case res of
1574 FL.Partial fs1 -> go SPEC st fs1
1575 FL.Done b -> return $ Yield b (GroupingInit st)
1577 where
1579 -- XXX code duplicated from the previous equation
1580 go !_ stt !acc = do
1581 res <- step (adaptState gst) stt
1582 case res of
1583 Yield x s -> do
1584 if cmp x prev
1585 then do
1586 r <- fstep acc x
1587 case r of
1588 FL.Partial fs1 -> go SPEC s fs1
1589 FL.Done b -> return $ Yield b (GroupingInit s)
1590 else do
1591 r <- done acc
1592 return $ Yield r (GroupingInitWith s x)
1593 Skip s -> go SPEC s acc
1594 Stop -> done acc >>= \r -> return $ Yield r GroupingDone
1595 stepOuter _ (GroupingYield _ _) = error "groupsBy: Unreachable"
1596 stepOuter _ GroupingDone = return Stop
1598{-# INLINE_NORMAL groupsRollingBy #-}
1599groupsRollingBy :: Monad m
1600 => (a -> a -> Bool)
1601 -> Fold m a b
1602 -> Stream m a
1603 -> Stream m b
1605groupsRollingBy eq fld = parseMany (PRD.groupByRolling eq fld)
1607groupsRollingBy cmp (Fold fstep initial done) (Stream step state) =
1608 Stream stepOuter (GroupingInit state)
1610 where
1612 {-# INLINE_LATE stepOuter #-}
1613 stepOuter _ (GroupingInit st) = do
1614 -- XXX Note that if the stream stops without yielding a single element
1615 -- in the group we discard the "initial" effect.
1616 res <- initial
1617 return
1618 $ case res of
1619 FL.Partial fs -> Skip $ GroupingDo st fs
1620 FL.Done fb -> Yield fb $ GroupingInit st
1621 stepOuter gst (GroupingDo st fs) = do
1622 res <- step (adaptState gst) st
1623 case res of
1624 Yield x s -> do
1625 r <- fstep fs x
1626 case r of
1627 FL.Partial fs1 -> go SPEC x s fs1
1628 FL.Done fb -> return $ Yield fb (GroupingInit s)
1629 Skip s -> return $ Skip $ GroupingDo s fs
1630 Stop -> return Stop
1632 where
1634 go !_ prev stt !acc = do
1635 res <- step (adaptState gst) stt
1636 case res of
1637 Yield x s -> do
1638 if cmp prev x
1639 then do
1640 r <- fstep acc x
1641 case r of
1642 FL.Partial fs1 -> go SPEC x s fs1
1643 FL.Done b -> return $ Yield b (GroupingInit s)
1644 else do
1645 r <- done acc
1646 return $ Yield r (GroupingInitWith s x)
1647 Skip s -> go SPEC prev s acc
1648 Stop -> done acc >>= \r -> return $ Yield r GroupingDone
1649 stepOuter _ (GroupingInitWith st x) = do
1650 res <- initial
1651 return
1652 $ case res of
1653 FL.Partial s -> Skip $ GroupingDoWith st s x
1654 FL.Done b -> Yield b $ GroupingInitWith st x
1655 stepOuter gst (GroupingDoWith st fs previous) = do
1656 res <- fstep fs previous
1657 case res of
1658 FL.Partial s -> go SPEC previous st s
1659 FL.Done b -> return $ Yield b (GroupingInit st)
1661 where
1663 -- XXX GHC: groupsBy has one less parameter in this go loop and it
1664 -- fuses. However, groupsRollingBy does not fuse, removing the prev
1665 -- parameter makes it fuse. Something needs to be fixed in GHC. The
1666 -- workaround for this is noted in the comments below.
1667 go !_ prev !stt !acc = do
1668 res <- step (adaptState gst) stt
1669 case res of
1670 Yield x s -> do
1671 if cmp prev x
1672 then do
1673 r <- fstep acc x
1674 case r of
1675 FL.Partial fs1 -> go SPEC x s fs1
1676 FL.Done b -> return $ Yield b (GroupingInit st)
1677 else do
1678 {-
1679 r <- done acc
1680 return $ Yield r (GroupingInitWith s x)
1681 -}
1682 -- The code above does not let groupBy fuse. We use the
1683 -- alternative code below instead. Instead of jumping
1684 -- to GroupingInitWith state, we unroll the code of
1685 -- GroupingInitWith state here to help GHC with stream
1686 -- fusion.
1687 result <- initial
1688 r <- done acc
1689 return
1690 $ Yield r
1691 $ case result of
1692 FL.Partial fsi -> GroupingDoWith s fsi x
1693 FL.Done b -> GroupingYield b (GroupingInit s)
1694 Skip s -> go SPEC prev s acc
1695 Stop -> done acc >>= \r -> return $ Yield r GroupingDone
1696 stepOuter _ (GroupingYield r next) = return $ Yield r next
1697 stepOuter _ GroupingDone = return Stop
1699------------------------------------------------------------------------------
1700-- Splitting - by a predicate
1701------------------------------------------------------------------------------
1703data WordsByState st fs b
1704 = WordsByInit st
1705 | WordsByDo st !fs
1706 | WordsByDone
1707 | WordsByYield !b (WordsByState st fs b)
1709{-# INLINE_NORMAL wordsBy #-}
1710wordsBy :: Monad m => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1711wordsBy predicate (Fold fstep initial done) (Stream step state) =
1712 Stream stepOuter (WordsByInit state)
1714 where
1716 {-# INLINE_LATE stepOuter #-}
1717 stepOuter _ (WordsByInit st) = do
1718 res <- initial
1719 return
1720 $ case res of
1721 FL.Partial s -> Skip $ WordsByDo st s
1722 FL.Done b -> Yield b (WordsByInit st)
1724 stepOuter gst (WordsByDo st fs) = do
1725 res <- step (adaptState gst) st
1726 case res of
1727 Yield x s -> do
1728 if predicate x
1729 then do
1730 resi <- initial
1731 return
1732 $ case resi of
1733 FL.Partial fs1 -> Skip $ WordsByDo s fs1
1734 FL.Done b -> Yield b (WordsByInit s)
1735 else do
1736 r <- fstep fs x
1737 case r of
1738 FL.Partial fs1 -> go SPEC s fs1
1739 FL.Done b -> return $ Yield b (WordsByInit s)
1740 Skip s -> return $ Skip $ WordsByDo s fs
1741 Stop -> return Stop
1743 where
1745 go !_ stt !acc = do
1746 res <- step (adaptState gst) stt
1747 case res of
1748 Yield x s -> do
1749 if predicate x
1750 then do
1751 {-
1752 r <- done acc
1753 return $ Yield r (WordsByInit s)
1754 -}
1755 -- The above code does not fuse well. Need to check why
1756 -- GHC is not able to simplify it well. Using the code
1757 -- below, instead of jumping through the WordsByInit
1758 -- state always, we directly go to WordsByDo state in
1759 -- the common case of Partial.
1760 resi <- initial
1761 r <- done acc
1762 return
1763 $ Yield r
1764 $ case resi of
1765 FL.Partial fs1 -> WordsByDo s fs1
1766 FL.Done b -> WordsByYield b (WordsByInit s)
1767 else do
1768 r <- fstep acc x
1769 case r of
1770 FL.Partial fs1 -> go SPEC s fs1
1771 FL.Done b -> return $ Yield b (WordsByInit s)
1772 Skip s -> go SPEC s acc
1773 Stop -> done acc >>= \r -> return $ Yield r WordsByDone
1775 stepOuter _ WordsByDone = return Stop
1777 stepOuter _ (WordsByYield b next) = return $ Yield b next
1779------------------------------------------------------------------------------
1780-- Splitting on a sequence
1781------------------------------------------------------------------------------
1783-- String search algorithms:
1784-- http://www-igm.univ-mlv.fr/~lecroq/string/index.html
1787-- TODO can we unify the splitting operations using a splitting configuration
1788-- like in the split package.
1790data SplitStyle = Infix | Suffix | Prefix deriving (Eq, Show)
1791data SplitOptions = SplitOptions
1792 { style :: SplitStyle
1793 , withSep :: Bool -- ^ keep the separators in output
1794 -- , compact :: Bool -- ^ treat multiple consecutive separators as one
1795 -- , trimHead :: Bool -- ^ drop blank at head
1796 -- , trimTail :: Bool -- ^ drop blank at tail
1800-- XXX using "fs" as the last arg in Constructors may simplify the code a bit,
1801-- because we can use the constructor directly without having to create "jump"
1802-- functions.
1803{-# ANN type SplitOnSeqState Fuse #-}
1804data SplitOnSeqState rb rh ck w fs s b x =
1805 SplitOnSeqInit
1806 | SplitOnSeqYield b (SplitOnSeqState rb rh ck w fs s b x)
1807 | SplitOnSeqDone
1809 | SplitOnSeqEmpty !fs s
1811 | SplitOnSeqSingle !fs s x
1813 | SplitOnSeqWordInit !fs s
1814 | SplitOnSeqWordLoop !w s !fs
1815 | SplitOnSeqWordDone Int !fs !w
1817 | SplitOnSeqKRInit Int !fs s rb !rh
1818 | SplitOnSeqKRLoop fs s rb !rh !ck
1819 | SplitOnSeqKRCheck fs s rb !rh
1820 | SplitOnSeqKRDone Int !fs rb !rh
1822 | SplitOnSeqReinit (fs -> SplitOnSeqState rb rh ck w fs s b x)
1824{-# INLINE_NORMAL splitOnSeq #-}
1825splitOnSeq
1826 :: forall m a b. (MonadIO m, Unboxed a, Enum a, Eq a)
1827 => Array a
1828 -> Fold m a b
1829 -> Stream m a
1830 -> Stream m b
1831splitOnSeq patArr (Fold fstep initial done) (Stream step state) =
1832 Stream stepOuter SplitOnSeqInit
1834 where
1836 patLen = A.length patArr
1837 maxIndex = patLen - 1
1838 elemBits = SIZE_OF(a) * 8
1840 -- For word pattern case
1841 wordMask :: Word
1842 wordMask = (1 `shiftL` (elemBits * patLen)) - 1
1844 elemMask :: Word
1845 elemMask = (1 `shiftL` elemBits) - 1
1847 wordPat :: Word
1848 wordPat = wordMask .&. A.foldl' addToWord 0 patArr
1850 addToWord wd a = (wd `shiftL` elemBits) .|. fromIntegral (fromEnum a)
1852 -- For Rabin-Karp search
1853 k = 2891336453 :: Word32
1854 coeff = k ^ patLen
1856 addCksum cksum a = cksum * k + fromIntegral (fromEnum a)
1858 deltaCksum cksum old new =
1859 addCksum cksum new - coeff * fromIntegral (fromEnum old)
1861 -- XXX shall we use a random starting hash or 1 instead of 0?
1862 patHash = A.foldl' addCksum 0 patArr
1864 skip = return . Skip
1866 nextAfterInit nextGen stepRes =
1867 case stepRes of
1868 FL.Partial s -> nextGen s
1869 FL.Done b -> SplitOnSeqYield b (SplitOnSeqReinit nextGen)
1871 {-# INLINE yieldProceed #-}
1872 yieldProceed nextGen fs =
1873 initial >>= skip . SplitOnSeqYield fs . nextAfterInit nextGen
1875 {-# INLINE_LATE stepOuter #-}
1876 stepOuter _ SplitOnSeqInit = do
1877 res <- initial
1878 case res of
1879 FL.Partial acc ->
1880 if patLen == 0
1881 then return $ Skip $ SplitOnSeqEmpty acc state
1882 else if patLen == 1
1883 then do
1884 pat <- liftIO $ A.unsafeIndexIO 0 patArr
1885 return $ Skip $ SplitOnSeqSingle acc state pat
1886 else if SIZE_OF(a) * patLen
1887 <= sizeOf (undefined :: Word)
1888 then return $ Skip $ SplitOnSeqWordInit acc state
1889 else do
1890 (rb, rhead) <- liftIO $ RB.new patLen
1891 skip $ SplitOnSeqKRInit 0 acc state rb rhead
1892 FL.Done b -> skip $ SplitOnSeqYield b SplitOnSeqInit
1894 stepOuter _ (SplitOnSeqYield x next) = return $ Yield x next
1896 ---------------------------
1897 -- Checkpoint
1898 ---------------------------
1900 stepOuter _ (SplitOnSeqReinit nextGen) =
1901 initial >>= skip . nextAfterInit nextGen
1903 ---------------------------
1904 -- Empty pattern
1905 ---------------------------
1907 stepOuter gst (SplitOnSeqEmpty acc st) = do
1908 res <- step (adaptState gst) st
1909 case res of
1910 Yield x s -> do
1911 r <- fstep acc x
1912 b1 <-
1913 case r of
1914 FL.Partial acc1 -> done acc1
1915 FL.Done b -> return b
1916 let jump c = SplitOnSeqEmpty c s
1917 in yieldProceed jump b1
1918 Skip s -> skip (SplitOnSeqEmpty acc s)
1919 Stop -> return Stop
1921 -----------------
1922 -- Done
1923 -----------------
1925 stepOuter _ SplitOnSeqDone = return Stop
1927 -----------------
1928 -- Single Pattern
1929 -----------------
1931 stepOuter gst (SplitOnSeqSingle fs st pat) = do
1932 res <- step (adaptState gst) st
1933 case res of
1934 Yield x s -> do
1935 let jump c = SplitOnSeqSingle c s pat
1936 if pat == x
1937 then done fs >>= yieldProceed jump
1938 else do
1939 r <- fstep fs x
1940 case r of
1941 FL.Partial fs1 -> skip $ jump fs1
1942 FL.Done b -> yieldProceed jump b
1943 Skip s -> return $ Skip $ SplitOnSeqSingle fs s pat
1944 Stop -> do
1945 r <- done fs
1946 return $ Skip $ SplitOnSeqYield r SplitOnSeqDone
1948 ---------------------------
1949 -- Short Pattern - Shift Or
1950 ---------------------------
1952 stepOuter _ (SplitOnSeqWordDone 0 fs _) = do
1953 r <- done fs
1954 skip $ SplitOnSeqYield r SplitOnSeqDone
1955 stepOuter _ (SplitOnSeqWordDone n fs wrd) = do
1956 let old = elemMask .&. (wrd `shiftR` (elemBits * (n - 1)))
1957 r <- fstep fs (toEnum $ fromIntegral old)
1958 case r of
1959 FL.Partial fs1 -> skip $ SplitOnSeqWordDone (n - 1) fs1 wrd
1960 FL.Done b -> do
1961 let jump c = SplitOnSeqWordDone (n - 1) c wrd
1962 yieldProceed jump b
1964 stepOuter gst (SplitOnSeqWordInit fs st0) =
1965 go SPEC 0 0 st0
1967 where
1969 {-# INLINE go #-}
1970 go !_ !idx !wrd !st = do
1971 res <- step (adaptState gst) st
1972 case res of
1973 Yield x s -> do
1974 let wrd1 = addToWord wrd x
1975 if idx == maxIndex
1976 then do
1977 if wrd1 .&. wordMask == wordPat
1978 then do
1979 let jump c = SplitOnSeqWordInit c s
1980 done fs >>= yieldProceed jump
1981 else skip $ SplitOnSeqWordLoop wrd1 s fs
1982 else go SPEC (idx + 1) wrd1 s
1983 Skip s -> go SPEC idx wrd s
1984 Stop -> do
1985 if idx /= 0
1986 then skip $ SplitOnSeqWordDone idx fs wrd
1987 else do
1988 r <- done fs
1989 skip $ SplitOnSeqYield r SplitOnSeqDone
1991 stepOuter gst (SplitOnSeqWordLoop wrd0 st0 fs0) =
1992 go SPEC wrd0 st0 fs0
1994 where
1996 {-# INLINE go #-}
1997 go !_ !wrd !st !fs = do
1998 res <- step (adaptState gst) st
1999 case res of
2000 Yield x s -> do
2001 let jump c = SplitOnSeqWordInit c s
2002 wrd1 = addToWord wrd x
2003 old = (wordMask .&. wrd)
2004 `shiftR` (elemBits * (patLen - 1))
2005 r <- fstep fs (toEnum $ fromIntegral old)
2006 case r of
2007 FL.Partial fs1 -> do
2008 if wrd1 .&. wordMask == wordPat
2009 then done fs1 >>= yieldProceed jump
2010 else go SPEC wrd1 s fs1
2011 FL.Done b -> yieldProceed jump b
2012 Skip s -> go SPEC wrd s fs
2013 Stop -> skip $ SplitOnSeqWordDone patLen fs wrd
2015 -------------------------------
2016 -- General Pattern - Karp Rabin
2017 -------------------------------
2019 stepOuter gst (SplitOnSeqKRInit idx fs st rb rh) = do
2020 res <- step (adaptState gst) st
2021 case res of
2022 Yield x s -> do
2023 rh1 <- liftIO $ RB.unsafeInsert rb rh x
2024 if idx == maxIndex
2025 then do
2026 let fld = RB.unsafeFoldRing (RB.ringBound rb)
2027 let !ringHash = fld addCksum 0 rb
2028 if ringHash == patHash
2029 then skip $ SplitOnSeqKRCheck fs s rb rh1
2030 else skip $ SplitOnSeqKRLoop fs s rb rh1 ringHash
2031 else skip $ SplitOnSeqKRInit (idx + 1) fs s rb rh1
2032 Skip s -> skip $ SplitOnSeqKRInit idx fs s rb rh
2033 Stop -> do
2034 skip $ SplitOnSeqKRDone idx fs rb (RB.startOf rb)
2036 -- XXX The recursive "go" is more efficient than the state based recursion
2037 -- code commented out below. Perhaps its more efficient because of
2038 -- factoring out "rb" outside the loop.
2039 --
2040 stepOuter gst (SplitOnSeqKRLoop fs0 st0 rb rh0 cksum0) =
2041 go SPEC fs0 st0 rh0 cksum0
2043 where
2045 go !_ !fs !st !rh !cksum = do
2046 res <- step (adaptState gst) st
2047 case res of
2048 Yield x s -> do
2049 old <- liftIO $ peek rh
2050 let cksum1 = deltaCksum cksum old x
2051 r <- fstep fs old
2052 case r of
2053 FL.Partial fs1 -> do
2054 rh1 <- liftIO (RB.unsafeInsert rb rh x)
2055 if cksum1 == patHash
2056 then skip $ SplitOnSeqKRCheck fs1 s rb rh1
2057 else go SPEC fs1 s rh1 cksum1
2058 FL.Done b -> do
2059 let rst = RB.startOf rb
2060 jump c = SplitOnSeqKRInit 0 c s rb rst
2061 yieldProceed jump b
2062 Skip s -> go SPEC fs s rh cksum
2063 Stop -> skip $ SplitOnSeqKRDone patLen fs rb rh
2065 -- XXX The following code is 5 times slower compared to the recursive loop
2066 -- based code above. Need to investigate why. One possibility is that the
2067 -- go loop above does not thread around the ring buffer (rb). This code may
2068 -- be causing the state to bloat and getting allocated on each iteration.
2069 -- We can check the cmm/asm code to confirm. If so a good GHC solution to
2070 -- such problem is needed. One way to avoid this could be to use unboxed
2071 -- mutable state?
2072 {-
2073 stepOuter gst (SplitOnSeqKRLoop fs st rb rh cksum) = do
2074 res <- step (adaptState gst) st
2075 case res of
2076 Yield x s -> do
2077 old <- liftIO $ peek rh
2078 let cksum1 = deltaCksum cksum old x
2079 fs1 <- fstep fs old
2080 if (cksum1 == patHash)
2081 then do
2082 r <- done fs1
2083 skip $ SplitOnSeqYield r $ SplitOnSeqKRInit 0 s rb rh
2084 else do
2085 rh1 <- liftIO (RB.unsafeInsert rb rh x)
2086 skip $ SplitOnSeqKRLoop fs1 s rb rh1 cksum1
2087 Skip s -> skip $ SplitOnSeqKRLoop fs s rb rh cksum
2088 Stop -> skip $ SplitOnSeqKRDone patLen fs rb rh
2089 -}
2091 stepOuter _ (SplitOnSeqKRCheck fs st rb rh) = do
2092 if RB.unsafeEqArray rb rh patArr
2093 then do
2094 r <- done fs
2095 let rst = RB.startOf rb
2096 jump c = SplitOnSeqKRInit 0 c st rb rst
2097 yieldProceed jump r
2098 else skip $ SplitOnSeqKRLoop fs st rb rh patHash
2100 stepOuter _ (SplitOnSeqKRDone 0 fs _ _) = do
2101 r <- done fs
2102 skip $ SplitOnSeqYield r SplitOnSeqDone
2103 stepOuter _ (SplitOnSeqKRDone n fs rb rh) = do
2104 old <- liftIO $ peek rh
2105 let rh1 = RB.advance rb rh
2106 r <- fstep fs old
2107 case r of
2108 FL.Partial fs1 -> skip $ SplitOnSeqKRDone (n - 1) fs1 rb rh1
2109 FL.Done b -> do
2110 let jump c = SplitOnSeqKRDone (n - 1) c rb rh1
2111 yieldProceed jump b
2113{-# ANN type SplitOnSuffixSeqState Fuse #-}
2114data SplitOnSuffixSeqState rb rh ck w fs s b x =
2115 SplitOnSuffixSeqInit
2116 | SplitOnSuffixSeqYield b (SplitOnSuffixSeqState rb rh ck w fs s b x)
2117 | SplitOnSuffixSeqDone
2119 | SplitOnSuffixSeqEmpty !fs s
2121 | SplitOnSuffixSeqSingleInit !fs s x
2122 | SplitOnSuffixSeqSingle !fs s x
2124 | SplitOnSuffixSeqWordInit !fs s
2125 | SplitOnSuffixSeqWordLoop !w s !fs
2126 | SplitOnSuffixSeqWordDone Int !fs !w
2128 | SplitOnSuffixSeqKRInit Int !fs s rb !rh
2129 | SplitOnSuffixSeqKRInit1 !fs s rb !rh
2130 | SplitOnSuffixSeqKRLoop fs s rb !rh !ck
2131 | SplitOnSuffixSeqKRCheck fs s rb !rh
2132 | SplitOnSuffixSeqKRDone Int !fs rb !rh
2134 | SplitOnSuffixSeqReinit
2135 (fs -> SplitOnSuffixSeqState rb rh ck w fs s b x)
2137{-# INLINE_NORMAL splitOnSuffixSeq #-}
2138splitOnSuffixSeq
2139 :: forall m a b. (MonadIO m, Unboxed a, Enum a, Eq a)
2140 => Bool
2141 -> Array a
2142 -> Fold m a b
2143 -> Stream m a
2144 -> Stream m b
2145splitOnSuffixSeq withSep patArr (Fold fstep initial done) (Stream step state) =
2146 Stream stepOuter SplitOnSuffixSeqInit
2148 where
2150 patLen = A.length patArr
2151 maxIndex = patLen - 1
2152 elemBits = SIZE_OF(a) * 8
2154 -- For word pattern case
2155 wordMask :: Word
2156 wordMask = (1 `shiftL` (elemBits * patLen)) - 1
2158 elemMask :: Word
2159 elemMask = (1 `shiftL` elemBits) - 1
2161 wordPat :: Word
2162 wordPat = wordMask .&. A.foldl' addToWord 0 patArr
2164 addToWord wd a = (wd `shiftL` elemBits) .|. fromIntegral (fromEnum a)
2166 nextAfterInit nextGen stepRes =
2167 case stepRes of
2168 FL.Partial s -> nextGen s
2169 FL.Done b ->
2170 SplitOnSuffixSeqYield b (SplitOnSuffixSeqReinit nextGen)
2172 {-# INLINE yieldProceed #-}
2173 yieldProceed nextGen fs =
2174 initial >>= skip . SplitOnSuffixSeqYield fs . nextAfterInit nextGen
2176 -- For single element pattern case
2177 {-# INLINE processYieldSingle #-}
2178 processYieldSingle pat x s fs = do
2179 let jump c = SplitOnSuffixSeqSingleInit c s pat
2180 if pat == x
2181 then do
2182 r <- if withSep then fstep fs x else return $ FL.Partial fs
2183 b1 <-
2184 case r of
2185 FL.Partial fs1 -> done fs1
2186 FL.Done b -> return b
2187 yieldProceed jump b1
2188 else do
2189 r <- fstep fs x
2190 case r of
2191 FL.Partial fs1 -> skip $ SplitOnSuffixSeqSingle fs1 s pat
2192 FL.Done b -> yieldProceed jump b
2194 -- For Rabin-Karp search
2195 k = 2891336453 :: Word32
2196 coeff = k ^ patLen
2198 addCksum cksum a = cksum * k + fromIntegral (fromEnum a)
2200 deltaCksum cksum old new =
2201 addCksum cksum new - coeff * fromIntegral (fromEnum old)
2203 -- XXX shall we use a random starting hash or 1 instead of 0?
2204 patHash = A.foldl' addCksum 0 patArr
2206 skip = return . Skip
2208 {-# INLINE_LATE stepOuter #-}
2209 stepOuter _ SplitOnSuffixSeqInit = do
2210 res <- initial
2211 case res of
2212 FL.Partial fs ->
2213 if patLen == 0
2214 then skip $ SplitOnSuffixSeqEmpty fs state
2215 else if patLen == 1
2216 then do
2217 pat <- liftIO $ A.unsafeIndexIO 0 patArr
2218 skip $ SplitOnSuffixSeqSingleInit fs state pat
2219 else if SIZE_OF(a) * patLen
2220 <= sizeOf (undefined :: Word)
2221 then skip $ SplitOnSuffixSeqWordInit fs state
2222 else do
2223 (rb, rhead) <- liftIO $ RB.new patLen
2224 skip $ SplitOnSuffixSeqKRInit 0 fs state rb rhead
2225 FL.Done fb -> skip $ SplitOnSuffixSeqYield fb SplitOnSuffixSeqInit
2227 stepOuter _ (SplitOnSuffixSeqYield x next) = return $ Yield x next
2229 ---------------------------
2230 -- Reinit
2231 ---------------------------
2233 stepOuter _ (SplitOnSuffixSeqReinit nextGen) =
2234 initial >>= skip . nextAfterInit nextGen
2236 ---------------------------
2237 -- Empty pattern
2238 ---------------------------
2240 stepOuter gst (SplitOnSuffixSeqEmpty acc st) = do
2241 res <- step (adaptState gst) st
2242 case res of
2243 Yield x s -> do
2244 let jump c = SplitOnSuffixSeqEmpty c s
2245 r <- fstep acc x
2246 b1 <-
2247 case r of
2248 FL.Partial fs -> done fs
2249 FL.Done b -> return b
2250 yieldProceed jump b1
2251 Skip s -> skip (SplitOnSuffixSeqEmpty acc s)
2252 Stop -> return Stop
2254 -----------------
2255 -- Done
2256 -----------------
2258 stepOuter _ SplitOnSuffixSeqDone = return Stop
2260 -----------------
2261 -- Single Pattern
2262 -----------------
2264 stepOuter gst (SplitOnSuffixSeqSingleInit fs st pat) = do
2265 res <- step (adaptState gst) st
2266 case res of
2267 Yield x s -> processYieldSingle pat x s fs
2268 Skip s -> skip $ SplitOnSuffixSeqSingleInit fs s pat
2269 Stop -> return Stop
2271 stepOuter gst (SplitOnSuffixSeqSingle fs st pat) = do
2272 res <- step (adaptState gst) st
2273 case res of
2274 Yield x s -> processYieldSingle pat x s fs
2275 Skip s -> skip $ SplitOnSuffixSeqSingle fs s pat
2276 Stop -> do
2277 r <- done fs
2278 skip $ SplitOnSuffixSeqYield r SplitOnSuffixSeqDone
2280 ---------------------------
2281 -- Short Pattern - Shift Or
2282 ---------------------------
2284 stepOuter _ (SplitOnSuffixSeqWordDone 0 fs _) = do
2285 r <- done fs
2286 skip $ SplitOnSuffixSeqYield r SplitOnSuffixSeqDone
2287 stepOuter _ (SplitOnSuffixSeqWordDone n fs wrd) = do
2288 let old = elemMask .&. (wrd `shiftR` (elemBits * (n - 1)))
2289 r <- fstep fs (toEnum $ fromIntegral old)
2290 case r of
2291 FL.Partial fs1 -> skip $ SplitOnSuffixSeqWordDone (n - 1) fs1 wrd
2292 FL.Done b -> do
2293 let jump c = SplitOnSuffixSeqWordDone (n - 1) c wrd
2294 yieldProceed jump b
2296 stepOuter gst (SplitOnSuffixSeqWordInit fs0 st0) = do
2297 res <- step (adaptState gst) st0
2298 case res of
2299 Yield x s -> do
2300 let wrd = addToWord 0 x
2301 r <- if withSep then fstep fs0 x else return $ FL.Partial fs0
2302 case r of
2303 FL.Partial fs1 -> go SPEC 1 wrd s fs1
2304 FL.Done b -> do
2305 let jump c = SplitOnSuffixSeqWordInit c s
2306 yieldProceed jump b
2307 Skip s -> skip (SplitOnSuffixSeqWordInit fs0 s)
2308 Stop -> return Stop
2310 where
2312 {-# INLINE go #-}
2313 go !_ !idx !wrd !st !fs = do
2314 res <- step (adaptState gst) st
2315 case res of
2316 Yield x s -> do
2317 let jump c = SplitOnSuffixSeqWordInit c s
2318 let wrd1 = addToWord wrd x
2319 r <- if withSep then fstep fs x else return $ FL.Partial fs
2320 case r of
2321 FL.Partial fs1 ->
2322 if idx /= maxIndex
2323 then go SPEC (idx + 1) wrd1 s fs1
2324 else if wrd1 .&. wordMask /= wordPat
2325 then skip $ SplitOnSuffixSeqWordLoop wrd1 s fs1
2326 else do done fs >>= yieldProceed jump
2327 FL.Done b -> yieldProceed jump b
2328 Skip s -> go SPEC idx wrd s fs
2329 Stop -> skip $ SplitOnSuffixSeqWordDone idx fs wrd
2331 stepOuter gst (SplitOnSuffixSeqWordLoop wrd0 st0 fs0) =
2332 go SPEC wrd0 st0 fs0
2334 where
2336 {-# INLINE go #-}
2337 go !_ !wrd !st !fs = do
2338 res <- step (adaptState gst) st
2339 case res of
2340 Yield x s -> do
2341 let jump c = SplitOnSuffixSeqWordInit c s
2342 wrd1 = addToWord wrd x
2343 old = (wordMask .&. wrd)
2344 `shiftR` (elemBits * (patLen - 1))
2345 r <-
2346 if withSep
2347 then fstep fs x
2348 else fstep fs (toEnum $ fromIntegral old)
2349 case r of
2350 FL.Partial fs1 ->
2351 if wrd1 .&. wordMask == wordPat
2352 then done fs1 >>= yieldProceed jump
2353 else go SPEC wrd1 s fs1
2354 FL.Done b -> yieldProceed jump b
2355 Skip s -> go SPEC wrd s fs
2356 Stop ->
2357 if wrd .&. wordMask == wordPat
2358 then return Stop
2359 else if withSep
2360 then do
2361 r <- done fs
2362 skip $ SplitOnSuffixSeqYield r SplitOnSuffixSeqDone
2363 else skip $ SplitOnSuffixSeqWordDone patLen fs wrd
2365 -------------------------------
2366 -- General Pattern - Karp Rabin
2367 -------------------------------
2369 stepOuter gst (SplitOnSuffixSeqKRInit idx0 fs st0 rb rh0) = do
2370 res <- step (adaptState gst) st0
2371 case res of
2372 Yield x s -> do
2373 rh1 <- liftIO $ RB.unsafeInsert rb rh0 x
2374 r <- if withSep then fstep fs x else return $ FL.Partial fs
2375 case r of
2376 FL.Partial fs1 ->
2377 skip $ SplitOnSuffixSeqKRInit1 fs1 s rb rh1
2378 FL.Done b -> do
2379 let rst = RB.startOf rb
2380 jump c = SplitOnSuffixSeqKRInit 0 c s rb rst
2381 yieldProceed jump b
2382 Skip s -> skip $ SplitOnSuffixSeqKRInit idx0 fs s rb rh0
2383 Stop -> return Stop
2385 stepOuter gst (SplitOnSuffixSeqKRInit1 fs0 st0 rb rh0) = do
2386 go SPEC 1 rh0 st0 fs0
2388 where
2390 go !_ !idx !rh st !fs = do
2391 res <- step (adaptState gst) st
2392 case res of
2393 Yield x s -> do
2394 rh1 <- liftIO (RB.unsafeInsert rb rh x)
2395 r <- if withSep then fstep fs x else return $ FL.Partial fs
2396 case r of
2397 FL.Partial fs1 ->
2398 if idx /= maxIndex
2399 then go SPEC (idx + 1) rh1 s fs1
2400 else skip $
2401 let fld = RB.unsafeFoldRing (RB.ringBound rb)
2402 !ringHash = fld addCksum 0 rb
2403 in if ringHash == patHash
2404 then SplitOnSuffixSeqKRCheck fs1 s rb rh1
2405 else SplitOnSuffixSeqKRLoop
2406 fs1 s rb rh1 ringHash
2407 FL.Done b -> do
2408 let rst = RB.startOf rb
2409 jump c = SplitOnSuffixSeqKRInit 0 c s rb rst
2410 yieldProceed jump b
2411 Skip s -> go SPEC idx rh s fs
2412 Stop -> do
2413 -- do not issue a blank segment when we end at pattern
2414 if (idx == maxIndex) && RB.unsafeEqArray rb rh patArr
2415 then return Stop
2416 else if withSep
2417 then do
2418 r <- done fs
2419 skip $ SplitOnSuffixSeqYield r SplitOnSuffixSeqDone
2420 else skip $ SplitOnSuffixSeqKRDone idx fs rb (RB.startOf rb)
2422 stepOuter gst (SplitOnSuffixSeqKRLoop fs0 st0 rb rh0 cksum0) =
2423 go SPEC fs0 st0 rh0 cksum0
2425 where
2427 go !_ !fs !st !rh !cksum = do
2428 res <- step (adaptState gst) st
2429 case res of
2430 Yield x s -> do
2431 old <- liftIO $ peek rh
2432 rh1 <- liftIO (RB.unsafeInsert rb rh x)
2433 let cksum1 = deltaCksum cksum old x
2434 r <- if withSep then fstep fs x else fstep fs old
2435 case r of
2436 FL.Partial fs1 ->
2437 if cksum1 /= patHash
2438 then go SPEC fs1 s rh1 cksum1
2439 else skip $ SplitOnSuffixSeqKRCheck fs1 s rb rh1
2440 FL.Done b -> do
2441 let rst = RB.startOf rb
2442 jump c = SplitOnSuffixSeqKRInit 0 c s rb rst
2443 yieldProceed jump b
2444 Skip s -> go SPEC fs s rh cksum
2445 Stop ->
2446 if RB.unsafeEqArray rb rh patArr
2447 then return Stop
2448 else if withSep
2449 then do
2450 r <- done fs
2451 skip $ SplitOnSuffixSeqYield r SplitOnSuffixSeqDone
2452 else skip $ SplitOnSuffixSeqKRDone patLen fs rb rh
2454 stepOuter _ (SplitOnSuffixSeqKRCheck fs st rb rh) = do
2455 if RB.unsafeEqArray rb rh patArr
2456 then do
2457 r <- done fs
2458 let rst = RB.startOf rb
2459 jump c = SplitOnSuffixSeqKRInit 0 c st rb rst
2460 yieldProceed jump r
2461 else skip $ SplitOnSuffixSeqKRLoop fs st rb rh patHash
2463 stepOuter _ (SplitOnSuffixSeqKRDone 0 fs _ _) = do
2464 r <- done fs
2465 skip $ SplitOnSuffixSeqYield r SplitOnSuffixSeqDone
2466 stepOuter _ (SplitOnSuffixSeqKRDone n fs rb rh) = do
2467 old <- liftIO $ peek rh
2468 let rh1 = RB.advance rb rh
2469 r <- fstep fs old
2470 case r of
2471 FL.Partial fs1 -> skip $ SplitOnSuffixSeqKRDone (n - 1) fs1 rb rh1
2472 FL.Done b -> do
2473 let jump c = SplitOnSuffixSeqKRDone (n - 1) c rb rh1
2474 yieldProceed jump b
2476------------------------------------------------------------------------------
2477-- Nested Container Transformation
2478------------------------------------------------------------------------------
2480{-# ANN type SplitState Fuse #-}
2481data SplitState s arr
2482 = SplitInitial s
2483 | SplitBuffering s arr
2484 | SplitSplitting s arr
2485 | SplitYielding arr (SplitState s arr)
2486 | SplitFinishing
2488-- XXX An alternative approach would be to use a partial fold (Fold m a b) to
2489-- split using a splitBy like combinator. The Fold would consume upto the
2490-- separator and return any leftover which can then be fed to the next fold.
2492-- We can revisit this once we have partial folds/parsers.
2494-- | Performs infix separator style splitting.
2495{-# INLINE_NORMAL splitInnerBy #-}
2496splitInnerBy
2497 :: Monad m
2498 => (f a -> m (f a, Maybe (f a))) -- splitter
2499 -> (f a -> f a -> m (f a)) -- joiner
2500 -> Stream m (f a)
2501 -> Stream m (f a)
2502splitInnerBy splitter joiner (Stream step1 state1) =
2503 Stream step (SplitInitial state1)
2505 where
2507 {-# INLINE_LATE step #-}
2508 step gst (SplitInitial st) = do
2509 r <- step1 gst st
2510 case r of
2511 Yield x s -> do
2512 (x1, mx2) <- splitter x
2513 return $ case mx2 of
2514 Nothing -> Skip (SplitBuffering s x1)
2515 Just x2 -> Skip (SplitYielding x1 (SplitSplitting s x2))
2516 Skip s -> return $ Skip (SplitInitial s)
2517 Stop -> return Stop
2519 step gst (SplitBuffering st buf) = do
2520 r <- step1 gst st
2521 case r of
2522 Yield x s -> do
2523 (x1, mx2) <- splitter x
2524 buf' <- joiner buf x1
2525 return $ case mx2 of
2526 Nothing -> Skip (SplitBuffering s buf')
2527 Just x2 -> Skip (SplitYielding buf' (SplitSplitting s x2))
2528 Skip s -> return $ Skip (SplitBuffering s buf)
2529 Stop -> return $ Skip (SplitYielding buf SplitFinishing)
2531 step _ (SplitSplitting st buf) = do
2532 (x1, mx2) <- splitter buf
2533 return $ case mx2 of
2534 Nothing -> Skip $ SplitBuffering st x1
2535 Just x2 -> Skip $ SplitYielding x1 (SplitSplitting st x2)
2537 step _ (SplitYielding x next) = return $ Yield x next
2538 step _ SplitFinishing = return Stop
2540-- | Performs infix separator style splitting.
2541{-# INLINE_NORMAL splitInnerBySuffix #-}
2542splitInnerBySuffix
2543 :: (Monad m, Eq (f a), Monoid (f a))
2544 => (f a -> m (f a, Maybe (f a))) -- splitter
2545 -> (f a -> f a -> m (f a)) -- joiner
2546 -> Stream m (f a)
2547 -> Stream m (f a)
2548splitInnerBySuffix splitter joiner (Stream step1 state1) =
2549 Stream step (SplitInitial state1)
2551 where
2553 {-# INLINE_LATE step #-}
2554 step gst (SplitInitial st) = do
2555 r <- step1 gst st
2556 case r of
2557 Yield x s -> do
2558 (x1, mx2) <- splitter x
2559 return $ case mx2 of
2560 Nothing -> Skip (SplitBuffering s x1)
2561 Just x2 -> Skip (SplitYielding x1 (SplitSplitting s x2))
2562 Skip s -> return $ Skip (SplitInitial s)
2563 Stop -> return Stop
2565 step gst (SplitBuffering st buf) = do
2566 r <- step1 gst st
2567 case r of
2568 Yield x s -> do
2569 (x1, mx2) <- splitter x
2570 buf' <- joiner buf x1
2571 return $ case mx2 of
2572 Nothing -> Skip (SplitBuffering s buf')
2573 Just x2 -> Skip (SplitYielding buf' (SplitSplitting s x2))
2574 Skip s -> return $ Skip (SplitBuffering s buf)
2575 Stop -> return $
2576 if buf == mempty
2577 then Stop
2578 else Skip (SplitYielding buf SplitFinishing)
2580 step _ (SplitSplitting st buf) = do
2581 (x1, mx2) <- splitter buf
2582 return $ case mx2 of
2583 Nothing -> Skip $ SplitBuffering st x1
2584 Just x2 -> Skip $ SplitYielding x1 (SplitSplitting st x2)
2586 step _ (SplitYielding x next) = return $ Yield x next
2587 step _ SplitFinishing = return Stop