git.iko.soy/streamly
on a296fac help sign in
streamly/src/Streamly/Internal/Data/Stream/StreamD/Nesting.hs
2326 lines 87 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 , foldIterateM
117 -- ** Parsing
118 -- | Parsing is opposite to flattening. 'parseMany' is dual to concatMap or
119 -- unfoldMany. concatMap generates a stream from single values in a
120 -- stream and flattens, parseMany does the opposite of flattening by
121 -- splitting the stream and then folds each such split to single value in
122 -- the output stream.
123 , parseMany
124 , parseIterate
126 -- ** Grouping
127 -- | Group segments of a stream and fold. Special case of parsing.
128 , chunksOf
129 , groupsOf2
130 , groupsBy
131 , groupsRollingBy
133 -- ** Splitting
134 -- | A special case of parsing.
135 , wordsBy
136 , splitOnSeq
137 , splitOnSuffixSeq
139 -- * Transform (Nested Containers)
140 -- | Opposite to compact in ArrayStream
141 , splitInnerBy
142 , splitInnerBySuffix
143 )
144where
146#include "inline.hs"
148import Control.Exception (assert)
149import Control.Monad.Catch (MonadThrow, throwM)
150import Control.Monad.IO.Class (MonadIO(..))
151import Data.Bits (shiftR, shiftL, (.|.), (.&.))
152#if __GLASGOW_HASKELL__ >= 801
153import Data.Functor.Identity ( Identity )
154#endif
155import Data.Word (Word32)
156import Foreign.Storable (Storable(..))
157import Fusion.Plugin.Types (Fuse(..))
158import GHC.Types (SPEC(..))
160import Streamly.Internal.Data.Array.Foreign.Type (Array(..))
161import Streamly.Internal.Data.Fold.Type (Fold(..))
162import Streamly.Internal.Data.Parser (ParseError(..))
163import Streamly.Internal.Data.SVar.Type (adaptState)
164import Streamly.Internal.Data.Unfold.Type (Unfold(..))
166import qualified Streamly.Internal.Data.Array.Foreign.Type as A
167import qualified Streamly.Internal.Data.Fold.Type as FL
168import qualified Streamly.Internal.Data.Parser as PR
169import qualified Streamly.Internal.Data.Parser.ParserD as PRD
170import qualified Streamly.Internal.Ring.Foreign as RB
172import Streamly.Internal.Data.Stream.StreamD.Type
174import Prelude hiding (concatMap, mapM, zipWith)
176------------------------------------------------------------------------------
177-- Appending
178------------------------------------------------------------------------------
180data AppendState s1 s2 = AppendFirst s1 | AppendSecond s2
182-- Note that this could be much faster compared to the CPS stream. However, as
183-- the number of streams being composed increases this may become expensive.
184-- Need to see where the breaking point is between the two.
185--
186{-# INLINE_NORMAL append #-}
187append :: Monad m => Stream m a -> Stream m a -> Stream m a
188append (Stream step1 state1) (Stream step2 state2) =
189 Stream step (AppendFirst state1)
191 where
193 {-# INLINE_LATE step #-}
194 step gst (AppendFirst st) = do
195 r <- step1 gst st
196 return $ case r of
197 Yield a s -> Yield a (AppendFirst s)
198 Skip s -> Skip (AppendFirst s)
199 Stop -> Skip (AppendSecond state2)
201 step gst (AppendSecond st) = do
202 r <- step2 gst st
203 return $ case r of
204 Yield a s -> Yield a (AppendSecond s)
205 Skip s -> Skip (AppendSecond s)
206 Stop -> Stop
208------------------------------------------------------------------------------
209-- Interleaving
210------------------------------------------------------------------------------
212data InterleaveState s1 s2 = InterleaveFirst s1 s2 | InterleaveSecond s1 s2
213 | InterleaveSecondOnly s2 | InterleaveFirstOnly s1
215{-# INLINE_NORMAL interleave #-}
216interleave :: Monad m => Stream m a -> Stream m a -> Stream m a
217interleave (Stream step1 state1) (Stream step2 state2) =
218 Stream step (InterleaveFirst state1 state2)
220 where
222 {-# INLINE_LATE step #-}
223 step gst (InterleaveFirst st1 st2) = do
224 r <- step1 gst st1
225 return $ case r of
226 Yield a s -> Yield a (InterleaveSecond s st2)
227 Skip s -> Skip (InterleaveFirst s st2)
228 Stop -> Skip (InterleaveSecondOnly st2)
230 step gst (InterleaveSecond st1 st2) = do
231 r <- step2 gst st2
232 return $ case r of
233 Yield a s -> Yield a (InterleaveFirst st1 s)
234 Skip s -> Skip (InterleaveSecond st1 s)
235 Stop -> Skip (InterleaveFirstOnly st1)
237 step gst (InterleaveFirstOnly st1) = do
238 r <- step1 gst st1
239 return $ case r of
240 Yield a s -> Yield a (InterleaveFirstOnly s)
241 Skip s -> Skip (InterleaveFirstOnly s)
242 Stop -> Stop
244 step gst (InterleaveSecondOnly st2) = do
245 r <- step2 gst st2
246 return $ case r of
247 Yield a s -> Yield a (InterleaveSecondOnly s)
248 Skip s -> Skip (InterleaveSecondOnly s)
249 Stop -> Stop
251{-# INLINE_NORMAL interleaveMin #-}
252interleaveMin :: Monad m => Stream m a -> Stream m a -> Stream m a
253interleaveMin (Stream step1 state1) (Stream step2 state2) =
254 Stream step (InterleaveFirst state1 state2)
256 where
258 {-# INLINE_LATE step #-}
259 step gst (InterleaveFirst st1 st2) = do
260 r <- step1 gst st1
261 return $ case r of
262 Yield a s -> Yield a (InterleaveSecond s st2)
263 Skip s -> Skip (InterleaveFirst s st2)
264 Stop -> Stop
266 step gst (InterleaveSecond st1 st2) = do
267 r <- step2 gst st2
268 return $ case r of
269 Yield a s -> Yield a (InterleaveFirst st1 s)
270 Skip s -> Skip (InterleaveSecond st1 s)
271 Stop -> Stop
273 step _ (InterleaveFirstOnly _) = undefined
274 step _ (InterleaveSecondOnly _) = undefined
276{-# INLINE_NORMAL interleaveSuffix #-}
277interleaveSuffix :: Monad m => Stream m a -> Stream m a -> Stream m a
278interleaveSuffix (Stream step1 state1) (Stream step2 state2) =
279 Stream step (InterleaveFirst state1 state2)
281 where
283 {-# INLINE_LATE step #-}
284 step gst (InterleaveFirst st1 st2) = do
285 r <- step1 gst st1
286 return $ case r of
287 Yield a s -> Yield a (InterleaveSecond s st2)
288 Skip s -> Skip (InterleaveFirst s st2)
289 Stop -> Stop
291 step gst (InterleaveSecond st1 st2) = do
292 r <- step2 gst st2
293 return $ case r of
294 Yield a s -> Yield a (InterleaveFirst st1 s)
295 Skip s -> Skip (InterleaveSecond st1 s)
296 Stop -> Skip (InterleaveFirstOnly st1)
298 step gst (InterleaveFirstOnly st1) = do
299 r <- step1 gst st1
300 return $ case r of
301 Yield a s -> Yield a (InterleaveFirstOnly s)
302 Skip s -> Skip (InterleaveFirstOnly s)
303 Stop -> Stop
305 step _ (InterleaveSecondOnly _) = undefined
307data InterleaveInfixState s1 s2 a
308 = InterleaveInfixFirst s1 s2
309 | InterleaveInfixSecondBuf s1 s2
310 | InterleaveInfixSecondYield s1 s2 a
311 | InterleaveInfixFirstYield s1 s2 a
312 | InterleaveInfixFirstOnly s1
314{-# INLINE_NORMAL interleaveInfix #-}
315interleaveInfix :: Monad m => Stream m a -> Stream m a -> Stream m a
316interleaveInfix (Stream step1 state1) (Stream step2 state2) =
317 Stream step (InterleaveInfixFirst state1 state2)
319 where
321 {-# INLINE_LATE step #-}
322 step gst (InterleaveInfixFirst st1 st2) = do
323 r <- step1 gst st1
324 return $ case r of
325 Yield a s -> Yield a (InterleaveInfixSecondBuf s st2)
326 Skip s -> Skip (InterleaveInfixFirst s st2)
327 Stop -> Stop
329 step gst (InterleaveInfixSecondBuf st1 st2) = do
330 r <- step2 gst st2
331 return $ case r of
332 Yield a s -> Skip (InterleaveInfixSecondYield st1 s a)
333 Skip s -> Skip (InterleaveInfixSecondBuf st1 s)
334 Stop -> Skip (InterleaveInfixFirstOnly st1)
336 step gst (InterleaveInfixSecondYield st1 st2 x) = do
337 r <- step1 gst st1
338 return $ case r of
339 Yield a s -> Yield x (InterleaveInfixFirstYield s st2 a)
340 Skip s -> Skip (InterleaveInfixSecondYield s st2 x)
341 Stop -> Stop
343 step _ (InterleaveInfixFirstYield st1 st2 x) = do
344 return $ Yield x (InterleaveInfixSecondBuf st1 st2)
346 step gst (InterleaveInfixFirstOnly st1) = do
347 r <- step1 gst st1
348 return $ case r of
349 Yield a s -> Yield a (InterleaveInfixFirstOnly s)
350 Skip s -> Skip (InterleaveInfixFirstOnly s)
351 Stop -> Stop
353------------------------------------------------------------------------------
354-- Scheduling
355------------------------------------------------------------------------------
357{-# INLINE_NORMAL roundRobin #-}
358roundRobin :: Monad m => Stream m a -> Stream m a -> Stream m a
359roundRobin (Stream step1 state1) (Stream step2 state2) =
360 Stream step (InterleaveFirst state1 state2)
362 where
364 {-# INLINE_LATE step #-}
365 step gst (InterleaveFirst st1 st2) = do
366 r <- step1 gst st1
367 return $ case r of
368 Yield a s -> Yield a (InterleaveSecond s st2)
369 Skip s -> Skip (InterleaveSecond s st2)
370 Stop -> Skip (InterleaveSecondOnly st2)
372 step gst (InterleaveSecond st1 st2) = do
373 r <- step2 gst st2
374 return $ case r of
375 Yield a s -> Yield a (InterleaveFirst st1 s)
376 Skip s -> Skip (InterleaveFirst st1 s)
377 Stop -> Skip (InterleaveFirstOnly st1)
379 step gst (InterleaveSecondOnly st2) = do
380 r <- step2 gst st2
381 return $ case r of
382 Yield a s -> Yield a (InterleaveSecondOnly s)
383 Skip s -> Skip (InterleaveSecondOnly s)
384 Stop -> Stop
386 step gst (InterleaveFirstOnly st1) = do
387 r <- step1 gst st1
388 return $ case r of
389 Yield a s -> Yield a (InterleaveFirstOnly s)
390 Skip s -> Skip (InterleaveFirstOnly s)
391 Stop -> Stop
393------------------------------------------------------------------------------
394-- Zipping
395------------------------------------------------------------------------------
397{-# INLINE_NORMAL zipWithM #-}
398zipWithM :: Monad m
399 => (a -> b -> m c) -> Stream m a -> Stream m b -> Stream m c
400zipWithM f (Stream stepa ta) (Stream stepb tb) = Stream step (ta, tb, Nothing)
401 where
402 {-# INLINE_LATE step #-}
403 step gst (sa, sb, Nothing) = do
404 r <- stepa (adaptState gst) sa
405 return $
406 case r of
407 Yield x sa' -> Skip (sa', sb, Just x)
408 Skip sa' -> Skip (sa', sb, Nothing)
409 Stop -> Stop
411 step gst (sa, sb, Just x) = do
412 r <- stepb (adaptState gst) sb
413 case r of
414 Yield y sb' -> do
415 z <- f x y
416 return $ Yield z (sa, sb', Nothing)
417 Skip sb' -> return $ Skip (sa, sb', Just x)
418 Stop -> return Stop
420#if __GLASGOW_HASKELL__ >= 801
421{-# RULES "zipWithM xs xs"
422 forall f xs. zipWithM @Identity f xs xs = mapM (\x -> f x x) xs #-}
423#endif
425{-# INLINE zipWith #-}
426zipWith :: Monad m => (a -> b -> c) -> Stream m a -> Stream m b -> Stream m c
427zipWith f = zipWithM (\a b -> return (f a b))
429------------------------------------------------------------------------------
430-- Merging
431------------------------------------------------------------------------------
433{-# INLINE_NORMAL mergeByM #-}
434mergeByM
435 :: (Monad m)
436 => (a -> a -> m Ordering) -> Stream m a -> Stream m a -> Stream m a
437mergeByM cmp (Stream stepa ta) (Stream stepb tb) =
438 Stream step (Just ta, Just tb, Nothing, Nothing)
439 where
440 {-# INLINE_LATE step #-}
442 -- one of the values is missing, and the corresponding stream is running
443 step gst (Just sa, sb, Nothing, b) = do
444 r <- stepa gst sa
445 return $ case r of
446 Yield a sa' -> Skip (Just sa', sb, Just a, b)
447 Skip sa' -> Skip (Just sa', sb, Nothing, b)
448 Stop -> Skip (Nothing, sb, Nothing, b)
450 step gst (sa, Just sb, a, Nothing) = do
451 r <- stepb gst sb
452 return $ case r of
453 Yield b sb' -> Skip (sa, Just sb', a, Just b)
454 Skip sb' -> Skip (sa, Just sb', a, Nothing)
455 Stop -> Skip (sa, Nothing, a, Nothing)
457 -- both the values are available
458 step _ (sa, sb, Just a, Just b) = do
459 res <- cmp a b
460 return $ case res of
461 GT -> Yield b (sa, sb, Just a, Nothing)
462 _ -> Yield a (sa, sb, Nothing, Just b)
464 -- one of the values is missing, corresponding stream is done
465 step _ (Nothing, sb, Nothing, Just b) =
466 return $ Yield b (Nothing, sb, Nothing, Nothing)
468 step _ (sa, Nothing, Just a, Nothing) =
469 return $ Yield a (sa, Nothing, Nothing, Nothing)
471 step _ (Nothing, Nothing, Nothing, Nothing) = return Stop
473{-# INLINE mergeBy #-}
474mergeBy
475 :: (Monad m)
476 => (a -> a -> Ordering) -> Stream m a -> Stream m a -> Stream m a
477mergeBy cmp = mergeByM (\a b -> return $ cmp a b)
479------------------------------------------------------------------------------
480-- Combine N Streams - unfoldMany
481------------------------------------------------------------------------------
483data ConcatUnfoldInterleaveState o i =
484 ConcatUnfoldInterleaveOuter o [i]
485 | ConcatUnfoldInterleaveInner o [i]
486 | ConcatUnfoldInterleaveInnerL [i] [i]
487 | ConcatUnfoldInterleaveInnerR [i] [i]
489-- XXX use arrays to store state instead of lists.
490-- XXX In general we can use different scheduling strategies e.g. how to
491-- schedule the outer vs inner loop or assigning weights to different streams
492-- or outer and inner loops.
494-- After a yield, switch to the next stream. Do not switch streams on Skip.
495-- Yield from outer stream switches to the inner stream.
496--
497-- There are two choices here, (1) exhaust the outer stream first and then
498-- start yielding from the inner streams, this is much simpler to implement,
499-- (2) yield at least one element from an inner stream before going back to
500-- outer stream and opening the next stream from it.
501--
502-- Ideally, we need some scheduling bias to inner streams vs outer stream.
503-- Maybe we can configure the behavior.
504--
505{-# INLINE_NORMAL unfoldManyInterleave #-}
506unfoldManyInterleave :: Monad m => Unfold m a b -> Stream m a -> Stream m b
507unfoldManyInterleave (Unfold istep inject) (Stream ostep ost) =
508 Stream step (ConcatUnfoldInterleaveOuter ost [])
509 where
510 {-# INLINE_LATE step #-}
511 step gst (ConcatUnfoldInterleaveOuter o ls) = do
512 r <- ostep (adaptState gst) o
513 case r of
514 Yield a o' -> do
515 i <- inject a
516 i `seq` return (Skip (ConcatUnfoldInterleaveInner o' (i : ls)))
517 Skip o' -> return $ Skip (ConcatUnfoldInterleaveOuter o' ls)
518 Stop -> return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
520 step _ (ConcatUnfoldInterleaveInner _ []) = undefined
521 step _ (ConcatUnfoldInterleaveInner o (st:ls)) = do
522 r <- istep st
523 return $ case r of
524 Yield x s -> Yield x (ConcatUnfoldInterleaveOuter o (s:ls))
525 Skip s -> Skip (ConcatUnfoldInterleaveInner o (s:ls))
526 Stop -> Skip (ConcatUnfoldInterleaveOuter o ls)
528 step _ (ConcatUnfoldInterleaveInnerL [] []) = return Stop
529 step _ (ConcatUnfoldInterleaveInnerL [] rs) =
530 return $ Skip (ConcatUnfoldInterleaveInnerR [] rs)
532 step _ (ConcatUnfoldInterleaveInnerL (st:ls) rs) = do
533 r <- istep st
534 return $ case r of
535 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerL ls (s:rs))
536 Skip s -> Skip (ConcatUnfoldInterleaveInnerL (s:ls) rs)
537 Stop -> Skip (ConcatUnfoldInterleaveInnerL ls rs)
539 step _ (ConcatUnfoldInterleaveInnerR [] []) = return Stop
540 step _ (ConcatUnfoldInterleaveInnerR ls []) =
541 return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
543 step _ (ConcatUnfoldInterleaveInnerR ls (st:rs)) = do
544 r <- istep st
545 return $ case r of
546 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerR (s:ls) rs)
547 Skip s -> Skip (ConcatUnfoldInterleaveInnerR ls (s:rs))
548 Stop -> Skip (ConcatUnfoldInterleaveInnerR ls rs)
550-- XXX In general we can use different scheduling strategies e.g. how to
551-- schedule the outer vs inner loop or assigning weights to different streams
552-- or outer and inner loops.
553--
554-- This could be inefficient if the tasks are too small.
555--
556-- Compared to unfoldManyInterleave this one switches streams on Skips.
557--
558{-# INLINE_NORMAL unfoldManyRoundRobin #-}
559unfoldManyRoundRobin :: Monad m => Unfold m a b -> Stream m a -> Stream m b
560unfoldManyRoundRobin (Unfold istep inject) (Stream ostep ost) =
561 Stream step (ConcatUnfoldInterleaveOuter ost [])
562 where
563 {-# INLINE_LATE step #-}
564 step gst (ConcatUnfoldInterleaveOuter o ls) = do
565 r <- ostep (adaptState gst) o
566 case r of
567 Yield a o' -> do
568 i <- inject a
569 i `seq` return (Skip (ConcatUnfoldInterleaveInner o' (i : ls)))
570 Skip o' -> return $ Skip (ConcatUnfoldInterleaveInner o' ls)
571 Stop -> return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
573 step _ (ConcatUnfoldInterleaveInner o []) =
574 return $ Skip (ConcatUnfoldInterleaveOuter o [])
576 step _ (ConcatUnfoldInterleaveInner o (st:ls)) = do
577 r <- istep st
578 return $ case r of
579 Yield x s -> Yield x (ConcatUnfoldInterleaveOuter o (s:ls))
580 Skip s -> Skip (ConcatUnfoldInterleaveOuter o (s:ls))
581 Stop -> Skip (ConcatUnfoldInterleaveOuter o ls)
583 step _ (ConcatUnfoldInterleaveInnerL [] []) = return Stop
584 step _ (ConcatUnfoldInterleaveInnerL [] rs) =
585 return $ Skip (ConcatUnfoldInterleaveInnerR [] rs)
587 step _ (ConcatUnfoldInterleaveInnerL (st:ls) rs) = do
588 r <- istep st
589 return $ case r of
590 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerL ls (s:rs))
591 Skip s -> Skip (ConcatUnfoldInterleaveInnerL ls (s:rs))
592 Stop -> Skip (ConcatUnfoldInterleaveInnerL ls rs)
594 step _ (ConcatUnfoldInterleaveInnerR [] []) = return Stop
595 step _ (ConcatUnfoldInterleaveInnerR ls []) =
596 return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
598 step _ (ConcatUnfoldInterleaveInnerR ls (st:rs)) = do
599 r <- istep st
600 return $ case r of
601 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerR (s:ls) rs)
602 Skip s -> Skip (ConcatUnfoldInterleaveInnerR (s:ls) rs)
603 Stop -> Skip (ConcatUnfoldInterleaveInnerR ls rs)
605------------------------------------------------------------------------------
606-- Combine N Streams - interpose
607------------------------------------------------------------------------------
609{-# ANN type InterposeSuffixState Fuse #-}
610data InterposeSuffixState s1 i1 =
611 InterposeSuffixFirst s1
612 -- | InterposeSuffixFirstYield s1 i1
613 | InterposeSuffixFirstInner s1 i1
614 | InterposeSuffixSecond s1
616-- Note that if an unfolded layer turns out to be nil we still emit the
617-- separator effect. An alternate behavior could be to emit the separator
618-- effect only if at least one element has been yielded by the unfolding.
619-- However, that becomes a bit complicated, so we have chosen the former
620-- behvaior for now.
621{-# INLINE_NORMAL interposeSuffix #-}
622interposeSuffix
623 :: Monad m
624 => m c -> Unfold m b c -> Stream m b -> Stream m c
625interposeSuffix
626 action
627 (Unfold istep1 inject1) (Stream step1 state1) =
628 Stream step (InterposeSuffixFirst state1)
630 where
632 {-# INLINE_LATE step #-}
633 step gst (InterposeSuffixFirst s1) = do
634 r <- step1 (adaptState gst) s1
635 case r of
636 Yield a s -> do
637 i <- inject1 a
638 i `seq` return (Skip (InterposeSuffixFirstInner s i))
639 -- i `seq` return (Skip (InterposeSuffixFirstYield s i))
640 Skip s -> return $ Skip (InterposeSuffixFirst s)
641 Stop -> return Stop
643 {-
644 step _ (InterposeSuffixFirstYield s1 i1) = do
645 r <- istep1 i1
646 return $ case r of
647 Yield x i' -> Yield x (InterposeSuffixFirstInner s1 i')
648 Skip i' -> Skip (InterposeSuffixFirstYield s1 i')
649 Stop -> Skip (InterposeSuffixFirst s1)
650 -}
652 step _ (InterposeSuffixFirstInner s1 i1) = do
653 r <- istep1 i1
654 return $ case r of
655 Yield x i' -> Yield x (InterposeSuffixFirstInner s1 i')
656 Skip i' -> Skip (InterposeSuffixFirstInner s1 i')
657 Stop -> Skip (InterposeSuffixSecond s1)
659 step _ (InterposeSuffixSecond s1) = do
660 r <- action
661 return $ Yield r (InterposeSuffixFirst s1)
663{-# ANN type InterposeState Fuse #-}
664data InterposeState s1 i1 a =
665 InterposeFirst s1
666 -- | InterposeFirstYield s1 i1
667 | InterposeFirstInner s1 i1
668 | InterposeFirstInject s1
669 -- | InterposeFirstBuf s1 i1
670 | InterposeSecondYield s1 i1
671 -- -- | InterposeSecondYield s1 i1 a
672 -- -- | InterposeFirstResume s1 i1 a
674-- Note that this only interposes the pure values, we may run many effects to
675-- generate those values as some effects may not generate anything (Skip).
676{-# INLINE_NORMAL interpose #-}
677interpose :: Monad m => m c -> Unfold m b c -> Stream m b -> Stream m c
678interpose
679 action
680 (Unfold istep1 inject1) (Stream step1 state1) =
681 Stream step (InterposeFirst state1)
683 where
685 {-# INLINE_LATE step #-}
686 step gst (InterposeFirst s1) = do
687 r <- step1 (adaptState gst) s1
688 case r of
689 Yield a s -> do
690 i <- inject1 a
691 i `seq` return (Skip (InterposeFirstInner s i))
692 -- i `seq` return (Skip (InterposeFirstYield s i))
693 Skip s -> return $ Skip (InterposeFirst s)
694 Stop -> return Stop
696 {-
697 step _ (InterposeFirstYield s1 i1) = do
698 r <- istep1 i1
699 return $ case r of
700 Yield x i' -> Yield x (InterposeFirstInner s1 i')
701 Skip i' -> Skip (InterposeFirstYield s1 i')
702 Stop -> Skip (InterposeFirst s1)
703 -}
705 step _ (InterposeFirstInner s1 i1) = do
706 r <- istep1 i1
707 return $ case r of
708 Yield x i' -> Yield x (InterposeFirstInner s1 i')
709 Skip i' -> Skip (InterposeFirstInner s1 i')
710 Stop -> Skip (InterposeFirstInject s1)
712 step gst (InterposeFirstInject s1) = do
713 r <- step1 (adaptState gst) s1
714 case r of
715 Yield a s -> do
716 i <- inject1 a
717 -- i `seq` return (Skip (InterposeFirstBuf s i))
718 i `seq` return (Skip (InterposeSecondYield s i))
719 Skip s -> return $ Skip (InterposeFirstInject s)
720 Stop -> return Stop
722 {-
723 step _ (InterposeFirstBuf s1 i1) = do
724 r <- istep1 i1
725 return $ case r of
726 Yield x i' -> Skip (InterposeSecondYield s1 i' x)
727 Skip i' -> Skip (InterposeFirstBuf s1 i')
728 Stop -> Stop
729 -}
731 {-
732 step _ (InterposeSecondYield s1 i1 v) = do
733 r <- action
734 return $ Yield r (InterposeFirstResume s1 i1 v)
735 -}
736 step _ (InterposeSecondYield s1 i1) = do
737 r <- action
738 return $ Yield r (InterposeFirstInner s1 i1)
740 {-
741 step _ (InterposeFirstResume s1 i1 v) = do
742 return $ Yield v (InterposeFirstInner s1 i1)
743 -}
745------------------------------------------------------------------------------
746-- Combine N Streams - intercalate
747------------------------------------------------------------------------------
749data ICUState s1 s2 i1 i2 =
750 ICUFirst s1 s2
751 | ICUSecond s1 s2
752 | ICUSecondOnly s2
753 | ICUFirstOnly s1
754 | ICUFirstInner s1 s2 i1
755 | ICUSecondInner s1 s2 i2
756 | ICUFirstOnlyInner s1 i1
757 | ICUSecondOnlyInner s2 i2
759-- | Interleave streams (full streams, not the elements) unfolded from two
760-- input streams and concat. Stop when the first stream stops. If the second
761-- stream ends before the first one then first stream still keeps running alone
762-- without any interleaving with the second stream.
763--
764-- [a1, a2, ... an] [b1, b2 ...]
765-- => [streamA1, streamA2, ... streamAn] [streamB1, streamB2, ...]
766-- => [streamA1, streamB1, streamA2...StreamAn, streamBn]
767-- => [a11, a12, ...a1j, b11, b12, ...b1k, a21, a22, ...]
768--
769{-# INLINE_NORMAL gintercalateSuffix #-}
770gintercalateSuffix
771 :: Monad m
772 => Unfold m a c -> Stream m a -> Unfold m b c -> Stream m b -> Stream m c
773gintercalateSuffix
774 (Unfold istep1 inject1) (Stream step1 state1)
775 (Unfold istep2 inject2) (Stream step2 state2) =
776 Stream step (ICUFirst state1 state2)
778 where
780 {-# INLINE_LATE step #-}
781 step gst (ICUFirst s1 s2) = do
782 r <- step1 (adaptState gst) s1
783 case r of
784 Yield a s -> do
785 i <- inject1 a
786 i `seq` return (Skip (ICUFirstInner s s2 i))
787 Skip s -> return $ Skip (ICUFirst s s2)
788 Stop -> return Stop
790 step gst (ICUFirstOnly s1) = do
791 r <- step1 (adaptState gst) s1
792 case r of
793 Yield a s -> do
794 i <- inject1 a
795 i `seq` return (Skip (ICUFirstOnlyInner s i))
796 Skip s -> return $ Skip (ICUFirstOnly s)
797 Stop -> return Stop
799 step _ (ICUFirstInner s1 s2 i1) = do
800 r <- istep1 i1
801 return $ case r of
802 Yield x i' -> Yield x (ICUFirstInner s1 s2 i')
803 Skip i' -> Skip (ICUFirstInner s1 s2 i')
804 Stop -> Skip (ICUSecond s1 s2)
806 step _ (ICUFirstOnlyInner s1 i1) = do
807 r <- istep1 i1
808 return $ case r of
809 Yield x i' -> Yield x (ICUFirstOnlyInner s1 i')
810 Skip i' -> Skip (ICUFirstOnlyInner s1 i')
811 Stop -> Skip (ICUFirstOnly s1)
813 step gst (ICUSecond s1 s2) = do
814 r <- step2 (adaptState gst) s2
815 case r of
816 Yield a s -> do
817 i <- inject2 a
818 i `seq` return (Skip (ICUSecondInner s1 s i))
819 Skip s -> return $ Skip (ICUSecond s1 s)
820 Stop -> return $ Skip (ICUFirstOnly s1)
822 step _ (ICUSecondInner s1 s2 i2) = do
823 r <- istep2 i2
824 return $ case r of
825 Yield x i' -> Yield x (ICUSecondInner s1 s2 i')
826 Skip i' -> Skip (ICUSecondInner s1 s2 i')
827 Stop -> Skip (ICUFirst s1 s2)
829 step _ (ICUSecondOnly _s2) = undefined
830 step _ (ICUSecondOnlyInner _s2 _i2) = undefined
832data ICALState s1 s2 i1 i2 a =
833 ICALFirst s1 s2
834 -- | ICALFirstYield s1 s2 i1
835 | ICALFirstInner s1 s2 i1
836 | ICALFirstOnly s1
837 | ICALFirstOnlyInner s1 i1
838 | ICALSecondInject s1 s2
839 | ICALFirstInject s1 s2 i2
840 -- | ICALFirstBuf s1 s2 i1 i2
841 | ICALSecondInner s1 s2 i1 i2
842 -- -- | ICALSecondInner s1 s2 i1 i2 a
843 -- -- | ICALFirstResume s1 s2 i1 i2 a
845-- | Interleave streams (full streams, not the elements) unfolded from two
846-- input streams and concat. Stop when the first stream stops. If the second
847-- stream ends before the first one then first stream still keeps running alone
848-- without any interleaving with the second stream.
849--
850-- [a1, a2, ... an] [b1, b2 ...]
851-- => [streamA1, streamA2, ... streamAn] [streamB1, streamB2, ...]
852-- => [streamA1, streamB1, streamA2...StreamAn, streamBn]
853-- => [a11, a12, ...a1j, b11, b12, ...b1k, a21, a22, ...]
854--
855{-# INLINE_NORMAL gintercalate #-}
856gintercalate
857 :: Monad m
858 => Unfold m a c -> Stream m a -> Unfold m b c -> Stream m b -> Stream m c
859gintercalate
860 (Unfold istep1 inject1) (Stream step1 state1)
861 (Unfold istep2 inject2) (Stream step2 state2) =
862 Stream step (ICALFirst state1 state2)
864 where
866 {-# INLINE_LATE step #-}
867 step gst (ICALFirst s1 s2) = do
868 r <- step1 (adaptState gst) s1
869 case r of
870 Yield a s -> do
871 i <- inject1 a
872 i `seq` return (Skip (ICALFirstInner s s2 i))
873 -- i `seq` return (Skip (ICALFirstYield s s2 i))
874 Skip s -> return $ Skip (ICALFirst s s2)
875 Stop -> return Stop
877 {-
878 step _ (ICALFirstYield s1 s2 i1) = do
879 r <- istep1 i1
880 return $ case r of
881 Yield x i' -> Yield x (ICALFirstInner s1 s2 i')
882 Skip i' -> Skip (ICALFirstYield s1 s2 i')
883 Stop -> Skip (ICALFirst s1 s2)
884 -}
886 step _ (ICALFirstInner s1 s2 i1) = do
887 r <- istep1 i1
888 return $ case r of
889 Yield x i' -> Yield x (ICALFirstInner s1 s2 i')
890 Skip i' -> Skip (ICALFirstInner s1 s2 i')
891 Stop -> Skip (ICALSecondInject s1 s2)
893 step gst (ICALFirstOnly s1) = do
894 r <- step1 (adaptState gst) s1
895 case r of
896 Yield a s -> do
897 i <- inject1 a
898 i `seq` return (Skip (ICALFirstOnlyInner s i))
899 Skip s -> return $ Skip (ICALFirstOnly s)
900 Stop -> return Stop
902 step _ (ICALFirstOnlyInner s1 i1) = do
903 r <- istep1 i1
904 return $ case r of
905 Yield x i' -> Yield x (ICALFirstOnlyInner s1 i')
906 Skip i' -> Skip (ICALFirstOnlyInner s1 i')
907 Stop -> Skip (ICALFirstOnly s1)
909 -- We inject the second stream even before checking if the first stream
910 -- would yield any more elements. There is no clear choice whether we
911 -- should do this before or after that. Doing it after may make the state
912 -- machine a bit simpler though.
913 step gst (ICALSecondInject s1 s2) = do
914 r <- step2 (adaptState gst) s2
915 case r of
916 Yield a s -> do
917 i <- inject2 a
918 i `seq` return (Skip (ICALFirstInject s1 s i))
919 Skip s -> return $ Skip (ICALSecondInject s1 s)
920 Stop -> return $ Skip (ICALFirstOnly s1)
922 step gst (ICALFirstInject s1 s2 i2) = do
923 r <- step1 (adaptState gst) s1
924 case r of
925 Yield a s -> do
926 i <- inject1 a
927 i `seq` return (Skip (ICALSecondInner s s2 i i2))
928 -- i `seq` return (Skip (ICALFirstBuf s s2 i i2))
929 Skip s -> return $ Skip (ICALFirstInject s s2 i2)
930 Stop -> return Stop
932 {-
933 step _ (ICALFirstBuf s1 s2 i1 i2) = do
934 r <- istep1 i1
935 return $ case r of
936 Yield x i' -> Skip (ICALSecondInner s1 s2 i' i2 x)
937 Skip i' -> Skip (ICALFirstBuf s1 s2 i' i2)
938 Stop -> Stop
940 step _ (ICALSecondInner s1 s2 i1 i2 v) = do
941 r <- istep2 i2
942 return $ case r of
943 Yield x i' -> Yield x (ICALSecondInner s1 s2 i1 i' v)
944 Skip i' -> Skip (ICALSecondInner s1 s2 i1 i' v)
945 Stop -> Skip (ICALFirstResume s1 s2 i1 i2 v)
946 -}
948 step _ (ICALSecondInner s1 s2 i1 i2) = do
949 r <- istep2 i2
950 return $ case r of
951 Yield x i' -> Yield x (ICALSecondInner s1 s2 i1 i')
952 Skip i' -> Skip (ICALSecondInner s1 s2 i1 i')
953 Stop -> Skip (ICALFirstInner s1 s2 i1)
954 -- Stop -> Skip (ICALFirstResume s1 s2 i1 i2)
956 {-
957 step _ (ICALFirstResume s1 s2 i1 i2 x) = do
958 return $ Yield x (ICALFirstInner s1 s2 i1 i2)
959 -}
961------------------------------------------------------------------------------
962-- Folding
963------------------------------------------------------------------------------
965{-# ANN type FIterState Fuse #-}
966data FIterState s f m a b
967 = FIterInit s f
968 | forall fs. FIterStream s (fs -> a -> m (FL.Step fs b)) fs (fs -> m b)
969 | FIterYield b (FIterState s f m a b)
970 | FIterStop
972{-# INLINE_NORMAL foldIterateM #-}
973foldIterateM ::
974 Monad m => (b -> m (FL.Fold m a b)) -> b -> Stream m a -> Stream m b
975foldIterateM func seed0 (Stream step state) =
976 Stream stepOuter (FIterInit state seed0)
978 where
980 {-# INLINE iterStep #-}
981 iterStep from st fstep extract = do
982 res <- from
983 return
984 $ Skip
985 $ case res of
986 FL.Partial fs -> FIterStream st fstep fs extract
987 FL.Done fb -> FIterYield fb $ FIterInit st fb
989 {-# INLINE_LATE stepOuter #-}
990 stepOuter _ (FIterInit st seed) = do
991 (FL.Fold fstep initial extract) <- func seed
992 iterStep initial st fstep extract
993 stepOuter gst (FIterStream st fstep fs extract) = do
994 r <- step (adaptState gst) st
995 case r of
996 Yield x s -> do
997 iterStep (fstep fs x) s fstep extract
998 Skip s -> return $ Skip $ FIterStream s fstep fs extract
999 Stop -> do
1000 b <- extract fs
1001 return $ Skip $ FIterYield b FIterStop
1002 stepOuter _ (FIterYield a next) = return $ Yield a next
1003 stepOuter _ FIterStop = return Stop
1005------------------------------------------------------------------------------
1006-- Parsing
1007------------------------------------------------------------------------------
1009{-# ANN type ParseChunksState Fuse #-}
1010data ParseChunksState x inpBuf st pst =
1011 ParseChunksInit inpBuf st
1012 | ParseChunksInitLeftOver inpBuf
1013 | ParseChunksStream st inpBuf !pst
1014 | ParseChunksBuf inpBuf st inpBuf !pst
1015 | ParseChunksYield x (ParseChunksState x inpBuf st pst)
1017{-# INLINE_NORMAL parseMany #-}
1018parseMany
1019 :: MonadThrow m
1020 => PRD.Parser m a b
1021 -> Stream m a
1022 -> Stream m b
1023parseMany (PRD.Parser pstep initial extract) (Stream step state) =
1024 Stream stepOuter (ParseChunksInit [] state)
1026 where
1028 {-# INLINE_LATE stepOuter #-}
1029 -- Buffer is empty, get the first element from the stream, initialize the
1030 -- fold and then go to stream processing loop.
1031 stepOuter gst (ParseChunksInit [] st) = do
1032 r <- step (adaptState gst) st
1033 case r of
1034 Yield x s -> do
1035 res <- initial
1036 case res of
1037 PRD.IPartial ps ->
1038 return $ Skip $ ParseChunksBuf [x] s [] ps
1039 PRD.IDone pb ->
1040 let next = ParseChunksInit [x] s
1041 in return $ Skip $ ParseChunksYield pb next
1042 PRD.IError err -> throwM $ ParseError err
1043 Skip s -> return $ Skip $ ParseChunksInit [] s
1044 Stop -> return Stop
1046 -- Buffer is not empty, go to buffered processing loop
1047 stepOuter _ (ParseChunksInit src st) = do
1048 res <- initial
1049 case res of
1050 PRD.IPartial ps ->
1051 return $ Skip $ ParseChunksBuf src st [] ps
1052 PRD.IDone pb ->
1053 let next = ParseChunksInit src st
1054 in return $ Skip $ ParseChunksYield pb next
1055 PRD.IError err -> throwM $ ParseError err
1057 -- XXX we just discard any leftover input at the end
1058 stepOuter _ (ParseChunksInitLeftOver _) = return Stop
1060 -- Buffer is empty, process elements from the stream
1061 stepOuter gst (ParseChunksStream st buf pst) = do
1062 r <- step (adaptState gst) st
1063 case r of
1064 Yield x s -> do
1065 pRes <- pstep pst x
1066 case pRes of
1067 PR.Partial 0 pst1 ->
1068 return $ Skip $ ParseChunksStream s [] pst1
1069 PR.Partial n pst1 -> do
1070 assert (n <= length (x:buf)) (return ())
1071 let src0 = Prelude.take n (x:buf)
1072 src = Prelude.reverse src0
1073 return $ Skip $ ParseChunksBuf src s [] pst1
1074 PR.Continue 0 pst1 ->
1075 return $ Skip $ ParseChunksStream s (x:buf) pst1
1076 PR.Continue n pst1 -> do
1077 assert (n <= length (x:buf)) (return ())
1078 let (src0, buf1) = splitAt n (x:buf)
1079 src = Prelude.reverse src0
1080 return $ Skip $ ParseChunksBuf src s buf1 pst1
1081 PR.Done 0 b -> do
1082 return $ Skip $
1083 ParseChunksYield b (ParseChunksInit [] s)
1084 PR.Done n b -> do
1085 assert (n <= length (x:buf)) (return ())
1086 let src = Prelude.reverse (Prelude.take n (x:buf))
1087 return $ Skip $
1088 ParseChunksYield b (ParseChunksInit src s)
1089 PR.Error err -> throwM $ ParseError err
1090 Skip s -> return $ Skip $ ParseChunksStream s buf pst
1091 Stop -> do
1092 b <- extract pst
1093 let src = Prelude.reverse buf
1094 return $ Skip $
1095 ParseChunksYield b (ParseChunksInitLeftOver src)
1097 -- go back to stream processing mode
1098 stepOuter _ (ParseChunksBuf [] s buf pst) =
1099 return $ Skip $ ParseChunksStream s buf pst
1101 -- buffered processing loop
1102 stepOuter _ (ParseChunksBuf (x:xs) s buf pst) = do
1103 pRes <- pstep pst x
1104 case pRes of
1105 PR.Partial 0 pst1 ->
1106 return $ Skip $ ParseChunksBuf xs s [] pst1
1107 PR.Partial n pst1 -> do
1108 assert (n <= length (x:buf)) (return ())
1109 let src0 = Prelude.take n (x:buf)
1110 src = Prelude.reverse src0 ++ xs
1111 return $ Skip $ ParseChunksBuf src s [] pst1
1112 PR.Continue 0 pst1 ->
1113 return $ Skip $ ParseChunksBuf xs s (x:buf) pst1
1114 PR.Continue n pst1 -> do
1115 assert (n <= length (x:buf)) (return ())
1116 let (src0, buf1) = splitAt n (x:buf)
1117 src = Prelude.reverse src0 ++ xs
1118 return $ Skip $ ParseChunksBuf src s buf1 pst1
1119 PR.Done 0 b ->
1120 return $ Skip $ ParseChunksYield b (ParseChunksInit xs s)
1121 PR.Done n b -> do
1122 assert (n <= length (x:buf)) (return ())
1123 let src = Prelude.reverse (Prelude.take n (x:buf)) ++ xs
1124 return $ Skip $ ParseChunksYield b (ParseChunksInit src s)
1125 PR.Error err -> throwM $ ParseError err
1127 stepOuter _ (ParseChunksYield a next) = return $ Yield a next
1129{-# ANN type ConcatParseState Fuse #-}
1130data ConcatParseState b inpBuf st p m a =
1131 ConcatParseInit inpBuf st p
1132 | ConcatParseInitLeftOver inpBuf
1133 | forall s. ConcatParseStream st inpBuf (s -> a -> m (PRD.Step s b)) s (s -> m b)
1134 | forall s. ConcatParseBuf inpBuf st inpBuf (s -> a -> m (PRD.Step s b)) s (s -> m b)
1135 | ConcatParseYield b (ConcatParseState b inpBuf st p m a)
1137{-# INLINE_NORMAL parseIterate #-}
1138parseIterate
1139 :: MonadThrow m
1140 => (b -> PRD.Parser m a b)
1141 -> b
1142 -> Stream m a
1143 -> Stream m b
1144parseIterate func seed (Stream step state) =
1145 Stream stepOuter (ConcatParseInit [] state (func seed))
1147 where
1149 {-# INLINE_LATE stepOuter #-}
1150 -- Buffer is empty, go to stream processing loop
1151 stepOuter _ (ConcatParseInit [] st (PRD.Parser pstep initial extract)) = do
1152 res <- initial
1153 case res of
1154 PRD.IPartial ps ->
1155 return $ Skip $ ConcatParseStream st [] pstep ps extract
1156 PRD.IDone pb ->
1157 let next = ConcatParseInit [] st (func pb)
1158 in return $ Skip $ ConcatParseYield pb next
1159 PRD.IError err -> throwM $ ParseError err
1161 -- Buffer is not empty, go to buffered processing loop
1162 stepOuter _ (ConcatParseInit src st
1163 (PRD.Parser pstep initial extract)) = do
1164 res <- initial
1165 case res of
1166 PRD.IPartial ps ->
1167 return $ Skip $ ConcatParseBuf src st [] pstep ps extract
1168 PRD.IDone pb ->
1169 let next = ConcatParseInit src st (func pb)
1170 in return $ Skip $ ConcatParseYield pb next
1171 PRD.IError err -> throwM $ ParseError err
1173 -- XXX we just discard any leftover input at the end
1174 stepOuter _ (ConcatParseInitLeftOver _) = return Stop
1176 -- Buffer is empty process elements from the stream
1177 stepOuter gst (ConcatParseStream st buf pstep pst extract) = do
1178 r <- step (adaptState gst) st
1179 case r of
1180 Yield x s -> do
1181 pRes <- pstep pst x
1182 case pRes of
1183 PR.Partial 0 pst1 ->
1184 return $ Skip $ ConcatParseStream s [] pstep pst1 extract
1185 PR.Partial n pst1 -> do
1186 assert (n <= length (x:buf)) (return ())
1187 let src0 = Prelude.take n (x:buf)
1188 src = Prelude.reverse src0
1189 return $ Skip $ ConcatParseBuf src s [] pstep pst1 extract
1190 -- PR.Continue 0 pst1 ->
1191 -- return $ Skip $ ConcatParseStream s (x:buf) pst1
1192 PR.Continue n pst1 -> do
1193 assert (n <= length (x:buf)) (return ())
1194 let (src0, buf1) = splitAt n (x:buf)
1195 src = Prelude.reverse src0
1196 return $ Skip $ ConcatParseBuf src s buf1 pstep pst1 extract
1197 -- XXX Specialize for Stop 0 common case?
1198 PR.Done n b -> do
1199 assert (n <= length (x:buf)) (return ())
1200 let src = Prelude.reverse (Prelude.take n (x:buf))
1201 return $ Skip $
1202 ConcatParseYield b (ConcatParseInit src s (func b))
1203 PR.Error err -> throwM $ ParseError err
1204 Skip s -> return $ Skip $ ConcatParseStream s buf pstep pst extract
1205 Stop -> do
1206 b <- extract pst
1207 let src = Prelude.reverse buf
1208 return $ Skip $ ConcatParseYield b (ConcatParseInitLeftOver src)
1210 -- go back to stream processing mode
1211 stepOuter _ (ConcatParseBuf [] s buf pstep ps extract) =
1212 return $ Skip $ ConcatParseStream s buf pstep ps extract
1214 -- buffered processing loop
1215 stepOuter _ (ConcatParseBuf (x:xs) s buf pstep pst extract) = do
1216 pRes <- pstep pst x
1217 case pRes of
1218 PR.Partial 0 pst1 ->
1219 return $ Skip $ ConcatParseBuf xs s [] pstep pst1 extract
1220 PR.Partial n pst1 -> do
1221 assert (n <= length (x:buf)) (return ())
1222 let src0 = Prelude.take n (x:buf)
1223 src = Prelude.reverse src0 ++ xs
1224 return $ Skip $ ConcatParseBuf src s [] pstep pst1 extract
1225 -- PR.Continue 0 pst1 -> return $ Skip $ ConcatParseBuf xs s (x:buf) pst1
1226 PR.Continue n pst1 -> do
1227 assert (n <= length (x:buf)) (return ())
1228 let (src0, buf1) = splitAt n (x:buf)
1229 src = Prelude.reverse src0 ++ xs
1230 return $ Skip $ ConcatParseBuf src s buf1 pstep pst1 extract
1231 -- XXX Specialize for Stop 0 common case?
1232 PR.Done n b -> do
1233 assert (n <= length (x:buf)) (return ())
1234 let src = Prelude.reverse (Prelude.take n (x:buf)) ++ xs
1235 return $ Skip $ ConcatParseYield b
1236 (ConcatParseInit src s (func b))
1237 PR.Error err -> throwM $ ParseError err
1239 stepOuter _ (ConcatParseYield a next) = return $ Yield a next
1241------------------------------------------------------------------------------
1242-- Grouping
1243------------------------------------------------------------------------------
1245data GroupByState st fs a b
1246 = GroupingInit st
1247 | GroupingDo st !fs
1248 | GroupingInitWith st !a
1249 | GroupingDoWith st !fs !a
1250 | GroupingYield !b (GroupByState st fs a b)
1251 | GroupingDone
1253{-# INLINE_NORMAL groupsBy #-}
1254groupsBy :: Monad m
1255 => (a -> a -> Bool)
1256 -> Fold m a b
1257 -> Stream m a
1258 -> Stream m b
1260groupsBy eq fld = parseMany (PRD.groupBy eq fld)
1262groupsBy cmp (Fold fstep initial done) (Stream step state) =
1263 Stream stepOuter (GroupingInit state)
1265 where
1267 {-# INLINE_LATE stepOuter #-}
1268 stepOuter _ (GroupingInit st) = do
1269 -- XXX Note that if the stream stops without yielding a single element
1270 -- in the group we discard the "initial" effect.
1271 res <- initial
1272 return
1273 $ case res of
1274 FL.Partial s -> Skip $ GroupingDo st s
1275 FL.Done b -> Yield b $ GroupingInit st
1276 stepOuter gst (GroupingDo st fs) = do
1277 res <- step (adaptState gst) st
1278 case res of
1279 Yield x s -> do
1280 r <- fstep fs x
1281 case r of
1282 FL.Partial fs1 -> go SPEC x s fs1
1283 FL.Done b -> return $ Yield b (GroupingInit s)
1284 Skip s -> return $ Skip $ GroupingDo s fs
1285 Stop -> return Stop
1287 where
1289 go !_ prev stt !acc = do
1290 res <- step (adaptState gst) stt
1291 case res of
1292 Yield x s -> do
1293 if cmp x prev
1294 then do
1295 r <- fstep acc x
1296 case r of
1297 FL.Partial fs1 -> go SPEC prev s fs1
1298 FL.Done b -> return $ Yield b (GroupingInit s)
1299 else do
1300 r <- done acc
1301 return $ Yield r (GroupingInitWith s x)
1302 Skip s -> go SPEC prev s acc
1303 Stop -> done acc >>= \r -> return $ Yield r GroupingDone
1304 stepOuter _ (GroupingInitWith st x) = do
1305 res <- initial
1306 return
1307 $ case res of
1308 FL.Partial s -> Skip $ GroupingDoWith st s x
1309 FL.Done b -> Yield b $ GroupingInitWith st x
1310 stepOuter gst (GroupingDoWith st fs prev) = do
1311 res <- fstep fs prev
1312 case res of
1313 FL.Partial fs1 -> go SPEC st fs1
1314 FL.Done b -> return $ Yield b (GroupingInit st)
1316 where
1318 -- XXX code duplicated from the previous equation
1319 go !_ stt !acc = do
1320 res <- step (adaptState gst) stt
1321 case res of
1322 Yield x s -> do
1323 if cmp x prev
1324 then do
1325 r <- fstep acc x
1326 case r of
1327 FL.Partial fs1 -> go SPEC s fs1
1328 FL.Done b -> return $ Yield b (GroupingInit s)
1329 else do
1330 r <- done acc
1331 return $ Yield r (GroupingInitWith s x)
1332 Skip s -> go SPEC s acc
1333 Stop -> done acc >>= \r -> return $ Yield r GroupingDone
1334 stepOuter _ (GroupingYield _ _) = error "groupsBy: Unreachable"
1335 stepOuter _ GroupingDone = return Stop
1337{-# INLINE_NORMAL groupsRollingBy #-}
1338groupsRollingBy :: Monad m
1339 => (a -> a -> Bool)
1340 -> Fold m a b
1341 -> Stream m a
1342 -> Stream m b
1344groupsRollingBy eq fld = parseMany (PRD.groupByRolling eq fld)
1346groupsRollingBy cmp (Fold fstep initial done) (Stream step state) =
1347 Stream stepOuter (GroupingInit state)
1349 where
1351 {-# INLINE_LATE stepOuter #-}
1352 stepOuter _ (GroupingInit st) = do
1353 -- XXX Note that if the stream stops without yielding a single element
1354 -- in the group we discard the "initial" effect.
1355 res <- initial
1356 return
1357 $ case res of
1358 FL.Partial fs -> Skip $ GroupingDo st fs
1359 FL.Done fb -> Yield fb $ GroupingInit st
1360 stepOuter gst (GroupingDo st fs) = do
1361 res <- step (adaptState gst) st
1362 case res of
1363 Yield x s -> do
1364 r <- fstep fs x
1365 case r of
1366 FL.Partial fs1 -> go SPEC x s fs1
1367 FL.Done fb -> return $ Yield fb (GroupingInit s)
1368 Skip s -> return $ Skip $ GroupingDo s fs
1369 Stop -> return Stop
1371 where
1373 go !_ prev stt !acc = do
1374 res <- step (adaptState gst) stt
1375 case res of
1376 Yield x s -> do
1377 if cmp prev x
1378 then do
1379 r <- fstep acc x
1380 case r of
1381 FL.Partial fs1 -> go SPEC x s fs1
1382 FL.Done b -> return $ Yield b (GroupingInit s)
1383 else do
1384 r <- done acc
1385 return $ Yield r (GroupingInitWith s x)
1386 Skip s -> go SPEC prev s acc
1387 Stop -> done acc >>= \r -> return $ Yield r GroupingDone
1388 stepOuter _ (GroupingInitWith st x) = do
1389 res <- initial
1390 return
1391 $ case res of
1392 FL.Partial s -> Skip $ GroupingDoWith st s x
1393 FL.Done b -> Yield b $ GroupingInitWith st x
1394 stepOuter gst (GroupingDoWith st fs previous) = do
1395 res <- fstep fs previous
1396 case res of
1397 FL.Partial s -> go SPEC previous st s
1398 FL.Done b -> return $ Yield b (GroupingInit st)
1400 where
1402 -- XXX GHC: groupsBy has one less parameter in this go loop and it
1403 -- fuses. However, groupsRollingBy does not fuse, removing the prev
1404 -- parameter makes it fuse. Something needs to be fixed in GHC. The
1405 -- workaround for this is noted in the comments below.
1406 go !_ prev !stt !acc = do
1407 res <- step (adaptState gst) stt
1408 case res of
1409 Yield x s -> do
1410 if cmp prev x
1411 then do
1412 r <- fstep acc x
1413 case r of
1414 FL.Partial fs1 -> go SPEC x s fs1
1415 FL.Done b -> return $ Yield b (GroupingInit st)
1416 else do
1417 {-
1418 r <- done acc
1419 return $ Yield r (GroupingInitWith s x)
1420 -}
1421 -- The code above does not let groupBy fuse. We use the
1422 -- alternative code below instead. Instead of jumping
1423 -- to GroupingInitWith state, we unroll the code of
1424 -- GroupingInitWith state here to help GHC with stream
1425 -- fusion.
1426 result <- initial
1427 r <- done acc
1428 return
1429 $ Yield r
1430 $ case result of
1431 FL.Partial fsi -> GroupingDoWith s fsi x
1432 FL.Done b -> GroupingYield b (GroupingInit s)
1433 Skip s -> go SPEC prev s acc
1434 Stop -> done acc >>= \r -> return $ Yield r GroupingDone
1435 stepOuter _ (GroupingYield r next) = return $ Yield r next
1436 stepOuter _ GroupingDone = return Stop
1438------------------------------------------------------------------------------
1439-- Splitting - by a predicate
1440------------------------------------------------------------------------------
1442data WordsByState st fs b
1443 = WordsByInit st
1444 | WordsByDo st !fs
1445 | WordsByDone
1446 | WordsByYield !b (WordsByState st fs b)
1448{-# INLINE_NORMAL wordsBy #-}
1449wordsBy :: Monad m => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1450wordsBy predicate (Fold fstep initial done) (Stream step state) =
1451 Stream stepOuter (WordsByInit state)
1453 where
1455 {-# INLINE_LATE stepOuter #-}
1456 stepOuter _ (WordsByInit st) = do
1457 res <- initial
1458 return
1459 $ case res of
1460 FL.Partial s -> Skip $ WordsByDo st s
1461 FL.Done b -> Yield b (WordsByInit st)
1463 stepOuter gst (WordsByDo st fs) = do
1464 res <- step (adaptState gst) st
1465 case res of
1466 Yield x s -> do
1467 if predicate x
1468 then do
1469 resi <- initial
1470 return
1471 $ case resi of
1472 FL.Partial fs1 -> Skip $ WordsByDo s fs1
1473 FL.Done b -> Yield b (WordsByInit s)
1474 else do
1475 r <- fstep fs x
1476 case r of
1477 FL.Partial fs1 -> go SPEC s fs1
1478 FL.Done b -> return $ Yield b (WordsByInit s)
1479 Skip s -> return $ Skip $ WordsByDo s fs
1480 Stop -> return Stop
1482 where
1484 go !_ stt !acc = do
1485 res <- step (adaptState gst) stt
1486 case res of
1487 Yield x s -> do
1488 if predicate x
1489 then do
1490 {-
1491 r <- done acc
1492 return $ Yield r (WordsByInit s)
1493 -}
1494 -- The above code does not fuse well. Need to check why
1495 -- GHC is not able to simplify it well. Using the code
1496 -- below, instead of jumping through the WordsByInit
1497 -- state always, we directly go to WordsByDo state in
1498 -- the common case of Partial.
1499 resi <- initial
1500 r <- done acc
1501 return
1502 $ Yield r
1503 $ case resi of
1504 FL.Partial fs1 -> WordsByDo s fs1
1505 FL.Done b -> WordsByYield b (WordsByInit s)
1506 else do
1507 r <- fstep acc x
1508 case r of
1509 FL.Partial fs1 -> go SPEC s fs1
1510 FL.Done b -> return $ Yield b (WordsByInit s)
1511 Skip s -> go SPEC s acc
1512 Stop -> done acc >>= \r -> return $ Yield r WordsByDone
1514 stepOuter _ WordsByDone = return Stop
1516 stepOuter _ (WordsByYield b next) = return $ Yield b next
1518------------------------------------------------------------------------------
1519-- Splitting on a sequence
1520------------------------------------------------------------------------------
1522-- String search algorithms:
1523-- http://www-igm.univ-mlv.fr/~lecroq/string/index.html
1526-- TODO can we unify the splitting operations using a splitting configuration
1527-- like in the split package.
1529data SplitStyle = Infix | Suffix | Prefix deriving (Eq, Show)
1530data SplitOptions = SplitOptions
1531 { style :: SplitStyle
1532 , withSep :: Bool -- ^ keep the separators in output
1533 -- , compact :: Bool -- ^ treat multiple consecutive separators as one
1534 -- , trimHead :: Bool -- ^ drop blank at head
1535 -- , trimTail :: Bool -- ^ drop blank at tail
1539-- XXX using "fs" as the last arg in Constructors may simplify the code a bit,
1540-- because we can use the constructor directly without having to create "jump"
1541-- functions.
1542{-# ANN type SplitOnSeqState Fuse #-}
1543data SplitOnSeqState rb rh ck w fs s b x =
1544 SplitOnSeqInit
1545 | SplitOnSeqYield b (SplitOnSeqState rb rh ck w fs s b x)
1546 | SplitOnSeqDone
1548 | SplitOnSeqEmpty !fs s
1550 | SplitOnSeqSingle !fs s x
1552 | SplitOnSeqWordInit !fs s
1553 | SplitOnSeqWordLoop !w s !fs
1554 | SplitOnSeqWordDone Int !fs !w
1556 | SplitOnSeqKRInit Int !fs s rb !rh
1557 | SplitOnSeqKRLoop fs s rb !rh !ck
1558 | SplitOnSeqKRCheck fs s rb !rh
1559 | SplitOnSeqKRDone Int !fs rb !rh
1561 | SplitOnSeqReinit (fs -> SplitOnSeqState rb rh ck w fs s b x)
1563{-# INLINE_NORMAL splitOnSeq #-}
1564splitOnSeq
1565 :: forall m a b. (MonadIO m, Storable a, Enum a, Eq a)
1566 => Array a
1567 -> Fold m a b
1568 -> Stream m a
1569 -> Stream m b
1570splitOnSeq patArr (Fold fstep initial done) (Stream step state) =
1571 Stream stepOuter SplitOnSeqInit
1573 where
1575 patLen = A.length patArr
1576 maxIndex = patLen - 1
1577 elemBits = sizeOf (undefined :: a) * 8
1579 -- For word pattern case
1580 wordMask :: Word
1581 wordMask = (1 `shiftL` (elemBits * patLen)) - 1
1583 elemMask :: Word
1584 elemMask = (1 `shiftL` elemBits) - 1
1586 wordPat :: Word
1587 wordPat = wordMask .&. A.foldl' addToWord 0 patArr
1589 addToWord wd a = (wd `shiftL` elemBits) .|. fromIntegral (fromEnum a)
1591 -- For Rabin-Karp search
1592 k = 2891336453 :: Word32
1593 coeff = k ^ patLen
1595 addCksum cksum a = cksum * k + fromIntegral (fromEnum a)
1597 deltaCksum cksum old new =
1598 addCksum cksum new - coeff * fromIntegral (fromEnum old)
1600 -- XXX shall we use a random starting hash or 1 instead of 0?
1601 patHash = A.foldl' addCksum 0 patArr
1603 skip = return . Skip
1605 nextAfterInit nextGen stepRes =
1606 case stepRes of
1607 FL.Partial s -> nextGen s
1608 FL.Done b -> SplitOnSeqYield b (SplitOnSeqReinit nextGen)
1610 {-# INLINE yieldProceed #-}
1611 yieldProceed nextGen fs =
1612 initial >>= skip . SplitOnSeqYield fs . nextAfterInit nextGen
1614 {-# INLINE_LATE stepOuter #-}
1615 stepOuter _ SplitOnSeqInit = do
1616 res <- initial
1617 case res of
1618 FL.Partial acc ->
1619 if patLen == 0
1620 then return $ Skip $ SplitOnSeqEmpty acc state
1621 else if patLen == 1
1622 then do
1623 pat <- liftIO $ A.unsafeIndexIO patArr 0
1624 return $ Skip $ SplitOnSeqSingle acc state pat
1625 else if sizeOf (undefined :: a) * patLen
1626 <= sizeOf (undefined :: Word)
1627 then return $ Skip $ SplitOnSeqWordInit acc state
1628 else do
1629 (rb, rhead) <- liftIO $ RB.new patLen
1630 skip $ SplitOnSeqKRInit 0 acc state rb rhead
1631 FL.Done b -> skip $ SplitOnSeqYield b SplitOnSeqInit
1633 stepOuter _ (SplitOnSeqYield x next) = return $ Yield x next
1635 ---------------------------
1636 -- Checkpoint
1637 ---------------------------
1639 stepOuter _ (SplitOnSeqReinit nextGen) =
1640 initial >>= skip . nextAfterInit nextGen
1642 ---------------------------
1643 -- Empty pattern
1644 ---------------------------
1646 stepOuter gst (SplitOnSeqEmpty acc st) = do
1647 res <- step (adaptState gst) st
1648 case res of
1649 Yield x s -> do
1650 r <- fstep acc x
1651 b1 <-
1652 case r of
1653 FL.Partial acc1 -> done acc1
1654 FL.Done b -> return b
1655 let jump c = SplitOnSeqEmpty c s
1656 in yieldProceed jump b1
1657 Skip s -> skip (SplitOnSeqEmpty acc s)
1658 Stop -> return Stop
1660 -----------------
1661 -- Done
1662 -----------------
1664 stepOuter _ SplitOnSeqDone = return Stop
1666 -----------------
1667 -- Single Pattern
1668 -----------------
1670 stepOuter gst (SplitOnSeqSingle fs st pat) = do
1671 res <- step (adaptState gst) st
1672 case res of
1673 Yield x s -> do
1674 let jump c = SplitOnSeqSingle c s pat
1675 if pat == x
1676 then done fs >>= yieldProceed jump
1677 else do
1678 r <- fstep fs x
1679 case r of
1680 FL.Partial fs1 -> skip $ jump fs1
1681 FL.Done b -> yieldProceed jump b
1682 Skip s -> return $ Skip $ SplitOnSeqSingle fs s pat
1683 Stop -> do
1684 r <- done fs
1685 return $ Skip $ SplitOnSeqYield r SplitOnSeqDone
1687 ---------------------------
1688 -- Short Pattern - Shift Or
1689 ---------------------------
1691 stepOuter _ (SplitOnSeqWordDone 0 fs _) = do
1692 r <- done fs
1693 skip $ SplitOnSeqYield r SplitOnSeqDone
1694 stepOuter _ (SplitOnSeqWordDone n fs wrd) = do
1695 let old = elemMask .&. (wrd `shiftR` (elemBits * (n - 1)))
1696 r <- fstep fs (toEnum $ fromIntegral old)
1697 case r of
1698 FL.Partial fs1 -> skip $ SplitOnSeqWordDone (n - 1) fs1 wrd
1699 FL.Done b -> do
1700 let jump c = SplitOnSeqWordDone (n - 1) c wrd
1701 yieldProceed jump b
1703 stepOuter gst (SplitOnSeqWordInit fs st0) =
1704 go SPEC 0 0 st0
1706 where
1708 {-# INLINE go #-}
1709 go !_ !idx !wrd !st = do
1710 res <- step (adaptState gst) st
1711 case res of
1712 Yield x s -> do
1713 let wrd1 = addToWord wrd x
1714 if idx == maxIndex
1715 then do
1716 if wrd1 .&. wordMask == wordPat
1717 then do
1718 let jump c = SplitOnSeqWordInit c s
1719 done fs >>= yieldProceed jump
1720 else skip $ SplitOnSeqWordLoop wrd1 s fs
1721 else go SPEC (idx + 1) wrd1 s
1722 Skip s -> go SPEC idx wrd s
1723 Stop -> do
1724 if idx /= 0
1725 then skip $ SplitOnSeqWordDone idx fs wrd
1726 else do
1727 r <- done fs
1728 skip $ SplitOnSeqYield r SplitOnSeqDone
1730 stepOuter gst (SplitOnSeqWordLoop wrd0 st0 fs0) =
1731 go SPEC wrd0 st0 fs0
1733 where
1735 {-# INLINE go #-}
1736 go !_ !wrd !st !fs = do
1737 res <- step (adaptState gst) st
1738 case res of
1739 Yield x s -> do
1740 let jump c = SplitOnSeqWordInit c s
1741 wrd1 = addToWord wrd x
1742 old = (wordMask .&. wrd)
1743 `shiftR` (elemBits * (patLen - 1))
1744 r <- fstep fs (toEnum $ fromIntegral old)
1745 case r of
1746 FL.Partial fs1 -> do
1747 if wrd1 .&. wordMask == wordPat
1748 then done fs1 >>= yieldProceed jump
1749 else go SPEC wrd1 s fs1
1750 FL.Done b -> yieldProceed jump b
1751 Skip s -> go SPEC wrd s fs
1752 Stop -> skip $ SplitOnSeqWordDone patLen fs wrd
1754 -------------------------------
1755 -- General Pattern - Karp Rabin
1756 -------------------------------
1758 stepOuter gst (SplitOnSeqKRInit idx fs st rb rh) = do
1759 res <- step (adaptState gst) st
1760 case res of
1761 Yield x s -> do
1762 rh1 <- liftIO $ RB.unsafeInsert rb rh x
1763 if idx == maxIndex
1764 then do
1765 let fld = RB.unsafeFoldRing (RB.ringBound rb)
1766 let !ringHash = fld addCksum 0 rb
1767 if ringHash == patHash
1768 then skip $ SplitOnSeqKRCheck fs s rb rh1
1769 else skip $ SplitOnSeqKRLoop fs s rb rh1 ringHash
1770 else skip $ SplitOnSeqKRInit (idx + 1) fs s rb rh1
1771 Skip s -> skip $ SplitOnSeqKRInit idx fs s rb rh
1772 Stop -> do
1773 skip $ SplitOnSeqKRDone idx fs rb (RB.startOf rb)
1775 -- XXX The recursive "go" is more efficient than the state based recursion
1776 -- code commented out below. Perhaps its more efficient because of
1777 -- factoring out "rb" outside the loop.
1778 --
1779 stepOuter gst (SplitOnSeqKRLoop fs0 st0 rb rh0 cksum0) =
1780 go SPEC fs0 st0 rh0 cksum0
1782 where
1784 go !_ !fs !st !rh !cksum = do
1785 res <- step (adaptState gst) st
1786 case res of
1787 Yield x s -> do
1788 old <- liftIO $ peek rh
1789 let cksum1 = deltaCksum cksum old x
1790 r <- fstep fs old
1791 case r of
1792 FL.Partial fs1 -> do
1793 rh1 <- liftIO (RB.unsafeInsert rb rh x)
1794 if cksum1 == patHash
1795 then skip $ SplitOnSeqKRCheck fs1 s rb rh1
1796 else go SPEC fs1 s rh1 cksum1
1797 FL.Done b -> do
1798 let rst = RB.startOf rb
1799 jump c = SplitOnSeqKRInit 0 c s rb rst
1800 yieldProceed jump b
1801 Skip s -> go SPEC fs s rh cksum
1802 Stop -> skip $ SplitOnSeqKRDone patLen fs rb rh
1804 -- XXX The following code is 5 times slower compared to the recursive loop
1805 -- based code above. Need to investigate why. One possibility is that the
1806 -- go loop above does not thread around the ring buffer (rb). This code may
1807 -- be causing the state to bloat and getting allocated on each iteration.
1808 -- We can check the cmm/asm code to confirm. If so a good GHC solution to
1809 -- such problem is needed. One way to avoid this could be to use unboxed
1810 -- mutable state?
1811 {-
1812 stepOuter gst (SplitOnSeqKRLoop fs st rb rh cksum) = do
1813 res <- step (adaptState gst) st
1814 case res of
1815 Yield x s -> do
1816 old <- liftIO $ peek rh
1817 let cksum1 = deltaCksum cksum old x
1818 fs1 <- fstep fs old
1819 if (cksum1 == patHash)
1820 then do
1821 r <- done fs1
1822 skip $ SplitOnSeqYield r $ SplitOnSeqKRInit 0 s rb rh
1823 else do
1824 rh1 <- liftIO (RB.unsafeInsert rb rh x)
1825 skip $ SplitOnSeqKRLoop fs1 s rb rh1 cksum1
1826 Skip s -> skip $ SplitOnSeqKRLoop fs s rb rh cksum
1827 Stop -> skip $ SplitOnSeqKRDone patLen fs rb rh
1828 -}
1830 stepOuter _ (SplitOnSeqKRCheck fs st rb rh) = do
1831 if RB.unsafeEqArray rb rh patArr
1832 then do
1833 r <- done fs
1834 let rst = RB.startOf rb
1835 jump c = SplitOnSeqKRInit 0 c st rb rst
1836 yieldProceed jump r
1837 else skip $ SplitOnSeqKRLoop fs st rb rh patHash
1839 stepOuter _ (SplitOnSeqKRDone 0 fs _ _) = do
1840 r <- done fs
1841 skip $ SplitOnSeqYield r SplitOnSeqDone
1842 stepOuter _ (SplitOnSeqKRDone n fs rb rh) = do
1843 old <- liftIO $ peek rh
1844 let rh1 = RB.advance rb rh
1845 r <- fstep fs old
1846 case r of
1847 FL.Partial fs1 -> skip $ SplitOnSeqKRDone (n - 1) fs1 rb rh1
1848 FL.Done b -> do
1849 let jump c = SplitOnSeqKRDone (n - 1) c rb rh1
1850 yieldProceed jump b
1852{-# ANN type SplitOnSuffixSeqState Fuse #-}
1853data SplitOnSuffixSeqState rb rh ck w fs s b x =
1854 SplitOnSuffixSeqInit
1855 | SplitOnSuffixSeqYield b (SplitOnSuffixSeqState rb rh ck w fs s b x)
1856 | SplitOnSuffixSeqDone
1858 | SplitOnSuffixSeqEmpty !fs s
1860 | SplitOnSuffixSeqSingleInit !fs s x
1861 | SplitOnSuffixSeqSingle !fs s x
1863 | SplitOnSuffixSeqWordInit !fs s
1864 | SplitOnSuffixSeqWordLoop !w s !fs
1865 | SplitOnSuffixSeqWordDone Int !fs !w
1867 | SplitOnSuffixSeqKRInit Int !fs s rb !rh
1868 | SplitOnSuffixSeqKRInit1 !fs s rb !rh
1869 | SplitOnSuffixSeqKRLoop fs s rb !rh !ck
1870 | SplitOnSuffixSeqKRCheck fs s rb !rh
1871 | SplitOnSuffixSeqKRDone Int !fs rb !rh
1873 | SplitOnSuffixSeqReinit
1874 (fs -> SplitOnSuffixSeqState rb rh ck w fs s b x)
1876{-# INLINE_NORMAL splitOnSuffixSeq #-}
1877splitOnSuffixSeq
1878 :: forall m a b. (MonadIO m, Storable a, Enum a, Eq a)
1879 => Bool
1880 -> Array a
1881 -> Fold m a b
1882 -> Stream m a
1883 -> Stream m b
1884splitOnSuffixSeq withSep patArr (Fold fstep initial done) (Stream step state) =
1885 Stream stepOuter SplitOnSuffixSeqInit
1887 where
1889 patLen = A.length patArr
1890 maxIndex = patLen - 1
1891 elemBits = sizeOf (undefined :: a) * 8
1893 -- For word pattern case
1894 wordMask :: Word
1895 wordMask = (1 `shiftL` (elemBits * patLen)) - 1
1897 elemMask :: Word
1898 elemMask = (1 `shiftL` elemBits) - 1
1900 wordPat :: Word
1901 wordPat = wordMask .&. A.foldl' addToWord 0 patArr
1903 addToWord wd a = (wd `shiftL` elemBits) .|. fromIntegral (fromEnum a)
1905 nextAfterInit nextGen stepRes =
1906 case stepRes of
1907 FL.Partial s -> nextGen s
1908 FL.Done b ->
1909 SplitOnSuffixSeqYield b (SplitOnSuffixSeqReinit nextGen)
1911 {-# INLINE yieldProceed #-}
1912 yieldProceed nextGen fs =
1913 initial >>= skip . SplitOnSuffixSeqYield fs . nextAfterInit nextGen
1915 -- For single element pattern case
1916 {-# INLINE processYieldSingle #-}
1917 processYieldSingle pat x s fs = do
1918 let jump c = SplitOnSuffixSeqSingleInit c s pat
1919 if pat == x
1920 then do
1921 r <- if withSep then fstep fs x else return $ FL.Partial fs
1922 b1 <-
1923 case r of
1924 FL.Partial fs1 -> done fs1
1925 FL.Done b -> return b
1926 yieldProceed jump b1
1927 else do
1928 r <- fstep fs x
1929 case r of
1930 FL.Partial fs1 -> skip $ SplitOnSuffixSeqSingle fs1 s pat
1931 FL.Done b -> yieldProceed jump b
1933 -- For Rabin-Karp search
1934 k = 2891336453 :: Word32
1935 coeff = k ^ patLen
1937 addCksum cksum a = cksum * k + fromIntegral (fromEnum a)
1939 deltaCksum cksum old new =
1940 addCksum cksum new - coeff * fromIntegral (fromEnum old)
1942 -- XXX shall we use a random starting hash or 1 instead of 0?
1943 patHash = A.foldl' addCksum 0 patArr
1945 skip = return . Skip
1947 {-# INLINE_LATE stepOuter #-}
1948 stepOuter _ SplitOnSuffixSeqInit = do
1949 res <- initial
1950 case res of
1951 FL.Partial fs ->
1952 if patLen == 0
1953 then skip $ SplitOnSuffixSeqEmpty fs state
1954 else if patLen == 1
1955 then do
1956 pat <- liftIO $ A.unsafeIndexIO patArr 0
1957 skip $ SplitOnSuffixSeqSingleInit fs state pat
1958 else if sizeOf (undefined :: a) * patLen
1959 <= sizeOf (undefined :: Word)
1960 then skip $ SplitOnSuffixSeqWordInit fs state
1961 else do
1962 (rb, rhead) <- liftIO $ RB.new patLen
1963 skip $ SplitOnSuffixSeqKRInit 0 fs state rb rhead
1964 FL.Done fb -> skip $ SplitOnSuffixSeqYield fb SplitOnSuffixSeqInit
1966 stepOuter _ (SplitOnSuffixSeqYield x next) = return $ Yield x next
1968 ---------------------------
1969 -- Reinit
1970 ---------------------------
1972 stepOuter _ (SplitOnSuffixSeqReinit nextGen) =
1973 initial >>= skip . nextAfterInit nextGen
1975 ---------------------------
1976 -- Empty pattern
1977 ---------------------------
1979 stepOuter gst (SplitOnSuffixSeqEmpty acc st) = do
1980 res <- step (adaptState gst) st
1981 case res of
1982 Yield x s -> do
1983 let jump c = SplitOnSuffixSeqEmpty c s
1984 r <- fstep acc x
1985 b1 <-
1986 case r of
1987 FL.Partial fs -> done fs
1988 FL.Done b -> return b
1989 yieldProceed jump b1
1990 Skip s -> skip (SplitOnSuffixSeqEmpty acc s)
1991 Stop -> return Stop
1993 -----------------
1994 -- Done
1995 -----------------
1997 stepOuter _ SplitOnSuffixSeqDone = return Stop
1999 -----------------
2000 -- Single Pattern
2001 -----------------
2003 stepOuter gst (SplitOnSuffixSeqSingleInit fs st pat) = do
2004 res <- step (adaptState gst) st
2005 case res of
2006 Yield x s -> processYieldSingle pat x s fs
2007 Skip s -> skip $ SplitOnSuffixSeqSingleInit fs s pat
2008 Stop -> return Stop
2010 stepOuter gst (SplitOnSuffixSeqSingle fs st pat) = do
2011 res <- step (adaptState gst) st
2012 case res of
2013 Yield x s -> processYieldSingle pat x s fs
2014 Skip s -> skip $ SplitOnSuffixSeqSingle fs s pat
2015 Stop -> do
2016 r <- done fs
2017 skip $ SplitOnSuffixSeqYield r SplitOnSuffixSeqDone
2019 ---------------------------
2020 -- Short Pattern - Shift Or
2021 ---------------------------
2023 stepOuter _ (SplitOnSuffixSeqWordDone 0 fs _) = do
2024 r <- done fs
2025 skip $ SplitOnSuffixSeqYield r SplitOnSuffixSeqDone
2026 stepOuter _ (SplitOnSuffixSeqWordDone n fs wrd) = do
2027 let old = elemMask .&. (wrd `shiftR` (elemBits * (n - 1)))
2028 r <- fstep fs (toEnum $ fromIntegral old)
2029 case r of
2030 FL.Partial fs1 -> skip $ SplitOnSuffixSeqWordDone (n - 1) fs1 wrd
2031 FL.Done b -> do
2032 let jump c = SplitOnSuffixSeqWordDone (n - 1) c wrd
2033 yieldProceed jump b
2035 stepOuter gst (SplitOnSuffixSeqWordInit fs0 st0) = do
2036 res <- step (adaptState gst) st0
2037 case res of
2038 Yield x s -> do
2039 let wrd = addToWord 0 x
2040 r <- if withSep then fstep fs0 x else return $ FL.Partial fs0
2041 case r of
2042 FL.Partial fs1 -> go SPEC 1 wrd s fs1
2043 FL.Done b -> do
2044 let jump c = SplitOnSuffixSeqWordInit c s
2045 yieldProceed jump b
2046 Skip s -> skip (SplitOnSuffixSeqWordInit fs0 s)
2047 Stop -> return Stop
2049 where
2051 {-# INLINE go #-}
2052 go !_ !idx !wrd !st !fs = do
2053 res <- step (adaptState gst) st
2054 case res of
2055 Yield x s -> do
2056 let jump c = SplitOnSuffixSeqWordInit c s
2057 let wrd1 = addToWord wrd x
2058 r <- if withSep then fstep fs x else return $ FL.Partial fs
2059 case r of
2060 FL.Partial fs1 ->
2061 if idx /= maxIndex
2062 then go SPEC (idx + 1) wrd1 s fs1
2063 else if wrd1 .&. wordMask /= wordPat
2064 then skip $ SplitOnSuffixSeqWordLoop wrd1 s fs1
2065 else do done fs >>= yieldProceed jump
2066 FL.Done b -> yieldProceed jump b
2067 Skip s -> go SPEC idx wrd s fs
2068 Stop -> skip $ SplitOnSuffixSeqWordDone idx fs wrd
2070 stepOuter gst (SplitOnSuffixSeqWordLoop wrd0 st0 fs0) =
2071 go SPEC wrd0 st0 fs0
2073 where
2075 {-# INLINE go #-}
2076 go !_ !wrd !st !fs = do
2077 res <- step (adaptState gst) st
2078 case res of
2079 Yield x s -> do
2080 let jump c = SplitOnSuffixSeqWordInit c s
2081 wrd1 = addToWord wrd x
2082 old = (wordMask .&. wrd)
2083 `shiftR` (elemBits * (patLen - 1))
2084 r <-
2085 if withSep
2086 then fstep fs x
2087 else fstep fs (toEnum $ fromIntegral old)
2088 case r of
2089 FL.Partial fs1 ->
2090 if wrd1 .&. wordMask == wordPat
2091 then done fs1 >>= yieldProceed jump
2092 else go SPEC wrd1 s fs1
2093 FL.Done b -> yieldProceed jump b
2094 Skip s -> go SPEC wrd s fs
2095 Stop ->
2096 if wrd .&. wordMask == wordPat
2097 then return Stop
2098 else if withSep
2099 then do
2100 r <- done fs
2101 skip $ SplitOnSuffixSeqYield r SplitOnSuffixSeqDone
2102 else skip $ SplitOnSuffixSeqWordDone patLen fs wrd
2104 -------------------------------
2105 -- General Pattern - Karp Rabin
2106 -------------------------------
2108 stepOuter gst (SplitOnSuffixSeqKRInit idx0 fs st0 rb rh0) = do
2109 res <- step (adaptState gst) st0
2110 case res of
2111 Yield x s -> do
2112 rh1 <- liftIO $ RB.unsafeInsert rb rh0 x
2113 r <- if withSep then fstep fs x else return $ FL.Partial fs
2114 case r of
2115 FL.Partial fs1 ->
2116 skip $ SplitOnSuffixSeqKRInit1 fs1 s rb rh1
2117 FL.Done b -> do
2118 let rst = RB.startOf rb
2119 jump c = SplitOnSuffixSeqKRInit 0 c s rb rst
2120 yieldProceed jump b
2121 Skip s -> skip $ SplitOnSuffixSeqKRInit idx0 fs s rb rh0
2122 Stop -> return Stop
2124 stepOuter gst (SplitOnSuffixSeqKRInit1 fs0 st0 rb rh0) = do
2125 go SPEC 1 rh0 st0 fs0
2127 where
2129 go !_ !idx !rh st !fs = do
2130 res <- step (adaptState gst) st
2131 case res of
2132 Yield x s -> do
2133 rh1 <- liftIO (RB.unsafeInsert rb rh x)
2134 r <- if withSep then fstep fs x else return $ FL.Partial fs
2135 case r of
2136 FL.Partial fs1 ->
2137 if idx /= maxIndex
2138 then go SPEC (idx + 1) rh1 s fs1
2139 else skip $
2140 let fld = RB.unsafeFoldRing (RB.ringBound rb)
2141 !ringHash = fld addCksum 0 rb
2142 in if ringHash == patHash
2143 then SplitOnSuffixSeqKRCheck fs1 s rb rh1
2144 else SplitOnSuffixSeqKRLoop
2145 fs1 s rb rh1 ringHash
2146 FL.Done b -> do
2147 let rst = RB.startOf rb
2148 jump c = SplitOnSuffixSeqKRInit 0 c s rb rst
2149 yieldProceed jump b
2150 Skip s -> go SPEC idx rh s fs
2151 Stop -> do
2152 -- do not issue a blank segment when we end at pattern
2153 if (idx == maxIndex) && RB.unsafeEqArray rb rh patArr
2154 then return Stop
2155 else if withSep
2156 then do
2157 r <- done fs
2158 skip $ SplitOnSuffixSeqYield r SplitOnSuffixSeqDone
2159 else skip $ SplitOnSuffixSeqKRDone idx fs rb (RB.startOf rb)
2161 stepOuter gst (SplitOnSuffixSeqKRLoop fs0 st0 rb rh0 cksum0) =
2162 go SPEC fs0 st0 rh0 cksum0
2164 where
2166 go !_ !fs !st !rh !cksum = do
2167 res <- step (adaptState gst) st
2168 case res of
2169 Yield x s -> do
2170 old <- liftIO $ peek rh
2171 rh1 <- liftIO (RB.unsafeInsert rb rh x)
2172 let cksum1 = deltaCksum cksum old x
2173 r <- if withSep then fstep fs x else fstep fs old
2174 case r of
2175 FL.Partial fs1 ->
2176 if (cksum1 /= patHash)
2177 then go SPEC fs1 s rh1 cksum1
2178 else skip $ SplitOnSuffixSeqKRCheck fs1 s rb rh1
2179 FL.Done b -> do
2180 let rst = RB.startOf rb
2181 jump c = SplitOnSuffixSeqKRInit 0 c s rb rst
2182 yieldProceed jump b
2183 Skip s -> go SPEC fs s rh cksum
2184 Stop ->
2185 if RB.unsafeEqArray rb rh patArr
2186 then return Stop
2187 else if withSep
2188 then do
2189 r <- done fs
2190 skip $ SplitOnSuffixSeqYield r SplitOnSuffixSeqDone
2191 else skip $ SplitOnSuffixSeqKRDone patLen fs rb rh
2193 stepOuter _ (SplitOnSuffixSeqKRCheck fs st rb rh) = do
2194 if RB.unsafeEqArray rb rh patArr
2195 then do
2196 r <- done fs
2197 let rst = RB.startOf rb
2198 jump c = SplitOnSuffixSeqKRInit 0 c st rb rst
2199 yieldProceed jump r
2200 else skip $ SplitOnSuffixSeqKRLoop fs st rb rh patHash
2202 stepOuter _ (SplitOnSuffixSeqKRDone 0 fs _ _) = do
2203 r <- done fs
2204 skip $ SplitOnSuffixSeqYield r SplitOnSuffixSeqDone
2205 stepOuter _ (SplitOnSuffixSeqKRDone n fs rb rh) = do
2206 old <- liftIO $ peek rh
2207 let rh1 = RB.advance rb rh
2208 r <- fstep fs old
2209 case r of
2210 FL.Partial fs1 -> skip $ SplitOnSuffixSeqKRDone (n - 1) fs1 rb rh1
2211 FL.Done b -> do
2212 let jump c = SplitOnSuffixSeqKRDone (n - 1) c rb rh1
2213 yieldProceed jump b
2215------------------------------------------------------------------------------
2216-- Nested Container Transformation
2217------------------------------------------------------------------------------
2219{-# ANN type SplitState Fuse #-}
2220data SplitState s arr
2221 = SplitInitial s
2222 | SplitBuffering s arr
2223 | SplitSplitting s arr
2224 | SplitYielding arr (SplitState s arr)
2225 | SplitFinishing
2227-- XXX An alternative approach would be to use a partial fold (Fold m a b) to
2228-- split using a splitBy like combinator. The Fold would consume upto the
2229-- separator and return any leftover which can then be fed to the next fold.
2231-- We can revisit this once we have partial folds/parsers.
2233-- | Performs infix separator style splitting.
2234{-# INLINE_NORMAL splitInnerBy #-}
2235splitInnerBy
2236 :: Monad m
2237 => (f a -> m (f a, Maybe (f a))) -- splitter
2238 -> (f a -> f a -> m (f a)) -- joiner
2239 -> Stream m (f a)
2240 -> Stream m (f a)
2241splitInnerBy splitter joiner (Stream step1 state1) =
2242 (Stream step (SplitInitial state1))
2244 where
2246 {-# INLINE_LATE step #-}
2247 step gst (SplitInitial st) = do
2248 r <- step1 gst st
2249 case r of
2250 Yield x s -> do
2251 (x1, mx2) <- splitter x
2252 return $ case mx2 of
2253 Nothing -> Skip (SplitBuffering s x1)
2254 Just x2 -> Skip (SplitYielding x1 (SplitSplitting s x2))
2255 Skip s -> return $ Skip (SplitInitial s)
2256 Stop -> return $ Stop
2258 step gst (SplitBuffering st buf) = do
2259 r <- step1 gst st
2260 case r of
2261 Yield x s -> do
2262 (x1, mx2) <- splitter x
2263 buf' <- joiner buf x1
2264 return $ case mx2 of
2265 Nothing -> Skip (SplitBuffering s buf')
2266 Just x2 -> Skip (SplitYielding buf' (SplitSplitting s x2))
2267 Skip s -> return $ Skip (SplitBuffering s buf)
2268 Stop -> return $ Skip (SplitYielding buf SplitFinishing)
2270 step _ (SplitSplitting st buf) = do
2271 (x1, mx2) <- splitter buf
2272 return $ case mx2 of
2273 Nothing -> Skip $ SplitBuffering st x1
2274 Just x2 -> Skip $ SplitYielding x1 (SplitSplitting st x2)
2276 step _ (SplitYielding x next) = return $ Yield x next
2277 step _ SplitFinishing = return $ Stop
2279-- | Performs infix separator style splitting.
2280{-# INLINE_NORMAL splitInnerBySuffix #-}
2281splitInnerBySuffix
2282 :: (Monad m, Eq (f a), Monoid (f a))
2283 => (f a -> m (f a, Maybe (f a))) -- splitter
2284 -> (f a -> f a -> m (f a)) -- joiner
2285 -> Stream m (f a)
2286 -> Stream m (f a)
2287splitInnerBySuffix splitter joiner (Stream step1 state1) =
2288 (Stream step (SplitInitial state1))
2290 where
2292 {-# INLINE_LATE step #-}
2293 step gst (SplitInitial st) = do
2294 r <- step1 gst st
2295 case r of
2296 Yield x s -> do
2297 (x1, mx2) <- splitter x
2298 return $ case mx2 of
2299 Nothing -> Skip (SplitBuffering s x1)
2300 Just x2 -> Skip (SplitYielding x1 (SplitSplitting s x2))
2301 Skip s -> return $ Skip (SplitInitial s)
2302 Stop -> return $ Stop
2304 step gst (SplitBuffering st buf) = do
2305 r <- step1 gst st
2306 case r of
2307 Yield x s -> do
2308 (x1, mx2) <- splitter x
2309 buf' <- joiner buf x1
2310 return $ case mx2 of
2311 Nothing -> Skip (SplitBuffering s buf')
2312 Just x2 -> Skip (SplitYielding buf' (SplitSplitting s x2))
2313 Skip s -> return $ Skip (SplitBuffering s buf)
2314 Stop -> return $
2315 if buf == mempty
2316 then Stop
2317 else Skip (SplitYielding buf SplitFinishing)
2319 step _ (SplitSplitting st buf) = do
2320 (x1, mx2) <- splitter buf
2321 return $ case mx2 of
2322 Nothing -> Skip $ SplitBuffering st x1
2323 Just x2 -> Skip $ SplitYielding x1 (SplitSplitting st x2)
2325 step _ (SplitYielding x next) = return $ Yield x next
2326 step _ SplitFinishing = return Stop