git.iko.soy/streamly
on 523f606 help sign in
streamly/src/Streamly/Streams/StreamD.hs
2312 lines 78 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-- Converting folds
1672------------------------------------------------------------------------------
1674{-# INLINE fromStreamD #-}
1675fromStreamD :: (K.IsStream t, Monad m) => Stream m a -> t m a
1676fromStreamD = K.fromStream . toStreamK
1678------------------------------------------------------------------------------
1679-- Exceptions
1680------------------------------------------------------------------------------
1682-- | Run a side effect before the stream yields its first element.
1683before :: Monad m => m b -> Stream m a -> Stream m a
1684before action (Stream step state) = Stream step' Nothing
1686 where
1688 {-# INLINE_LATE step' #-}
1689 step' _ Nothing = action >> return (Skip (Just state))
1691 step' gst (Just st) = do
1692 res <- step gst st
1693 case res of
1694 Yield x s -> return $ Yield x (Just s)
1695 Skip s -> return $ Skip (Just s)
1696 Stop -> return Stop
1698-- | Run a side effect whenever the stream stops normally.
1699after :: Monad m => m b -> Stream m a -> Stream m a
1700after action (Stream step state) = Stream step' state
1702 where
1704 {-# INLINE_LATE step' #-}
1705 step' gst st = do
1706 res <- step gst st
1707 case res of
1708 Yield x s -> return $ Yield x s
1709 Skip s -> return $ Skip s
1710 Stop -> action >> return Stop
1712-- | Run a side effect whenever the stream aborts due to an exception. The
1713-- exception is not caught simply rethrown.
1714onException :: MonadCatch m => m b -> Stream m a -> Stream m a
1715onException action (Stream step state) = Stream step' state
1717 where
1719 {-# INLINE_LATE step' #-}
1720 step' gst st = do
1721 res <- step gst st `MC.onException` action
1722 case res of
1723 Yield x s -> return $ Yield x s
1724 Skip s -> return $ Skip s
1725 Stop -> action >> return Stop
1727-- XXX check if a monolithic implementation would be more performant.
1729-- | Run a side effect whenever the stream stops normally or aborts due to an
1730-- exception.
1731finally :: MonadCatch m => m b -> Stream m a -> Stream m a
1732finally action xs = after action $ onException action xs
1734-- | Run the first action before the stream starts and remember its output,
1735-- generate a stream using the output, run the second action using the
1736-- remembered value as an argument whenever the stream ends normally or due to
1737-- an exception.
1738bracket :: MonadCatch m => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
1739bracket bef aft bet = Stream step' Nothing
1741 where
1743 {-# INLINE_LATE step' #-}
1744 step' _ Nothing = bef >>= \x -> return (Skip (Just (bet x, x)))
1746 step' gst (Just (Stream step state, v)) = do
1747 res <- step gst state `MC.onException` aft v
1748 case res of
1749 Yield x s -> return $ Yield x (Just (Stream step s, v))
1750 Skip s -> return $ Skip (Just (Stream step s, v))
1751 Stop -> aft v >> return Stop
1753-- | When evaluating a stream if an exception occurs, stream evaluation aborts
1754-- and the specified exception handler is run with the exception as argument.
1755handle :: (MonadCatch m, Exception e) => (e -> m a) -> Stream m a -> Stream m a
1756handle f (Stream step state) = Stream step' (Just state)
1758 where
1760 {-# INLINE_LATE step' #-}
1761 step' gst (Just st) = do
1762 res <- MC.try $ step gst st
1763 case res of
1764 Left e -> f e >>= \x -> return (Yield x Nothing)
1765 Right r -> case r of
1766 Yield x s -> return $ Yield x (Just s)
1767 Skip s -> return $ Skip (Just s)
1768 Stop -> return Stop
1770 step' _ Nothing = return Stop
1772------------------------------------------------------------------------------
1773-- Transformation by Folding (Scans)
1774------------------------------------------------------------------------------
1776------------------------------------------------------------------------------
1777-- Prescans
1778------------------------------------------------------------------------------
1780-- XXX Is a prescan useful, discarding the last step does not sound useful? I
1781-- am not sure about the utility of this function, so this is implemented but
1782-- not exposed. We can expose it if someone provides good reasons why this is
1783-- useful.
1785-- XXX We have to execute the stream one step ahead to know that we are at the
1786-- last step. The vector implementation of prescan executes the last fold step
1787-- but does not yield the result. This means we have executed the effect but
1788-- discarded value. This does not sound right. In this implementation we are
1789-- not executing the last fold step.
1790{-# INLINE_NORMAL prescanlM' #-}
1791prescanlM' :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
1792prescanlM' f mz (Stream step state) = Stream step' (state, mz)
1793 where
1794 {-# INLINE_LATE step' #-}
1795 step' gst (st, prev) = do
1796 r <- step (adaptState gst) st
1797 case r of
1798 Yield x s -> do
1799 acc <- prev
1800 return $ Yield acc (s, f acc x)
1801 Skip s -> return $ Skip (s, prev)
1802 Stop -> return Stop
1804{-# INLINE prescanl' #-}
1805prescanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
1806prescanl' f z = prescanlM' (\a b -> return (f a b)) (return z)
1808------------------------------------------------------------------------------
1809-- Monolithic postscans (postscan followed by a map)
1810------------------------------------------------------------------------------
1812-- The performance of a modular postscan followed by a map seems to be
1813-- equivalent to this monolithic scan followed by map therefore we may not need
1814-- this implementation. We just have it for performance comparison and in case
1815-- modular version does not perform well in some situation.
1817{-# INLINE_NORMAL postscanlMx' #-}
1818postscanlMx' :: Monad m
1819 => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> Stream m b
1820postscanlMx' fstep begin done (Stream step state) = do
1821 Stream step' (state, begin)
1822 where
1823 {-# INLINE_LATE step' #-}
1824 step' gst (st, acc) = do
1825 r <- step (adaptState gst) st
1826 case r of
1827 Yield x s -> do
1828 old <- acc
1829 y <- fstep old x
1830 v <- done y
1831 v `seq` y `seq` return (Yield v (s, return y))
1832 Skip s -> return $ Skip (s, acc)
1833 Stop -> return Stop
1835{-# INLINE_NORMAL postscanlx' #-}
1836postscanlx' :: Monad m
1837 => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> Stream m b
1838postscanlx' fstep begin done s =
1839 postscanlMx' (\b a -> return (fstep b a)) (return begin) (return . done) s
1841-- XXX do we need consM strict to evaluate the begin value?
1842{-# INLINE scanlMx' #-}
1843scanlMx' :: Monad m
1844 => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> Stream m b
1845scanlMx' fstep begin done s =
1846 (begin >>= \x -> x `seq` done x) `consM` postscanlMx' fstep begin done s
1848{-# INLINE scanlx' #-}
1849scanlx' :: Monad m
1850 => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> Stream m b
1851scanlx' fstep begin done s =
1852 scanlMx' (\b a -> return (fstep b a)) (return begin) (return . done) s
1854------------------------------------------------------------------------------
1855-- postscans
1856------------------------------------------------------------------------------
1858-- XXX if we make the initial value of the accumulator monadic then should we
1859-- execute it even if the stream is empty? In that case we would have generated
1860-- the effect but discarded the value, but that is what a fold does when the
1861-- stream is empty. Whatever we decide, need to reconcile this with prescan.
1862-- If we execute the initial value here without even using it then it is ok to
1863-- execute the last step there as well without using the value.
1864-- Looking at the duality with right fold, in case of right fold we always
1865-- perform the action when the construction terminates, so in case of left fold
1866-- we should perform it only when the reduction starts.
1867{-# INLINE_NORMAL postscanlM' #-}
1868postscanlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
1869postscanlM' fstep begin (Stream step state) =
1870 begin `seq` Stream step' (state, begin)
1871 where
1872 {-# INLINE_LATE step' #-}
1873 step' gst (st, acc) = acc `seq` do
1874 r <- step (adaptState gst) st
1875 case r of
1876 Yield x s -> do
1877 y <- fstep acc x
1878 y `seq` return (Yield y (s, y))
1879 Skip s -> return $ Skip (s, acc)
1880 Stop -> return Stop
1882{-# INLINE_NORMAL postscanl' #-}
1883postscanl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
1884postscanl' f = postscanlM' (\a b -> return (f a b))
1886{-# INLINE_NORMAL postscanlM #-}
1887postscanlM :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
1888postscanlM fstep begin (Stream step state) = Stream step' (state, begin)
1889 where
1890 {-# INLINE_LATE step' #-}
1891 step' gst (st, acc) = do
1892 r <- step (adaptState gst) st
1893 case r of
1894 Yield x s -> do
1895 y <- fstep acc x
1896 return (Yield y (s, y))
1897 Skip s -> return $ Skip (s, acc)
1898 Stop -> return Stop
1900{-# INLINE_NORMAL postscanl #-}
1901postscanl :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
1902postscanl f = postscanlM (\a b -> return (f a b))
1904{-# INLINE_NORMAL scanlM' #-}
1905scanlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
1906scanlM' fstep begin s = begin `seq` (begin `cons` postscanlM' fstep begin s)
1908{-# INLINE scanl' #-}
1909scanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
1910scanl' f = scanlM' (\a b -> return (f a b))
1912{-# INLINE_NORMAL scanlM #-}
1913scanlM :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
1914scanlM fstep begin s = begin `cons` postscanlM fstep begin s
1916{-# INLINE scanl #-}
1917scanl :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
1918scanl f = scanlM (\a b -> return (f a b))
1920{-# INLINE_NORMAL scanl1M #-}
1921scanl1M :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
1922scanl1M fstep (Stream step state) = Stream step' (state, Nothing)
1923 where
1924 {-# INLINE_LATE step' #-}
1925 step' gst (st, Nothing) = do
1926 r <- step gst st
1927 case r of
1928 Yield x s -> return $ Yield x (s, Just x)
1929 Skip s -> return $ Skip (s, Nothing)
1930 Stop -> return Stop
1932 step' gst (st, Just acc) = do
1933 r <- step gst st
1934 case r of
1935 Yield y s -> do
1936 z <- fstep acc y
1937 return $ Yield z (s, Just z)
1938 Skip s -> return $ Skip (s, Just acc)
1939 Stop -> return Stop
1941{-# INLINE scanl1 #-}
1942scanl1 :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
1943scanl1 f = scanl1M (\x y -> return (f x y))
1945{-# INLINE_NORMAL scanl1M' #-}
1946scanl1M' :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
1947scanl1M' fstep (Stream step state) = Stream step' (state, Nothing)
1948 where
1949 {-# INLINE_LATE step' #-}
1950 step' gst (st, Nothing) = do
1951 r <- step gst st
1952 case r of
1953 Yield x s -> x `seq` return $ Yield x (s, Just x)
1954 Skip s -> return $ Skip (s, Nothing)
1955 Stop -> return Stop
1957 step' gst (st, Just acc) = acc `seq` do
1958 r <- step gst st
1959 case r of
1960 Yield y s -> do
1961 z <- fstep acc y
1962 z `seq` return $ Yield z (s, Just z)
1963 Skip s -> return $ Skip (s, Just acc)
1964 Stop -> return Stop
1966{-# INLINE scanl1' #-}
1967scanl1' :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
1968scanl1' f = scanl1M' (\x y -> return (f x y))
1970{-# INLINE tap #-}
1971tap :: Monad m => Sink m a -> Stream m a -> Stream m a
1972tap (Sink fstep) (Stream step state) = Stream step' state
1974 where
1976 step' gst st = do
1977 r <- step gst st
1978 case r of
1979 Yield x s -> do
1980 fstep x
1981 return $ Yield x s
1982 Skip s -> return $ Skip s
1983 Stop -> return $ Stop
1985-------------------------------------------------------------------------------
1986-- Filtering
1987-------------------------------------------------------------------------------
1989{-# INLINE_NORMAL takeWhileM #-}
1990takeWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
1991takeWhileM f (Stream step state) = Stream step' state
1992 where
1993 {-# INLINE_LATE step' #-}
1994 step' gst st = do
1995 r <- step gst st
1996 case r of
1997 Yield x s -> do
1998 b <- f x
1999 return $ if b then Yield x s else Stop
2000 Skip s -> return $ Skip s
2001 Stop -> return Stop
2003{-# INLINE takeWhile #-}
2004takeWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
2005takeWhile f = takeWhileM (return . f)
2007{-# INLINE_NORMAL drop #-}
2008drop :: Monad m => Int -> Stream m a -> Stream m a
2009drop n (Stream step state) = Stream step' (state, Just n)
2010 where
2011 {-# INLINE_LATE step' #-}
2012 step' gst (st, Just i)
2013 | i > 0 = do
2014 r <- step gst st
2015 return $
2016 case r of
2017 Yield _ s -> Skip (s, Just (i - 1))
2018 Skip s -> Skip (s, Just i)
2019 Stop -> Stop
2020 | otherwise = return $ Skip (st, Nothing)
2022 step' gst (st, Nothing) = do
2023 r <- step gst st
2024 return $
2025 case r of
2026 Yield x s -> Yield x (s, Nothing)
2027 Skip s -> Skip (s, Nothing)
2028 Stop -> Stop
2030data DropWhileState s a
2031 = DropWhileDrop s
2032 | DropWhileYield a s
2033 | DropWhileNext s
2035{-# INLINE_NORMAL dropWhileM #-}
2036dropWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
2037dropWhileM f (Stream step state) = Stream step' (DropWhileDrop state)
2038 where
2039 {-# INLINE_LATE step' #-}
2040 step' gst (DropWhileDrop st) = do
2041 r <- step gst st
2042 case r of
2043 Yield x s -> do
2044 b <- f x
2045 if b
2046 then return $ Skip (DropWhileDrop s)
2047 else return $ Skip (DropWhileYield x s)
2048 Skip s -> return $ Skip (DropWhileDrop s)
2049 Stop -> return Stop
2051 step' gst (DropWhileNext st) = do
2052 r <- step gst st
2053 case r of
2054 Yield x s -> return $ Skip (DropWhileYield x s)
2055 Skip s -> return $ Skip (DropWhileNext s)
2056 Stop -> return Stop
2058 step' _ (DropWhileYield x st) = return $ Yield x (DropWhileNext st)
2060{-# INLINE dropWhile #-}
2061dropWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
2062dropWhile f = dropWhileM (return . f)
2064{-# INLINE_NORMAL filterM #-}
2065filterM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
2066filterM f (Stream step state) = Stream step' state
2067 where
2068 {-# INLINE_LATE step' #-}
2069 step' gst st = do
2070 r <- step gst st
2071 case r of
2072 Yield x s -> do
2073 b <- f x
2074 return $ if b
2075 then Yield x s
2076 else Skip s
2077 Skip s -> return $ Skip s
2078 Stop -> return Stop
2080{-# INLINE filter #-}
2081filter :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
2082filter f = filterM (return . f)
2084{-# INLINE_NORMAL uniq #-}
2085uniq :: (Eq a, Monad m) => Stream m a -> Stream m a
2086uniq (Stream step state) = Stream step' (Nothing, state)
2087 where
2088 {-# INLINE_LATE step' #-}
2089 step' gst (Nothing, st) = do
2090 r <- step gst st
2091 case r of
2092 Yield x s -> return $ Yield x (Just x, s)
2093 Skip s -> return $ Skip (Nothing, s)
2094 Stop -> return Stop
2095 step' gst (Just x, st) = do
2096 r <- step gst st
2097 case r of
2098 Yield y s | x == y -> return $ Skip (Just x, s)
2099 | otherwise -> return $ Yield x (Just y, s)
2100 Skip s -> return $ Skip (Just x, s)
2101 Stop -> return Stop
2103------------------------------------------------------------------------------
2104-- Transformation by Mapping
2105------------------------------------------------------------------------------
2107{-# INLINE_NORMAL sequence #-}
2108sequence :: Monad m => Stream m (m a) -> Stream m a
2109sequence (Stream step state) = Stream step' state
2110 where
2111 {-# INLINE_LATE step' #-}
2112 step' gst st = do
2113 r <- step (adaptState gst) st
2114 case r of
2115 Yield x s -> x >>= \a -> return (Yield a s)
2116 Skip s -> return $ Skip s
2117 Stop -> return Stop
2119------------------------------------------------------------------------------
2120-- Inserting
2121------------------------------------------------------------------------------
2123{-# INLINE_NORMAL insertBy #-}
2124insertBy :: Monad m => (a -> a -> Ordering) -> a -> Stream m a -> Stream m a
2125insertBy cmp a (Stream step state) = Stream step' (state, False, Nothing)
2126 where
2127 {-# INLINE_LATE step' #-}
2128 step' gst (st, False, _) = do
2129 r <- step gst st
2130 case r of
2131 Yield x s -> case cmp a x of
2132 GT -> return $ Yield x (s, False, Nothing)
2133 _ -> return $ Yield a (s, True, Just x)
2134 Skip s -> return $ Skip (s, False, Nothing)
2135 Stop -> return $ Yield a (st, True, Nothing)
2137 step' _ (_, True, Nothing) = return Stop
2139 step' gst (st, True, Just prev) = do
2140 r <- step gst st
2141 case r of
2142 Yield x s -> return $ Yield prev (s, True, Just x)
2143 Skip s -> return $ Skip (s, True, Just prev)
2144 Stop -> return $ Yield prev (st, True, Nothing)
2146------------------------------------------------------------------------------
2147-- Deleting
2148------------------------------------------------------------------------------
2150{-# INLINE_NORMAL deleteBy #-}
2151deleteBy :: Monad m => (a -> a -> Bool) -> a -> Stream m a -> Stream m a
2152deleteBy eq x (Stream step state) = Stream step' (state, False)
2153 where
2154 {-# INLINE_LATE step' #-}
2155 step' gst (st, False) = do
2156 r <- step gst st
2157 case r of
2158 Yield y s -> return $
2159 if eq x y then Skip (s, True) else Yield y (s, False)
2160 Skip s -> return $ Skip (s, False)
2161 Stop -> return Stop
2163 step' gst (st, True) = do
2164 r <- step gst st
2165 case r of
2166 Yield y s -> return $ Yield y (s, True)
2167 Skip s -> return $ Skip (s, True)
2168 Stop -> return Stop
2170------------------------------------------------------------------------------
2171-- Transformation by Map and Filter
2172------------------------------------------------------------------------------
2174-- XXX Will this always fuse properly?
2175{-# INLINE_NORMAL mapMaybe #-}
2176mapMaybe :: Monad m => (a -> Maybe b) -> Stream m a -> Stream m b
2177mapMaybe f = fmap fromJust . filter isJust . map f
2179{-# INLINE_NORMAL mapMaybeM #-}
2180mapMaybeM :: Monad m => (a -> m (Maybe b)) -> Stream m a -> Stream m b
2181mapMaybeM f = fmap fromJust . filter isJust . mapM f
2183------------------------------------------------------------------------------
2184-- Zipping
2185------------------------------------------------------------------------------
2187{-# INLINE_NORMAL indexed #-}
2188indexed :: Monad m => Stream m a -> Stream m (Int, a)
2189indexed (Stream step state) = Stream step' (state, 0)
2190 where
2191 {-# INLINE_LATE step' #-}
2192 step' gst (st, i) = i `seq` do
2193 r <- step (adaptState gst) st
2194 case r of
2195 Yield x s -> return $ Yield (i, x) (s, i+1)
2196 Skip s -> return $ Skip (s, i)
2197 Stop -> return Stop
2199{-# INLINE_NORMAL indexedR #-}
2200indexedR :: Monad m => Int -> Stream m a -> Stream m (Int, a)
2201indexedR m (Stream step state) = Stream step' (state, m)
2202 where
2203 {-# INLINE_LATE step' #-}
2204 step' gst (st, i) = i `seq` do
2205 r <- step (adaptState gst) st
2206 case r of
2207 Yield x s -> let i' = i - 1
2208 in return $ Yield (i, x) (s, i')
2209 Skip s -> return $ Skip (s, i)
2210 Stop -> return Stop
2212{-# INLINE_NORMAL zipWithM #-}
2213zipWithM :: Monad m
2214 => (a -> b -> m c) -> Stream m a -> Stream m b -> Stream m c
2215zipWithM f (Stream stepa ta) (Stream stepb tb) = Stream step (ta, tb, Nothing)
2216 where
2217 {-# INLINE_LATE step #-}
2218 step gst (sa, sb, Nothing) = do
2219 r <- stepa (adaptState gst) sa
2220 return $
2221 case r of
2222 Yield x sa' -> Skip (sa', sb, Just x)
2223 Skip sa' -> Skip (sa', sb, Nothing)
2224 Stop -> Stop
2226 step gst (sa, sb, Just x) = do
2227 r <- stepb (adaptState gst) sb
2228 case r of
2229 Yield y sb' -> do
2230 z <- f x y
2231 return $ Yield z (sa, sb', Nothing)
2232 Skip sb' -> return $ Skip (sa, sb', Just x)
2233 Stop -> return Stop
2235{-# RULES "zipWithM xs xs"
2236 forall f xs. zipWithM f xs xs = mapM (\x -> f x x) xs #-}
2238{-# INLINE zipWith #-}
2239zipWith :: Monad m => (a -> b -> c) -> Stream m a -> Stream m b -> Stream m c
2240zipWith f = zipWithM (\a b -> return (f a b))
2242------------------------------------------------------------------------------
2243-- Merging
2244------------------------------------------------------------------------------
2246{-# INLINE_NORMAL mergeByM #-}
2247mergeByM
2248 :: (Monad m)
2249 => (a -> a -> m Ordering) -> Stream m a -> Stream m a -> Stream m a
2250mergeByM cmp (Stream stepa ta) (Stream stepb tb) =
2251 Stream step (Just ta, Just tb, Nothing, Nothing)
2252 where
2253 {-# INLINE_LATE step #-}
2255 -- one of the values is missing, and the corresponding stream is running
2256 step gst (Just sa, sb, Nothing, b) = do
2257 r <- stepa gst sa
2258 return $ case r of
2259 Yield a sa' -> Skip (Just sa', sb, Just a, b)
2260 Skip sa' -> Skip (Just sa', sb, Nothing, b)
2261 Stop -> Skip (Nothing, sb, Nothing, b)
2263 step gst (sa, Just sb, a, Nothing) = do
2264 r <- stepb gst sb
2265 return $ case r of
2266 Yield b sb' -> Skip (sa, Just sb', a, Just b)
2267 Skip sb' -> Skip (sa, Just sb', a, Nothing)
2268 Stop -> Skip (sa, Nothing, a, Nothing)
2270 -- both the values are available
2271 step _ (sa, sb, Just a, Just b) = do
2272 res <- cmp a b
2273 return $ case res of
2274 GT -> Yield b (sa, sb, Just a, Nothing)
2275 _ -> Yield a (sa, sb, Nothing, Just b)
2277 -- one of the values is missing, corresponding stream is done
2278 step _ (Nothing, sb, Nothing, Just b) =
2279 return $ Yield b (Nothing, sb, Nothing, Nothing)
2281 step _ (sa, Nothing, Just a, Nothing) =
2282 return $ Yield a (sa, Nothing, Nothing, Nothing)
2284 step _ (Nothing, Nothing, Nothing, Nothing) = return Stop
2286{-# INLINE mergeBy #-}
2287mergeBy
2288 :: (Monad m)
2289 => (a -> a -> Ordering) -> Stream m a -> Stream m a -> Stream m a
2290mergeBy cmp = mergeByM (\a b -> return $ cmp a b)
2292------------------------------------------------------------------------------
2293-- Transformation comprehensions
2294------------------------------------------------------------------------------
2296{-# INLINE_NORMAL the #-}
2297the :: (Eq a, Monad m) => Stream m a -> m (Maybe a)
2298the (Stream step state) = go state
2299 where
2300 go st = do
2301 r <- step defState st
2302 case r of
2303 Yield x s -> go' x s
2304 Skip s -> go s
2305 Stop -> return Nothing
2306 go' n st = do
2307 r <- step defState st
2308 case r of
2309 Yield x s | x == n -> go' n s
2310 | otherwise -> return Nothing
2311 Skip s -> go' n s
2312 Stop -> return (Just n)