git.iko.soy/streamly
on bffdd45 help sign in
streamly/src/Streamly/Internal/Data/Stream/StreamD/Nesting.hs
2260 lines 84 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 -- concatUnfold.
39 AppendState(..)
40 , append
42 -- *** Interleaving
43 -- | Interleave elements from two streams alternately. A special case of
44 -- concatUnfoldInterleave.
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 concatUnfoldRoundrobin.
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 -- concatUnfold: 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 , concatUnfold
87 , ConcatUnfoldInterleaveState (..)
88 , concatUnfoldInterleave
89 , concatUnfoldRoundrobin
91 -- *** Interpose
92 -- | Like concatUnfold but intersperses an effect between the streams. A
93 -- special case of gintercalate.
94 , interpose
95 , interposeSuffix
97 -- *** Intercalate
98 -- | Like concatUnfold 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
116 -- ** Parsing
117 -- | Parsing is opposite to flattening. 'parseMany' is dual to concatMap or
118 -- concatUnfold. concatMap generates a stream from single values in a
119 -- stream and flattens, parseMany does the opposite of flattening by
120 -- splitting the stream and then folds each such split to single value in
121 -- the output stream.
122 , parseMany
123 , parseIterate
125 -- ** Grouping
126 -- | Group segments of a stream and fold. Special case of parsing.
127 , chunksOf
128 , groupsOf2
129 , groupsBy
130 , groupsRollingBy
132 -- ** Splitting
133 -- | A special case of parsing.
134 , wordsBy
135 , splitOnSeq
136 , splitOnSuffixSeq
138 -- * Transform (Nested Containers)
139 -- | Opposite to compact in ArrayStream
140 , splitInnerBy
141 , splitInnerBySuffix
142 )
143where
145#include "inline.hs"
147import Control.Exception (assert)
148import Control.Monad.Catch (MonadThrow, throwM)
149import Control.Monad.IO.Class (MonadIO(..))
150import Data.Bits (shiftR, shiftL, (.|.), (.&.))
151#if __GLASGOW_HASKELL__ >= 801
152import Data.Functor.Identity ( Identity )
153#endif
154import Data.Word (Word32)
155import Foreign.Storable (Storable(..))
156import Fusion.Plugin.Types (Fuse(..))
157import GHC.Types (SPEC(..))
159import Streamly.Internal.Data.Array.Foreign.Types (Array(..))
160import Streamly.Internal.Data.Fold.Types (Fold(..))
161import Streamly.Internal.Data.Parser (ParseError(..))
162import Streamly.Internal.Data.Unfold.Types (Unfold(..))
164import qualified Streamly.Internal.Data.Array.Foreign.Types as A
165import qualified Streamly.Internal.Data.Fold as FL
166import qualified Streamly.Internal.Data.Parser as PR
167import qualified Streamly.Internal.Data.Parser.ParserD as PRD
168import qualified Streamly.Internal.Ring.Foreign as RB
170import Streamly.Internal.Data.Stream.StreamD.Type
171import Streamly.Internal.Data.SVar
173import Prelude hiding (concatMap, mapM, zipWith)
175------------------------------------------------------------------------------
176-- Appending
177------------------------------------------------------------------------------
179data AppendState s1 s2 = AppendFirst s1 | AppendSecond s2
181-- Note that this could be much faster compared to the CPS stream. However, as
182-- the number of streams being composed increases this may become expensive.
183-- Need to see where the breaking point is between the two.
184--
185{-# INLINE_NORMAL append #-}
186append :: Monad m => Stream m a -> Stream m a -> Stream m a
187append (Stream step1 state1) (Stream step2 state2) =
188 Stream step (AppendFirst state1)
190 where
192 {-# INLINE_LATE step #-}
193 step gst (AppendFirst st) = do
194 r <- step1 gst st
195 return $ case r of
196 Yield a s -> Yield a (AppendFirst s)
197 Skip s -> Skip (AppendFirst s)
198 Stop -> Skip (AppendSecond state2)
200 step gst (AppendSecond st) = do
201 r <- step2 gst st
202 return $ case r of
203 Yield a s -> Yield a (AppendSecond s)
204 Skip s -> Skip (AppendSecond s)
205 Stop -> Stop
207------------------------------------------------------------------------------
208-- Interleaving
209------------------------------------------------------------------------------
211data InterleaveState s1 s2 = InterleaveFirst s1 s2 | InterleaveSecond s1 s2
212 | InterleaveSecondOnly s2 | InterleaveFirstOnly s1
214{-# INLINE_NORMAL interleave #-}
215interleave :: Monad m => Stream m a -> Stream m a -> Stream m a
216interleave (Stream step1 state1) (Stream step2 state2) =
217 Stream step (InterleaveFirst state1 state2)
219 where
221 {-# INLINE_LATE step #-}
222 step gst (InterleaveFirst st1 st2) = do
223 r <- step1 gst st1
224 return $ case r of
225 Yield a s -> Yield a (InterleaveSecond s st2)
226 Skip s -> Skip (InterleaveFirst s st2)
227 Stop -> Skip (InterleaveSecondOnly st2)
229 step gst (InterleaveSecond st1 st2) = do
230 r <- step2 gst st2
231 return $ case r of
232 Yield a s -> Yield a (InterleaveFirst st1 s)
233 Skip s -> Skip (InterleaveSecond st1 s)
234 Stop -> Skip (InterleaveFirstOnly st1)
236 step gst (InterleaveFirstOnly st1) = do
237 r <- step1 gst st1
238 return $ case r of
239 Yield a s -> Yield a (InterleaveFirstOnly s)
240 Skip s -> Skip (InterleaveFirstOnly s)
241 Stop -> Stop
243 step gst (InterleaveSecondOnly st2) = do
244 r <- step2 gst st2
245 return $ case r of
246 Yield a s -> Yield a (InterleaveSecondOnly s)
247 Skip s -> Skip (InterleaveSecondOnly s)
248 Stop -> Stop
250{-# INLINE_NORMAL interleaveMin #-}
251interleaveMin :: Monad m => Stream m a -> Stream m a -> Stream m a
252interleaveMin (Stream step1 state1) (Stream step2 state2) =
253 Stream step (InterleaveFirst state1 state2)
255 where
257 {-# INLINE_LATE step #-}
258 step gst (InterleaveFirst st1 st2) = do
259 r <- step1 gst st1
260 return $ case r of
261 Yield a s -> Yield a (InterleaveSecond s st2)
262 Skip s -> Skip (InterleaveFirst s st2)
263 Stop -> Stop
265 step gst (InterleaveSecond st1 st2) = do
266 r <- step2 gst st2
267 return $ case r of
268 Yield a s -> Yield a (InterleaveFirst st1 s)
269 Skip s -> Skip (InterleaveSecond st1 s)
270 Stop -> Stop
272 step _ (InterleaveFirstOnly _) = undefined
273 step _ (InterleaveSecondOnly _) = undefined
275{-# INLINE_NORMAL interleaveSuffix #-}
276interleaveSuffix :: Monad m => Stream m a -> Stream m a -> Stream m a
277interleaveSuffix (Stream step1 state1) (Stream step2 state2) =
278 Stream step (InterleaveFirst state1 state2)
280 where
282 {-# INLINE_LATE step #-}
283 step gst (InterleaveFirst st1 st2) = do
284 r <- step1 gst st1
285 return $ case r of
286 Yield a s -> Yield a (InterleaveSecond s st2)
287 Skip s -> Skip (InterleaveFirst s st2)
288 Stop -> Stop
290 step gst (InterleaveSecond st1 st2) = do
291 r <- step2 gst st2
292 return $ case r of
293 Yield a s -> Yield a (InterleaveFirst st1 s)
294 Skip s -> Skip (InterleaveSecond st1 s)
295 Stop -> Skip (InterleaveFirstOnly st1)
297 step gst (InterleaveFirstOnly st1) = do
298 r <- step1 gst st1
299 return $ case r of
300 Yield a s -> Yield a (InterleaveFirstOnly s)
301 Skip s -> Skip (InterleaveFirstOnly s)
302 Stop -> Stop
304 step _ (InterleaveSecondOnly _) = undefined
306data InterleaveInfixState s1 s2 a
307 = InterleaveInfixFirst s1 s2
308 | InterleaveInfixSecondBuf s1 s2
309 | InterleaveInfixSecondYield s1 s2 a
310 | InterleaveInfixFirstYield s1 s2 a
311 | InterleaveInfixFirstOnly s1
313{-# INLINE_NORMAL interleaveInfix #-}
314interleaveInfix :: Monad m => Stream m a -> Stream m a -> Stream m a
315interleaveInfix (Stream step1 state1) (Stream step2 state2) =
316 Stream step (InterleaveInfixFirst state1 state2)
318 where
320 {-# INLINE_LATE step #-}
321 step gst (InterleaveInfixFirst st1 st2) = do
322 r <- step1 gst st1
323 return $ case r of
324 Yield a s -> Yield a (InterleaveInfixSecondBuf s st2)
325 Skip s -> Skip (InterleaveInfixFirst s st2)
326 Stop -> Stop
328 step gst (InterleaveInfixSecondBuf st1 st2) = do
329 r <- step2 gst st2
330 return $ case r of
331 Yield a s -> Skip (InterleaveInfixSecondYield st1 s a)
332 Skip s -> Skip (InterleaveInfixSecondBuf st1 s)
333 Stop -> Skip (InterleaveInfixFirstOnly st1)
335 step gst (InterleaveInfixSecondYield st1 st2 x) = do
336 r <- step1 gst st1
337 return $ case r of
338 Yield a s -> Yield x (InterleaveInfixFirstYield s st2 a)
339 Skip s -> Skip (InterleaveInfixSecondYield s st2 x)
340 Stop -> Stop
342 step _ (InterleaveInfixFirstYield st1 st2 x) = do
343 return $ Yield x (InterleaveInfixSecondBuf st1 st2)
345 step gst (InterleaveInfixFirstOnly st1) = do
346 r <- step1 gst st1
347 return $ case r of
348 Yield a s -> Yield a (InterleaveInfixFirstOnly s)
349 Skip s -> Skip (InterleaveInfixFirstOnly s)
350 Stop -> Stop
352------------------------------------------------------------------------------
353-- Scheduling
354------------------------------------------------------------------------------
356{-# INLINE_NORMAL roundRobin #-}
357roundRobin :: Monad m => Stream m a -> Stream m a -> Stream m a
358roundRobin (Stream step1 state1) (Stream step2 state2) =
359 Stream step (InterleaveFirst state1 state2)
361 where
363 {-# INLINE_LATE step #-}
364 step gst (InterleaveFirst st1 st2) = do
365 r <- step1 gst st1
366 return $ case r of
367 Yield a s -> Yield a (InterleaveSecond s st2)
368 Skip s -> Skip (InterleaveSecond s st2)
369 Stop -> Skip (InterleaveSecondOnly st2)
371 step gst (InterleaveSecond st1 st2) = do
372 r <- step2 gst st2
373 return $ case r of
374 Yield a s -> Yield a (InterleaveFirst st1 s)
375 Skip s -> Skip (InterleaveFirst st1 s)
376 Stop -> Skip (InterleaveFirstOnly st1)
378 step gst (InterleaveSecondOnly st2) = do
379 r <- step2 gst st2
380 return $ case r of
381 Yield a s -> Yield a (InterleaveSecondOnly s)
382 Skip s -> Skip (InterleaveSecondOnly s)
383 Stop -> Stop
385 step gst (InterleaveFirstOnly st1) = do
386 r <- step1 gst st1
387 return $ case r of
388 Yield a s -> Yield a (InterleaveFirstOnly s)
389 Skip s -> Skip (InterleaveFirstOnly s)
390 Stop -> Stop
392------------------------------------------------------------------------------
393-- Zipping
394------------------------------------------------------------------------------
396{-# INLINE_NORMAL zipWithM #-}
397zipWithM :: Monad m
398 => (a -> b -> m c) -> Stream m a -> Stream m b -> Stream m c
399zipWithM f (Stream stepa ta) (Stream stepb tb) = Stream step (ta, tb, Nothing)
400 where
401 {-# INLINE_LATE step #-}
402 step gst (sa, sb, Nothing) = do
403 r <- stepa (adaptState gst) sa
404 return $
405 case r of
406 Yield x sa' -> Skip (sa', sb, Just x)
407 Skip sa' -> Skip (sa', sb, Nothing)
408 Stop -> Stop
410 step gst (sa, sb, Just x) = do
411 r <- stepb (adaptState gst) sb
412 case r of
413 Yield y sb' -> do
414 z <- f x y
415 return $ Yield z (sa, sb', Nothing)
416 Skip sb' -> return $ Skip (sa, sb', Just x)
417 Stop -> return Stop
419#if __GLASGOW_HASKELL__ >= 801
420{-# RULES "zipWithM xs xs"
421 forall f xs. zipWithM @Identity f xs xs = mapM (\x -> f x x) xs #-}
422#endif
424{-# INLINE zipWith #-}
425zipWith :: Monad m => (a -> b -> c) -> Stream m a -> Stream m b -> Stream m c
426zipWith f = zipWithM (\a b -> return (f a b))
428------------------------------------------------------------------------------
429-- Merging
430------------------------------------------------------------------------------
432{-# INLINE_NORMAL mergeByM #-}
433mergeByM
434 :: (Monad m)
435 => (a -> a -> m Ordering) -> Stream m a -> Stream m a -> Stream m a
436mergeByM cmp (Stream stepa ta) (Stream stepb tb) =
437 Stream step (Just ta, Just tb, Nothing, Nothing)
438 where
439 {-# INLINE_LATE step #-}
441 -- one of the values is missing, and the corresponding stream is running
442 step gst (Just sa, sb, Nothing, b) = do
443 r <- stepa gst sa
444 return $ case r of
445 Yield a sa' -> Skip (Just sa', sb, Just a, b)
446 Skip sa' -> Skip (Just sa', sb, Nothing, b)
447 Stop -> Skip (Nothing, sb, Nothing, b)
449 step gst (sa, Just sb, a, Nothing) = do
450 r <- stepb gst sb
451 return $ case r of
452 Yield b sb' -> Skip (sa, Just sb', a, Just b)
453 Skip sb' -> Skip (sa, Just sb', a, Nothing)
454 Stop -> Skip (sa, Nothing, a, Nothing)
456 -- both the values are available
457 step _ (sa, sb, Just a, Just b) = do
458 res <- cmp a b
459 return $ case res of
460 GT -> Yield b (sa, sb, Just a, Nothing)
461 _ -> Yield a (sa, sb, Nothing, Just b)
463 -- one of the values is missing, corresponding stream is done
464 step _ (Nothing, sb, Nothing, Just b) =
465 return $ Yield b (Nothing, sb, Nothing, Nothing)
467 step _ (sa, Nothing, Just a, Nothing) =
468 return $ Yield a (sa, Nothing, Nothing, Nothing)
470 step _ (Nothing, Nothing, Nothing, Nothing) = return Stop
472{-# INLINE mergeBy #-}
473mergeBy
474 :: (Monad m)
475 => (a -> a -> Ordering) -> Stream m a -> Stream m a -> Stream m a
476mergeBy cmp = mergeByM (\a b -> return $ cmp a b)
478------------------------------------------------------------------------------
479-- Combine N Streams - concatUnfold
480------------------------------------------------------------------------------
482data ConcatUnfoldInterleaveState o i =
483 ConcatUnfoldInterleaveOuter o [i]
484 | ConcatUnfoldInterleaveInner o [i]
485 | ConcatUnfoldInterleaveInnerL [i] [i]
486 | ConcatUnfoldInterleaveInnerR [i] [i]
488-- XXX use arrays to store state instead of lists.
489-- XXX In general we can use different scheduling strategies e.g. how to
490-- schedule the outer vs inner loop or assigning weights to different streams
491-- or outer and inner loops.
493-- After a yield, switch to the next stream. Do not switch streams on Skip.
494-- Yield from outer stream switches to the inner stream.
495--
496-- There are two choices here, (1) exhaust the outer stream first and then
497-- start yielding from the inner streams, this is much simpler to implement,
498-- (2) yield at least one element from an inner stream before going back to
499-- outer stream and opening the next stream from it.
500--
501-- Ideally, we need some scheduling bias to inner streams vs outer stream.
502-- Maybe we can configure the behavior.
503--
504{-# INLINE_NORMAL concatUnfoldInterleave #-}
505concatUnfoldInterleave :: Monad m => Unfold m a b -> Stream m a -> Stream m b
506concatUnfoldInterleave (Unfold istep inject) (Stream ostep ost) =
507 Stream step (ConcatUnfoldInterleaveOuter ost [])
508 where
509 {-# INLINE_LATE step #-}
510 step gst (ConcatUnfoldInterleaveOuter o ls) = do
511 r <- ostep (adaptState gst) o
512 case r of
513 Yield a o' -> do
514 i <- inject a
515 i `seq` return (Skip (ConcatUnfoldInterleaveInner o' (i : ls)))
516 Skip o' -> return $ Skip (ConcatUnfoldInterleaveOuter o' ls)
517 Stop -> return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
519 step _ (ConcatUnfoldInterleaveInner _ []) = undefined
520 step _ (ConcatUnfoldInterleaveInner o (st:ls)) = do
521 r <- istep st
522 return $ case r of
523 Yield x s -> Yield x (ConcatUnfoldInterleaveOuter o (s:ls))
524 Skip s -> Skip (ConcatUnfoldInterleaveInner o (s:ls))
525 Stop -> Skip (ConcatUnfoldInterleaveOuter o ls)
527 step _ (ConcatUnfoldInterleaveInnerL [] []) = return Stop
528 step _ (ConcatUnfoldInterleaveInnerL [] rs) =
529 return $ Skip (ConcatUnfoldInterleaveInnerR [] rs)
531 step _ (ConcatUnfoldInterleaveInnerL (st:ls) rs) = do
532 r <- istep st
533 return $ case r of
534 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerL ls (s:rs))
535 Skip s -> Skip (ConcatUnfoldInterleaveInnerL (s:ls) rs)
536 Stop -> Skip (ConcatUnfoldInterleaveInnerL ls rs)
538 step _ (ConcatUnfoldInterleaveInnerR [] []) = return Stop
539 step _ (ConcatUnfoldInterleaveInnerR ls []) =
540 return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
542 step _ (ConcatUnfoldInterleaveInnerR ls (st:rs)) = do
543 r <- istep st
544 return $ case r of
545 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerR (s:ls) rs)
546 Skip s -> Skip (ConcatUnfoldInterleaveInnerR ls (s:rs))
547 Stop -> Skip (ConcatUnfoldInterleaveInnerR ls rs)
549-- XXX In general we can use different scheduling strategies e.g. how to
550-- schedule the outer vs inner loop or assigning weights to different streams
551-- or outer and inner loops.
552--
553-- This could be inefficient if the tasks are too small.
554--
555-- Compared to concatUnfoldInterleave this one switches streams on Skips.
556--
557{-# INLINE_NORMAL concatUnfoldRoundrobin #-}
558concatUnfoldRoundrobin :: Monad m => Unfold m a b -> Stream m a -> Stream m b
559concatUnfoldRoundrobin (Unfold istep inject) (Stream ostep ost) =
560 Stream step (ConcatUnfoldInterleaveOuter ost [])
561 where
562 {-# INLINE_LATE step #-}
563 step gst (ConcatUnfoldInterleaveOuter o ls) = do
564 r <- ostep (adaptState gst) o
565 case r of
566 Yield a o' -> do
567 i <- inject a
568 i `seq` return (Skip (ConcatUnfoldInterleaveInner o' (i : ls)))
569 Skip o' -> return $ Skip (ConcatUnfoldInterleaveInner o' ls)
570 Stop -> return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
572 step _ (ConcatUnfoldInterleaveInner o []) =
573 return $ Skip (ConcatUnfoldInterleaveOuter o [])
575 step _ (ConcatUnfoldInterleaveInner o (st:ls)) = do
576 r <- istep st
577 return $ case r of
578 Yield x s -> Yield x (ConcatUnfoldInterleaveOuter o (s:ls))
579 Skip s -> Skip (ConcatUnfoldInterleaveOuter o (s:ls))
580 Stop -> Skip (ConcatUnfoldInterleaveOuter o ls)
582 step _ (ConcatUnfoldInterleaveInnerL [] []) = return Stop
583 step _ (ConcatUnfoldInterleaveInnerL [] rs) =
584 return $ Skip (ConcatUnfoldInterleaveInnerR [] rs)
586 step _ (ConcatUnfoldInterleaveInnerL (st:ls) rs) = do
587 r <- istep st
588 return $ case r of
589 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerL ls (s:rs))
590 Skip s -> Skip (ConcatUnfoldInterleaveInnerL ls (s:rs))
591 Stop -> Skip (ConcatUnfoldInterleaveInnerL ls rs)
593 step _ (ConcatUnfoldInterleaveInnerR [] []) = return Stop
594 step _ (ConcatUnfoldInterleaveInnerR ls []) =
595 return $ Skip (ConcatUnfoldInterleaveInnerL ls [])
597 step _ (ConcatUnfoldInterleaveInnerR ls (st:rs)) = do
598 r <- istep st
599 return $ case r of
600 Yield x s -> Yield x (ConcatUnfoldInterleaveInnerR (s:ls) rs)
601 Skip s -> Skip (ConcatUnfoldInterleaveInnerR (s:ls) rs)
602 Stop -> Skip (ConcatUnfoldInterleaveInnerR ls rs)
604------------------------------------------------------------------------------
605-- Combine N Streams - interpose
606------------------------------------------------------------------------------
608{-# ANN type InterposeSuffixState Fuse #-}
609data InterposeSuffixState s1 i1 =
610 InterposeSuffixFirst s1
611 -- | InterposeSuffixFirstYield s1 i1
612 | InterposeSuffixFirstInner s1 i1
613 | InterposeSuffixSecond s1
615-- Note that if an unfolded layer turns out to be nil we still emit the
616-- separator effect. An alternate behavior could be to emit the separator
617-- effect only if at least one element has been yielded by the unfolding.
618-- However, that becomes a bit complicated, so we have chosen the former
619-- behvaior for now.
620{-# INLINE_NORMAL interposeSuffix #-}
621interposeSuffix
622 :: Monad m
623 => m c -> Unfold m b c -> Stream m b -> Stream m c
624interposeSuffix
625 action
626 (Unfold istep1 inject1) (Stream step1 state1) =
627 Stream step (InterposeSuffixFirst state1)
629 where
631 {-# INLINE_LATE step #-}
632 step gst (InterposeSuffixFirst s1) = do
633 r <- step1 (adaptState gst) s1
634 case r of
635 Yield a s -> do
636 i <- inject1 a
637 i `seq` return (Skip (InterposeSuffixFirstInner s i))
638 -- i `seq` return (Skip (InterposeSuffixFirstYield s i))
639 Skip s -> return $ Skip (InterposeSuffixFirst s)
640 Stop -> return Stop
642 {-
643 step _ (InterposeSuffixFirstYield s1 i1) = do
644 r <- istep1 i1
645 return $ case r of
646 Yield x i' -> Yield x (InterposeSuffixFirstInner s1 i')
647 Skip i' -> Skip (InterposeSuffixFirstYield s1 i')
648 Stop -> Skip (InterposeSuffixFirst s1)
649 -}
651 step _ (InterposeSuffixFirstInner s1 i1) = do
652 r <- istep1 i1
653 return $ case r of
654 Yield x i' -> Yield x (InterposeSuffixFirstInner s1 i')
655 Skip i' -> Skip (InterposeSuffixFirstInner s1 i')
656 Stop -> Skip (InterposeSuffixSecond s1)
658 step _ (InterposeSuffixSecond s1) = do
659 r <- action
660 return $ Yield r (InterposeSuffixFirst s1)
662{-# ANN type InterposeState Fuse #-}
663data InterposeState s1 i1 a =
664 InterposeFirst s1
665 -- | InterposeFirstYield s1 i1
666 | InterposeFirstInner s1 i1
667 | InterposeFirstInject s1
668 -- | InterposeFirstBuf s1 i1
669 | InterposeSecondYield s1 i1
670 -- -- | InterposeSecondYield s1 i1 a
671 -- -- | InterposeFirstResume s1 i1 a
673-- Note that this only interposes the pure values, we may run many effects to
674-- generate those values as some effects may not generate anything (Skip).
675{-# INLINE_NORMAL interpose #-}
676interpose :: Monad m => m c -> Unfold m b c -> Stream m b -> Stream m c
677interpose
678 action
679 (Unfold istep1 inject1) (Stream step1 state1) =
680 Stream step (InterposeFirst state1)
682 where
684 {-# INLINE_LATE step #-}
685 step gst (InterposeFirst s1) = do
686 r <- step1 (adaptState gst) s1
687 case r of
688 Yield a s -> do
689 i <- inject1 a
690 i `seq` return (Skip (InterposeFirstInner s i))
691 -- i `seq` return (Skip (InterposeFirstYield s i))
692 Skip s -> return $ Skip (InterposeFirst s)
693 Stop -> return Stop
695 {-
696 step _ (InterposeFirstYield s1 i1) = do
697 r <- istep1 i1
698 return $ case r of
699 Yield x i' -> Yield x (InterposeFirstInner s1 i')
700 Skip i' -> Skip (InterposeFirstYield s1 i')
701 Stop -> Skip (InterposeFirst s1)
702 -}
704 step _ (InterposeFirstInner s1 i1) = do
705 r <- istep1 i1
706 return $ case r of
707 Yield x i' -> Yield x (InterposeFirstInner s1 i')
708 Skip i' -> Skip (InterposeFirstInner s1 i')
709 Stop -> Skip (InterposeFirstInject s1)
711 step gst (InterposeFirstInject s1) = do
712 r <- step1 (adaptState gst) s1
713 case r of
714 Yield a s -> do
715 i <- inject1 a
716 -- i `seq` return (Skip (InterposeFirstBuf s i))
717 i `seq` return (Skip (InterposeSecondYield s i))
718 Skip s -> return $ Skip (InterposeFirstInject s)
719 Stop -> return Stop
721 {-
722 step _ (InterposeFirstBuf s1 i1) = do
723 r <- istep1 i1
724 return $ case r of
725 Yield x i' -> Skip (InterposeSecondYield s1 i' x)
726 Skip i' -> Skip (InterposeFirstBuf s1 i')
727 Stop -> Stop
728 -}
730 {-
731 step _ (InterposeSecondYield s1 i1 v) = do
732 r <- action
733 return $ Yield r (InterposeFirstResume s1 i1 v)
734 -}
735 step _ (InterposeSecondYield s1 i1) = do
736 r <- action
737 return $ Yield r (InterposeFirstInner s1 i1)
739 {-
740 step _ (InterposeFirstResume s1 i1 v) = do
741 return $ Yield v (InterposeFirstInner s1 i1)
742 -}
744------------------------------------------------------------------------------
745-- Combine N Streams - intercalate
746------------------------------------------------------------------------------
748data ICUState s1 s2 i1 i2 =
749 ICUFirst s1 s2
750 | ICUSecond s1 s2
751 | ICUSecondOnly s2
752 | ICUFirstOnly s1
753 | ICUFirstInner s1 s2 i1
754 | ICUSecondInner s1 s2 i2
755 | ICUFirstOnlyInner s1 i1
756 | ICUSecondOnlyInner s2 i2
758-- | Interleave streams (full streams, not the elements) unfolded from two
759-- input streams and concat. Stop when the first stream stops. If the second
760-- stream ends before the first one then first stream still keeps running alone
761-- without any interleaving with the second stream.
762--
763-- [a1, a2, ... an] [b1, b2 ...]
764-- => [streamA1, streamA2, ... streamAn] [streamB1, streamB2, ...]
765-- => [streamA1, streamB1, streamA2...StreamAn, streamBn]
766-- => [a11, a12, ...a1j, b11, b12, ...b1k, a21, a22, ...]
767--
768{-# INLINE_NORMAL gintercalateSuffix #-}
769gintercalateSuffix
770 :: Monad m
771 => Unfold m a c -> Stream m a -> Unfold m b c -> Stream m b -> Stream m c
772gintercalateSuffix
773 (Unfold istep1 inject1) (Stream step1 state1)
774 (Unfold istep2 inject2) (Stream step2 state2) =
775 Stream step (ICUFirst state1 state2)
777 where
779 {-# INLINE_LATE step #-}
780 step gst (ICUFirst s1 s2) = do
781 r <- step1 (adaptState gst) s1
782 case r of
783 Yield a s -> do
784 i <- inject1 a
785 i `seq` return (Skip (ICUFirstInner s s2 i))
786 Skip s -> return $ Skip (ICUFirst s s2)
787 Stop -> return Stop
789 step gst (ICUFirstOnly s1) = do
790 r <- step1 (adaptState gst) s1
791 case r of
792 Yield a s -> do
793 i <- inject1 a
794 i `seq` return (Skip (ICUFirstOnlyInner s i))
795 Skip s -> return $ Skip (ICUFirstOnly s)
796 Stop -> return Stop
798 step _ (ICUFirstInner s1 s2 i1) = do
799 r <- istep1 i1
800 return $ case r of
801 Yield x i' -> Yield x (ICUFirstInner s1 s2 i')
802 Skip i' -> Skip (ICUFirstInner s1 s2 i')
803 Stop -> Skip (ICUSecond s1 s2)
805 step _ (ICUFirstOnlyInner s1 i1) = do
806 r <- istep1 i1
807 return $ case r of
808 Yield x i' -> Yield x (ICUFirstOnlyInner s1 i')
809 Skip i' -> Skip (ICUFirstOnlyInner s1 i')
810 Stop -> Skip (ICUFirstOnly s1)
812 step gst (ICUSecond s1 s2) = do
813 r <- step2 (adaptState gst) s2
814 case r of
815 Yield a s -> do
816 i <- inject2 a
817 i `seq` return (Skip (ICUSecondInner s1 s i))
818 Skip s -> return $ Skip (ICUSecond s1 s)
819 Stop -> return $ Skip (ICUFirstOnly s1)
821 step _ (ICUSecondInner s1 s2 i2) = do
822 r <- istep2 i2
823 return $ case r of
824 Yield x i' -> Yield x (ICUSecondInner s1 s2 i')
825 Skip i' -> Skip (ICUSecondInner s1 s2 i')
826 Stop -> Skip (ICUFirst s1 s2)
828 step _ (ICUSecondOnly _s2) = undefined
829 step _ (ICUSecondOnlyInner _s2 _i2) = undefined
831data ICALState s1 s2 i1 i2 a =
832 ICALFirst s1 s2
833 -- | ICALFirstYield s1 s2 i1
834 | ICALFirstInner s1 s2 i1
835 | ICALFirstOnly s1
836 | ICALFirstOnlyInner s1 i1
837 | ICALSecondInject s1 s2
838 | ICALFirstInject s1 s2 i2
839 -- | ICALFirstBuf s1 s2 i1 i2
840 | ICALSecondInner s1 s2 i1 i2
841 -- -- | ICALSecondInner s1 s2 i1 i2 a
842 -- -- | ICALFirstResume s1 s2 i1 i2 a
844-- | Interleave streams (full streams, not the elements) unfolded from two
845-- input streams and concat. Stop when the first stream stops. If the second
846-- stream ends before the first one then first stream still keeps running alone
847-- without any interleaving with the second stream.
848--
849-- [a1, a2, ... an] [b1, b2 ...]
850-- => [streamA1, streamA2, ... streamAn] [streamB1, streamB2, ...]
851-- => [streamA1, streamB1, streamA2...StreamAn, streamBn]
852-- => [a11, a12, ...a1j, b11, b12, ...b1k, a21, a22, ...]
853--
854{-# INLINE_NORMAL gintercalate #-}
855gintercalate
856 :: Monad m
857 => Unfold m a c -> Stream m a -> Unfold m b c -> Stream m b -> Stream m c
858gintercalate
859 (Unfold istep1 inject1) (Stream step1 state1)
860 (Unfold istep2 inject2) (Stream step2 state2) =
861 Stream step (ICALFirst state1 state2)
863 where
865 {-# INLINE_LATE step #-}
866 step gst (ICALFirst s1 s2) = do
867 r <- step1 (adaptState gst) s1
868 case r of
869 Yield a s -> do
870 i <- inject1 a
871 i `seq` return (Skip (ICALFirstInner s s2 i))
872 -- i `seq` return (Skip (ICALFirstYield s s2 i))
873 Skip s -> return $ Skip (ICALFirst s s2)
874 Stop -> return Stop
876 {-
877 step _ (ICALFirstYield s1 s2 i1) = do
878 r <- istep1 i1
879 return $ case r of
880 Yield x i' -> Yield x (ICALFirstInner s1 s2 i')
881 Skip i' -> Skip (ICALFirstYield s1 s2 i')
882 Stop -> Skip (ICALFirst s1 s2)
883 -}
885 step _ (ICALFirstInner s1 s2 i1) = do
886 r <- istep1 i1
887 return $ case r of
888 Yield x i' -> Yield x (ICALFirstInner s1 s2 i')
889 Skip i' -> Skip (ICALFirstInner s1 s2 i')
890 Stop -> Skip (ICALSecondInject s1 s2)
892 step gst (ICALFirstOnly s1) = do
893 r <- step1 (adaptState gst) s1
894 case r of
895 Yield a s -> do
896 i <- inject1 a
897 i `seq` return (Skip (ICALFirstOnlyInner s i))
898 Skip s -> return $ Skip (ICALFirstOnly s)
899 Stop -> return Stop
901 step _ (ICALFirstOnlyInner s1 i1) = do
902 r <- istep1 i1
903 return $ case r of
904 Yield x i' -> Yield x (ICALFirstOnlyInner s1 i')
905 Skip i' -> Skip (ICALFirstOnlyInner s1 i')
906 Stop -> Skip (ICALFirstOnly s1)
908 -- We inject the second stream even before checking if the first stream
909 -- would yield any more elements. There is no clear choice whether we
910 -- should do this before or after that. Doing it after may make the state
911 -- machine a bit simpler though.
912 step gst (ICALSecondInject s1 s2) = do
913 r <- step2 (adaptState gst) s2
914 case r of
915 Yield a s -> do
916 i <- inject2 a
917 i `seq` return (Skip (ICALFirstInject s1 s i))
918 Skip s -> return $ Skip (ICALSecondInject s1 s)
919 Stop -> return $ Skip (ICALFirstOnly s1)
921 step gst (ICALFirstInject s1 s2 i2) = do
922 r <- step1 (adaptState gst) s1
923 case r of
924 Yield a s -> do
925 i <- inject1 a
926 i `seq` return (Skip (ICALSecondInner s s2 i i2))
927 -- i `seq` return (Skip (ICALFirstBuf s s2 i i2))
928 Skip s -> return $ Skip (ICALFirstInject s s2 i2)
929 Stop -> return Stop
931 {-
932 step _ (ICALFirstBuf s1 s2 i1 i2) = do
933 r <- istep1 i1
934 return $ case r of
935 Yield x i' -> Skip (ICALSecondInner s1 s2 i' i2 x)
936 Skip i' -> Skip (ICALFirstBuf s1 s2 i' i2)
937 Stop -> Stop
939 step _ (ICALSecondInner s1 s2 i1 i2 v) = do
940 r <- istep2 i2
941 return $ case r of
942 Yield x i' -> Yield x (ICALSecondInner s1 s2 i1 i' v)
943 Skip i' -> Skip (ICALSecondInner s1 s2 i1 i' v)
944 Stop -> Skip (ICALFirstResume s1 s2 i1 i2 v)
945 -}
947 step _ (ICALSecondInner s1 s2 i1 i2) = do
948 r <- istep2 i2
949 return $ case r of
950 Yield x i' -> Yield x (ICALSecondInner s1 s2 i1 i')
951 Skip i' -> Skip (ICALSecondInner s1 s2 i1 i')
952 Stop -> Skip (ICALFirstInner s1 s2 i1)
953 -- Stop -> Skip (ICALFirstResume s1 s2 i1 i2)
955 {-
956 step _ (ICALFirstResume s1 s2 i1 i2 x) = do
957 return $ Yield x (ICALFirstInner s1 s2 i1 i2)
958 -}
960------------------------------------------------------------------------------
961-- Parsing
962------------------------------------------------------------------------------
964{-# ANN type ParseChunksState Fuse #-}
965data ParseChunksState x inpBuf st pst =
966 ParseChunksInit inpBuf st
967 | ParseChunksInitLeftOver inpBuf
968 | ParseChunksStream st inpBuf !pst
969 | ParseChunksBuf inpBuf st inpBuf !pst
970 | ParseChunksYield x (ParseChunksState x inpBuf st pst)
972{-# INLINE_NORMAL parseMany #-}
973parseMany
974 :: MonadThrow m
975 => PRD.Parser m a b
976 -> Stream m a
977 -> Stream m b
978parseMany (PRD.Parser pstep initial extract) (Stream step state) =
979 Stream stepOuter (ParseChunksInit [] state)
981 where
983 {-# INLINE_LATE stepOuter #-}
984 -- Buffer is empty, get the first element from the stream, initialize the
985 -- fold and then go to stream processing loop.
986 stepOuter gst (ParseChunksInit [] st) = do
987 r <- step (adaptState gst) st
988 case r of
989 Yield x s -> initial >>= return . Skip . ParseChunksBuf [x] s []
990 Skip s -> return $ Skip $ ParseChunksInit [] s
991 Stop -> return Stop
993 -- Buffer is not empty, go to buffered processing loop
994 stepOuter _ (ParseChunksInit src st) = do
995 initial >>= return . Skip . ParseChunksBuf src st []
997 -- XXX we just discard any leftover input at the end
998 stepOuter _ (ParseChunksInitLeftOver _) = return Stop
1000 -- Buffer is empty, process elements from the stream
1001 stepOuter gst (ParseChunksStream st buf pst) = do
1002 r <- step (adaptState gst) st
1003 case r of
1004 Yield x s -> do
1005 pRes <- pstep pst x
1006 case pRes of
1007 PR.Partial 0 pst1 ->
1008 return $ Skip $ ParseChunksStream s [] pst1
1009 PR.Partial n pst1 -> do
1010 assert (n <= length (x:buf)) (return ())
1011 let src0 = Prelude.take n (x:buf)
1012 src = Prelude.reverse src0
1013 return $ Skip $ ParseChunksBuf src s [] pst1
1014 PR.Continue 0 pst1 ->
1015 return $ Skip $ ParseChunksStream s (x:buf) pst1
1016 PR.Continue n pst1 -> do
1017 assert (n <= length (x:buf)) (return ())
1018 let (src0, buf1) = splitAt n (x:buf)
1019 src = Prelude.reverse src0
1020 return $ Skip $ ParseChunksBuf src s buf1 pst1
1021 PR.Done 0 b -> do
1022 return $ Skip $
1023 ParseChunksYield b (ParseChunksInit [] s)
1024 PR.Done n b -> do
1025 assert (n <= length (x:buf)) (return ())
1026 let src = Prelude.reverse (Prelude.take n (x:buf))
1027 return $ Skip $
1028 ParseChunksYield b (ParseChunksInit src s)
1029 PR.Error err -> throwM $ ParseError err
1030 Skip s -> return $ Skip $ ParseChunksStream s buf pst
1031 Stop -> do
1032 b <- extract pst
1033 let src = Prelude.reverse buf
1034 return $ Skip $
1035 ParseChunksYield b (ParseChunksInitLeftOver src)
1037 -- go back to stream processing mode
1038 stepOuter _ (ParseChunksBuf [] s buf pst) =
1039 return $ Skip $ ParseChunksStream s buf pst
1041 -- buffered processing loop
1042 stepOuter _ (ParseChunksBuf (x:xs) s buf pst) = do
1043 pRes <- pstep pst x
1044 case pRes of
1045 PR.Partial 0 pst1 ->
1046 return $ Skip $ ParseChunksBuf xs s [] pst1
1047 PR.Partial n pst1 -> do
1048 assert (n <= length (x:buf)) (return ())
1049 let src0 = Prelude.take n (x:buf)
1050 src = Prelude.reverse src0 ++ xs
1051 return $ Skip $ ParseChunksBuf src s [] pst1
1052 PR.Continue 0 pst1 ->
1053 return $ Skip $ ParseChunksBuf xs s (x:buf) pst1
1054 PR.Continue n pst1 -> do
1055 assert (n <= length (x:buf)) (return ())
1056 let (src0, buf1) = splitAt n (x:buf)
1057 src = Prelude.reverse src0 ++ xs
1058 return $ Skip $ ParseChunksBuf src s buf1 pst1
1059 PR.Done 0 b ->
1060 return $ Skip $ ParseChunksYield b (ParseChunksInit xs s)
1061 PR.Done n b -> do
1062 assert (n <= length (x:buf)) (return ())
1063 let src = Prelude.reverse (Prelude.take n (x:buf)) ++ xs
1064 return $ Skip $ ParseChunksYield b (ParseChunksInit src s)
1065 PR.Error err -> throwM $ ParseError err
1067 stepOuter _ (ParseChunksYield a next) = return $ Yield a next
1069{-# ANN type ConcatParseState Fuse #-}
1070data ConcatParseState x inpBuf st p =
1071 ConcatParseInit inpBuf st p
1072 | ConcatParseInitLeftOver inpBuf
1073 | ConcatParseStream st inpBuf p
1074 | ConcatParseBuf inpBuf st inpBuf p
1075 | ConcatParseYield x (ConcatParseState x inpBuf st p)
1077{-# INLINE_NORMAL parseIterate #-}
1078parseIterate
1079 :: MonadThrow m
1080 => (b -> PRD.Parser m a b)
1081 -> b
1082 -> Stream m a
1083 -> Stream m b
1084parseIterate func seed (Stream step state) =
1085 Stream stepOuter (ConcatParseInit [] state (func seed))
1087 where
1089 {-# INLINE_LATE stepOuter #-}
1090 -- Buffer is empty, go to stream processing loop
1091 stepOuter _ (ConcatParseInit [] st (PRD.Parser pstep initial extract)) = do
1092 initial >>= \r -> return $ Skip $ ConcatParseStream st []
1093 (PRD.Parser pstep (return r) extract)
1095 -- Buffer is not empty, go to buffered processing loop
1096 stepOuter _ (ConcatParseInit src st
1097 (PRD.Parser pstep initial extract)) = do
1098 initial >>= \r -> return $ Skip $ ConcatParseBuf src st []
1099 (PRD.Parser pstep (return r) extract)
1101 -- XXX we just discard any leftover input at the end
1102 stepOuter _ (ConcatParseInitLeftOver _) = return Stop
1104 -- Buffer is empty process elements from the stream
1105 stepOuter gst (ConcatParseStream st buf
1106 p@(PRD.Parser pstep initial extract)) = do
1107 r <- step (adaptState gst) st
1108 case r of
1109 Yield x s -> do
1110 pst <- initial
1111 pRes <- pstep pst x
1112 case pRes of
1113 PR.Partial 0 pst1 ->
1114 return $ Skip $ ConcatParseStream s []
1115 (PRD.Parser pstep (return pst1) extract)
1116 PR.Partial n pst1 -> do
1117 assert (n <= length (x:buf)) (return ())
1118 let src0 = Prelude.take n (x:buf)
1119 src = Prelude.reverse src0
1120 return $ Skip $ ConcatParseBuf src s []
1121 (PRD.Parser pstep (return pst1) extract)
1122 -- PR.Continue 0 pst1 ->
1123 -- return $ Skip $ ConcatParseStream s (x:buf) pst1
1124 PR.Continue n pst1 -> do
1125 assert (n <= length (x:buf)) (return ())
1126 let (src0, buf1) = splitAt n (x:buf)
1127 src = Prelude.reverse src0
1128 return $ Skip $ ConcatParseBuf src s buf1
1129 (PRD.Parser pstep (return pst1) extract)
1130 -- XXX Specialize for Stop 0 common case?
1131 PR.Done n b -> do
1132 assert (n <= length (x:buf)) (return ())
1133 let src = Prelude.reverse (Prelude.take n (x:buf))
1134 return $ Skip $
1135 ConcatParseYield b (ConcatParseInit src s (func b))
1136 PR.Error err -> throwM $ ParseError err
1137 Skip s -> return $ Skip $ ConcatParseStream s buf p
1138 Stop -> do
1139 pst <- initial
1140 b <- extract pst
1141 let src = Prelude.reverse buf
1142 return $ Skip $
1143 ConcatParseYield b (ConcatParseInitLeftOver src)
1145 -- go back to stream processing mode
1146 stepOuter _ (ConcatParseBuf [] s buf p) =
1147 return $ Skip $ ConcatParseStream s buf p
1149 -- buffered processing loop
1150 stepOuter _ (ConcatParseBuf (x:xs) s buf
1151 (PRD.Parser pstep initial extract)) = do
1152 pst <- initial
1153 pRes <- pstep pst x
1154 case pRes of
1155 PR.Partial 0 pst1 ->
1156 return $ Skip $ ConcatParseBuf xs s []
1157 (PRD.Parser pstep (return pst1) extract)
1158 PR.Partial n pst1 -> do
1159 assert (n <= length (x:buf)) (return ())
1160 let src0 = Prelude.take n (x:buf)
1161 src = Prelude.reverse src0 ++ xs
1162 return $ Skip $ ConcatParseBuf src s []
1163 (PRD.Parser pstep (return pst1) extract)
1164 -- PR.Continue 0 pst1 -> return $ Skip $ ConcatParseBuf xs s (x:buf) pst1
1165 PR.Continue n pst1 -> do
1166 assert (n <= length (x:buf)) (return ())
1167 let (src0, buf1) = splitAt n (x:buf)
1168 src = Prelude.reverse src0 ++ xs
1169 return $ Skip $ ConcatParseBuf src s buf1
1170 (PRD.Parser pstep (return pst1) extract)
1171 -- XXX Specialize for Stop 0 common case?
1172 PR.Done n b -> do
1173 assert (n <= length (x:buf)) (return ())
1174 let src = Prelude.reverse (Prelude.take n (x:buf)) ++ xs
1175 return $ Skip $ ConcatParseYield b
1176 (ConcatParseInit src s (func b))
1177 PR.Error err -> throwM $ ParseError err
1179 stepOuter _ (ConcatParseYield a next) = return $ Yield a next
1181------------------------------------------------------------------------------
1182-- Grouping
1183------------------------------------------------------------------------------
1185data GroupByState st fs a b
1186 = GroupingInit st
1187 | GroupingDo st !fs
1188 | GroupingInitWith st !a
1189 | GroupingDoWith st !fs !a
1190 | GroupingYield !b (GroupByState st fs a b)
1191 | GroupingDone
1193{-# INLINE_NORMAL groupsBy #-}
1194groupsBy :: Monad m
1195 => (a -> a -> Bool)
1196 -> Fold m a b
1197 -> Stream m a
1198 -> Stream m b
1199groupsBy cmp (Fold fstep initial done) (Stream step state) =
1200 Stream stepOuter (GroupingInit state)
1202 where
1204 {-# INLINE_LATE stepOuter #-}
1205 stepOuter _ (GroupingInit st) = do
1206 -- XXX Note that if the stream stops without yielding a single element
1207 -- in the group we discard the "initial" effect.
1208 res <- initial
1209 return
1210 $ case res of
1211 FL.Partial s -> Skip $ GroupingDo st s
1212 FL.Done b -> Yield b $ GroupingInit st
1213 stepOuter gst (GroupingDo st fs) = do
1214 res <- step (adaptState gst) st
1215 case res of
1216 Yield x s -> do
1217 r <- fstep fs x
1218 case r of
1219 FL.Partial fs1 -> go SPEC x s fs1
1220 FL.Done b -> return $ Yield b (GroupingInit s)
1221 Skip s -> return $ Skip $ GroupingDo s fs
1222 Stop -> return Stop
1224 where
1226 go !_ prev stt !acc = do
1227 res <- step (adaptState gst) stt
1228 case res of
1229 Yield x s -> do
1230 if cmp x prev
1231 then do
1232 r <- fstep acc x
1233 case r of
1234 FL.Partial fs1 -> go SPEC prev s fs1
1235 FL.Done b -> return $ Yield b (GroupingInit s)
1236 else do
1237 r <- done acc
1238 return $ Yield r (GroupingInitWith s x)
1239 Skip s -> go SPEC prev s acc
1240 Stop -> done acc >>= \r -> return $ Yield r GroupingDone
1241 stepOuter _ (GroupingInitWith st x) = do
1242 res <- initial
1243 return
1244 $ case res of
1245 FL.Partial s -> Skip $ GroupingDoWith st s x
1246 FL.Done b -> Yield b $ GroupingInitWith st x
1247 stepOuter gst (GroupingDoWith st fs prev) = do
1248 res <- fstep fs prev
1249 case res of
1250 FL.Partial fs1 -> go SPEC st fs1
1251 FL.Done b -> return $ Yield b (GroupingInit st)
1253 where
1255 -- XXX code duplicated from the previous equation
1256 go !_ stt !acc = do
1257 res <- step (adaptState gst) stt
1258 case res of
1259 Yield x s -> do
1260 if cmp x prev
1261 then do
1262 r <- fstep acc x
1263 case r of
1264 FL.Partial fs1 -> go SPEC s fs1
1265 FL.Done b -> return $ Yield b (GroupingInit s)
1266 else do
1267 r <- done acc
1268 return $ Yield r (GroupingInitWith s x)
1269 Skip s -> go SPEC s acc
1270 Stop -> done acc >>= \r -> return $ Yield r GroupingDone
1271 stepOuter _ (GroupingYield _ _) = error "groupsBy: Unreachable"
1272 stepOuter _ GroupingDone = return Stop
1274{-# INLINE_NORMAL groupsRollingBy #-}
1275groupsRollingBy :: Monad m
1276 => (a -> a -> Bool)
1277 -> Fold m a b
1278 -> Stream m a
1279 -> Stream m b
1280groupsRollingBy cmp (Fold fstep initial done) (Stream step state) =
1281 Stream stepOuter (GroupingInit state)
1283 where
1285 {-# INLINE_LATE stepOuter #-}
1286 stepOuter _ (GroupingInit st) = do
1287 -- XXX Note that if the stream stops without yielding a single element
1288 -- in the group we discard the "initial" effect.
1289 res <- initial
1290 return
1291 $ case res of
1292 FL.Partial fs -> Skip $ GroupingDo st fs
1293 FL.Done fb -> Yield fb $ GroupingInit st
1294 stepOuter gst (GroupingDo st fs) = do
1295 res <- step (adaptState gst) st
1296 case res of
1297 Yield x s -> do
1298 r <- fstep fs x
1299 case r of
1300 FL.Partial fs1 -> go SPEC x s fs1
1301 FL.Done fb -> return $ Yield fb (GroupingInit s)
1302 Skip s -> return $ Skip $ GroupingDo s fs
1303 Stop -> return Stop
1305 where
1307 go !_ prev stt !acc = do
1308 res <- step (adaptState gst) stt
1309 case res of
1310 Yield x s -> do
1311 if cmp prev x
1312 then do
1313 r <- fstep acc x
1314 case r of
1315 FL.Partial fs1 -> go SPEC x s fs1
1316 FL.Done b -> return $ Yield b (GroupingInit s)
1317 else do
1318 r <- done acc
1319 return $ Yield r (GroupingInitWith s x)
1320 Skip s -> go SPEC prev s acc
1321 Stop -> done acc >>= \r -> return $ Yield r GroupingDone
1322 stepOuter _ (GroupingInitWith st x) = do
1323 res <- initial
1324 return
1325 $ case res of
1326 FL.Partial s -> Skip $ GroupingDoWith st s x
1327 FL.Done b -> Yield b $ GroupingInitWith st x
1328 stepOuter gst (GroupingDoWith st fs previous) = do
1329 res <- fstep fs previous
1330 case res of
1331 FL.Partial s -> go SPEC previous st s
1332 FL.Done b -> return $ Yield b (GroupingInit st)
1334 where
1336 -- XXX GHC: groupsBy has one less parameter in this go loop and it
1337 -- fuses. However, groupsRollingBy does not fuse, removing the prev
1338 -- parameter makes it fuse. Something needs to be fixed in GHC. The
1339 -- workaround for this is noted in the comments below.
1340 go !_ prev !stt !acc = do
1341 res <- step (adaptState gst) stt
1342 case res of
1343 Yield x s -> do
1344 if cmp prev x
1345 then do
1346 r <- fstep acc x
1347 case r of
1348 FL.Partial fs1 -> go SPEC x s fs1
1349 FL.Done b -> return $ Yield b (GroupingInit st)
1350 else do
1351 {-
1352 r <- done acc
1353 return $ Yield r (GroupingInitWith s x)
1354 -}
1355 -- The code above does not let groupBy fuse. We use the
1356 -- alternative code below instead. Instead of jumping
1357 -- to GroupingInitWith state, we unroll the code of
1358 -- GroupingInitWith state here to help GHC with stream
1359 -- fusion.
1360 result <- initial
1361 r <- done acc
1362 return
1363 $ Yield r
1364 $ case result of
1365 FL.Partial fsi -> GroupingDoWith s fsi x
1366 FL.Done b -> GroupingYield b (GroupingInit s)
1367 Skip s -> go SPEC prev s acc
1368 Stop -> done acc >>= \r -> return $ Yield r GroupingDone
1369 stepOuter _ (GroupingYield r next) = return $ Yield r next
1370 stepOuter _ GroupingDone = return Stop
1372------------------------------------------------------------------------------
1373-- Splitting - by a predicate
1374------------------------------------------------------------------------------
1376data WordsByState st fs b
1377 = WordsByInit st
1378 | WordsByDo st !fs
1379 | WordsByDone
1380 | WordsByYield !b (WordsByState st fs b)
1382{-# INLINE_NORMAL wordsBy #-}
1383wordsBy :: Monad m => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1384wordsBy predicate (Fold fstep initial done) (Stream step state) =
1385 Stream stepOuter (WordsByInit state)
1387 where
1389 {-# INLINE_LATE stepOuter #-}
1390 stepOuter _ (WordsByInit st) = do
1391 res <- initial
1392 return
1393 $ case res of
1394 FL.Partial s -> Skip $ WordsByDo st s
1395 FL.Done b -> Yield b (WordsByInit st)
1397 stepOuter gst (WordsByDo st fs) = do
1398 res <- step (adaptState gst) st
1399 case res of
1400 Yield x s -> do
1401 if predicate x
1402 then do
1403 resi <- initial
1404 return
1405 $ case resi of
1406 FL.Partial fs1 -> Skip $ WordsByDo s fs1
1407 FL.Done b -> Yield b (WordsByInit s)
1408 else do
1409 r <- fstep fs x
1410 case r of
1411 FL.Partial fs1 -> go SPEC s fs1
1412 FL.Done b -> return $ Yield b (WordsByInit s)
1413 Skip s -> return $ Skip $ WordsByDo s fs
1414 Stop -> return Stop
1416 where
1418 go !_ stt !acc = do
1419 res <- step (adaptState gst) stt
1420 case res of
1421 Yield x s -> do
1422 if predicate x
1423 then do
1424 {-
1425 r <- done acc
1426 return $ Yield r (WordsByInit s)
1427 -}
1428 -- The above code does not fuse well. Need to check why
1429 -- GHC is not able to simplify it well. Using the code
1430 -- below, instead of jumping through the WordsByInit
1431 -- state always, we directly go to WordsByDo state in
1432 -- the common case of Partial.
1433 resi <- initial
1434 r <- done acc
1435 return
1436 $ Yield r
1437 $ case resi of
1438 FL.Partial fs1 -> WordsByDo s fs1
1439 FL.Done b -> WordsByYield b (WordsByInit s)
1440 else do
1441 r <- fstep acc x
1442 case r of
1443 FL.Partial fs1 -> go SPEC s fs1
1444 FL.Done b -> return $ Yield b (WordsByInit s)
1445 Skip s -> go SPEC s acc
1446 Stop -> done acc >>= \r -> return $ Yield r WordsByDone
1448 stepOuter _ WordsByDone = return Stop
1450 stepOuter _ (WordsByYield b next) = return $ Yield b next
1452------------------------------------------------------------------------------
1453-- Splitting on a sequence
1454------------------------------------------------------------------------------
1456-- String search algorithms:
1457-- http://www-igm.univ-mlv.fr/~lecroq/string/index.html
1460-- TODO can we unify the splitting operations using a splitting configuration
1461-- like in the split package.
1463data SplitStyle = Infix | Suffix | Prefix deriving (Eq, Show)
1464data SplitOptions = SplitOptions
1465 { style :: SplitStyle
1466 , withSep :: Bool -- ^ keep the separators in output
1467 -- , compact :: Bool -- ^ treat multiple consecutive separators as one
1468 -- , trimHead :: Bool -- ^ drop blank at head
1469 -- , trimTail :: Bool -- ^ drop blank at tail
1473-- XXX using "fs" as the last arg in Constructors may simplify the code a bit,
1474-- because we can use the constructor directly without having to create "jump"
1475-- functions.
1476{-# ANN type SplitOnSeqState Fuse #-}
1477data SplitOnSeqState rb rh ck w fs s b x =
1478 SplitOnSeqInit
1479 | SplitOnSeqYield b (SplitOnSeqState rb rh ck w fs s b x)
1480 | SplitOnSeqDone
1482 | SplitOnSeqEmpty !fs s
1484 | SplitOnSeqSingle !fs s x
1486 | SplitOnSeqWordInit !fs s
1487 | SplitOnSeqWordLoop !w s !fs
1488 | SplitOnSeqWordDone Int !fs !w
1490 | SplitOnSeqKRInit Int !fs s rb !rh
1491 | SplitOnSeqKRLoop fs s rb !rh !ck
1492 | SplitOnSeqKRCheck fs s rb !rh
1493 | SplitOnSeqKRDone Int !fs rb !rh
1495 | SplitOnSeqReinit (fs -> SplitOnSeqState rb rh ck w fs s b x)
1497{-# INLINE_NORMAL splitOnSeq #-}
1498splitOnSeq
1499 :: forall m a b. (MonadIO m, Storable a, Enum a, Eq a)
1500 => Array a
1501 -> Fold m a b
1502 -> Stream m a
1503 -> Stream m b
1504splitOnSeq patArr (Fold fstep initial done) (Stream step state) =
1505 Stream stepOuter SplitOnSeqInit
1507 where
1509 patLen = A.length patArr
1510 maxIndex = patLen - 1
1511 elemBits = sizeOf (undefined :: a) * 8
1513 -- For word pattern case
1514 wordMask :: Word
1515 wordMask = (1 `shiftL` (elemBits * patLen)) - 1
1517 elemMask :: Word
1518 elemMask = (1 `shiftL` elemBits) - 1
1520 wordPat :: Word
1521 wordPat = wordMask .&. A.foldl' addToWord 0 patArr
1523 addToWord wd a = (wd `shiftL` elemBits) .|. fromIntegral (fromEnum a)
1525 -- For Rabin-Karp search
1526 k = 2891336453 :: Word32
1527 coeff = k ^ patLen
1529 addCksum cksum a = cksum * k + fromIntegral (fromEnum a)
1531 deltaCksum cksum old new =
1532 addCksum cksum new - coeff * fromIntegral (fromEnum old)
1534 -- XXX shall we use a random starting hash or 1 instead of 0?
1535 patHash = A.foldl' addCksum 0 patArr
1537 skip = return . Skip
1539 nextAfterInit nextGen stepRes =
1540 case stepRes of
1541 FL.Partial s -> nextGen s
1542 FL.Done b -> SplitOnSeqYield b (SplitOnSeqReinit nextGen)
1544 {-# INLINE yieldProceed #-}
1545 yieldProceed nextGen fs =
1546 initial >>= skip . SplitOnSeqYield fs . nextAfterInit nextGen
1548 {-# INLINE_LATE stepOuter #-}
1549 stepOuter _ SplitOnSeqInit = do
1550 res <- initial
1551 case res of
1552 FL.Partial acc ->
1553 if patLen == 0
1554 then return $ Skip $ SplitOnSeqEmpty acc state
1555 else if patLen == 1
1556 then do
1557 pat <- liftIO $ A.unsafeIndexIO patArr 0
1558 return $ Skip $ SplitOnSeqSingle acc state pat
1559 else if sizeOf (undefined :: a) * patLen
1560 <= sizeOf (undefined :: Word)
1561 then return $ Skip $ SplitOnSeqWordInit acc state
1562 else do
1563 (rb, rhead) <- liftIO $ RB.new patLen
1564 skip $ SplitOnSeqKRInit 0 acc state rb rhead
1565 FL.Done b -> skip $ SplitOnSeqYield b SplitOnSeqInit
1567 stepOuter _ (SplitOnSeqYield x next) = return $ Yield x next
1569 ---------------------------
1570 -- Checkpoint
1571 ---------------------------
1573 stepOuter _ (SplitOnSeqReinit nextGen) =
1574 initial >>= skip . nextAfterInit nextGen
1576 ---------------------------
1577 -- Empty pattern
1578 ---------------------------
1580 stepOuter gst (SplitOnSeqEmpty acc st) = do
1581 res <- step (adaptState gst) st
1582 case res of
1583 Yield x s -> do
1584 r <- fstep acc x
1585 b1 <-
1586 case r of
1587 FL.Partial acc1 -> done acc1
1588 FL.Done b -> return b
1589 let jump c = SplitOnSeqEmpty c s
1590 in yieldProceed jump b1
1591 Skip s -> skip (SplitOnSeqEmpty acc s)
1592 Stop -> return Stop
1594 -----------------
1595 -- Done
1596 -----------------
1598 stepOuter _ SplitOnSeqDone = return Stop
1600 -----------------
1601 -- Single Pattern
1602 -----------------
1604 stepOuter gst (SplitOnSeqSingle fs st pat) = do
1605 res <- step (adaptState gst) st
1606 case res of
1607 Yield x s -> do
1608 let jump c = SplitOnSeqSingle c s pat
1609 if pat == x
1610 then done fs >>= yieldProceed jump
1611 else do
1612 r <- fstep fs x
1613 case r of
1614 FL.Partial fs1 -> skip $ jump fs1
1615 FL.Done b -> yieldProceed jump b
1616 Skip s -> return $ Skip $ SplitOnSeqSingle fs s pat
1617 Stop -> do
1618 r <- done fs
1619 return $ Skip $ SplitOnSeqYield r SplitOnSeqDone
1621 ---------------------------
1622 -- Short Pattern - Shift Or
1623 ---------------------------
1625 stepOuter _ (SplitOnSeqWordDone 0 fs _) = do
1626 r <- done fs
1627 skip $ SplitOnSeqYield r SplitOnSeqDone
1628 stepOuter _ (SplitOnSeqWordDone n fs wrd) = do
1629 let old = elemMask .&. (wrd `shiftR` (elemBits * (n - 1)))
1630 r <- fstep fs (toEnum $ fromIntegral old)
1631 case r of
1632 FL.Partial fs1 -> skip $ SplitOnSeqWordDone (n - 1) fs1 wrd
1633 FL.Done b -> do
1634 let jump c = SplitOnSeqWordDone (n - 1) c wrd
1635 yieldProceed jump b
1637 stepOuter gst (SplitOnSeqWordInit fs st0) =
1638 go SPEC 0 0 st0
1640 where
1642 {-# INLINE go #-}
1643 go !_ !idx !wrd !st = do
1644 res <- step (adaptState gst) st
1645 case res of
1646 Yield x s -> do
1647 let wrd1 = addToWord wrd x
1648 if idx == maxIndex
1649 then do
1650 if wrd1 .&. wordMask == wordPat
1651 then do
1652 let jump c = SplitOnSeqWordInit c s
1653 done fs >>= yieldProceed jump
1654 else skip $ SplitOnSeqWordLoop wrd1 s fs
1655 else go SPEC (idx + 1) wrd1 s
1656 Skip s -> go SPEC idx wrd s
1657 Stop -> do
1658 if idx /= 0
1659 then skip $ SplitOnSeqWordDone idx fs wrd
1660 else do
1661 r <- done fs
1662 skip $ SplitOnSeqYield r SplitOnSeqDone
1664 stepOuter gst (SplitOnSeqWordLoop wrd0 st0 fs0) =
1665 go SPEC wrd0 st0 fs0
1667 where
1669 {-# INLINE go #-}
1670 go !_ !wrd !st !fs = do
1671 res <- step (adaptState gst) st
1672 case res of
1673 Yield x s -> do
1674 let jump c = SplitOnSeqWordInit c s
1675 wrd1 = addToWord wrd x
1676 old = (wordMask .&. wrd)
1677 `shiftR` (elemBits * (patLen - 1))
1678 r <- fstep fs (toEnum $ fromIntegral old)
1679 case r of
1680 FL.Partial fs1 -> do
1681 if wrd1 .&. wordMask == wordPat
1682 then done fs1 >>= yieldProceed jump
1683 else go SPEC wrd1 s fs1
1684 FL.Done b -> yieldProceed jump b
1685 Skip s -> go SPEC wrd s fs
1686 Stop -> skip $ SplitOnSeqWordDone patLen fs wrd
1688 -------------------------------
1689 -- General Pattern - Karp Rabin
1690 -------------------------------
1692 stepOuter gst (SplitOnSeqKRInit idx fs st rb rh) = do
1693 res <- step (adaptState gst) st
1694 case res of
1695 Yield x s -> do
1696 rh1 <- liftIO $ RB.unsafeInsert rb rh x
1697 if idx == maxIndex
1698 then do
1699 let fold = RB.unsafeFoldRing (RB.ringBound rb)
1700 let !ringHash = fold addCksum 0 rb
1701 if ringHash == patHash
1702 then skip $ SplitOnSeqKRCheck fs s rb rh1
1703 else skip $ SplitOnSeqKRLoop fs s rb rh1 ringHash
1704 else skip $ SplitOnSeqKRInit (idx + 1) fs s rb rh1
1705 Skip s -> skip $ SplitOnSeqKRInit idx fs s rb rh
1706 Stop -> do
1707 skip $ SplitOnSeqKRDone idx fs rb (RB.startOf rb)
1709 -- XXX The recursive "go" is more efficient than the state based recursion
1710 -- code commented out below. Perhaps its more efficient because of
1711 -- factoring out "rb" outside the loop.
1712 --
1713 stepOuter gst (SplitOnSeqKRLoop fs0 st0 rb rh0 cksum0) =
1714 go SPEC fs0 st0 rh0 cksum0
1716 where
1718 go !_ !fs !st !rh !cksum = do
1719 res <- step (adaptState gst) st
1720 case res of
1721 Yield x s -> do
1722 old <- liftIO $ peek rh
1723 let cksum1 = deltaCksum cksum old x
1724 r <- fstep fs old
1725 case r of
1726 FL.Partial fs1 -> do
1727 rh1 <- liftIO (RB.unsafeInsert rb rh x)
1728 if cksum1 == patHash
1729 then skip $ SplitOnSeqKRCheck fs1 s rb rh1
1730 else go SPEC fs1 s rh1 cksum1
1731 FL.Done b -> do
1732 let rst = RB.startOf rb
1733 jump c = SplitOnSeqKRInit 0 c s rb rst
1734 yieldProceed jump b
1735 Skip s -> go SPEC fs s rh cksum
1736 Stop -> skip $ SplitOnSeqKRDone patLen fs rb rh
1738 -- XXX The following code is 5 times slower compared to the recursive loop
1739 -- based code above. Need to investigate why. One possibility is that the
1740 -- go loop above does not thread around the ring buffer (rb). This code may
1741 -- be causing the state to bloat and getting allocated on each iteration.
1742 -- We can check the cmm/asm code to confirm. If so a good GHC solution to
1743 -- such problem is needed. One way to avoid this could be to use unboxed
1744 -- mutable state?
1745 {-
1746 stepOuter gst (SplitOnSeqKRLoop fs st rb rh cksum) = do
1747 res <- step (adaptState gst) st
1748 case res of
1749 Yield x s -> do
1750 old <- liftIO $ peek rh
1751 let cksum1 = deltaCksum cksum old x
1752 fs1 <- fstep fs old
1753 if (cksum1 == patHash)
1754 then do
1755 r <- done fs1
1756 skip $ SplitOnSeqYield r $ SplitOnSeqKRInit 0 s rb rh
1757 else do
1758 rh1 <- liftIO (RB.unsafeInsert rb rh x)
1759 skip $ SplitOnSeqKRLoop fs1 s rb rh1 cksum1
1760 Skip s -> skip $ SplitOnSeqKRLoop fs s rb rh cksum
1761 Stop -> skip $ SplitOnSeqKRDone patLen fs rb rh
1762 -}
1764 stepOuter _ (SplitOnSeqKRCheck fs st rb rh) = do
1765 if RB.unsafeEqArray rb rh patArr
1766 then do
1767 r <- done fs
1768 let rst = RB.startOf rb
1769 jump c = SplitOnSeqKRInit 0 c st rb rst
1770 yieldProceed jump r
1771 else skip $ SplitOnSeqKRLoop fs st rb rh patHash
1773 stepOuter _ (SplitOnSeqKRDone 0 fs _ _) = do
1774 r <- done fs
1775 skip $ SplitOnSeqYield r SplitOnSeqDone
1776 stepOuter _ (SplitOnSeqKRDone n fs rb rh) = do
1777 old <- liftIO $ peek rh
1778 let rh1 = RB.advance rb rh
1779 r <- fstep fs old
1780 case r of
1781 FL.Partial fs1 -> skip $ SplitOnSeqKRDone (n - 1) fs1 rb rh1
1782 FL.Done b -> do
1783 let jump c = SplitOnSeqKRDone (n - 1) c rb rh1
1784 yieldProceed jump b
1786{-# ANN type SplitOnSuffixSeqState Fuse #-}
1787data SplitOnSuffixSeqState rb rh ck w fs s b x =
1788 SplitOnSuffixSeqInit
1789 | SplitOnSuffixSeqYield b (SplitOnSuffixSeqState rb rh ck w fs s b x)
1790 | SplitOnSuffixSeqDone
1792 | SplitOnSuffixSeqEmpty !fs s
1794 | SplitOnSuffixSeqSingleInit !fs s x
1795 | SplitOnSuffixSeqSingle !fs s x
1797 | SplitOnSuffixSeqWordInit !fs s
1798 | SplitOnSuffixSeqWordLoop !w s !fs
1799 | SplitOnSuffixSeqWordDone Int !fs !w
1801 | SplitOnSuffixSeqKRInit Int !fs s rb !rh
1802 | SplitOnSuffixSeqKRInit1 !fs s rb !rh
1803 | SplitOnSuffixSeqKRLoop fs s rb !rh !ck
1804 | SplitOnSuffixSeqKRCheck fs s rb !rh
1805 | SplitOnSuffixSeqKRDone Int !fs rb !rh
1807 | SplitOnSuffixSeqReinit
1808 (fs -> SplitOnSuffixSeqState rb rh ck w fs s b x)
1810{-# INLINE_NORMAL splitOnSuffixSeq #-}
1811splitOnSuffixSeq
1812 :: forall m a b. (MonadIO m, Storable a, Enum a, Eq a)
1813 => Bool
1814 -> Array a
1815 -> Fold m a b
1816 -> Stream m a
1817 -> Stream m b
1818splitOnSuffixSeq withSep patArr (Fold fstep initial done) (Stream step state) =
1819 Stream stepOuter SplitOnSuffixSeqInit
1821 where
1823 patLen = A.length patArr
1824 maxIndex = patLen - 1
1825 elemBits = sizeOf (undefined :: a) * 8
1827 -- For word pattern case
1828 wordMask :: Word
1829 wordMask = (1 `shiftL` (elemBits * patLen)) - 1
1831 elemMask :: Word
1832 elemMask = (1 `shiftL` elemBits) - 1
1834 wordPat :: Word
1835 wordPat = wordMask .&. A.foldl' addToWord 0 patArr
1837 addToWord wd a = (wd `shiftL` elemBits) .|. fromIntegral (fromEnum a)
1839 nextAfterInit nextGen stepRes =
1840 case stepRes of
1841 FL.Partial s -> nextGen s
1842 FL.Done b ->
1843 SplitOnSuffixSeqYield b (SplitOnSuffixSeqReinit nextGen)
1845 {-# INLINE yieldProceed #-}
1846 yieldProceed nextGen fs =
1847 initial >>= skip . SplitOnSuffixSeqYield fs . nextAfterInit nextGen
1849 -- For single element pattern case
1850 {-# INLINE processYieldSingle #-}
1851 processYieldSingle pat x s fs = do
1852 let jump c = SplitOnSuffixSeqSingleInit c s pat
1853 if pat == x
1854 then do
1855 r <- if withSep then fstep fs x else return $ FL.Partial fs
1856 b1 <-
1857 case r of
1858 FL.Partial fs1 -> done fs1
1859 FL.Done b -> return b
1860 yieldProceed jump b1
1861 else do
1862 r <- fstep fs x
1863 case r of
1864 FL.Partial fs1 -> skip $ SplitOnSuffixSeqSingle fs1 s pat
1865 FL.Done b -> yieldProceed jump b
1867 -- For Rabin-Karp search
1868 k = 2891336453 :: Word32
1869 coeff = k ^ patLen
1871 addCksum cksum a = cksum * k + fromIntegral (fromEnum a)
1873 deltaCksum cksum old new =
1874 addCksum cksum new - coeff * fromIntegral (fromEnum old)
1876 -- XXX shall we use a random starting hash or 1 instead of 0?
1877 patHash = A.foldl' addCksum 0 patArr
1879 skip = return . Skip
1881 {-# INLINE_LATE stepOuter #-}
1882 stepOuter _ SplitOnSuffixSeqInit = do
1883 res <- initial
1884 case res of
1885 FL.Partial fs ->
1886 if patLen == 0
1887 then skip $ SplitOnSuffixSeqEmpty fs state
1888 else if patLen == 1
1889 then do
1890 pat <- liftIO $ A.unsafeIndexIO patArr 0
1891 skip $ SplitOnSuffixSeqSingleInit fs state pat
1892 else if sizeOf (undefined :: a) * patLen
1893 <= sizeOf (undefined :: Word)
1894 then skip $ SplitOnSuffixSeqWordInit fs state
1895 else do
1896 (rb, rhead) <- liftIO $ RB.new patLen
1897 skip $ SplitOnSuffixSeqKRInit 0 fs state rb rhead
1898 FL.Done fb -> skip $ SplitOnSuffixSeqYield fb SplitOnSuffixSeqInit
1900 stepOuter _ (SplitOnSuffixSeqYield x next) = return $ Yield x next
1902 ---------------------------
1903 -- Reinit
1904 ---------------------------
1906 stepOuter _ (SplitOnSuffixSeqReinit nextGen) =
1907 initial >>= skip . nextAfterInit nextGen
1909 ---------------------------
1910 -- Empty pattern
1911 ---------------------------
1913 stepOuter gst (SplitOnSuffixSeqEmpty acc st) = do
1914 res <- step (adaptState gst) st
1915 case res of
1916 Yield x s -> do
1917 let jump c = SplitOnSuffixSeqEmpty c s
1918 r <- fstep acc x
1919 b1 <-
1920 case r of
1921 FL.Partial fs -> done fs
1922 FL.Done b -> return b
1923 yieldProceed jump b1
1924 Skip s -> skip (SplitOnSuffixSeqEmpty acc s)
1925 Stop -> return Stop
1927 -----------------
1928 -- Done
1929 -----------------
1931 stepOuter _ SplitOnSuffixSeqDone = return Stop
1933 -----------------
1934 -- Single Pattern
1935 -----------------
1937 stepOuter gst (SplitOnSuffixSeqSingleInit fs st pat) = do
1938 res <- step (adaptState gst) st
1939 case res of
1940 Yield x s -> processYieldSingle pat x s fs
1941 Skip s -> skip $ SplitOnSuffixSeqSingleInit fs s pat
1942 Stop -> return Stop
1944 stepOuter gst (SplitOnSuffixSeqSingle fs st pat) = do
1945 res <- step (adaptState gst) st
1946 case res of
1947 Yield x s -> processYieldSingle pat x s fs
1948 Skip s -> skip $ SplitOnSuffixSeqSingle fs s pat
1949 Stop -> do
1950 r <- done fs
1951 skip $ SplitOnSuffixSeqYield r SplitOnSuffixSeqDone
1953 ---------------------------
1954 -- Short Pattern - Shift Or
1955 ---------------------------
1957 stepOuter _ (SplitOnSuffixSeqWordDone 0 fs _) = do
1958 r <- done fs
1959 skip $ SplitOnSuffixSeqYield r SplitOnSuffixSeqDone
1960 stepOuter _ (SplitOnSuffixSeqWordDone n fs wrd) = do
1961 let old = elemMask .&. (wrd `shiftR` (elemBits * (n - 1)))
1962 r <- fstep fs (toEnum $ fromIntegral old)
1963 case r of
1964 FL.Partial fs1 -> skip $ SplitOnSuffixSeqWordDone (n - 1) fs1 wrd
1965 FL.Done b -> do
1966 let jump c = SplitOnSuffixSeqWordDone (n - 1) c wrd
1967 yieldProceed jump b
1969 stepOuter gst (SplitOnSuffixSeqWordInit fs0 st0) = do
1970 res <- step (adaptState gst) st0
1971 case res of
1972 Yield x s -> do
1973 let wrd = addToWord 0 x
1974 r <- if withSep then fstep fs0 x else return $ FL.Partial fs0
1975 case r of
1976 FL.Partial fs1 -> go SPEC 1 wrd s fs1
1977 FL.Done b -> do
1978 let jump c = SplitOnSuffixSeqWordInit c s
1979 yieldProceed jump b
1980 Skip s -> skip (SplitOnSuffixSeqWordInit fs0 s)
1981 Stop -> return Stop
1983 where
1985 {-# INLINE go #-}
1986 go !_ !idx !wrd !st !fs = do
1987 res <- step (adaptState gst) st
1988 case res of
1989 Yield x s -> do
1990 let jump c = SplitOnSuffixSeqWordInit c s
1991 let wrd1 = addToWord wrd x
1992 r <- if withSep then fstep fs x else return $ FL.Partial fs
1993 case r of
1994 FL.Partial fs1 ->
1995 if idx /= maxIndex
1996 then go SPEC (idx + 1) wrd1 s fs1
1997 else if wrd1 .&. wordMask /= wordPat
1998 then skip $ SplitOnSuffixSeqWordLoop wrd1 s fs1
1999 else do done fs >>= yieldProceed jump
2000 FL.Done b -> yieldProceed jump b
2001 Skip s -> go SPEC idx wrd s fs
2002 Stop -> skip $ SplitOnSuffixSeqWordDone idx fs wrd
2004 stepOuter gst (SplitOnSuffixSeqWordLoop wrd0 st0 fs0) =
2005 go SPEC wrd0 st0 fs0
2007 where
2009 {-# INLINE go #-}
2010 go !_ !wrd !st !fs = do
2011 res <- step (adaptState gst) st
2012 case res of
2013 Yield x s -> do
2014 let jump c = SplitOnSuffixSeqWordInit c s
2015 wrd1 = addToWord wrd x
2016 old = (wordMask .&. wrd)
2017 `shiftR` (elemBits * (patLen - 1))
2018 r <-
2019 if withSep
2020 then fstep fs x
2021 else fstep fs (toEnum $ fromIntegral old)
2022 case r of
2023 FL.Partial fs1 ->
2024 if wrd1 .&. wordMask == wordPat
2025 then done fs1 >>= yieldProceed jump
2026 else go SPEC wrd1 s fs1
2027 FL.Done b -> yieldProceed jump b
2028 Skip s -> go SPEC wrd s fs
2029 Stop ->
2030 if wrd .&. wordMask == wordPat
2031 then return Stop
2032 else if withSep
2033 then do
2034 r <- done fs
2035 skip $ SplitOnSuffixSeqYield r SplitOnSuffixSeqDone
2036 else skip $ SplitOnSuffixSeqWordDone patLen fs wrd
2038 -------------------------------
2039 -- General Pattern - Karp Rabin
2040 -------------------------------
2042 stepOuter gst (SplitOnSuffixSeqKRInit idx0 fs st0 rb rh0) = do
2043 res <- step (adaptState gst) st0
2044 case res of
2045 Yield x s -> do
2046 rh1 <- liftIO $ RB.unsafeInsert rb rh0 x
2047 r <- if withSep then fstep fs x else return $ FL.Partial fs
2048 case r of
2049 FL.Partial fs1 ->
2050 skip $ SplitOnSuffixSeqKRInit1 fs1 s rb rh1
2051 FL.Done b -> do
2052 let rst = RB.startOf rb
2053 jump c = SplitOnSuffixSeqKRInit 0 c s rb rst
2054 yieldProceed jump b
2055 Skip s -> skip $ SplitOnSuffixSeqKRInit idx0 fs s rb rh0
2056 Stop -> return Stop
2058 stepOuter gst (SplitOnSuffixSeqKRInit1 fs0 st0 rb rh0) = do
2059 go SPEC 1 rh0 st0 fs0
2061 where
2063 go !_ !idx !rh st !fs = do
2064 res <- step (adaptState gst) st
2065 case res of
2066 Yield x s -> do
2067 rh1 <- liftIO (RB.unsafeInsert rb rh x)
2068 r <- if withSep then fstep fs x else return $ FL.Partial fs
2069 case r of
2070 FL.Partial fs1 ->
2071 if idx /= maxIndex
2072 then go SPEC (idx + 1) rh1 s fs1
2073 else skip $
2074 let fold = RB.unsafeFoldRing (RB.ringBound rb)
2075 !ringHash = fold addCksum 0 rb
2076 in if ringHash == patHash
2077 then SplitOnSuffixSeqKRCheck fs1 s rb rh1
2078 else SplitOnSuffixSeqKRLoop
2079 fs1 s rb rh1 ringHash
2080 FL.Done b -> do
2081 let rst = RB.startOf rb
2082 jump c = SplitOnSuffixSeqKRInit 0 c s rb rst
2083 yieldProceed jump b
2084 Skip s -> go SPEC idx rh s fs
2085 Stop -> do
2086 -- do not issue a blank segment when we end at pattern
2087 if (idx == maxIndex) && RB.unsafeEqArray rb rh patArr
2088 then return Stop
2089 else if withSep
2090 then do
2091 r <- done fs
2092 skip $ SplitOnSuffixSeqYield r SplitOnSuffixSeqDone
2093 else skip $ SplitOnSuffixSeqKRDone idx fs rb (RB.startOf rb)
2095 stepOuter gst (SplitOnSuffixSeqKRLoop fs0 st0 rb rh0 cksum0) =
2096 go SPEC fs0 st0 rh0 cksum0
2098 where
2100 go !_ !fs !st !rh !cksum = do
2101 res <- step (adaptState gst) st
2102 case res of
2103 Yield x s -> do
2104 old <- liftIO $ peek rh
2105 rh1 <- liftIO (RB.unsafeInsert rb rh x)
2106 let cksum1 = deltaCksum cksum old x
2107 r <- if withSep then fstep fs x else fstep fs old
2108 case r of
2109 FL.Partial fs1 ->
2110 if (cksum1 /= patHash)
2111 then go SPEC fs1 s rh1 cksum1
2112 else skip $ SplitOnSuffixSeqKRCheck fs1 s rb rh1
2113 FL.Done b -> do
2114 let rst = RB.startOf rb
2115 jump c = SplitOnSuffixSeqKRInit 0 c s rb rst
2116 yieldProceed jump b
2117 Skip s -> go SPEC fs s rh cksum
2118 Stop ->
2119 if RB.unsafeEqArray rb rh patArr
2120 then return Stop
2121 else if withSep
2122 then do
2123 r <- done fs
2124 skip $ SplitOnSuffixSeqYield r SplitOnSuffixSeqDone
2125 else skip $ SplitOnSuffixSeqKRDone patLen fs rb rh
2127 stepOuter _ (SplitOnSuffixSeqKRCheck fs st rb rh) = do
2128 if RB.unsafeEqArray rb rh patArr
2129 then do
2130 r <- done fs
2131 let rst = RB.startOf rb
2132 jump c = SplitOnSuffixSeqKRInit 0 c st rb rst
2133 yieldProceed jump r
2134 else skip $ SplitOnSuffixSeqKRLoop fs st rb rh patHash
2136 stepOuter _ (SplitOnSuffixSeqKRDone 0 fs _ _) = do
2137 r <- done fs
2138 skip $ SplitOnSuffixSeqYield r SplitOnSuffixSeqDone
2139 stepOuter _ (SplitOnSuffixSeqKRDone n fs rb rh) = do
2140 old <- liftIO $ peek rh
2141 let rh1 = RB.advance rb rh
2142 r <- fstep fs old
2143 case r of
2144 FL.Partial fs1 -> skip $ SplitOnSuffixSeqKRDone (n - 1) fs1 rb rh1
2145 FL.Done b -> do
2146 let jump c = SplitOnSuffixSeqKRDone (n - 1) c rb rh1
2147 yieldProceed jump b
2149------------------------------------------------------------------------------
2150-- Nested Container Transformation
2151------------------------------------------------------------------------------
2153{-# ANN type SplitState Fuse #-}
2154data SplitState s arr
2155 = SplitInitial s
2156 | SplitBuffering s arr
2157 | SplitSplitting s arr
2158 | SplitYielding arr (SplitState s arr)
2159 | SplitFinishing
2161-- XXX An alternative approach would be to use a partial fold (Fold m a b) to
2162-- split using a splitBy like combinator. The Fold would consume upto the
2163-- separator and return any leftover which can then be fed to the next fold.
2165-- We can revisit this once we have partial folds/parsers.
2167-- | Performs infix separator style splitting.
2168{-# INLINE_NORMAL splitInnerBy #-}
2169splitInnerBy
2170 :: Monad m
2171 => (f a -> m (f a, Maybe (f a))) -- splitter
2172 -> (f a -> f a -> m (f a)) -- joiner
2173 -> Stream m (f a)
2174 -> Stream m (f a)
2175splitInnerBy splitter joiner (Stream step1 state1) =
2176 (Stream step (SplitInitial state1))
2178 where
2180 {-# INLINE_LATE step #-}
2181 step gst (SplitInitial st) = do
2182 r <- step1 gst st
2183 case r of
2184 Yield x s -> do
2185 (x1, mx2) <- splitter x
2186 return $ case mx2 of
2187 Nothing -> Skip (SplitBuffering s x1)
2188 Just x2 -> Skip (SplitYielding x1 (SplitSplitting s x2))
2189 Skip s -> return $ Skip (SplitInitial s)
2190 Stop -> return $ Stop
2192 step gst (SplitBuffering st buf) = do
2193 r <- step1 gst st
2194 case r of
2195 Yield x s -> do
2196 (x1, mx2) <- splitter x
2197 buf' <- joiner buf x1
2198 return $ case mx2 of
2199 Nothing -> Skip (SplitBuffering s buf')
2200 Just x2 -> Skip (SplitYielding buf' (SplitSplitting s x2))
2201 Skip s -> return $ Skip (SplitBuffering s buf)
2202 Stop -> return $ Skip (SplitYielding buf SplitFinishing)
2204 step _ (SplitSplitting st buf) = do
2205 (x1, mx2) <- splitter buf
2206 return $ case mx2 of
2207 Nothing -> Skip $ SplitBuffering st x1
2208 Just x2 -> Skip $ SplitYielding x1 (SplitSplitting st x2)
2210 step _ (SplitYielding x next) = return $ Yield x next
2211 step _ SplitFinishing = return $ Stop
2213-- | Performs infix separator style splitting.
2214{-# INLINE_NORMAL splitInnerBySuffix #-}
2215splitInnerBySuffix
2216 :: (Monad m, Eq (f a), Monoid (f a))
2217 => (f a -> m (f a, Maybe (f a))) -- splitter
2218 -> (f a -> f a -> m (f a)) -- joiner
2219 -> Stream m (f a)
2220 -> Stream m (f a)
2221splitInnerBySuffix splitter joiner (Stream step1 state1) =
2222 (Stream step (SplitInitial state1))
2224 where
2226 {-# INLINE_LATE step #-}
2227 step gst (SplitInitial st) = do
2228 r <- step1 gst st
2229 case r of
2230 Yield x s -> do
2231 (x1, mx2) <- splitter x
2232 return $ case mx2 of
2233 Nothing -> Skip (SplitBuffering s x1)
2234 Just x2 -> Skip (SplitYielding x1 (SplitSplitting s x2))
2235 Skip s -> return $ Skip (SplitInitial s)
2236 Stop -> return $ Stop
2238 step gst (SplitBuffering st buf) = do
2239 r <- step1 gst st
2240 case r of
2241 Yield x s -> do
2242 (x1, mx2) <- splitter x
2243 buf' <- joiner buf x1
2244 return $ case mx2 of
2245 Nothing -> Skip (SplitBuffering s buf')
2246 Just x2 -> Skip (SplitYielding buf' (SplitSplitting s x2))
2247 Skip s -> return $ Skip (SplitBuffering s buf)
2248 Stop -> return $
2249 if buf == mempty
2250 then Stop
2251 else Skip (SplitYielding buf SplitFinishing)
2253 step _ (SplitSplitting st buf) = do
2254 (x1, mx2) <- splitter buf
2255 return $ case mx2 of
2256 Nothing -> Skip $ SplitBuffering st x1
2257 Just x2 -> Skip $ SplitYielding x1 (SplitSplitting st x2)
2259 step _ (SplitYielding x next) = return $ Yield x next
2260 step _ SplitFinishing = return Stop