git.iko.soy/streamly
on 260e31d help sign in
streamly/src/Streamly/Streams/StreamD.hs
2333 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-- 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.
1683{-# INLINE_NORMAL before #-}
1684before :: Monad m => m b -> Stream m a -> Stream m a
1685before action (Stream step state) = Stream step' Nothing
1687 where
1689 {-# INLINE_LATE step' #-}
1690 step' _ Nothing = action >> return (Skip (Just state))
1692 step' gst (Just st) = do
1693 res <- step gst st
1694 case res of
1695 Yield x s -> return $ Yield x (Just s)
1696 Skip s -> return $ Skip (Just s)
1697 Stop -> return Stop
1699-- | Run a side effect whenever the stream stops normally.
1700{-# INLINE_NORMAL after #-}
1701after :: Monad m => m b -> Stream m a -> Stream m a
1702after action (Stream step state) = Stream step' state
1704 where
1706 {-# INLINE_LATE step' #-}
1707 step' gst st = do
1708 res <- step gst st
1709 case res of
1710 Yield x s -> return $ Yield x s
1711 Skip s -> return $ Skip s
1712 Stop -> action >> return Stop
1714-- XXX These combinators are expensive due to the call to
1715-- onException/handle/try on each step. Therefore, they should called in an
1716-- outer loop where we perform less iterations. For example, we cannot call
1717-- them on each iteration in a char stream, instead we can call them when doing
1718-- an IO on an array.
1720-- XXX For high performance error checks in busy streams we may need another
1721-- Error constructor in step.
1723-- | Run a side effect whenever the stream aborts due to an exception. The
1724-- exception is not caught simply rethrown.
1725{-# INLINE_NORMAL onException #-}
1726onException :: MonadCatch m => m b -> Stream m a -> Stream m a
1727onException action (Stream step state) = Stream step' state
1729 where
1731 {-# INLINE_LATE step' #-}
1732 step' gst st = do
1733 res <- step gst st `MC.onException` action
1734 case res of
1735 Yield x s -> return $ Yield x s
1736 Skip s -> return $ Skip s
1737 Stop -> action >> return Stop
1739-- XXX check if a monolithic implementation would be more performant.
1741-- | Run a side effect whenever the stream stops normally or aborts due to an
1742-- exception.
1743{-# INLINE finally #-}
1744finally :: MonadCatch m => m b -> Stream m a -> Stream m a
1745finally action xs = after action $ onException action xs
1747-- XXX bracket is like concatMap, it generates a stream and then flattens it.
1748-- Like concatMap it has 10x worse performance compared to linear fused
1749-- compositions.
1751-- | Run the first action before the stream starts and remember its output,
1752-- generate a stream using the output, run the second action using the
1753-- remembered value as an argument whenever the stream ends normally or due to
1754-- an exception.
1755{-# INLINE_NORMAL bracket #-}
1756bracket :: MonadCatch m => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
1757bracket bef aft bet = Stream step' Nothing
1759 where
1761 {-# INLINE_LATE step' #-}
1762 step' _ Nothing = bef >>= \x -> return (Skip (Just (bet x, x)))
1764 -- NOTE: It is important to use UnStream instead of the Stream pattern
1765 -- here, otherwise we get huge perf degradation, see note in concatMap.
1766 step' gst (Just (UnStream step state, v)) = do
1767 res <- step gst state `MC.onException` aft v
1768 case res of
1769 Yield x s -> return $ Yield x (Just (Stream step s, v))
1770 Skip s -> return $ Skip (Just (Stream step s, v))
1771 Stop -> aft v >> return Stop
1773-- | When evaluating a stream if an exception occurs, stream evaluation aborts
1774-- and the specified exception handler is run with the exception as argument.
1775{-# INLINE_NORMAL handle #-}
1776handle :: (MonadCatch m, Exception e) => (e -> m a) -> Stream m a -> Stream m a
1777handle f (Stream step state) = Stream step' (Just state)
1779 where
1781 {-# INLINE_LATE step' #-}
1782 step' gst (Just st) = do
1783 res <- MC.try $ step gst st
1784 case res of
1785 Left e -> f e >>= \x -> return (Yield x Nothing)
1786 Right r -> case r of
1787 Yield x s -> return $ Yield x (Just s)
1788 Skip s -> return $ Skip (Just s)
1789 Stop -> return Stop
1791 step' _ Nothing = return Stop
1793------------------------------------------------------------------------------
1794-- Transformation by Folding (Scans)
1795------------------------------------------------------------------------------
1797------------------------------------------------------------------------------
1798-- Prescans
1799------------------------------------------------------------------------------
1801-- XXX Is a prescan useful, discarding the last step does not sound useful? I
1802-- am not sure about the utility of this function, so this is implemented but
1803-- not exposed. We can expose it if someone provides good reasons why this is
1804-- useful.
1806-- XXX We have to execute the stream one step ahead to know that we are at the
1807-- last step. The vector implementation of prescan executes the last fold step
1808-- but does not yield the result. This means we have executed the effect but
1809-- discarded value. This does not sound right. In this implementation we are
1810-- not executing the last fold step.
1811{-# INLINE_NORMAL prescanlM' #-}
1812prescanlM' :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
1813prescanlM' f mz (Stream step state) = Stream step' (state, mz)
1814 where
1815 {-# INLINE_LATE step' #-}
1816 step' gst (st, prev) = do
1817 r <- step (adaptState gst) st
1818 case r of
1819 Yield x s -> do
1820 acc <- prev
1821 return $ Yield acc (s, f acc x)
1822 Skip s -> return $ Skip (s, prev)
1823 Stop -> return Stop
1825{-# INLINE prescanl' #-}
1826prescanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
1827prescanl' f z = prescanlM' (\a b -> return (f a b)) (return z)
1829------------------------------------------------------------------------------
1830-- Monolithic postscans (postscan followed by a map)
1831------------------------------------------------------------------------------
1833-- The performance of a modular postscan followed by a map seems to be
1834-- equivalent to this monolithic scan followed by map therefore we may not need
1835-- this implementation. We just have it for performance comparison and in case
1836-- modular version does not perform well in some situation.
1838{-# INLINE_NORMAL postscanlMx' #-}
1839postscanlMx' :: Monad m
1840 => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> Stream m b
1841postscanlMx' fstep begin done (Stream step state) = do
1842 Stream step' (state, begin)
1843 where
1844 {-# INLINE_LATE step' #-}
1845 step' gst (st, acc) = do
1846 r <- step (adaptState gst) st
1847 case r of
1848 Yield x s -> do
1849 old <- acc
1850 y <- fstep old x
1851 v <- done y
1852 v `seq` y `seq` return (Yield v (s, return y))
1853 Skip s -> return $ Skip (s, acc)
1854 Stop -> return Stop
1856{-# INLINE_NORMAL postscanlx' #-}
1857postscanlx' :: Monad m
1858 => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> Stream m b
1859postscanlx' fstep begin done s =
1860 postscanlMx' (\b a -> return (fstep b a)) (return begin) (return . done) s
1862-- XXX do we need consM strict to evaluate the begin value?
1863{-# INLINE scanlMx' #-}
1864scanlMx' :: Monad m
1865 => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> Stream m b
1866scanlMx' fstep begin done s =
1867 (begin >>= \x -> x `seq` done x) `consM` postscanlMx' fstep begin done s
1869{-# INLINE scanlx' #-}
1870scanlx' :: Monad m
1871 => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> Stream m b
1872scanlx' fstep begin done s =
1873 scanlMx' (\b a -> return (fstep b a)) (return begin) (return . done) s
1875------------------------------------------------------------------------------
1876-- postscans
1877------------------------------------------------------------------------------
1879-- XXX if we make the initial value of the accumulator monadic then should we
1880-- execute it even if the stream is empty? In that case we would have generated
1881-- the effect but discarded the value, but that is what a fold does when the
1882-- stream is empty. Whatever we decide, need to reconcile this with prescan.
1883-- If we execute the initial value here without even using it then it is ok to
1884-- execute the last step there as well without using the value.
1885-- Looking at the duality with right fold, in case of right fold we always
1886-- perform the action when the construction terminates, so in case of left fold
1887-- we should perform it only when the reduction starts.
1888{-# INLINE_NORMAL postscanlM' #-}
1889postscanlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
1890postscanlM' fstep begin (Stream step state) =
1891 begin `seq` Stream step' (state, begin)
1892 where
1893 {-# INLINE_LATE step' #-}
1894 step' gst (st, acc) = acc `seq` do
1895 r <- step (adaptState gst) st
1896 case r of
1897 Yield x s -> do
1898 y <- fstep acc x
1899 y `seq` return (Yield y (s, y))
1900 Skip s -> return $ Skip (s, acc)
1901 Stop -> return Stop
1903{-# INLINE_NORMAL postscanl' #-}
1904postscanl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
1905postscanl' f = postscanlM' (\a b -> return (f a b))
1907{-# INLINE_NORMAL postscanlM #-}
1908postscanlM :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
1909postscanlM fstep begin (Stream step state) = Stream step' (state, begin)
1910 where
1911 {-# INLINE_LATE step' #-}
1912 step' gst (st, acc) = do
1913 r <- step (adaptState gst) st
1914 case r of
1915 Yield x s -> do
1916 y <- fstep acc x
1917 return (Yield y (s, y))
1918 Skip s -> return $ Skip (s, acc)
1919 Stop -> return Stop
1921{-# INLINE_NORMAL postscanl #-}
1922postscanl :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
1923postscanl f = postscanlM (\a b -> return (f a b))
1925{-# INLINE_NORMAL scanlM' #-}
1926scanlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
1927scanlM' fstep begin s = begin `seq` (begin `cons` postscanlM' fstep begin s)
1929{-# INLINE scanl' #-}
1930scanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
1931scanl' f = scanlM' (\a b -> return (f a b))
1933{-# INLINE_NORMAL scanlM #-}
1934scanlM :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
1935scanlM fstep begin s = begin `cons` postscanlM fstep begin s
1937{-# INLINE scanl #-}
1938scanl :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
1939scanl f = scanlM (\a b -> return (f a b))
1941{-# INLINE_NORMAL scanl1M #-}
1942scanl1M :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
1943scanl1M fstep (Stream step state) = Stream step' (state, Nothing)
1944 where
1945 {-# INLINE_LATE step' #-}
1946 step' gst (st, Nothing) = do
1947 r <- step gst st
1948 case r of
1949 Yield x s -> return $ Yield x (s, Just x)
1950 Skip s -> return $ Skip (s, Nothing)
1951 Stop -> return Stop
1953 step' gst (st, Just acc) = do
1954 r <- step gst st
1955 case r of
1956 Yield y s -> do
1957 z <- fstep acc y
1958 return $ Yield z (s, Just z)
1959 Skip s -> return $ Skip (s, Just acc)
1960 Stop -> return Stop
1962{-# INLINE scanl1 #-}
1963scanl1 :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
1964scanl1 f = scanl1M (\x y -> return (f x y))
1966{-# INLINE_NORMAL scanl1M' #-}
1967scanl1M' :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
1968scanl1M' fstep (Stream step state) = Stream step' (state, Nothing)
1969 where
1970 {-# INLINE_LATE step' #-}
1971 step' gst (st, Nothing) = do
1972 r <- step gst st
1973 case r of
1974 Yield x s -> x `seq` return $ Yield x (s, Just x)
1975 Skip s -> return $ Skip (s, Nothing)
1976 Stop -> return Stop
1978 step' gst (st, Just acc) = acc `seq` do
1979 r <- step gst st
1980 case r of
1981 Yield y s -> do
1982 z <- fstep acc y
1983 z `seq` return $ Yield z (s, Just z)
1984 Skip s -> return $ Skip (s, Just acc)
1985 Stop -> return Stop
1987{-# INLINE scanl1' #-}
1988scanl1' :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
1989scanl1' f = scanl1M' (\x y -> return (f x y))
1991{-# INLINE tap #-}
1992tap :: Monad m => Sink m a -> Stream m a -> Stream m a
1993tap (Sink fstep) (Stream step state) = Stream step' state
1995 where
1997 step' gst st = do
1998 r <- step gst st
1999 case r of
2000 Yield x s -> do
2001 fstep x
2002 return $ Yield x s
2003 Skip s -> return $ Skip s
2004 Stop -> return $ Stop
2006-------------------------------------------------------------------------------
2007-- Filtering
2008-------------------------------------------------------------------------------
2010{-# INLINE_NORMAL takeWhileM #-}
2011takeWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
2012takeWhileM f (Stream step state) = Stream step' state
2013 where
2014 {-# INLINE_LATE step' #-}
2015 step' gst st = do
2016 r <- step gst st
2017 case r of
2018 Yield x s -> do
2019 b <- f x
2020 return $ if b then Yield x s else Stop
2021 Skip s -> return $ Skip s
2022 Stop -> return Stop
2024{-# INLINE takeWhile #-}
2025takeWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
2026takeWhile f = takeWhileM (return . f)
2028{-# INLINE_NORMAL drop #-}
2029drop :: Monad m => Int -> Stream m a -> Stream m a
2030drop n (Stream step state) = Stream step' (state, Just n)
2031 where
2032 {-# INLINE_LATE step' #-}
2033 step' gst (st, Just i)
2034 | i > 0 = do
2035 r <- step gst st
2036 return $
2037 case r of
2038 Yield _ s -> Skip (s, Just (i - 1))
2039 Skip s -> Skip (s, Just i)
2040 Stop -> Stop
2041 | otherwise = return $ Skip (st, Nothing)
2043 step' gst (st, Nothing) = do
2044 r <- step gst st
2045 return $
2046 case r of
2047 Yield x s -> Yield x (s, Nothing)
2048 Skip s -> Skip (s, Nothing)
2049 Stop -> Stop
2051data DropWhileState s a
2052 = DropWhileDrop s
2053 | DropWhileYield a s
2054 | DropWhileNext s
2056{-# INLINE_NORMAL dropWhileM #-}
2057dropWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
2058dropWhileM f (Stream step state) = Stream step' (DropWhileDrop state)
2059 where
2060 {-# INLINE_LATE step' #-}
2061 step' gst (DropWhileDrop st) = do
2062 r <- step gst st
2063 case r of
2064 Yield x s -> do
2065 b <- f x
2066 if b
2067 then return $ Skip (DropWhileDrop s)
2068 else return $ Skip (DropWhileYield x s)
2069 Skip s -> return $ Skip (DropWhileDrop s)
2070 Stop -> return Stop
2072 step' gst (DropWhileNext st) = do
2073 r <- step gst st
2074 case r of
2075 Yield x s -> return $ Skip (DropWhileYield x s)
2076 Skip s -> return $ Skip (DropWhileNext s)
2077 Stop -> return Stop
2079 step' _ (DropWhileYield x st) = return $ Yield x (DropWhileNext st)
2081{-# INLINE dropWhile #-}
2082dropWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
2083dropWhile f = dropWhileM (return . f)
2085{-# INLINE_NORMAL filterM #-}
2086filterM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
2087filterM f (Stream step state) = Stream step' state
2088 where
2089 {-# INLINE_LATE step' #-}
2090 step' gst st = do
2091 r <- step gst st
2092 case r of
2093 Yield x s -> do
2094 b <- f x
2095 return $ if b
2096 then Yield x s
2097 else Skip s
2098 Skip s -> return $ Skip s
2099 Stop -> return Stop
2101{-# INLINE filter #-}
2102filter :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
2103filter f = filterM (return . f)
2105{-# INLINE_NORMAL uniq #-}
2106uniq :: (Eq a, Monad m) => Stream m a -> Stream m a
2107uniq (Stream step state) = Stream step' (Nothing, state)
2108 where
2109 {-# INLINE_LATE step' #-}
2110 step' gst (Nothing, st) = do
2111 r <- step gst st
2112 case r of
2113 Yield x s -> return $ Yield x (Just x, s)
2114 Skip s -> return $ Skip (Nothing, s)
2115 Stop -> return Stop
2116 step' gst (Just x, st) = do
2117 r <- step gst st
2118 case r of
2119 Yield y s | x == y -> return $ Skip (Just x, s)
2120 | otherwise -> return $ Yield x (Just y, s)
2121 Skip s -> return $ Skip (Just x, s)
2122 Stop -> return Stop
2124------------------------------------------------------------------------------
2125-- Transformation by Mapping
2126------------------------------------------------------------------------------
2128{-# INLINE_NORMAL sequence #-}
2129sequence :: Monad m => Stream m (m a) -> Stream m a
2130sequence (Stream step state) = Stream step' state
2131 where
2132 {-# INLINE_LATE step' #-}
2133 step' gst st = do
2134 r <- step (adaptState gst) st
2135 case r of
2136 Yield x s -> x >>= \a -> return (Yield a s)
2137 Skip s -> return $ Skip s
2138 Stop -> return Stop
2140------------------------------------------------------------------------------
2141-- Inserting
2142------------------------------------------------------------------------------
2144{-# INLINE_NORMAL insertBy #-}
2145insertBy :: Monad m => (a -> a -> Ordering) -> a -> Stream m a -> Stream m a
2146insertBy cmp a (Stream step state) = Stream step' (state, False, Nothing)
2147 where
2148 {-# INLINE_LATE step' #-}
2149 step' gst (st, False, _) = do
2150 r <- step gst st
2151 case r of
2152 Yield x s -> case cmp a x of
2153 GT -> return $ Yield x (s, False, Nothing)
2154 _ -> return $ Yield a (s, True, Just x)
2155 Skip s -> return $ Skip (s, False, Nothing)
2156 Stop -> return $ Yield a (st, True, Nothing)
2158 step' _ (_, True, Nothing) = return Stop
2160 step' gst (st, True, Just prev) = do
2161 r <- step gst st
2162 case r of
2163 Yield x s -> return $ Yield prev (s, True, Just x)
2164 Skip s -> return $ Skip (s, True, Just prev)
2165 Stop -> return $ Yield prev (st, True, Nothing)
2167------------------------------------------------------------------------------
2168-- Deleting
2169------------------------------------------------------------------------------
2171{-# INLINE_NORMAL deleteBy #-}
2172deleteBy :: Monad m => (a -> a -> Bool) -> a -> Stream m a -> Stream m a
2173deleteBy eq x (Stream step state) = Stream step' (state, False)
2174 where
2175 {-# INLINE_LATE step' #-}
2176 step' gst (st, False) = do
2177 r <- step gst st
2178 case r of
2179 Yield y s -> return $
2180 if eq x y then Skip (s, True) else Yield y (s, False)
2181 Skip s -> return $ Skip (s, False)
2182 Stop -> return Stop
2184 step' gst (st, True) = do
2185 r <- step gst st
2186 case r of
2187 Yield y s -> return $ Yield y (s, True)
2188 Skip s -> return $ Skip (s, True)
2189 Stop -> return Stop
2191------------------------------------------------------------------------------
2192-- Transformation by Map and Filter
2193------------------------------------------------------------------------------
2195-- XXX Will this always fuse properly?
2196{-# INLINE_NORMAL mapMaybe #-}
2197mapMaybe :: Monad m => (a -> Maybe b) -> Stream m a -> Stream m b
2198mapMaybe f = fmap fromJust . filter isJust . map f
2200{-# INLINE_NORMAL mapMaybeM #-}
2201mapMaybeM :: Monad m => (a -> m (Maybe b)) -> Stream m a -> Stream m b
2202mapMaybeM f = fmap fromJust . filter isJust . mapM f
2204------------------------------------------------------------------------------
2205-- Zipping
2206------------------------------------------------------------------------------
2208{-# INLINE_NORMAL indexed #-}
2209indexed :: Monad m => Stream m a -> Stream m (Int, a)
2210indexed (Stream step state) = Stream step' (state, 0)
2211 where
2212 {-# INLINE_LATE step' #-}
2213 step' gst (st, i) = i `seq` do
2214 r <- step (adaptState gst) st
2215 case r of
2216 Yield x s -> return $ Yield (i, x) (s, i+1)
2217 Skip s -> return $ Skip (s, i)
2218 Stop -> return Stop
2220{-# INLINE_NORMAL indexedR #-}
2221indexedR :: Monad m => Int -> Stream m a -> Stream m (Int, a)
2222indexedR m (Stream step state) = Stream step' (state, m)
2223 where
2224 {-# INLINE_LATE step' #-}
2225 step' gst (st, i) = i `seq` do
2226 r <- step (adaptState gst) st
2227 case r of
2228 Yield x s -> let i' = i - 1
2229 in return $ Yield (i, x) (s, i')
2230 Skip s -> return $ Skip (s, i)
2231 Stop -> return Stop
2233{-# INLINE_NORMAL zipWithM #-}
2234zipWithM :: Monad m
2235 => (a -> b -> m c) -> Stream m a -> Stream m b -> Stream m c
2236zipWithM f (Stream stepa ta) (Stream stepb tb) = Stream step (ta, tb, Nothing)
2237 where
2238 {-# INLINE_LATE step #-}
2239 step gst (sa, sb, Nothing) = do
2240 r <- stepa (adaptState gst) sa
2241 return $
2242 case r of
2243 Yield x sa' -> Skip (sa', sb, Just x)
2244 Skip sa' -> Skip (sa', sb, Nothing)
2245 Stop -> Stop
2247 step gst (sa, sb, Just x) = do
2248 r <- stepb (adaptState gst) sb
2249 case r of
2250 Yield y sb' -> do
2251 z <- f x y
2252 return $ Yield z (sa, sb', Nothing)
2253 Skip sb' -> return $ Skip (sa, sb', Just x)
2254 Stop -> return Stop
2256{-# RULES "zipWithM xs xs"
2257 forall f xs. zipWithM f xs xs = mapM (\x -> f x x) xs #-}
2259{-# INLINE zipWith #-}
2260zipWith :: Monad m => (a -> b -> c) -> Stream m a -> Stream m b -> Stream m c
2261zipWith f = zipWithM (\a b -> return (f a b))
2263------------------------------------------------------------------------------
2264-- Merging
2265------------------------------------------------------------------------------
2267{-# INLINE_NORMAL mergeByM #-}
2268mergeByM
2269 :: (Monad m)
2270 => (a -> a -> m Ordering) -> Stream m a -> Stream m a -> Stream m a
2271mergeByM cmp (Stream stepa ta) (Stream stepb tb) =
2272 Stream step (Just ta, Just tb, Nothing, Nothing)
2273 where
2274 {-# INLINE_LATE step #-}
2276 -- one of the values is missing, and the corresponding stream is running
2277 step gst (Just sa, sb, Nothing, b) = do
2278 r <- stepa gst sa
2279 return $ case r of
2280 Yield a sa' -> Skip (Just sa', sb, Just a, b)
2281 Skip sa' -> Skip (Just sa', sb, Nothing, b)
2282 Stop -> Skip (Nothing, sb, Nothing, b)
2284 step gst (sa, Just sb, a, Nothing) = do
2285 r <- stepb gst sb
2286 return $ case r of
2287 Yield b sb' -> Skip (sa, Just sb', a, Just b)
2288 Skip sb' -> Skip (sa, Just sb', a, Nothing)
2289 Stop -> Skip (sa, Nothing, a, Nothing)
2291 -- both the values are available
2292 step _ (sa, sb, Just a, Just b) = do
2293 res <- cmp a b
2294 return $ case res of
2295 GT -> Yield b (sa, sb, Just a, Nothing)
2296 _ -> Yield a (sa, sb, Nothing, Just b)
2298 -- one of the values is missing, corresponding stream is done
2299 step _ (Nothing, sb, Nothing, Just b) =
2300 return $ Yield b (Nothing, sb, Nothing, Nothing)
2302 step _ (sa, Nothing, Just a, Nothing) =
2303 return $ Yield a (sa, Nothing, Nothing, Nothing)
2305 step _ (Nothing, Nothing, Nothing, Nothing) = return Stop
2307{-# INLINE mergeBy #-}
2308mergeBy
2309 :: (Monad m)
2310 => (a -> a -> Ordering) -> Stream m a -> Stream m a -> Stream m a
2311mergeBy cmp = mergeByM (\a b -> return $ cmp a b)
2313------------------------------------------------------------------------------
2314-- Transformation comprehensions
2315------------------------------------------------------------------------------
2317{-# INLINE_NORMAL the #-}
2318the :: (Eq a, Monad m) => Stream m a -> m (Maybe a)
2319the (Stream step state) = go state
2320 where
2321 go st = do
2322 r <- step defState st
2323 case r of
2324 Yield x s -> go' x s
2325 Skip s -> go s
2326 Stop -> return Nothing
2327 go' n st = do
2328 r <- step defState st
2329 case r of
2330 Yield x s | x == n -> go' n s
2331 | otherwise -> return Nothing
2332 Skip s -> go' n s
2333 Stop -> return (Just n)