git.iko.soy/streamly
on 3ce574a help sign in
streamly/src/Streamly/Streams/StreamD.hs
2323 lines 79 kB history raw
1{-# LANGUAGE BangPatterns #-}
2{-# LANGUAGE CPP #-}
3{-# LANGUAGE ConstraintKinds #-}
4{-# LANGUAGE ExistentialQuantification #-}
5{-# LANGUAGE FlexibleContexts #-}
6{-# LANGUAGE FlexibleInstances #-}
7{-# LANGUAGE MultiParamTypeClasses #-}
8{-# LANGUAGE PatternSynonyms #-}
9{-# LANGUAGE RecordWildCards #-}
10{-# LANGUAGE ScopedTypeVariables #-}
11{-# LANGUAGE ViewPatterns #-}
12{-# LANGUAGE RankNTypes #-}
14#include "inline.hs"
16-- |
17-- Module : Streamly.Streams.StreamD
18-- Copyright : (c) 2018 Harendra Kumar
19-- Copyright : (c) Roman Leshchinskiy 2008-2010
20--
21-- License : BSD3
22-- Maintainer : harendra.kumar@gmail.com
23-- Stability : experimental
24-- Portability : GHC
25--
26-- Direct style re-implementation of CPS style stream in StreamK module. The
27-- symbol or suffix 'D' in this module denotes the "Direct" style. GHC is able
28-- to INLINE and fuse direct style better, providing better performance than
29-- CPS implementation.
30--
31-- @
32-- import qualified Streamly.Streams.StreamD as D
33-- @
35-- Some of the functions in this file have been adapted from the vector
36-- library, https://hackage.haskell.org/package/vector.
38module Streamly.Streams.StreamD
39 (
40 -- * The stream type
41 Step (..)
42 , Stream (..)
44 -- * Construction
45 , nil
46 , cons
48 -- * Deconstruction
49 , uncons
51 -- * Generation
52 -- ** Unfolds
53 , unfoldr
54 , unfoldrM
56 -- ** Specialized Generation
57 -- | Generate a monadic stream from a seed.
58 , repeat
59 , replicate
60 , replicateM
61 , fromIndices
62 , fromIndicesM
63 , generate
64 , generateM
66 -- ** Enumerations
67 , enumerateFromStepIntegral
68 , enumerateFromIntegral
69 , enumerateFromThenIntegral
70 , enumerateFromToIntegral
71 , enumerateFromThenToIntegral
73 , enumerateFromStepNum
74 , numFrom
75 , numFromThen
76 , enumerateFromToFractional
77 , enumerateFromThenToFractional
79 -- ** Conversions
80 -- | Transform an input structure into a stream.
81 -- | Direct style stream does not support @fromFoldable@.
82 , yield
83 , yieldM
84 , fromList
85 , fromListM
86 , fromStreamK
87 , fromStreamD
89 -- * Elimination
90 -- ** General Folds
91 , foldrS
92 , foldrT
93 , foldrM
94 , foldrMx
95 , foldr
96 , foldr1
98 , foldl'
99 , foldlM'
100 , foldlS
101 , foldlT
102 , reverse
103 , reverse'
105 , foldlx'
106 , foldlMx'
108 -- ** Specialized Folds
109 , tap
110 , drain
111 , null
112 , head
113 , tail
114 , last
115 , elem
116 , notElem
117 , all
118 , any
119 , maximum
120 , maximumBy
121 , minimum
122 , minimumBy
123 , findIndices
124 , lookup
125 , findM
126 , find
127 , (!!)
129 -- ** Flattening nested streams
130 , concatMapM
131 , concatMap
133 -- ** Grouping
134 , groupsOf
135 , groupsBy
136 , groupsRollingBy
138 -- ** Splitting
139 , splitBy
140 , splitSuffixBy
141 , wordsBy
142 , splitSuffixBy'
144 , splitOn
145 , splitSuffixOn
147 -- ** Substreams
148 , isPrefixOf
149 , isSubsequenceOf
150 , stripPrefix
152 -- ** Map and Fold
153 , mapM_
155 -- ** Conversions
156 -- | Transform a stream into another type.
157 , toList
158 , toListRev
159 , toStreamK
160 , toStreamD
162 -- * Transformation
163 -- ** By folding (scans)
164 , scanlM'
165 , scanl'
166 , scanlM
167 , scanl
168 , scanl1M'
169 , scanl1'
170 , scanl1M
171 , scanl1
173 , prescanl'
174 , prescanlM'
176 , postscanl
177 , postscanlM
178 , postscanl'
179 , postscanlM'
181 , postscanlx'
182 , postscanlMx'
183 , scanlMx'
184 , scanlx'
186 -- * Filtering
187 , filter
188 , filterM
189 , uniq
190 , take
191 , takeWhile
192 , takeWhileM
193 , drop
194 , dropWhile
195 , dropWhileM
197 -- * Mapping
198 , map
199 , mapM
200 , sequence
202 -- * Inserting
203 , insertBy
205 -- * Deleting
206 , deleteBy
208 -- ** Map and Filter
209 , mapMaybe
210 , mapMaybeM
212 -- * Zipping
213 , indexed
214 , indexedR
215 , zipWith
216 , zipWithM
218 -- * Comparisons
219 , eqBy
220 , cmpBy
222 -- * Merging
223 , mergeBy
224 , mergeByM
226 -- * Transformation comprehensions
227 , the
229 -- * Exceptions
230 , before
231 , after
232 , bracket
233 , onException
234 , finally
235 , handle
236 )
237where
239import Control.Exception (Exception)
240import Control.Monad.Catch (MonadCatch)
241import Control.Monad.IO.Class (MonadIO(..))
242import Control.Monad.Trans (MonadTrans(lift))
243import Data.Word (Word32)
244import Data.Bits (shiftR, shiftL, (.|.), (.&.))
245import Data.Maybe (fromJust, isJust)
246import Foreign.Ptr (Ptr)
247import Foreign.Storable (Storable(..))
248import GHC.Types (SPEC(..))
249import Prelude
250 hiding (map, mapM, mapM_, repeat, foldr, last, take, filter,
251 takeWhile, drop, dropWhile, all, any, maximum, minimum, elem,
252 notElem, null, head, tail, zipWith, lookup, foldr1, sequence,
253 (!!), scanl, scanl1, concatMap, replicate, enumFromTo, concat,
254 reverse)
256import qualified Control.Monad.Catch as MC
258import Streamly.Mem.Array.Types (Array(..))
259import Streamly.Fold.Types (Fold(..))
260import Streamly.SVar (MonadAsync, defState, adaptState)
261import Streamly.Sink.Types (Sink(..))
263import Streamly.Streams.StreamD.Type
265import qualified Streamly.Mem.Array.Types as A
266import qualified Streamly.Mem.Ring as RB
267import qualified Streamly.Streams.StreamK as K
269------------------------------------------------------------------------------
270-- Construction
271------------------------------------------------------------------------------
273-- | An empty 'Stream'.
274{-# INLINE_NORMAL nil #-}
275nil :: Monad m => Stream m a
276nil = Stream (\_ _ -> return Stop) ()
278{-# INLINE_NORMAL consM #-}
279consM :: Monad m => m a -> Stream m a -> Stream m a
280consM m (Stream step state) = Stream step1 Nothing
281 where
282 {-# INLINE_LATE step1 #-}
283 step1 _ Nothing = m >>= \x -> return $ Yield x (Just state)
284 step1 gst (Just st) = do
285 r <- step gst st
286 return $
287 case r of
288 Yield a s -> Yield a (Just s)
289 Skip s -> Skip (Just s)
290 Stop -> Stop
292-- XXX implement in terms of consM?
293-- cons x = consM (return x)
294--
295-- | Can fuse but has O(n^2) complexity.
296{-# INLINE_NORMAL cons #-}
297cons :: Monad m => a -> Stream m a -> Stream m a
298cons x (Stream step state) = Stream step1 Nothing
299 where
300 {-# INLINE_LATE step1 #-}
301 step1 _ Nothing = return $ Yield x (Just state)
302 step1 gst (Just st) = do
303 r <- step gst st
304 return $
305 case r of
306 Yield a s -> Yield a (Just s)
307 Skip s -> Skip (Just s)
308 Stop -> Stop
310-------------------------------------------------------------------------------
311-- Deconstruction
312-------------------------------------------------------------------------------
314-- Does not fuse, has the same performance as the StreamK version.
315{-# INLINE_NORMAL uncons #-}
316uncons :: Monad m => Stream m a -> m (Maybe (a, Stream m a))
317uncons (UnStream step state) = go state
318 where
319 go st = do
320 r <- step defState st
321 case r of
322 Yield x s -> return $ Just (x, Stream step s)
323 Skip s -> go s
324 Stop -> return Nothing
326------------------------------------------------------------------------------
327-- Generation by unfold
328------------------------------------------------------------------------------
330{-# INLINE_NORMAL unfoldrM #-}
331unfoldrM :: Monad m => (s -> m (Maybe (a, s))) -> s -> Stream m a
332unfoldrM next state = Stream step state
333 where
334 {-# INLINE_LATE step #-}
335 step _ st = do
336 r <- next st
337 return $ case r of
338 Just (x, s) -> Yield x s
339 Nothing -> Stop
341{-# INLINE_LATE unfoldr #-}
342unfoldr :: Monad m => (s -> Maybe (a, s)) -> s -> Stream m a
343unfoldr f = unfoldrM (return . f)
345------------------------------------------------------------------------------
346-- Specialized Generation
347------------------------------------------------------------------------------
349repeat :: Monad m => a -> Stream m a
350repeat x = Stream (\_ _ -> return $ Yield x ()) ()
352{-# INLINE_NORMAL replicateM #-}
353replicateM :: Monad m => Int -> m a -> Stream m a
354replicateM n p = Stream step n
355 where
356 {-# INLINE_LATE step #-}
357 step _ i | i <= 0 = return Stop
358 | otherwise = do
359 x <- p
360 return $ Yield x (i - 1)
362{-# INLINE_NORMAL replicate #-}
363replicate :: Monad m => Int -> a -> Stream m a
364replicate n x = replicateM n (return x)
366-- This would not work properly for floats, therefore we put an Integral
367-- constraint.
368-- | Can be used to enumerate unbounded integrals. This does not check for
369-- overflow or underflow for bounded integrals.
370{-# INLINE_NORMAL enumerateFromStepIntegral #-}
371enumerateFromStepIntegral :: (Integral a, Monad m) => a -> a -> Stream m a
372enumerateFromStepIntegral from stride =
373 from `seq` stride `seq` Stream step from
374 where
375 {-# INLINE_LATE step #-}
376 step _ !x = return $ Yield x $! (x + stride)
378-- We are assuming that "to" is constrained by the type to be within
379-- max/min bounds.
380{-# INLINE enumerateFromToIntegral #-}
381enumerateFromToIntegral :: (Monad m, Integral a) => a -> a -> Stream m a
382enumerateFromToIntegral from to =
383 takeWhile (<= to) $ enumerateFromStepIntegral from 1
385{-# INLINE enumerateFromIntegral #-}
386enumerateFromIntegral :: (Monad m, Integral a, Bounded a) => a -> Stream m a
387enumerateFromIntegral from = enumerateFromToIntegral from maxBound
389data EnumState a = EnumInit | EnumYield a a a | EnumStop
391{-# INLINE_NORMAL enumerateFromThenToIntegralUp #-}
392enumerateFromThenToIntegralUp
393 :: (Monad m, Integral a)
394 => a -> a -> a -> Stream m a
395enumerateFromThenToIntegralUp from next to = Stream step EnumInit
396 where
397 {-# INLINE_LATE step #-}
398 step _ EnumInit =
399 return $
400 if to < next
401 then if to < from
402 then Stop
403 else Yield from EnumStop
404 else -- from <= next <= to
405 let stride = next - from
406 in Skip $ EnumYield from stride (to - stride)
408 step _ (EnumYield x stride toMinus) =
409 return $
410 if x > toMinus
411 then Yield x EnumStop
412 else Yield x $ EnumYield (x + stride) stride toMinus
414 step _ EnumStop = return Stop
416{-# INLINE_NORMAL enumerateFromThenToIntegralDn #-}
417enumerateFromThenToIntegralDn
418 :: (Monad m, Integral a)
419 => a -> a -> a -> Stream m a
420enumerateFromThenToIntegralDn from next to = Stream step EnumInit
421 where
422 {-# INLINE_LATE step #-}
423 step _ EnumInit =
424 return $ if to > next
425 then if to > from
426 then Stop
427 else Yield from EnumStop
428 else -- from >= next >= to
429 let stride = next - from
430 in Skip $ EnumYield from stride (to - stride)
432 step _ (EnumYield x stride toMinus) =
433 return $
434 if x < toMinus
435 then Yield x EnumStop
436 else Yield x $ EnumYield (x + stride) stride toMinus
438 step _ EnumStop = return Stop
440{-# INLINE_NORMAL enumerateFromThenToIntegral #-}
441enumerateFromThenToIntegral
442 :: (Monad m, Integral a)
443 => a -> a -> a -> Stream m a
444enumerateFromThenToIntegral from next to
445 | next >= from = enumerateFromThenToIntegralUp from next to
446 | otherwise = enumerateFromThenToIntegralDn from next to
448{-# INLINE_NORMAL enumerateFromThenIntegral #-}
449enumerateFromThenIntegral
450 :: (Monad m, Integral a, Bounded a)
451 => a -> a -> Stream m a
452enumerateFromThenIntegral from next =
453 if next > from
454 then enumerateFromThenToIntegralUp from next maxBound
455 else enumerateFromThenToIntegralDn from next minBound
457-- For floating point numbers if the increment is less than the precision then
458-- it just gets lost. Therefore we cannot always increment it correctly by just
459-- repeated addition.
460-- 9007199254740992 + 1 + 1 :: Double => 9.007199254740992e15
461-- 9007199254740992 + 2 :: Double => 9.007199254740994e15
463-- Instead we accumulate the increment counter and compute the increment
464-- everytime before adding it to the starting number.
465--
466-- This works for Integrals as well as floating point numbers, but
467-- enumerateFromStepIntegral is faster for integrals.
468{-# INLINE_NORMAL enumerateFromStepNum #-}
469enumerateFromStepNum :: (Monad m, Num a) => a -> a -> Stream m a
470enumerateFromStepNum from stride = Stream step 0
471 where
472 {-# INLINE_LATE step #-}
473 step _ !i = return $ (Yield $! (from + i * stride)) $! (i + 1)
475{-# INLINE_NORMAL numFrom #-}
476numFrom :: (Monad m, Num a) => a -> Stream m a
477numFrom from = enumerateFromStepNum from 1
479{-# INLINE_NORMAL numFromThen #-}
480numFromThen :: (Monad m, Num a) => a -> a -> Stream m a
481numFromThen from next = enumerateFromStepNum from (next - from)
483-- We cannot write a general function for Num. The only way to write code
484-- portable between the two is to use a 'Real' constraint and convert between
485-- Fractional and Integral using fromRational which is horribly slow.
486{-# INLINE_NORMAL enumerateFromToFractional #-}
487enumerateFromToFractional
488 :: (Monad m, Fractional a, Ord a)
489 => a -> a -> Stream m a
490enumerateFromToFractional from to =
491 takeWhile (<= to + 1 / 2) $ enumerateFromStepNum from 1
493{-# INLINE_NORMAL enumerateFromThenToFractional #-}
494enumerateFromThenToFractional
495 :: (Monad m, Fractional a, Ord a)
496 => a -> a -> a -> Stream m a
497enumerateFromThenToFractional from next to =
498 takeWhile predicate $ numFromThen from next
499 where
500 mid = (next - from) / 2
501 predicate | next >= from = (<= to + mid)
502 | otherwise = (>= to + mid)
504-------------------------------------------------------------------------------
505-- Generation by Conversion
506-------------------------------------------------------------------------------
508{-# INLINE_NORMAL fromIndicesM #-}
509fromIndicesM :: Monad m => (Int -> m a) -> Stream m a
510fromIndicesM gen = Stream step 0
511 where
512 {-# INLINE_LATE step #-}
513 step _ i = do
514 x <- gen i
515 return $ Yield x (i + 1)
517{-# INLINE fromIndices #-}
518fromIndices :: Monad m => (Int -> a) -> Stream m a
519fromIndices gen = fromIndicesM (return . gen)
521{-# INLINE_NORMAL generateM #-}
522generateM :: Monad m => Int -> (Int -> m a) -> Stream m a
523generateM n gen = n `seq` Stream step 0
524 where
525 {-# INLINE_LATE step #-}
526 step _ i | i < n = do
527 x <- gen i
528 return $ Yield x (i + 1)
529 | otherwise = return Stop
531{-# INLINE generate #-}
532generate :: Monad m => Int -> (Int -> a) -> Stream m a
533generate n gen = generateM n (return . gen)
535-- XXX we need the MonadAsync constraint because of a rewrite rule.
536-- | Convert a list of monadic actions to a 'Stream'
537{-# INLINE_LATE fromListM #-}
538fromListM :: MonadAsync m => [m a] -> Stream m a
539fromListM = Stream step
540 where
541 {-# INLINE_LATE step #-}
542 step _ (m:ms) = m >>= \x -> return $ Yield x ms
543 step _ [] = return Stop
545{-# INLINE toStreamD #-}
546toStreamD :: (K.IsStream t, Monad m) => t m a -> Stream m a
547toStreamD = fromStreamK . K.toStream
549------------------------------------------------------------------------------
550-- Elimination by Folds
551------------------------------------------------------------------------------
553------------------------------------------------------------------------------
554-- Right Folds
555------------------------------------------------------------------------------
557{-# INLINE_NORMAL foldr1 #-}
558foldr1 :: Monad m => (a -> a -> a) -> Stream m a -> m (Maybe a)
559foldr1 f m = do
560 r <- uncons m
561 case r of
562 Nothing -> return Nothing
563 Just (h, t) -> fmap Just (foldr f h t)
565------------------------------------------------------------------------------
566-- Left Folds
567------------------------------------------------------------------------------
569{-# INLINE_NORMAL foldlT #-}
570foldlT :: (Monad m, Monad (s m), MonadTrans s)
571 => (s m b -> a -> s m b) -> s m b -> Stream m a -> s m b
572foldlT fstep begin (Stream step state) = go SPEC begin state
573 where
574 go !_ acc st = do
575 r <- lift $ step defState st
576 case r of
577 Yield x s -> go SPEC (fstep acc x) s
578 Skip s -> go SPEC acc s
579 Stop -> acc
581-- Note, this is going to have horrible performance, because of the nature of
582-- the stream type (i.e. direct stream vs CPS). Its only for reference, it is
583-- likely be practically unusable.
584{-# INLINE_NORMAL foldlS #-}
585foldlS :: Monad m
586 => (Stream m b -> a -> Stream m b) -> Stream m b -> Stream m a -> Stream m b
587foldlS fstep begin (Stream step state) = Stream step' (Left (state, begin))
588 where
589 step' gst (Left (st, acc)) = do
590 r <- step (adaptState gst) st
591 return $ case r of
592 Yield x s -> Skip (Left (s, fstep acc x))
593 Skip s -> Skip (Left (s, acc))
594 Stop -> Skip (Right acc)
596 step' gst (Right (Stream stp stt)) = do
597 r <- stp (adaptState gst) stt
598 return $ case r of
599 Yield x s -> Yield x (Right (Stream stp s))
600 Skip s -> Skip (Right (Stream stp s))
601 Stop -> Stop
603------------------------------------------------------------------------------
604-- Specialized Folds
605------------------------------------------------------------------------------
607-- | Run a streaming composition, discard the results.
608{-# INLINE_LATE drain #-}
609drain :: Monad m => Stream m a -> m ()
610-- drain m = foldrM (\x xs -> x `seq` xs) (return ()) m
611drain (Stream step state) = go SPEC state
612 where
613 go !_ st = do
614 r <- step defState st
615 case r of
616 Yield _ s -> go SPEC s
617 Skip s -> go SPEC s
618 Stop -> return ()
620{-# INLINE_NORMAL null #-}
621null :: Monad m => Stream m a -> m Bool
622null m = foldrM (\_ _ -> return False) (return True) m
624-- XXX SPEC?
625{-# INLINE_NORMAL head #-}
626head :: Monad m => Stream m a -> m (Maybe a)
627head m = foldrM (\x _ -> return (Just x)) (return Nothing) m
629-- Does not fuse, has the same performance as the StreamK version.
630{-# INLINE_NORMAL tail #-}
631tail :: Monad m => Stream m a -> m (Maybe (Stream m a))
632tail (UnStream step state) = go state
633 where
634 go st = do
635 r <- step defState st
636 case r of
637 Yield _ s -> return (Just $ Stream step s)
638 Skip s -> go s
639 Stop -> return Nothing
641-- XXX will it fuse? need custom impl?
642{-# INLINE_NORMAL last #-}
643last :: Monad m => Stream m a -> m (Maybe a)
644last = foldl' (\_ y -> Just y) Nothing
646{-# INLINE_NORMAL elem #-}
647elem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
648-- elem e m = foldrM (\x xs -> if x == e then return True else xs) (return False) m
649elem e (Stream step state) = go state
650 where
651 go st = do
652 r <- step defState st
653 case r of
654 Yield x s
655 | x == e -> return True
656 | otherwise -> go s
657 Skip s -> go s
658 Stop -> return False
660{-# INLINE_NORMAL notElem #-}
661notElem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
662notElem e s = fmap not (elem e s)
664{-# INLINE_NORMAL all #-}
665all :: Monad m => (a -> Bool) -> Stream m a -> m Bool
666-- all p m = foldrM (\x xs -> if p x then xs else return False) (return True) m
667all p (Stream step state) = go state
668 where
669 go st = do
670 r <- step defState st
671 case r of
672 Yield x s
673 | p x -> go s
674 | otherwise -> return False
675 Skip s -> go s
676 Stop -> return True
678{-# INLINE_NORMAL any #-}
679any :: Monad m => (a -> Bool) -> Stream m a -> m Bool
680-- any p m = foldrM (\x xs -> if p x then return True else xs) (return False) m
681any p (Stream step state) = go state
682 where
683 go st = do
684 r <- step defState st
685 case r of
686 Yield x s
687 | p x -> return True
688 | otherwise -> go s
689 Skip s -> go s
690 Stop -> return False
692{-# INLINE_NORMAL maximum #-}
693maximum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
694maximum (Stream step state) = go Nothing state
695 where
696 go Nothing st = do
697 r <- step defState st
698 case r of
699 Yield x s -> go (Just x) s
700 Skip s -> go Nothing s
701 Stop -> return Nothing
702 go (Just acc) st = do
703 r <- step defState st
704 case r of
705 Yield x s
706 | acc <= x -> go (Just x) s
707 | otherwise -> go (Just acc) s
708 Skip s -> go (Just acc) s
709 Stop -> return (Just acc)
711{-# INLINE_NORMAL maximumBy #-}
712maximumBy :: Monad m => (a -> a -> Ordering) -> Stream m a -> m (Maybe a)
713maximumBy cmp (Stream step state) = go Nothing state
714 where
715 go Nothing st = do
716 r <- step defState st
717 case r of
718 Yield x s -> go (Just x) s
719 Skip s -> go Nothing s
720 Stop -> return Nothing
721 go (Just acc) st = do
722 r <- step defState st
723 case r of
724 Yield x s -> case cmp acc x of
725 GT -> go (Just acc) s
726 _ -> go (Just x) s
727 Skip s -> go (Just acc) s
728 Stop -> return (Just acc)
730{-# INLINE_NORMAL minimum #-}
731minimum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
732minimum (Stream step state) = go Nothing state
733 where
734 go Nothing st = do
735 r <- step defState st
736 case r of
737 Yield x s -> go (Just x) s
738 Skip s -> go Nothing s
739 Stop -> return Nothing
740 go (Just acc) st = do
741 r <- step defState st
742 case r of
743 Yield x s
744 | acc <= x -> go (Just acc) s
745 | otherwise -> go (Just x) s
746 Skip s -> go (Just acc) s
747 Stop -> return (Just acc)
749{-# INLINE_NORMAL minimumBy #-}
750minimumBy :: Monad m => (a -> a -> Ordering) -> Stream m a -> m (Maybe a)
751minimumBy cmp (Stream step state) = go Nothing state
752 where
753 go Nothing st = do
754 r <- step defState st
755 case r of
756 Yield x s -> go (Just x) s
757 Skip s -> go Nothing s
758 Stop -> return Nothing
759 go (Just acc) st = do
760 r <- step defState st
761 case r of
762 Yield x s -> case cmp acc x of
763 GT -> go (Just x) s
764 _ -> go (Just acc) s
765 Skip s -> go (Just acc) s
766 Stop -> return (Just acc)
768{-# INLINE_NORMAL (!!) #-}
769(!!) :: (Monad m) => Stream m a -> Int -> m (Maybe a)
770(Stream step state) !! i = go i state
771 where
772 go n st = do
773 r <- step defState st
774 case r of
775 Yield x s | n < 0 -> return Nothing
776 | n == 0 -> return $ Just x
777 | otherwise -> go (n - 1) s
778 Skip s -> go n s
779 Stop -> return Nothing
781{-# INLINE_NORMAL lookup #-}
782lookup :: (Monad m, Eq a) => a -> Stream m (a, b) -> m (Maybe b)
783lookup e m = foldrM (\(a, b) xs -> if e == a then return (Just b) else xs)
784 (return Nothing) m
786{-# INLINE_NORMAL findM #-}
787findM :: Monad m => (a -> m Bool) -> Stream m a -> m (Maybe a)
788findM p m = foldrM (\x xs -> p x >>= \r -> if r then return (Just x) else xs)
789 (return Nothing) m
791{-# INLINE find #-}
792find :: Monad m => (a -> Bool) -> Stream m a -> m (Maybe a)
793find p = findM (return . p)
795{-# INLINE_NORMAL findIndices #-}
796findIndices :: Monad m => (a -> Bool) -> Stream m a -> Stream m Int
797findIndices p (Stream step state) = Stream step' (state, 0)
798 where
799 {-# INLINE_LATE step' #-}
800 step' gst (st, i) = do
801 r <- step (adaptState gst) st
802 return $ case r of
803 Yield x s -> if p x then Yield i (s, i+1) else Skip (s, i+1)
804 Skip s -> Skip (s, i+1)
805 Stop -> Stop
807{-# INLINE toListRev #-}
808toListRev :: Monad m => Stream m a -> m [a]
809toListRev = foldl' (flip (:)) []
811-- We can implement reverse as:
812--
813-- > reverse = foldlS (flip cons) nil
814--
815-- However, this implementation is unusable because of the horrible performance
816-- of cons. So we just convert it to a list first and then stream from the
817-- list.
818--
819-- XXX Maybe we can use an Array instead of a list here?
820{-# INLINE_NORMAL reverse #-}
821reverse :: Monad m => Stream m a -> Stream m a
822reverse m = Stream step Nothing
823 where
824 {-# INLINE_LATE step #-}
825 step _ Nothing = do
826 xs <- toListRev m
827 return $ Skip (Just xs)
828 step _ (Just (x:xs)) = return $ Yield x (Just xs)
829 step _ (Just []) = return Stop
831-- Much faster reverse for Storables
832{-# INLINE_NORMAL reverse' #-}
833reverse' :: forall m a. (MonadIO m, Storable a) => Stream m a -> Stream m a
834{-
835-- This commented implementation copies the whole stream into one single array
836-- and then streams from that array, this is 3-4x faster than the chunked code
837-- that follows. Though this could be problematic due to unbounded large
838-- allocations. We need to figure out why the chunked code is slower and if we
839-- can optimize the chunked code to work as fast as this one. It may be a
840-- fusion issue?
841import Foreign.ForeignPtr (touchForeignPtr)
842import Foreign.ForeignPtr.Unsafe (unsafeForeignPtrToPtr)
843import Foreign.Ptr (Ptr, plusPtr)
844reverse' m = Stream step Nothing
845 where
846 {-# INLINE_LATE step #-}
847 step _ Nothing = do
848 arr <- A.fromStreamD m
849 let p = aEnd arr `plusPtr` negate (sizeOf (undefined :: a))
850 return $ Skip $ Just (aStart arr, p)
852 step _ (Just (start, p)) | p < unsafeForeignPtrToPtr start = return Stop
854 step _ (Just (start, p)) = do
855 let !x = A.unsafeInlineIO $ do
856 r <- peek p
857 touchForeignPtr start
858 return r
859 next = p `plusPtr` negate (sizeOf (undefined :: a))
860 return $ Yield x (Just (start, next))
861-}
862reverse' m =
863 A.flattenArraysRev
864 $ fromStreamK
865 $ K.reverse
866 $ toStreamK
867 $ A.fromStreamDArraysOf A.defaultChunkSize m
870------------------------------------------------------------------------------
871-- Grouping/Splitting
872------------------------------------------------------------------------------
874{-# INLINE_NORMAL splitSuffixBy' #-}
875splitSuffixBy' :: Monad m
876 => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
877splitSuffixBy' predicate f (Stream step state) =
878 Stream (stepOuter f) (Just state)
880 where
882 {-# INLINE_LATE stepOuter #-}
883 stepOuter (Fold fstep initial done) gst (Just st) = do
884 res <- step (adaptState gst) st
885 case res of
886 Yield x s -> do
887 acc <- initial
888 acc' <- fstep acc x
889 if (predicate x)
890 then done acc' >>= \val -> return $ Yield val (Just s)
891 else go SPEC s acc'
893 Skip s -> return $ Skip $ Just s
894 Stop -> return Stop
896 where
898 go !_ stt !acc = do
899 res <- step (adaptState gst) stt
900 case res of
901 Yield x s -> do
902 acc' <- fstep acc x
903 if (predicate x)
904 then done acc' >>= \val -> return $ Yield val (Just s)
905 else go SPEC s acc'
906 Skip s -> go SPEC s acc
907 Stop -> done acc >>= \val -> return $ Yield val Nothing
909 stepOuter _ _ Nothing = return Stop
911{-# INLINE_NORMAL groupsBy #-}
912groupsBy :: Monad m
913 => (a -> a -> Bool)
914 -> Fold m a b
915 -> Stream m a
916 -> Stream m b
917groupsBy cmp f (Stream step state) = Stream (stepOuter f) (Just state, Nothing)
919 where
921 {-# INLINE_LATE stepOuter #-}
922 stepOuter (Fold fstep initial done) gst (Just st, Nothing) = do
923 res <- step (adaptState gst) st
924 case res of
925 Yield x s -> do
926 acc <- initial
927 acc' <- fstep acc x
928 go SPEC x s acc'
930 Skip s -> return $ Skip $ (Just s, Nothing)
931 Stop -> return Stop
933 where
935 go !_ prev stt !acc = do
936 res <- step (adaptState gst) stt
937 case res of
938 Yield x s -> do
939 if cmp x prev
940 then do
941 acc' <- fstep acc x
942 go SPEC prev s acc'
943 else done acc >>= \r -> return $ Yield r (Just s, Just x)
944 Skip s -> go SPEC prev s acc
945 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
947 stepOuter (Fold fstep initial done) gst (Just st, Just prev) = do
948 acc <- initial
949 acc' <- fstep acc prev
950 go SPEC st acc'
952 where
954 -- XXX code duplicated from the previous equation
955 go !_ stt !acc = do
956 res <- step (adaptState gst) stt
957 case res of
958 Yield x s -> do
959 if cmp x prev
960 then do
961 acc' <- fstep acc x
962 go SPEC s acc'
963 else done acc >>= \r -> return $ Yield r (Just s, Just x)
964 Skip s -> go SPEC s acc
965 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
967 stepOuter _ _ (Nothing,_) = return Stop
969{-# INLINE_NORMAL groupsRollingBy #-}
970groupsRollingBy :: Monad m
971 => (a -> a -> Bool)
972 -> Fold m a b
973 -> Stream m a
974 -> Stream m b
975groupsRollingBy cmp f (Stream step state) =
976 Stream (stepOuter f) (Just state, Nothing)
977 where
979 {-# INLINE_LATE stepOuter #-}
980 stepOuter (Fold fstep initial done) gst (Just st, Nothing) = do
981 res <- step (adaptState gst) st
982 case res of
983 Yield x s -> do
984 acc <- initial
985 acc' <- fstep acc x
986 go SPEC x s acc'
988 Skip s -> return $ Skip $ (Just s, Nothing)
989 Stop -> return Stop
991 where
992 go !_ prev stt !acc = do
993 res <- step (adaptState gst) stt
994 case res of
995 Yield x s -> do
996 if cmp prev x
997 then do
998 acc' <- fstep acc x
999 go SPEC x s acc'
1000 else
1001 done acc >>= \r -> return $ Yield r (Just s, Just x)
1002 Skip s -> go SPEC prev s acc
1003 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1005 stepOuter (Fold fstep initial done) gst (Just st, Just prev') = do
1006 acc <- initial
1007 acc' <- fstep acc prev'
1008 go SPEC prev' st acc'
1010 where
1011 go !_ prevv stt !acc = do
1012 res <- step (adaptState gst) stt
1013 case res of
1014 Yield x s -> do
1015 if cmp prevv x
1016 then do
1017 acc' <- fstep acc x
1018 go SPEC x s acc'
1019 else done acc >>= \r -> return $ Yield r (Just s, Just x)
1020 Skip s -> go SPEC prevv s acc
1021 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1023 stepOuter _ _ (Nothing, _) = return Stop
1025{-# INLINE_NORMAL splitBy #-}
1026splitBy :: Monad m => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1027splitBy predicate f (Stream step state) = Stream (step' f) (Just state)
1029 where
1031 {-# INLINE_LATE step' #-}
1032 step' (Fold fstep initial done) gst (Just st) = initial >>= go SPEC st
1034 where
1036 go !_ stt !acc = do
1037 res <- step (adaptState gst) stt
1038 case res of
1039 Yield x s -> do
1040 if predicate x
1041 then done acc >>= \r -> return $ Yield r (Just s)
1042 else do
1043 acc' <- fstep acc x
1044 go SPEC s acc'
1045 Skip s -> go SPEC s acc
1046 Stop -> done acc >>= \r -> return $ Yield r Nothing
1048 step' _ _ Nothing = return Stop
1050{-# INLINE_NORMAL splitSuffixBy #-}
1051splitSuffixBy :: Monad m
1052 => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1053splitSuffixBy predicate f (Stream step state) = Stream (step' f) (Just state)
1055 where
1057 {-# INLINE_LATE step' #-}
1058 step' (Fold fstep initial done) gst (Just st) = do
1059 res <- step (adaptState gst) st
1060 case res of
1061 Yield x s -> do
1062 acc <- initial
1063 if predicate x
1064 then done acc >>= \val -> return $ Yield val (Just s)
1065 else do
1066 acc' <- fstep acc x
1067 go SPEC s acc'
1069 Skip s -> return $ Skip $ Just s
1070 Stop -> return Stop
1072 where
1074 go !_ stt !acc = do
1075 res <- step (adaptState gst) stt
1076 case res of
1077 Yield x s -> do
1078 if predicate x
1079 then done acc >>= \r -> return $ Yield r (Just s)
1080 else do
1081 acc' <- fstep acc x
1082 go SPEC s acc'
1083 Skip s -> go SPEC s acc
1084 Stop -> done acc >>= \r -> return $ Yield r Nothing
1086 step' _ _ Nothing = return Stop
1088{-# INLINE_NORMAL wordsBy #-}
1089wordsBy :: Monad m => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1090wordsBy predicate f (Stream step state) = Stream (stepOuter f) (Just state)
1092 where
1094 {-# INLINE_LATE stepOuter #-}
1095 stepOuter (Fold fstep initial done) gst (Just st) = do
1096 res <- step (adaptState gst) st
1097 case res of
1098 Yield x s -> do
1099 if predicate x
1100 then return $ Skip (Just s)
1101 else do
1102 acc <- initial
1103 acc' <- fstep acc x
1104 go SPEC s acc'
1106 Skip s -> return $ Skip $ Just s
1107 Stop -> return Stop
1109 where
1111 go !_ stt !acc = do
1112 res <- step (adaptState gst) stt
1113 case res of
1114 Yield x s -> do
1115 if predicate x
1116 then done acc >>= \r -> return $ Yield r (Just s)
1117 else do
1118 acc' <- fstep acc x
1119 go SPEC s acc'
1120 Skip s -> go SPEC s acc
1121 Stop -> done acc >>= \r -> return $ Yield r Nothing
1123 stepOuter _ _ Nothing = return Stop
1125-- String search algorithms:
1126-- http://www-igm.univ-mlv.fr/~lecroq/string/index.html
1129-- TODO can we unify the splitting operations using a splitting configuration
1130-- like in the split package.
1132data SplitStyle = Infix | Suffix | Prefix deriving (Eq, Show)
1134data SplitOptions = SplitOptions
1135 { style :: SplitStyle
1136 , withSep :: Bool -- ^ keep the separators in output
1137 -- , compact :: Bool -- ^ treat multiple consecutive separators as one
1138 -- , trimHead :: Bool -- ^ drop blank at head
1139 -- , trimTail :: Bool -- ^ drop blank at tail
1143data SplitOnState s a =
1144 GO_START
1145 | GO_EMPTY_PAT s
1146 | GO_SINGLE_PAT s a
1147 | GO_SHORT_PAT s
1148 | GO_KARP_RABIN s !(RB.Ring a) !(Ptr a)
1149 | GO_DONE
1151{-# INLINE_NORMAL splitOn #-}
1152splitOn
1153 :: forall m a b. (MonadIO m, Storable a, Enum a, Eq a)
1154 => Array a
1155 -> Fold m a b
1156 -> Stream m a
1157 -> Stream m b
1158splitOn patArr@Array{..} (Fold fstep initial done) (Stream step state) =
1159 Stream stepOuter GO_START
1161 where
1163 patLen = A.length patArr
1164 maxIndex = patLen - 1
1165 elemBits = sizeOf (undefined :: a) * 8
1167 {-# INLINE_LATE stepOuter #-}
1168 stepOuter _ GO_START =
1169 if patLen == 0
1170 then return $ Skip $ GO_EMPTY_PAT state
1171 else if patLen == 1
1172 then do
1173 r <- liftIO $ (A.unsafeIndexIO patArr 0)
1174 return $ Skip $ GO_SINGLE_PAT state r
1175 else if sizeOf (undefined :: a) * patLen
1176 <= sizeOf (undefined :: Word)
1177 then return $ Skip $ GO_SHORT_PAT state
1178 else do
1179 (rb, rhead) <- liftIO $ RB.new patLen
1180 return $ Skip $ GO_KARP_RABIN state rb rhead
1182 stepOuter gst (GO_SINGLE_PAT stt pat) = initial >>= go SPEC stt
1184 where
1186 go !_ st !acc = do
1187 res <- step (adaptState gst) st
1188 case res of
1189 Yield x s -> do
1190 if pat == x
1191 then do
1192 r <- done acc
1193 return $ Yield r (GO_SINGLE_PAT s pat)
1194 else fstep acc x >>= go SPEC s
1195 Skip s -> go SPEC s acc
1196 Stop -> done acc >>= \r -> return $ Yield r GO_DONE
1198 stepOuter gst (GO_SHORT_PAT stt) = initial >>= go0 SPEC 0 (0 :: Word) stt
1200 where
1202 mask :: Word
1203 mask = (1 `shiftL` (elemBits * patLen)) - 1
1205 addToWord wrd a = (wrd `shiftL` elemBits) .|. fromIntegral (fromEnum a)
1207 patWord :: Word
1208 patWord = mask .&. A.foldl' addToWord 0 patArr
1210 go0 !_ !idx wrd st !acc = do
1211 res <- step (adaptState gst) st
1212 case res of
1213 Yield x s -> do
1214 let wrd' = addToWord wrd x
1215 if idx == maxIndex
1216 then do
1217 if wrd' .&. mask == patWord
1218 then do
1219 r <- done acc
1220 return $ Yield r (GO_SHORT_PAT s)
1221 else go1 SPEC wrd' s acc
1222 else go0 SPEC (idx + 1) wrd' s acc
1223 Skip s -> go0 SPEC idx wrd s acc
1224 Stop -> do
1225 acc' <- if idx /= 0
1226 then go2 wrd idx acc
1227 else return acc
1228 done acc' >>= \r -> return $ Yield r GO_DONE
1230 {-# INLINE go1 #-}
1231 go1 !_ wrd st !acc = do
1232 res <- step (adaptState gst) st
1233 case res of
1234 Yield x s -> do
1235 let wrd' = addToWord wrd x
1236 old = (mask .&. wrd) `shiftR` (elemBits * (patLen - 1))
1237 acc' <- fstep acc (toEnum $ fromIntegral old)
1238 if wrd' .&. mask == patWord
1239 then done acc' >>= \r -> return $ Yield r (GO_SHORT_PAT s)
1240 else go1 SPEC wrd' s acc'
1241 Skip s -> go1 SPEC wrd s acc
1242 Stop -> do
1243 acc' <- go2 wrd patLen acc
1244 done acc' >>= \r -> return $ Yield r GO_DONE
1246 go2 !wrd !n !acc | n > 0 = do
1247 let old = (mask .&. wrd) `shiftR` (elemBits * (n - 1))
1248 fstep acc (toEnum $ fromIntegral old) >>= go2 wrd (n - 1)
1249 go2 _ _ acc = return acc
1251 stepOuter gst (GO_KARP_RABIN stt rb rhead) = do
1252 initial >>= go0 SPEC 0 rhead stt
1254 where
1256 k = 2891336453 :: Word32
1257 coeff = k ^ patLen
1258 addCksum cksum a = cksum * k + fromIntegral (fromEnum a)
1259 deltaCksum cksum old new =
1260 addCksum cksum new - coeff * fromIntegral (fromEnum old)
1262 -- XXX shall we use a random starting hash or 1 instead of 0?
1263 patHash = A.foldl' addCksum 0 patArr
1265 -- rh == ringHead
1266 go0 !_ !idx !rh st !acc = do
1267 res <- step (adaptState gst) st
1268 case res of
1269 Yield x s -> do
1270 rh' <- liftIO $ RB.unsafeInsert rb rh x
1271 if idx == maxIndex
1272 then do
1273 let fold = RB.unsafeFoldRing (RB.ringBound rb)
1274 let !ringHash = fold addCksum 0 rb
1275 if ringHash == patHash
1276 then go2 SPEC ringHash rh' s acc
1277 else go1 SPEC ringHash rh' s acc
1278 else go0 SPEC (idx + 1) rh' s acc
1279 Skip s -> go0 SPEC idx rh s acc
1280 Stop -> do
1281 !acc' <- if idx /= 0
1282 then RB.unsafeFoldRingM rh fstep acc rb
1283 else return acc
1284 done acc' >>= \r -> return $ Yield r GO_DONE
1286 -- XXX Theoretically this code can do 4 times faster if GHC generates
1287 -- optimal code. If we use just "(cksum' == patHash)" condition it goes
1288 -- 4x faster, as soon as we add the "RB.unsafeEqArray rb v" condition
1289 -- the generated code changes drastically and becomes 4x slower. Need
1290 -- to investigate what is going on with GHC.
1291 {-# INLINE go1 #-}
1292 go1 !_ !cksum !rh st !acc = do
1293 res <- step (adaptState gst) st
1294 case res of
1295 Yield x s -> do
1296 old <- liftIO $ peek rh
1297 let cksum' = deltaCksum cksum old x
1298 acc' <- fstep acc old
1300 if (cksum' == patHash)
1301 then do
1302 rh' <- liftIO (RB.unsafeInsert rb rh x)
1303 go2 SPEC cksum' rh' s acc'
1304 else do
1305 rh' <- liftIO (RB.unsafeInsert rb rh x)
1306 go1 SPEC cksum' rh' s acc'
1307 Skip s -> go1 SPEC cksum rh s acc
1308 Stop -> do
1309 acc' <- RB.unsafeFoldRingFullM rh fstep acc rb
1310 done acc' >>= \r -> return $ Yield r GO_DONE
1312 go2 !_ !cksum' !rh' s !acc' = do
1313 if RB.unsafeEqArray rb rh' patArr
1314 then do
1315 r <- done acc'
1316 return $ Yield r (GO_KARP_RABIN s rb rhead)
1317 else go1 SPEC cksum' rh' s acc'
1319 stepOuter gst (GO_EMPTY_PAT st) = do
1320 res <- step (adaptState gst) st
1321 case res of
1322 Yield x s -> do
1323 acc <- initial
1324 acc' <- fstep acc x
1325 done acc' >>= \r -> return $ Yield r (GO_EMPTY_PAT s)
1326 Skip s -> return $ Skip (GO_EMPTY_PAT s)
1327 Stop -> return Stop
1329 stepOuter _ GO_DONE = return Stop
1331{-# INLINE_NORMAL splitSuffixOn #-}
1332splitSuffixOn
1333 :: forall m a b. (MonadIO m, Storable a, Enum a, Eq a)
1334 => Bool
1335 -> Array a
1336 -> Fold m a b
1337 -> Stream m a
1338 -> Stream m b
1339splitSuffixOn withSep patArr@Array{..} (Fold fstep initial done)
1340 (Stream step state) =
1341 Stream stepOuter GO_START
1343 where
1345 patLen = A.length patArr
1346 maxIndex = patLen - 1
1347 elemBits = sizeOf (undefined :: a) * 8
1349 {-# INLINE_LATE stepOuter #-}
1350 stepOuter _ GO_START =
1351 if patLen == 0
1352 then return $ Skip $ GO_EMPTY_PAT state
1353 else if patLen == 1
1354 then do
1355 r <- liftIO $ (A.unsafeIndexIO patArr 0)
1356 return $ Skip $ GO_SINGLE_PAT state r
1357 else if sizeOf (undefined :: a) * patLen
1358 <= sizeOf (undefined :: Word)
1359 then return $ Skip $ GO_SHORT_PAT state
1360 else do
1361 (rb, rhead) <- liftIO $ RB.new patLen
1362 return $ Skip $ GO_KARP_RABIN state rb rhead
1364 stepOuter gst (GO_SINGLE_PAT stt pat) = do
1365 -- This first part is the only difference between splitOn and
1366 -- splitSuffixOn.
1367 -- If the last element is a separator do not issue a blank segment.
1368 res <- step (adaptState gst) stt
1369 case res of
1370 Yield x s -> do
1371 acc <- initial
1372 if pat == x
1373 then do
1374 acc' <- if withSep then fstep acc x else return acc
1375 done acc' >>= \r -> return $ Yield r (GO_SINGLE_PAT s pat)
1376 else fstep acc x >>= go SPEC s
1377 Skip s -> return $ Skip $ (GO_SINGLE_PAT s pat)
1378 Stop -> return Stop
1380 where
1382 -- This is identical for splitOn and splitSuffixOn
1383 go !_ st !acc = do
1384 res <- step (adaptState gst) st
1385 case res of
1386 Yield x s -> do
1387 if pat == x
1388 then do
1389 acc' <- if withSep then fstep acc x else return acc
1390 r <- done acc'
1391 return $ Yield r (GO_SINGLE_PAT s pat)
1392 else fstep acc x >>= go SPEC s
1393 Skip s -> go SPEC s acc
1394 Stop -> done acc >>= \r -> return $ Yield r GO_DONE
1396 stepOuter gst (GO_SHORT_PAT stt) = do
1398 -- Call "initial" only if the stream yields an element, otherwise we
1399 -- may call "initial" but never yield anything. initial may produce a
1400 -- side effect, therefore we will end up doing and discard a side
1401 -- effect.
1403 let idx = 0
1404 let wrd = 0
1405 res <- step (adaptState gst) stt
1406 case res of
1407 Yield x s -> do
1408 acc <- initial
1409 let wrd' = addToWord wrd x
1410 acc' <- if withSep then fstep acc x else return acc
1411 if idx == maxIndex
1412 then do
1413 if wrd' .&. mask == patWord
1414 then done acc' >>= \r -> return $ Yield r (GO_SHORT_PAT s)
1415 else go0 SPEC (idx + 1) wrd' s acc'
1416 else go0 SPEC (idx + 1) wrd' s acc'
1417 Skip s -> return $ Skip (GO_SHORT_PAT s)
1418 Stop -> return Stop
1420 where
1422 mask :: Word
1423 mask = (1 `shiftL` (elemBits * patLen)) - 1
1425 addToWord wrd a = (wrd `shiftL` elemBits) .|. fromIntegral (fromEnum a)
1427 patWord :: Word
1428 patWord = mask .&. A.foldl' addToWord 0 patArr
1430 go0 !_ !idx wrd st !acc = do
1431 res <- step (adaptState gst) st
1432 case res of
1433 Yield x s -> do
1434 let wrd' = addToWord wrd x
1435 acc' <- if withSep then fstep acc x else return acc
1436 if idx == maxIndex
1437 then do
1438 if wrd' .&. mask == patWord
1439 then do
1440 r <- done acc'
1441 return $ Yield r (GO_SHORT_PAT s)
1442 else go1 SPEC wrd' s acc'
1443 else go0 SPEC (idx + 1) wrd' s acc'
1444 Skip s -> go0 SPEC idx wrd s acc
1445 Stop -> do
1446 if (idx == maxIndex) && (wrd .&. mask == patWord)
1447 then return Stop
1448 else do
1449 acc' <- if idx /= 0 && not withSep
1450 then go2 wrd idx acc
1451 else return acc
1452 done acc' >>= \r -> return $ Yield r GO_DONE
1454 {-# INLINE go1 #-}
1455 go1 !_ wrd st !acc = do
1456 res <- step (adaptState gst) st
1457 case res of
1458 Yield x s -> do
1459 let wrd' = addToWord wrd x
1460 old = (mask .&. wrd) `shiftR` (elemBits * (patLen - 1))
1461 acc' <- if withSep
1462 then fstep acc x
1463 else fstep acc (toEnum $ fromIntegral old)
1464 if wrd' .&. mask == patWord
1465 then done acc' >>= \r -> return $ Yield r (GO_SHORT_PAT s)
1466 else go1 SPEC wrd' s acc'
1467 Skip s -> go1 SPEC wrd s acc
1468 Stop ->
1469 -- If the last sequence is a separator do not issue a blank
1470 -- segment.
1471 if wrd .&. mask == patWord
1472 then return Stop
1473 else do
1474 acc' <- if withSep
1475 then return acc
1476 else go2 wrd patLen acc
1477 done acc' >>= \r -> return $ Yield r GO_DONE
1479 go2 !wrd !n !acc | n > 0 = do
1480 let old = (mask .&. wrd) `shiftR` (elemBits * (n - 1))
1481 fstep acc (toEnum $ fromIntegral old) >>= go2 wrd (n - 1)
1482 go2 _ _ acc = return acc
1484 stepOuter gst (GO_KARP_RABIN stt rb rhead) = do
1485 let idx = 0
1486 res <- step (adaptState gst) stt
1487 case res of
1488 Yield x s -> do
1489 acc <- initial
1490 acc' <- if withSep then fstep acc x else return acc
1491 rh' <- liftIO (RB.unsafeInsert rb rhead x)
1492 if idx == maxIndex
1493 then do
1494 let fold = RB.unsafeFoldRing (RB.ringBound rb)
1495 let !ringHash = fold addCksum 0 rb
1496 if ringHash == patHash
1497 then go2 SPEC ringHash rh' s acc'
1498 else go0 SPEC (idx + 1) rh' s acc'
1499 else go0 SPEC (idx + 1) rh' s acc'
1500 Skip s -> return $ Skip (GO_KARP_RABIN s rb rhead)
1501 Stop -> return Stop
1503 where
1505 k = 2891336453 :: Word32
1506 coeff = k ^ patLen
1507 addCksum cksum a = cksum * k + fromIntegral (fromEnum a)
1508 deltaCksum cksum old new =
1509 addCksum cksum new - coeff * fromIntegral (fromEnum old)
1511 -- XXX shall we use a random starting hash or 1 instead of 0?
1512 patHash = A.foldl' addCksum 0 patArr
1514 -- rh == ringHead
1515 go0 !_ !idx !rh st !acc = do
1516 res <- step (adaptState gst) st
1517 case res of
1518 Yield x s -> do
1519 acc' <- if withSep then fstep acc x else return acc
1520 rh' <- liftIO (RB.unsafeInsert rb rh x)
1521 if idx == maxIndex
1522 then do
1523 let fold = RB.unsafeFoldRing (RB.ringBound rb)
1524 let !ringHash = fold addCksum 0 rb
1525 if ringHash == patHash
1526 then go2 SPEC ringHash rh' s acc'
1527 else go1 SPEC ringHash rh' s acc'
1528 else go0 SPEC (idx + 1) rh' s acc'
1529 Skip s -> go0 SPEC idx rh s acc
1530 Stop -> do
1531 -- do not issue a blank segment when we end at pattern
1532 if (idx == maxIndex) && RB.unsafeEqArray rb rh patArr
1533 then return Stop
1534 else do
1535 !acc' <- if idx /= 0 && not withSep
1536 then RB.unsafeFoldRingM rh fstep acc rb
1537 else return acc
1538 done acc' >>= \r -> return $ Yield r GO_DONE
1540 -- XXX Theoretically this code can do 4 times faster if GHC generates
1541 -- optimal code. If we use just "(cksum' == patHash)" condition it goes
1542 -- 4x faster, as soon as we add the "RB.unsafeEqArray rb v" condition
1543 -- the generated code changes drastically and becomes 4x slower. Need
1544 -- to investigate what is going on with GHC.
1545 {-# INLINE go1 #-}
1546 go1 !_ !cksum !rh st !acc = do
1547 res <- step (adaptState gst) st
1548 case res of
1549 Yield x s -> do
1550 old <- liftIO $ peek rh
1551 let cksum' = deltaCksum cksum old x
1552 acc' <- if withSep
1553 then fstep acc x
1554 else fstep acc old
1556 if (cksum' == patHash)
1557 then do
1558 rh' <- liftIO (RB.unsafeInsert rb rh x)
1559 go2 SPEC cksum' rh' s acc'
1560 else do
1561 rh' <- liftIO (RB.unsafeInsert rb rh x)
1562 go1 SPEC cksum' rh' s acc'
1563 Skip s -> go1 SPEC cksum rh s acc
1564 Stop -> do
1565 if RB.unsafeEqArray rb rh patArr
1566 then return Stop
1567 else do
1568 acc' <- if withSep
1569 then return acc
1570 else RB.unsafeFoldRingFullM rh fstep acc rb
1571 done acc' >>= \r -> return $ Yield r GO_DONE
1573 go2 !_ !cksum' !rh' s !acc' = do
1574 if RB.unsafeEqArray rb rh' patArr
1575 then do
1576 r <- done acc'
1577 return $ Yield r (GO_KARP_RABIN s rb rhead)
1578 else go1 SPEC cksum' rh' s acc'
1580 stepOuter gst (GO_EMPTY_PAT st) = do
1581 res <- step (adaptState gst) st
1582 case res of
1583 Yield x s -> do
1584 acc <- initial
1585 acc' <- fstep acc x
1586 done acc' >>= \r -> return $ Yield r (GO_EMPTY_PAT s)
1587 Skip s -> return $ Skip (GO_EMPTY_PAT s)
1588 Stop -> return Stop
1590 stepOuter _ GO_DONE = return Stop
1592------------------------------------------------------------------------------
1593-- Substreams
1594------------------------------------------------------------------------------
1596{-# INLINE_NORMAL isPrefixOf #-}
1597isPrefixOf :: (Eq a, Monad m) => Stream m a -> Stream m a -> m Bool
1598isPrefixOf (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
1599 where
1600 go (sa, sb, Nothing) = do
1601 r <- stepa defState sa
1602 case r of
1603 Yield x sa' -> go (sa', sb, Just x)
1604 Skip sa' -> go (sa', sb, Nothing)
1605 Stop -> return True
1607 go (sa, sb, Just x) = do
1608 r <- stepb defState sb
1609 case r of
1610 Yield y sb' ->
1611 if x == y
1612 then go (sa, sb', Nothing)
1613 else return False
1614 Skip sb' -> go (sa, sb', Just x)
1615 Stop -> return False
1617{-# INLINE_NORMAL isSubsequenceOf #-}
1618isSubsequenceOf :: (Eq a, Monad m) => Stream m a -> Stream m a -> m Bool
1619isSubsequenceOf (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
1620 where
1621 go (sa, sb, Nothing) = do
1622 r <- stepa defState sa
1623 case r of
1624 Yield x sa' -> go (sa', sb, Just x)
1625 Skip sa' -> go (sa', sb, Nothing)
1626 Stop -> return True
1628 go (sa, sb, Just x) = do
1629 r <- stepb defState sb
1630 case r of
1631 Yield y sb' ->
1632 if x == y
1633 then go (sa, sb', Nothing)
1634 else go (sa, sb', Just x)
1635 Skip sb' -> go (sa, sb', Just x)
1636 Stop -> return False
1638{-# INLINE_NORMAL stripPrefix #-}
1639stripPrefix
1640 :: (Eq a, Monad m)
1641 => Stream m a -> Stream m a -> m (Maybe (Stream m a))
1642stripPrefix (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
1643 where
1644 go (sa, sb, Nothing) = do
1645 r <- stepa defState sa
1646 case r of
1647 Yield x sa' -> go (sa', sb, Just x)
1648 Skip sa' -> go (sa', sb, Nothing)
1649 Stop -> return $ Just (Stream stepb sb)
1651 go (sa, sb, Just x) = do
1652 r <- stepb defState sb
1653 case r of
1654 Yield y sb' ->
1655 if x == y
1656 then go (sa, sb', Nothing)
1657 else return Nothing
1658 Skip sb' -> go (sa, sb', Just x)
1659 Stop -> return Nothing
1661------------------------------------------------------------------------------
1662-- Map and Fold
1663------------------------------------------------------------------------------
1665-- | Execute a monadic action for each element of the 'Stream'
1666{-# INLINE_NORMAL mapM_ #-}
1667mapM_ :: Monad m => (a -> m b) -> Stream m a -> m ()
1668mapM_ m = drain . mapM m
1670------------------------------------------------------------------------------
1671-- Exceptions
1672------------------------------------------------------------------------------
1674-- | Run a side effect before the stream yields its first element.
1675{-# INLINE_NORMAL before #-}
1676before :: Monad m => m b -> Stream m a -> Stream m a
1677before action (Stream step state) = Stream step' Nothing
1679 where
1681 {-# INLINE_LATE step' #-}
1682 step' _ Nothing = action >> return (Skip (Just state))
1684 step' gst (Just st) = do
1685 res <- step gst st
1686 case res of
1687 Yield x s -> return $ Yield x (Just s)
1688 Skip s -> return $ Skip (Just s)
1689 Stop -> return Stop
1691-- | Run a side effect whenever the stream stops normally.
1692{-# INLINE_NORMAL after #-}
1693after :: Monad m => m b -> Stream m a -> Stream m a
1694after action (Stream step state) = Stream step' state
1696 where
1698 {-# INLINE_LATE step' #-}
1699 step' gst st = do
1700 res <- step gst st
1701 case res of
1702 Yield x s -> return $ Yield x s
1703 Skip s -> return $ Skip s
1704 Stop -> action >> return Stop
1706-- XXX These combinators are expensive due to the call to
1707-- onException/handle/try on each step. Therefore, when possible, they should
1708-- be called in an outer loop where we perform less iterations. For example, we
1709-- cannot call them on each iteration in a char stream, instead we can call
1710-- them when doing an IO on an array.
1712-- XXX For high performance error checks in busy streams we may need another
1713-- Error constructor in step.
1715-- | Run a side effect whenever the stream aborts due to an exception. The
1716-- exception is not caught simply rethrown.
1717{-# INLINE_NORMAL onException #-}
1718onException :: MonadCatch m => m b -> Stream m a -> Stream m a
1719onException action (Stream step state) = Stream step' state
1721 where
1723 {-# INLINE_LATE step' #-}
1724 step' gst st = do
1725 res <- step gst st `MC.onException` action
1726 case res of
1727 Yield x s -> return $ Yield x s
1728 Skip s -> return $ Skip s
1729 Stop -> return Stop
1731-- | Run a side effect whenever the stream stops normally or aborts due to an
1732-- exception.
1733{-# INLINE finally #-}
1734finally :: MonadCatch m => m b -> Stream m a -> Stream m a
1735finally action xs = after action $ onException action xs
1737-- XXX bracket is like concatMap, it generates a stream and then flattens it.
1738-- Like concatMap it has 10x worse performance compared to linear fused
1739-- compositions.
1741-- | Run the first action before the stream starts and remember its output,
1742-- generate a stream using the output, run the second action using the
1743-- remembered value as an argument whenever the stream ends normally or due to
1744-- an exception.
1745{-# INLINE_NORMAL bracket #-}
1746bracket :: MonadCatch m => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
1747bracket bef aft bet = Stream step' Nothing
1749 where
1751 {-# INLINE_LATE step' #-}
1752 step' _ Nothing = bef >>= \x -> return (Skip (Just (bet x, x)))
1754 -- NOTE: It is important to use UnStream instead of the Stream pattern
1755 -- here, otherwise we get huge perf degradation, see note in concatMap.
1756 step' gst (Just (UnStream step state, v)) = do
1757 res <- step gst state `MC.onException` aft v
1758 case res of
1759 Yield x s -> return $ Yield x (Just (Stream step s, v))
1760 Skip s -> return $ Skip (Just (Stream step s, v))
1761 Stop -> aft v >> return Stop
1763-- | When evaluating a stream if an exception occurs, stream evaluation aborts
1764-- and the specified exception handler is run with the exception as argument.
1765{-# INLINE_NORMAL handle #-}
1766handle :: (MonadCatch m, Exception e) => (e -> m a) -> Stream m a -> Stream m a
1767handle f (Stream step state) = Stream step' (Just state)
1769 where
1771 {-# INLINE_LATE step' #-}
1772 step' gst (Just st) = do
1773 res <- MC.try $ step gst st
1774 case res of
1775 Left e -> f e >>= \x -> return (Yield x Nothing)
1776 Right r -> case r of
1777 Yield x s -> return $ Yield x (Just s)
1778 Skip s -> return $ Skip (Just s)
1779 Stop -> return Stop
1781 step' _ Nothing = return Stop
1783------------------------------------------------------------------------------
1784-- Transformation by Folding (Scans)
1785------------------------------------------------------------------------------
1787------------------------------------------------------------------------------
1788-- Prescans
1789------------------------------------------------------------------------------
1791-- XXX Is a prescan useful, discarding the last step does not sound useful? I
1792-- am not sure about the utility of this function, so this is implemented but
1793-- not exposed. We can expose it if someone provides good reasons why this is
1794-- useful.
1796-- XXX We have to execute the stream one step ahead to know that we are at the
1797-- last step. The vector implementation of prescan executes the last fold step
1798-- but does not yield the result. This means we have executed the effect but
1799-- discarded value. This does not sound right. In this implementation we are
1800-- not executing the last fold step.
1801{-# INLINE_NORMAL prescanlM' #-}
1802prescanlM' :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
1803prescanlM' f mz (Stream step state) = Stream step' (state, mz)
1804 where
1805 {-# INLINE_LATE step' #-}
1806 step' gst (st, prev) = do
1807 r <- step (adaptState gst) st
1808 case r of
1809 Yield x s -> do
1810 acc <- prev
1811 return $ Yield acc (s, f acc x)
1812 Skip s -> return $ Skip (s, prev)
1813 Stop -> return Stop
1815{-# INLINE prescanl' #-}
1816prescanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
1817prescanl' f z = prescanlM' (\a b -> return (f a b)) (return z)
1819------------------------------------------------------------------------------
1820-- Monolithic postscans (postscan followed by a map)
1821------------------------------------------------------------------------------
1823-- The performance of a modular postscan followed by a map seems to be
1824-- equivalent to this monolithic scan followed by map therefore we may not need
1825-- this implementation. We just have it for performance comparison and in case
1826-- modular version does not perform well in some situation.
1828{-# INLINE_NORMAL postscanlMx' #-}
1829postscanlMx' :: Monad m
1830 => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> Stream m b
1831postscanlMx' fstep begin done (Stream step state) = do
1832 Stream step' (state, begin)
1833 where
1834 {-# INLINE_LATE step' #-}
1835 step' gst (st, acc) = do
1836 r <- step (adaptState gst) st
1837 case r of
1838 Yield x s -> do
1839 old <- acc
1840 y <- fstep old x
1841 v <- done y
1842 v `seq` y `seq` return (Yield v (s, return y))
1843 Skip s -> return $ Skip (s, acc)
1844 Stop -> return Stop
1846{-# INLINE_NORMAL postscanlx' #-}
1847postscanlx' :: Monad m
1848 => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> Stream m b
1849postscanlx' fstep begin done s =
1850 postscanlMx' (\b a -> return (fstep b a)) (return begin) (return . done) s
1852-- XXX do we need consM strict to evaluate the begin value?
1853{-# INLINE scanlMx' #-}
1854scanlMx' :: Monad m
1855 => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> Stream m b
1856scanlMx' fstep begin done s =
1857 (begin >>= \x -> x `seq` done x) `consM` postscanlMx' fstep begin done s
1859{-# INLINE scanlx' #-}
1860scanlx' :: Monad m
1861 => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> Stream m b
1862scanlx' fstep begin done s =
1863 scanlMx' (\b a -> return (fstep b a)) (return begin) (return . done) s
1865------------------------------------------------------------------------------
1866-- postscans
1867------------------------------------------------------------------------------
1869-- XXX if we make the initial value of the accumulator monadic then should we
1870-- execute it even if the stream is empty? In that case we would have generated
1871-- the effect but discarded the value, but that is what a fold does when the
1872-- stream is empty. Whatever we decide, need to reconcile this with prescan.
1873-- If we execute the initial value here without even using it then it is ok to
1874-- execute the last step there as well without using the value.
1875-- Looking at the duality with right fold, in case of right fold we always
1876-- perform the action when the construction terminates, so in case of left fold
1877-- we should perform it only when the reduction starts.
1878{-# INLINE_NORMAL postscanlM' #-}
1879postscanlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
1880postscanlM' fstep begin (Stream step state) =
1881 begin `seq` Stream step' (state, begin)
1882 where
1883 {-# INLINE_LATE step' #-}
1884 step' gst (st, acc) = acc `seq` do
1885 r <- step (adaptState gst) st
1886 case r of
1887 Yield x s -> do
1888 y <- fstep acc x
1889 y `seq` return (Yield y (s, y))
1890 Skip s -> return $ Skip (s, acc)
1891 Stop -> return Stop
1893{-# INLINE_NORMAL postscanl' #-}
1894postscanl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
1895postscanl' f = postscanlM' (\a b -> return (f a b))
1897{-# INLINE_NORMAL postscanlM #-}
1898postscanlM :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
1899postscanlM fstep begin (Stream step state) = Stream step' (state, begin)
1900 where
1901 {-# INLINE_LATE step' #-}
1902 step' gst (st, acc) = do
1903 r <- step (adaptState gst) st
1904 case r of
1905 Yield x s -> do
1906 y <- fstep acc x
1907 return (Yield y (s, y))
1908 Skip s -> return $ Skip (s, acc)
1909 Stop -> return Stop
1911{-# INLINE_NORMAL postscanl #-}
1912postscanl :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
1913postscanl f = postscanlM (\a b -> return (f a b))
1915{-# INLINE_NORMAL scanlM' #-}
1916scanlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
1917scanlM' fstep begin s = begin `seq` (begin `cons` postscanlM' fstep begin s)
1919{-# INLINE scanl' #-}
1920scanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
1921scanl' f = scanlM' (\a b -> return (f a b))
1923{-# INLINE_NORMAL scanlM #-}
1924scanlM :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
1925scanlM fstep begin s = begin `cons` postscanlM fstep begin s
1927{-# INLINE scanl #-}
1928scanl :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
1929scanl f = scanlM (\a b -> return (f a b))
1931{-# INLINE_NORMAL scanl1M #-}
1932scanl1M :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
1933scanl1M fstep (Stream step state) = Stream step' (state, Nothing)
1934 where
1935 {-# INLINE_LATE step' #-}
1936 step' gst (st, Nothing) = do
1937 r <- step gst st
1938 case r of
1939 Yield x s -> return $ Yield x (s, Just x)
1940 Skip s -> return $ Skip (s, Nothing)
1941 Stop -> return Stop
1943 step' gst (st, Just acc) = do
1944 r <- step gst st
1945 case r of
1946 Yield y s -> do
1947 z <- fstep acc y
1948 return $ Yield z (s, Just z)
1949 Skip s -> return $ Skip (s, Just acc)
1950 Stop -> return Stop
1952{-# INLINE scanl1 #-}
1953scanl1 :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
1954scanl1 f = scanl1M (\x y -> return (f x y))
1956{-# INLINE_NORMAL scanl1M' #-}
1957scanl1M' :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
1958scanl1M' fstep (Stream step state) = Stream step' (state, Nothing)
1959 where
1960 {-# INLINE_LATE step' #-}
1961 step' gst (st, Nothing) = do
1962 r <- step gst st
1963 case r of
1964 Yield x s -> x `seq` return $ Yield x (s, Just x)
1965 Skip s -> return $ Skip (s, Nothing)
1966 Stop -> return Stop
1968 step' gst (st, Just acc) = acc `seq` do
1969 r <- step gst st
1970 case r of
1971 Yield y s -> do
1972 z <- fstep acc y
1973 z `seq` return $ Yield z (s, Just z)
1974 Skip s -> return $ Skip (s, Just acc)
1975 Stop -> return Stop
1977{-# INLINE scanl1' #-}
1978scanl1' :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
1979scanl1' f = scanl1M' (\x y -> return (f x y))
1981{-# INLINE tap #-}
1982tap :: Monad m => Sink m a -> Stream m a -> Stream m a
1983tap (Sink fstep) (Stream step state) = Stream step' state
1985 where
1987 step' gst st = do
1988 r <- step gst st
1989 case r of
1990 Yield x s -> do
1991 fstep x
1992 return $ Yield x s
1993 Skip s -> return $ Skip s
1994 Stop -> return $ Stop
1996-------------------------------------------------------------------------------
1997-- Filtering
1998-------------------------------------------------------------------------------
2000{-# INLINE_NORMAL takeWhileM #-}
2001takeWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
2002takeWhileM f (Stream step state) = Stream step' state
2003 where
2004 {-# INLINE_LATE step' #-}
2005 step' gst st = do
2006 r <- step gst st
2007 case r of
2008 Yield x s -> do
2009 b <- f x
2010 return $ if b then Yield x s else Stop
2011 Skip s -> return $ Skip s
2012 Stop -> return Stop
2014{-# INLINE takeWhile #-}
2015takeWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
2016takeWhile f = takeWhileM (return . f)
2018{-# INLINE_NORMAL drop #-}
2019drop :: Monad m => Int -> Stream m a -> Stream m a
2020drop n (Stream step state) = Stream step' (state, Just n)
2021 where
2022 {-# INLINE_LATE step' #-}
2023 step' gst (st, Just i)
2024 | i > 0 = do
2025 r <- step gst st
2026 return $
2027 case r of
2028 Yield _ s -> Skip (s, Just (i - 1))
2029 Skip s -> Skip (s, Just i)
2030 Stop -> Stop
2031 | otherwise = return $ Skip (st, Nothing)
2033 step' gst (st, Nothing) = do
2034 r <- step gst st
2035 return $
2036 case r of
2037 Yield x s -> Yield x (s, Nothing)
2038 Skip s -> Skip (s, Nothing)
2039 Stop -> Stop
2041data DropWhileState s a
2042 = DropWhileDrop s
2043 | DropWhileYield a s
2044 | DropWhileNext s
2046{-# INLINE_NORMAL dropWhileM #-}
2047dropWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
2048dropWhileM f (Stream step state) = Stream step' (DropWhileDrop state)
2049 where
2050 {-# INLINE_LATE step' #-}
2051 step' gst (DropWhileDrop st) = do
2052 r <- step gst st
2053 case r of
2054 Yield x s -> do
2055 b <- f x
2056 if b
2057 then return $ Skip (DropWhileDrop s)
2058 else return $ Skip (DropWhileYield x s)
2059 Skip s -> return $ Skip (DropWhileDrop s)
2060 Stop -> return Stop
2062 step' gst (DropWhileNext st) = do
2063 r <- step gst st
2064 case r of
2065 Yield x s -> return $ Skip (DropWhileYield x s)
2066 Skip s -> return $ Skip (DropWhileNext s)
2067 Stop -> return Stop
2069 step' _ (DropWhileYield x st) = return $ Yield x (DropWhileNext st)
2071{-# INLINE dropWhile #-}
2072dropWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
2073dropWhile f = dropWhileM (return . f)
2075{-# INLINE_NORMAL filterM #-}
2076filterM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
2077filterM f (Stream step state) = Stream step' state
2078 where
2079 {-# INLINE_LATE step' #-}
2080 step' gst st = do
2081 r <- step gst st
2082 case r of
2083 Yield x s -> do
2084 b <- f x
2085 return $ if b
2086 then Yield x s
2087 else Skip s
2088 Skip s -> return $ Skip s
2089 Stop -> return Stop
2091{-# INLINE filter #-}
2092filter :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
2093filter f = filterM (return . f)
2095{-# INLINE_NORMAL uniq #-}
2096uniq :: (Eq a, Monad m) => Stream m a -> Stream m a
2097uniq (Stream step state) = Stream step' (Nothing, state)
2098 where
2099 {-# INLINE_LATE step' #-}
2100 step' gst (Nothing, st) = do
2101 r <- step gst st
2102 case r of
2103 Yield x s -> return $ Yield x (Just x, s)
2104 Skip s -> return $ Skip (Nothing, s)
2105 Stop -> return Stop
2106 step' gst (Just x, st) = do
2107 r <- step gst st
2108 case r of
2109 Yield y s | x == y -> return $ Skip (Just x, s)
2110 | otherwise -> return $ Yield x (Just y, s)
2111 Skip s -> return $ Skip (Just x, s)
2112 Stop -> return Stop
2114------------------------------------------------------------------------------
2115-- Transformation by Mapping
2116------------------------------------------------------------------------------
2118{-# INLINE_NORMAL sequence #-}
2119sequence :: Monad m => Stream m (m a) -> Stream m a
2120sequence (Stream step state) = Stream step' state
2121 where
2122 {-# INLINE_LATE step' #-}
2123 step' gst st = do
2124 r <- step (adaptState gst) st
2125 case r of
2126 Yield x s -> x >>= \a -> return (Yield a s)
2127 Skip s -> return $ Skip s
2128 Stop -> return Stop
2130------------------------------------------------------------------------------
2131-- Inserting
2132------------------------------------------------------------------------------
2134{-# INLINE_NORMAL insertBy #-}
2135insertBy :: Monad m => (a -> a -> Ordering) -> a -> Stream m a -> Stream m a
2136insertBy cmp a (Stream step state) = Stream step' (state, False, Nothing)
2137 where
2138 {-# INLINE_LATE step' #-}
2139 step' gst (st, False, _) = do
2140 r <- step gst st
2141 case r of
2142 Yield x s -> case cmp a x of
2143 GT -> return $ Yield x (s, False, Nothing)
2144 _ -> return $ Yield a (s, True, Just x)
2145 Skip s -> return $ Skip (s, False, Nothing)
2146 Stop -> return $ Yield a (st, True, Nothing)
2148 step' _ (_, True, Nothing) = return Stop
2150 step' gst (st, True, Just prev) = do
2151 r <- step gst st
2152 case r of
2153 Yield x s -> return $ Yield prev (s, True, Just x)
2154 Skip s -> return $ Skip (s, True, Just prev)
2155 Stop -> return $ Yield prev (st, True, Nothing)
2157------------------------------------------------------------------------------
2158-- Deleting
2159------------------------------------------------------------------------------
2161{-# INLINE_NORMAL deleteBy #-}
2162deleteBy :: Monad m => (a -> a -> Bool) -> a -> Stream m a -> Stream m a
2163deleteBy eq x (Stream step state) = Stream step' (state, False)
2164 where
2165 {-# INLINE_LATE step' #-}
2166 step' gst (st, False) = do
2167 r <- step gst st
2168 case r of
2169 Yield y s -> return $
2170 if eq x y then Skip (s, True) else Yield y (s, False)
2171 Skip s -> return $ Skip (s, False)
2172 Stop -> return Stop
2174 step' gst (st, True) = do
2175 r <- step gst st
2176 case r of
2177 Yield y s -> return $ Yield y (s, True)
2178 Skip s -> return $ Skip (s, True)
2179 Stop -> return Stop
2181------------------------------------------------------------------------------
2182-- Transformation by Map and Filter
2183------------------------------------------------------------------------------
2185-- XXX Will this always fuse properly?
2186{-# INLINE_NORMAL mapMaybe #-}
2187mapMaybe :: Monad m => (a -> Maybe b) -> Stream m a -> Stream m b
2188mapMaybe f = fmap fromJust . filter isJust . map f
2190{-# INLINE_NORMAL mapMaybeM #-}
2191mapMaybeM :: Monad m => (a -> m (Maybe b)) -> Stream m a -> Stream m b
2192mapMaybeM f = fmap fromJust . filter isJust . mapM f
2194------------------------------------------------------------------------------
2195-- Zipping
2196------------------------------------------------------------------------------
2198{-# INLINE_NORMAL indexed #-}
2199indexed :: Monad m => Stream m a -> Stream m (Int, a)
2200indexed (Stream step state) = Stream step' (state, 0)
2201 where
2202 {-# INLINE_LATE step' #-}
2203 step' gst (st, i) = i `seq` do
2204 r <- step (adaptState gst) st
2205 case r of
2206 Yield x s -> return $ Yield (i, x) (s, i+1)
2207 Skip s -> return $ Skip (s, i)
2208 Stop -> return Stop
2210{-# INLINE_NORMAL indexedR #-}
2211indexedR :: Monad m => Int -> Stream m a -> Stream m (Int, a)
2212indexedR m (Stream step state) = Stream step' (state, m)
2213 where
2214 {-# INLINE_LATE step' #-}
2215 step' gst (st, i) = i `seq` do
2216 r <- step (adaptState gst) st
2217 case r of
2218 Yield x s -> let i' = i - 1
2219 in return $ Yield (i, x) (s, i')
2220 Skip s -> return $ Skip (s, i)
2221 Stop -> return Stop
2223{-# INLINE_NORMAL zipWithM #-}
2224zipWithM :: Monad m
2225 => (a -> b -> m c) -> Stream m a -> Stream m b -> Stream m c
2226zipWithM f (Stream stepa ta) (Stream stepb tb) = Stream step (ta, tb, Nothing)
2227 where
2228 {-# INLINE_LATE step #-}
2229 step gst (sa, sb, Nothing) = do
2230 r <- stepa (adaptState gst) sa
2231 return $
2232 case r of
2233 Yield x sa' -> Skip (sa', sb, Just x)
2234 Skip sa' -> Skip (sa', sb, Nothing)
2235 Stop -> Stop
2237 step gst (sa, sb, Just x) = do
2238 r <- stepb (adaptState gst) sb
2239 case r of
2240 Yield y sb' -> do
2241 z <- f x y
2242 return $ Yield z (sa, sb', Nothing)
2243 Skip sb' -> return $ Skip (sa, sb', Just x)
2244 Stop -> return Stop
2246{-# RULES "zipWithM xs xs"
2247 forall f xs. zipWithM f xs xs = mapM (\x -> f x x) xs #-}
2249{-# INLINE zipWith #-}
2250zipWith :: Monad m => (a -> b -> c) -> Stream m a -> Stream m b -> Stream m c
2251zipWith f = zipWithM (\a b -> return (f a b))
2253------------------------------------------------------------------------------
2254-- Merging
2255------------------------------------------------------------------------------
2257{-# INLINE_NORMAL mergeByM #-}
2258mergeByM
2259 :: (Monad m)
2260 => (a -> a -> m Ordering) -> Stream m a -> Stream m a -> Stream m a
2261mergeByM cmp (Stream stepa ta) (Stream stepb tb) =
2262 Stream step (Just ta, Just tb, Nothing, Nothing)
2263 where
2264 {-# INLINE_LATE step #-}
2266 -- one of the values is missing, and the corresponding stream is running
2267 step gst (Just sa, sb, Nothing, b) = do
2268 r <- stepa gst sa
2269 return $ case r of
2270 Yield a sa' -> Skip (Just sa', sb, Just a, b)
2271 Skip sa' -> Skip (Just sa', sb, Nothing, b)
2272 Stop -> Skip (Nothing, sb, Nothing, b)
2274 step gst (sa, Just sb, a, Nothing) = do
2275 r <- stepb gst sb
2276 return $ case r of
2277 Yield b sb' -> Skip (sa, Just sb', a, Just b)
2278 Skip sb' -> Skip (sa, Just sb', a, Nothing)
2279 Stop -> Skip (sa, Nothing, a, Nothing)
2281 -- both the values are available
2282 step _ (sa, sb, Just a, Just b) = do
2283 res <- cmp a b
2284 return $ case res of
2285 GT -> Yield b (sa, sb, Just a, Nothing)
2286 _ -> Yield a (sa, sb, Nothing, Just b)
2288 -- one of the values is missing, corresponding stream is done
2289 step _ (Nothing, sb, Nothing, Just b) =
2290 return $ Yield b (Nothing, sb, Nothing, Nothing)
2292 step _ (sa, Nothing, Just a, Nothing) =
2293 return $ Yield a (sa, Nothing, Nothing, Nothing)
2295 step _ (Nothing, Nothing, Nothing, Nothing) = return Stop
2297{-# INLINE mergeBy #-}
2298mergeBy
2299 :: (Monad m)
2300 => (a -> a -> Ordering) -> Stream m a -> Stream m a -> Stream m a
2301mergeBy cmp = mergeByM (\a b -> return $ cmp a b)
2303------------------------------------------------------------------------------
2304-- Transformation comprehensions
2305------------------------------------------------------------------------------
2307{-# INLINE_NORMAL the #-}
2308the :: (Eq a, Monad m) => Stream m a -> m (Maybe a)
2309the (Stream step state) = go state
2310 where
2311 go st = do
2312 r <- step defState st
2313 case r of
2314 Yield x s -> go' x s
2315 Skip s -> go s
2316 Stop -> return Nothing
2317 go' n st = do
2318 r <- step defState st
2319 case r of
2320 Yield x s | x == n -> go' n s
2321 | otherwise -> return Nothing
2322 Skip s -> go' n s
2323 Stop -> return (Just n)