git.iko.soy/streamly
on 2255dac help sign in
streamly/src/Streamly/Streams/StreamD.hs
2638 lines 90 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 #-}
13{-# LANGUAGE MagicHash #-}
15#include "inline.hs"
17-- |
18-- Module : Streamly.Streams.StreamD
19-- Copyright : (c) 2018 Harendra Kumar
20-- Copyright : (c) Roman Leshchinskiy 2008-2010
21--
22-- License : BSD3
23-- Maintainer : harendra.kumar@gmail.com
24-- Stability : experimental
25-- Portability : GHC
26--
27-- Direct style re-implementation of CPS style stream in StreamK module. The
28-- symbol or suffix 'D' in this module denotes the "Direct" style. GHC is able
29-- to INLINE and fuse direct style better, providing better performance than
30-- CPS implementation.
31--
32-- @
33-- import qualified Streamly.Streams.StreamD as D
34-- @
36-- Some of the functions in this file have been adapted from the vector
37-- library, https://hackage.haskell.org/package/vector.
39module Streamly.Streams.StreamD
40 (
41 -- * The stream type
42 Step (..)
43 , Stream (..)
45 -- * Construction
46 , nil
47 , cons
49 -- * Deconstruction
50 , uncons
52 -- * Generation
53 -- ** Unfolds
54 , unfoldr
55 , unfoldrM
57 -- ** Specialized Generation
58 -- | Generate a monadic stream from a seed.
59 , repeat
60 , replicate
61 , replicateM
62 , fromIndices
63 , fromIndicesM
64 , generate
65 , generateM
67 -- ** Enumerations
68 , enumerateFromStepIntegral
69 , enumerateFromIntegral
70 , enumerateFromThenIntegral
71 , enumerateFromToIntegral
72 , enumerateFromThenToIntegral
74 , enumerateFromStepNum
75 , numFrom
76 , numFromThen
77 , enumerateFromToFractional
78 , enumerateFromThenToFractional
80 -- ** Conversions
81 -- | Transform an input structure into a stream.
82 -- | Direct style stream does not support @fromFoldable@.
83 , yield
84 , yieldM
85 , fromList
86 , fromListM
87 , fromStreamK
88 , fromStreamD
90 -- * Elimination
91 -- ** General Folds
92 , foldrS
93 , foldrT
94 , foldrM
95 , foldrMx
96 , foldr
97 , foldr1
99 , foldl'
100 , foldlM'
101 , foldlS
102 , foldlT
103 , reverse
104 , reverse'
106 , foldlx'
107 , foldlMx'
109 -- ** Specialized Folds
110 , tap
111 , drain
112 , null
113 , head
114 , tail
115 , last
116 , elem
117 , notElem
118 , all
119 , any
120 , maximum
121 , maximumBy
122 , minimum
123 , minimumBy
124 , findIndices
125 , lookup
126 , findM
127 , find
128 , (!!)
130 -- ** Flattening nested streams
131 , concatMapM
132 , concatMap
134 -- ** Grouping
135 , groupsOf
136 , groupsBy
137 , groupsRollingBy
139 -- ** Splitting
140 , splitBy
141 , splitSuffixBy
142 , wordsBy
143 , splitSuffixBy'
145 , splitOn
146 , splitSuffixOn
148 -- ** Substreams
149 , isPrefixOf
150 , isSubsequenceOf
151 , stripPrefix
153 -- ** Map and Fold
154 , mapM_
156 -- ** Conversions
157 -- | Transform a stream into another type.
158 , toList
159 , toListRev
160 , toStreamK
161 , toStreamD
163 -- * Transformation
164 , transform
166 -- ** By folding (scans)
167 , scanlM'
168 , scanl'
169 , scanlM
170 , scanl
171 , scanl1M'
172 , scanl1'
173 , scanl1M
174 , scanl1
176 , prescanl'
177 , prescanlM'
179 , postscanl
180 , postscanlM
181 , postscanl'
182 , postscanlM'
184 , postscanlx'
185 , postscanlMx'
186 , scanlMx'
187 , scanlx'
189 -- * Filtering
190 , filter
191 , filterM
192 , uniq
193 , take
194 , takeWhile
195 , takeWhileM
196 , drop
197 , dropWhile
198 , dropWhileM
200 -- * Mapping
201 , map
202 , mapM
203 , sequence
205 -- * Inserting
206 , intersperseM
207 , intersperse
208 , insertBy
210 -- * Deleting
211 , deleteBy
213 -- ** Map and Filter
214 , mapMaybe
215 , mapMaybeM
217 -- * Zipping
218 , indexed
219 , indexedR
220 , zipWith
221 , zipWithM
223 -- * Comparisons
224 , eqBy
225 , cmpBy
227 -- * Merging
228 , mergeBy
229 , mergeByM
231 -- * Transformation comprehensions
232 , the
234 -- * Exceptions
235 , before
236 , after
237 , bracket
238 , onException
239 , finally
240 , handle
242 -- * UTF8 Encoding / Decoding transformations.
243 , decodeUtf8
244 , encodeUtf8
245 )
246where
248import Control.Exception (Exception)
249import Control.Monad.Catch (MonadCatch)
250import Control.Monad.IO.Class (MonadIO(..))
251import Control.Monad.Trans (MonadTrans(lift))
252import Data.Bits (shiftR, shiftL, (.|.), (.&.))
253import Data.Maybe (fromJust, isJust)
254import Data.Word (Word32)
255import Foreign.Ptr (Ptr)
256import Foreign.Storable (Storable(..))
257import GHC.Base (assert, Char(..), unsafeChr, (-#), (+#), uncheckedIShiftL#,
258 word2Int#, chr#, ord)
259import GHC.IO.Encoding.Failure (isSurrogate)
260import GHC.Types (SPEC(..))
261import GHC.Word (Word8(..))
262import Prelude
263 hiding (map, mapM, mapM_, repeat, foldr, last, take, filter,
264 takeWhile, drop, dropWhile, all, any, maximum, minimum, elem,
265 notElem, null, head, tail, zipWith, lookup, foldr1, sequence,
266 (!!), scanl, scanl1, concatMap, replicate, enumFromTo, concat,
267 reverse)
269import qualified Control.Monad.Catch as MC
271import Streamly.Mem.Array.Types (Array(..))
272import Streamly.Fold.Types (Fold(..))
273import Streamly.Pipe.Types (Pipe(..), PipeState(..))
274import Streamly.SVar (MonadAsync, defState, adaptState)
275import Streamly.Sink.Types (Sink(..))
277import Streamly.Streams.StreamD.Type
279import qualified Streamly.Pipe.Types as Pipe
280import qualified Streamly.Mem.Array.Types as A
281import qualified Streamly.Mem.Ring as RB
282import qualified Streamly.Streams.StreamK as K
284------------------------------------------------------------------------------
285-- Construction
286------------------------------------------------------------------------------
288-- | An empty 'Stream'.
289{-# INLINE_NORMAL nil #-}
290nil :: Monad m => Stream m a
291nil = Stream (\_ _ -> return Stop) ()
293{-# INLINE_NORMAL consM #-}
294consM :: Monad m => m a -> Stream m a -> Stream m a
295consM m (Stream step state) = Stream step1 Nothing
296 where
297 {-# INLINE_LATE step1 #-}
298 step1 _ Nothing = m >>= \x -> return $ Yield x (Just state)
299 step1 gst (Just st) = do
300 r <- step gst st
301 return $
302 case r of
303 Yield a s -> Yield a (Just s)
304 Skip s -> Skip (Just s)
305 Stop -> Stop
307-- XXX implement in terms of consM?
308-- cons x = consM (return x)
309--
310-- | Can fuse but has O(n^2) complexity.
311{-# INLINE_NORMAL cons #-}
312cons :: Monad m => a -> Stream m a -> Stream m a
313cons x (Stream step state) = Stream step1 Nothing
314 where
315 {-# INLINE_LATE step1 #-}
316 step1 _ Nothing = return $ Yield x (Just state)
317 step1 gst (Just st) = do
318 r <- step gst st
319 return $
320 case r of
321 Yield a s -> Yield a (Just s)
322 Skip s -> Skip (Just s)
323 Stop -> Stop
325-------------------------------------------------------------------------------
326-- Deconstruction
327-------------------------------------------------------------------------------
329-- Does not fuse, has the same performance as the StreamK version.
330{-# INLINE_NORMAL uncons #-}
331uncons :: Monad m => Stream m a -> m (Maybe (a, Stream m a))
332uncons (UnStream step state) = go state
333 where
334 go st = do
335 r <- step defState st
336 case r of
337 Yield x s -> return $ Just (x, Stream step s)
338 Skip s -> go s
339 Stop -> return Nothing
341------------------------------------------------------------------------------
342-- Generation by unfold
343------------------------------------------------------------------------------
345{-# INLINE_NORMAL unfoldrM #-}
346unfoldrM :: Monad m => (s -> m (Maybe (a, s))) -> s -> Stream m a
347unfoldrM next state = Stream step state
348 where
349 {-# INLINE_LATE step #-}
350 step _ st = do
351 r <- next st
352 return $ case r of
353 Just (x, s) -> Yield x s
354 Nothing -> Stop
356{-# INLINE_LATE unfoldr #-}
357unfoldr :: Monad m => (s -> Maybe (a, s)) -> s -> Stream m a
358unfoldr f = unfoldrM (return . f)
360------------------------------------------------------------------------------
361-- Specialized Generation
362------------------------------------------------------------------------------
364repeat :: Monad m => a -> Stream m a
365repeat x = Stream (\_ _ -> return $ Yield x ()) ()
367{-# INLINE_NORMAL replicateM #-}
368replicateM :: Monad m => Int -> m a -> Stream m a
369replicateM n p = Stream step n
370 where
371 {-# INLINE_LATE step #-}
372 step _ i | i <= 0 = return Stop
373 | otherwise = do
374 x <- p
375 return $ Yield x (i - 1)
377{-# INLINE_NORMAL replicate #-}
378replicate :: Monad m => Int -> a -> Stream m a
379replicate n x = replicateM n (return x)
381-- This would not work properly for floats, therefore we put an Integral
382-- constraint.
383-- | Can be used to enumerate unbounded integrals. This does not check for
384-- overflow or underflow for bounded integrals.
385{-# INLINE_NORMAL enumerateFromStepIntegral #-}
386enumerateFromStepIntegral :: (Integral a, Monad m) => a -> a -> Stream m a
387enumerateFromStepIntegral from stride =
388 from `seq` stride `seq` Stream step from
389 where
390 {-# INLINE_LATE step #-}
391 step _ !x = return $ Yield x $! (x + stride)
393-- We are assuming that "to" is constrained by the type to be within
394-- max/min bounds.
395{-# INLINE enumerateFromToIntegral #-}
396enumerateFromToIntegral :: (Monad m, Integral a) => a -> a -> Stream m a
397enumerateFromToIntegral from to =
398 takeWhile (<= to) $ enumerateFromStepIntegral from 1
400{-# INLINE enumerateFromIntegral #-}
401enumerateFromIntegral :: (Monad m, Integral a, Bounded a) => a -> Stream m a
402enumerateFromIntegral from = enumerateFromToIntegral from maxBound
404data EnumState a = EnumInit | EnumYield a a a | EnumStop
406{-# INLINE_NORMAL enumerateFromThenToIntegralUp #-}
407enumerateFromThenToIntegralUp
408 :: (Monad m, Integral a)
409 => a -> a -> a -> Stream m a
410enumerateFromThenToIntegralUp from next to = Stream step EnumInit
411 where
412 {-# INLINE_LATE step #-}
413 step _ EnumInit =
414 return $
415 if to < next
416 then if to < from
417 then Stop
418 else Yield from EnumStop
419 else -- from <= next <= to
420 let stride = next - from
421 in Skip $ EnumYield from stride (to - stride)
423 step _ (EnumYield x stride toMinus) =
424 return $
425 if x > toMinus
426 then Yield x EnumStop
427 else Yield x $ EnumYield (x + stride) stride toMinus
429 step _ EnumStop = return Stop
431{-# INLINE_NORMAL enumerateFromThenToIntegralDn #-}
432enumerateFromThenToIntegralDn
433 :: (Monad m, Integral a)
434 => a -> a -> a -> Stream m a
435enumerateFromThenToIntegralDn from next to = Stream step EnumInit
436 where
437 {-# INLINE_LATE step #-}
438 step _ EnumInit =
439 return $ if to > next
440 then if to > from
441 then Stop
442 else Yield from EnumStop
443 else -- from >= next >= to
444 let stride = next - from
445 in Skip $ EnumYield from stride (to - stride)
447 step _ (EnumYield x stride toMinus) =
448 return $
449 if x < toMinus
450 then Yield x EnumStop
451 else Yield x $ EnumYield (x + stride) stride toMinus
453 step _ EnumStop = return Stop
455{-# INLINE_NORMAL enumerateFromThenToIntegral #-}
456enumerateFromThenToIntegral
457 :: (Monad m, Integral a)
458 => a -> a -> a -> Stream m a
459enumerateFromThenToIntegral from next to
460 | next >= from = enumerateFromThenToIntegralUp from next to
461 | otherwise = enumerateFromThenToIntegralDn from next to
463{-# INLINE_NORMAL enumerateFromThenIntegral #-}
464enumerateFromThenIntegral
465 :: (Monad m, Integral a, Bounded a)
466 => a -> a -> Stream m a
467enumerateFromThenIntegral from next =
468 if next > from
469 then enumerateFromThenToIntegralUp from next maxBound
470 else enumerateFromThenToIntegralDn from next minBound
472-- For floating point numbers if the increment is less than the precision then
473-- it just gets lost. Therefore we cannot always increment it correctly by just
474-- repeated addition.
475-- 9007199254740992 + 1 + 1 :: Double => 9.007199254740992e15
476-- 9007199254740992 + 2 :: Double => 9.007199254740994e15
478-- Instead we accumulate the increment counter and compute the increment
479-- everytime before adding it to the starting number.
480--
481-- This works for Integrals as well as floating point numbers, but
482-- enumerateFromStepIntegral is faster for integrals.
483{-# INLINE_NORMAL enumerateFromStepNum #-}
484enumerateFromStepNum :: (Monad m, Num a) => a -> a -> Stream m a
485enumerateFromStepNum from stride = Stream step 0
486 where
487 {-# INLINE_LATE step #-}
488 step _ !i = return $ (Yield $! (from + i * stride)) $! (i + 1)
490{-# INLINE_NORMAL numFrom #-}
491numFrom :: (Monad m, Num a) => a -> Stream m a
492numFrom from = enumerateFromStepNum from 1
494{-# INLINE_NORMAL numFromThen #-}
495numFromThen :: (Monad m, Num a) => a -> a -> Stream m a
496numFromThen from next = enumerateFromStepNum from (next - from)
498-- We cannot write a general function for Num. The only way to write code
499-- portable between the two is to use a 'Real' constraint and convert between
500-- Fractional and Integral using fromRational which is horribly slow.
501{-# INLINE_NORMAL enumerateFromToFractional #-}
502enumerateFromToFractional
503 :: (Monad m, Fractional a, Ord a)
504 => a -> a -> Stream m a
505enumerateFromToFractional from to =
506 takeWhile (<= to + 1 / 2) $ enumerateFromStepNum from 1
508{-# INLINE_NORMAL enumerateFromThenToFractional #-}
509enumerateFromThenToFractional
510 :: (Monad m, Fractional a, Ord a)
511 => a -> a -> a -> Stream m a
512enumerateFromThenToFractional from next to =
513 takeWhile predicate $ numFromThen from next
514 where
515 mid = (next - from) / 2
516 predicate | next >= from = (<= to + mid)
517 | otherwise = (>= to + mid)
519-------------------------------------------------------------------------------
520-- Generation by Conversion
521-------------------------------------------------------------------------------
523{-# INLINE_NORMAL fromIndicesM #-}
524fromIndicesM :: Monad m => (Int -> m a) -> Stream m a
525fromIndicesM gen = Stream step 0
526 where
527 {-# INLINE_LATE step #-}
528 step _ i = do
529 x <- gen i
530 return $ Yield x (i + 1)
532{-# INLINE fromIndices #-}
533fromIndices :: Monad m => (Int -> a) -> Stream m a
534fromIndices gen = fromIndicesM (return . gen)
536{-# INLINE_NORMAL generateM #-}
537generateM :: Monad m => Int -> (Int -> m a) -> Stream m a
538generateM n gen = n `seq` Stream step 0
539 where
540 {-# INLINE_LATE step #-}
541 step _ i | i < n = do
542 x <- gen i
543 return $ Yield x (i + 1)
544 | otherwise = return Stop
546{-# INLINE generate #-}
547generate :: Monad m => Int -> (Int -> a) -> Stream m a
548generate n gen = generateM n (return . gen)
550-- XXX we need the MonadAsync constraint because of a rewrite rule.
551-- | Convert a list of monadic actions to a 'Stream'
552{-# INLINE_LATE fromListM #-}
553fromListM :: MonadAsync m => [m a] -> Stream m a
554fromListM = Stream step
555 where
556 {-# INLINE_LATE step #-}
557 step _ (m:ms) = m >>= \x -> return $ Yield x ms
558 step _ [] = return Stop
560{-# INLINE toStreamD #-}
561toStreamD :: (K.IsStream t, Monad m) => t m a -> Stream m a
562toStreamD = fromStreamK . K.toStream
564------------------------------------------------------------------------------
565-- Elimination by Folds
566------------------------------------------------------------------------------
568------------------------------------------------------------------------------
569-- Right Folds
570------------------------------------------------------------------------------
572{-# INLINE_NORMAL foldr1 #-}
573foldr1 :: Monad m => (a -> a -> a) -> Stream m a -> m (Maybe a)
574foldr1 f m = do
575 r <- uncons m
576 case r of
577 Nothing -> return Nothing
578 Just (h, t) -> fmap Just (foldr f h t)
580------------------------------------------------------------------------------
581-- Left Folds
582------------------------------------------------------------------------------
584{-# INLINE_NORMAL foldlT #-}
585foldlT :: (Monad m, Monad (s m), MonadTrans s)
586 => (s m b -> a -> s m b) -> s m b -> Stream m a -> s m b
587foldlT fstep begin (Stream step state) = go SPEC begin state
588 where
589 go !_ acc st = do
590 r <- lift $ step defState st
591 case r of
592 Yield x s -> go SPEC (fstep acc x) s
593 Skip s -> go SPEC acc s
594 Stop -> acc
596-- Note, this is going to have horrible performance, because of the nature of
597-- the stream type (i.e. direct stream vs CPS). Its only for reference, it is
598-- likely be practically unusable.
599{-# INLINE_NORMAL foldlS #-}
600foldlS :: Monad m
601 => (Stream m b -> a -> Stream m b) -> Stream m b -> Stream m a -> Stream m b
602foldlS fstep begin (Stream step state) = Stream step' (Left (state, begin))
603 where
604 step' gst (Left (st, acc)) = do
605 r <- step (adaptState gst) st
606 return $ case r of
607 Yield x s -> Skip (Left (s, fstep acc x))
608 Skip s -> Skip (Left (s, acc))
609 Stop -> Skip (Right acc)
611 step' gst (Right (Stream stp stt)) = do
612 r <- stp (adaptState gst) stt
613 return $ case r of
614 Yield x s -> Yield x (Right (Stream stp s))
615 Skip s -> Skip (Right (Stream stp s))
616 Stop -> Stop
618------------------------------------------------------------------------------
619-- Specialized Folds
620------------------------------------------------------------------------------
622-- | Run a streaming composition, discard the results.
623{-# INLINE_LATE drain #-}
624drain :: Monad m => Stream m a -> m ()
625-- drain m = foldrM (\x xs -> x `seq` xs) (return ()) m
626drain (Stream step state) = go SPEC state
627 where
628 go !_ st = do
629 r <- step defState st
630 case r of
631 Yield _ s -> go SPEC s
632 Skip s -> go SPEC s
633 Stop -> return ()
635{-# INLINE_NORMAL null #-}
636null :: Monad m => Stream m a -> m Bool
637null m = foldrM (\_ _ -> return False) (return True) m
639-- XXX SPEC?
640{-# INLINE_NORMAL head #-}
641head :: Monad m => Stream m a -> m (Maybe a)
642head m = foldrM (\x _ -> return (Just x)) (return Nothing) m
644-- Does not fuse, has the same performance as the StreamK version.
645{-# INLINE_NORMAL tail #-}
646tail :: Monad m => Stream m a -> m (Maybe (Stream m a))
647tail (UnStream step state) = go state
648 where
649 go st = do
650 r <- step defState st
651 case r of
652 Yield _ s -> return (Just $ Stream step s)
653 Skip s -> go s
654 Stop -> return Nothing
656-- XXX will it fuse? need custom impl?
657{-# INLINE_NORMAL last #-}
658last :: Monad m => Stream m a -> m (Maybe a)
659last = foldl' (\_ y -> Just y) Nothing
661{-# INLINE_NORMAL elem #-}
662elem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
663-- elem e m = foldrM (\x xs -> if x == e then return True else xs) (return False) m
664elem e (Stream step state) = go state
665 where
666 go st = do
667 r <- step defState st
668 case r of
669 Yield x s
670 | x == e -> return True
671 | otherwise -> go s
672 Skip s -> go s
673 Stop -> return False
675{-# INLINE_NORMAL notElem #-}
676notElem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
677notElem e s = fmap not (elem e s)
679{-# INLINE_NORMAL all #-}
680all :: Monad m => (a -> Bool) -> Stream m a -> m Bool
681-- all p m = foldrM (\x xs -> if p x then xs else return False) (return True) m
682all p (Stream step state) = go state
683 where
684 go st = do
685 r <- step defState st
686 case r of
687 Yield x s
688 | p x -> go s
689 | otherwise -> return False
690 Skip s -> go s
691 Stop -> return True
693{-# INLINE_NORMAL any #-}
694any :: Monad m => (a -> Bool) -> Stream m a -> m Bool
695-- any p m = foldrM (\x xs -> if p x then return True else xs) (return False) m
696any p (Stream step state) = go state
697 where
698 go st = do
699 r <- step defState st
700 case r of
701 Yield x s
702 | p x -> return True
703 | otherwise -> go s
704 Skip s -> go s
705 Stop -> return False
707{-# INLINE_NORMAL maximum #-}
708maximum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
709maximum (Stream step state) = go Nothing state
710 where
711 go Nothing st = do
712 r <- step defState st
713 case r of
714 Yield x s -> go (Just x) s
715 Skip s -> go Nothing s
716 Stop -> return Nothing
717 go (Just acc) st = do
718 r <- step defState st
719 case r of
720 Yield x s
721 | acc <= x -> go (Just x) s
722 | otherwise -> go (Just acc) s
723 Skip s -> go (Just acc) s
724 Stop -> return (Just acc)
726{-# INLINE_NORMAL maximumBy #-}
727maximumBy :: Monad m => (a -> a -> Ordering) -> Stream m a -> m (Maybe a)
728maximumBy cmp (Stream step state) = go Nothing state
729 where
730 go Nothing st = do
731 r <- step defState st
732 case r of
733 Yield x s -> go (Just x) s
734 Skip s -> go Nothing s
735 Stop -> return Nothing
736 go (Just acc) st = do
737 r <- step defState st
738 case r of
739 Yield x s -> case cmp acc x of
740 GT -> go (Just acc) s
741 _ -> go (Just x) s
742 Skip s -> go (Just acc) s
743 Stop -> return (Just acc)
745{-# INLINE_NORMAL minimum #-}
746minimum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
747minimum (Stream step state) = go Nothing state
748 where
749 go Nothing st = do
750 r <- step defState st
751 case r of
752 Yield x s -> go (Just x) s
753 Skip s -> go Nothing s
754 Stop -> return Nothing
755 go (Just acc) st = do
756 r <- step defState st
757 case r of
758 Yield x s
759 | acc <= x -> go (Just acc) s
760 | otherwise -> go (Just x) s
761 Skip s -> go (Just acc) s
762 Stop -> return (Just acc)
764{-# INLINE_NORMAL minimumBy #-}
765minimumBy :: Monad m => (a -> a -> Ordering) -> Stream m a -> m (Maybe a)
766minimumBy cmp (Stream step state) = go Nothing state
767 where
768 go Nothing st = do
769 r <- step defState st
770 case r of
771 Yield x s -> go (Just x) s
772 Skip s -> go Nothing s
773 Stop -> return Nothing
774 go (Just acc) st = do
775 r <- step defState st
776 case r of
777 Yield x s -> case cmp acc x of
778 GT -> go (Just x) s
779 _ -> go (Just acc) s
780 Skip s -> go (Just acc) s
781 Stop -> return (Just acc)
783{-# INLINE_NORMAL (!!) #-}
784(!!) :: (Monad m) => Stream m a -> Int -> m (Maybe a)
785(Stream step state) !! i = go i state
786 where
787 go n st = do
788 r <- step defState st
789 case r of
790 Yield x s | n < 0 -> return Nothing
791 | n == 0 -> return $ Just x
792 | otherwise -> go (n - 1) s
793 Skip s -> go n s
794 Stop -> return Nothing
796{-# INLINE_NORMAL lookup #-}
797lookup :: (Monad m, Eq a) => a -> Stream m (a, b) -> m (Maybe b)
798lookup e m = foldrM (\(a, b) xs -> if e == a then return (Just b) else xs)
799 (return Nothing) m
801{-# INLINE_NORMAL findM #-}
802findM :: Monad m => (a -> m Bool) -> Stream m a -> m (Maybe a)
803findM p m = foldrM (\x xs -> p x >>= \r -> if r then return (Just x) else xs)
804 (return Nothing) m
806{-# INLINE find #-}
807find :: Monad m => (a -> Bool) -> Stream m a -> m (Maybe a)
808find p = findM (return . p)
810{-# INLINE_NORMAL findIndices #-}
811findIndices :: Monad m => (a -> Bool) -> Stream m a -> Stream m Int
812findIndices p (Stream step state) = Stream step' (state, 0)
813 where
814 {-# INLINE_LATE step' #-}
815 step' gst (st, i) = do
816 r <- step (adaptState gst) st
817 return $ case r of
818 Yield x s -> if p x then Yield i (s, i+1) else Skip (s, i+1)
819 Skip s -> Skip (s, i+1)
820 Stop -> Stop
822{-# INLINE toListRev #-}
823toListRev :: Monad m => Stream m a -> m [a]
824toListRev = foldl' (flip (:)) []
826-- We can implement reverse as:
827--
828-- > reverse = foldlS (flip cons) nil
829--
830-- However, this implementation is unusable because of the horrible performance
831-- of cons. So we just convert it to a list first and then stream from the
832-- list.
833--
834-- XXX Maybe we can use an Array instead of a list here?
835{-# INLINE_NORMAL reverse #-}
836reverse :: Monad m => Stream m a -> Stream m a
837reverse m = Stream step Nothing
838 where
839 {-# INLINE_LATE step #-}
840 step _ Nothing = do
841 xs <- toListRev m
842 return $ Skip (Just xs)
843 step _ (Just (x:xs)) = return $ Yield x (Just xs)
844 step _ (Just []) = return Stop
846-- Much faster reverse for Storables
847{-# INLINE_NORMAL reverse' #-}
848reverse' :: forall m a. (MonadIO m, Storable a) => Stream m a -> Stream m a
849{-
850-- This commented implementation copies the whole stream into one single array
851-- and then streams from that array, this is 3-4x faster than the chunked code
852-- that follows. Though this could be problematic due to unbounded large
853-- allocations. We need to figure out why the chunked code is slower and if we
854-- can optimize the chunked code to work as fast as this one. It may be a
855-- fusion issue?
856import Foreign.ForeignPtr (touchForeignPtr)
857import Foreign.ForeignPtr.Unsafe (unsafeForeignPtrToPtr)
858import Foreign.Ptr (Ptr, plusPtr)
859reverse' m = Stream step Nothing
860 where
861 {-# INLINE_LATE step #-}
862 step _ Nothing = do
863 arr <- A.fromStreamD m
864 let p = aEnd arr `plusPtr` negate (sizeOf (undefined :: a))
865 return $ Skip $ Just (aStart arr, p)
867 step _ (Just (start, p)) | p < unsafeForeignPtrToPtr start = return Stop
869 step _ (Just (start, p)) = do
870 let !x = A.unsafeInlineIO $ do
871 r <- peek p
872 touchForeignPtr start
873 return r
874 next = p `plusPtr` negate (sizeOf (undefined :: a))
875 return $ Yield x (Just (start, next))
876-}
877reverse' m =
878 A.flattenArraysRev
879 $ fromStreamK
880 $ K.reverse
881 $ toStreamK
882 $ A.fromStreamDArraysOf A.defaultChunkSize m
885------------------------------------------------------------------------------
886-- Grouping/Splitting
887------------------------------------------------------------------------------
889{-# INLINE_NORMAL splitSuffixBy' #-}
890splitSuffixBy' :: Monad m
891 => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
892splitSuffixBy' predicate f (Stream step state) =
893 Stream (stepOuter f) (Just state)
895 where
897 {-# INLINE_LATE stepOuter #-}
898 stepOuter (Fold fstep initial done) gst (Just st) = do
899 res <- step (adaptState gst) st
900 case res of
901 Yield x s -> do
902 acc <- initial
903 acc' <- fstep acc x
904 if (predicate x)
905 then done acc' >>= \val -> return $ Yield val (Just s)
906 else go SPEC s acc'
908 Skip s -> return $ Skip $ Just s
909 Stop -> return Stop
911 where
913 go !_ stt !acc = do
914 res <- step (adaptState gst) stt
915 case res of
916 Yield x s -> do
917 acc' <- fstep acc x
918 if (predicate x)
919 then done acc' >>= \val -> return $ Yield val (Just s)
920 else go SPEC s acc'
921 Skip s -> go SPEC s acc
922 Stop -> done acc >>= \val -> return $ Yield val Nothing
924 stepOuter _ _ Nothing = return Stop
926{-# INLINE_NORMAL groupsBy #-}
927groupsBy :: Monad m
928 => (a -> a -> Bool)
929 -> Fold m a b
930 -> Stream m a
931 -> Stream m b
932groupsBy cmp f (Stream step state) = Stream (stepOuter f) (Just state, Nothing)
934 where
936 {-# INLINE_LATE stepOuter #-}
937 stepOuter (Fold fstep initial done) gst (Just st, Nothing) = do
938 res <- step (adaptState gst) st
939 case res of
940 Yield x s -> do
941 acc <- initial
942 acc' <- fstep acc x
943 go SPEC x s acc'
945 Skip s -> return $ Skip $ (Just s, Nothing)
946 Stop -> return Stop
948 where
950 go !_ prev stt !acc = do
951 res <- step (adaptState gst) stt
952 case res of
953 Yield x s -> do
954 if cmp x prev
955 then do
956 acc' <- fstep acc x
957 go SPEC prev s acc'
958 else done acc >>= \r -> return $ Yield r (Just s, Just x)
959 Skip s -> go SPEC prev s acc
960 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
962 stepOuter (Fold fstep initial done) gst (Just st, Just prev) = do
963 acc <- initial
964 acc' <- fstep acc prev
965 go SPEC st acc'
967 where
969 -- XXX code duplicated from the previous equation
970 go !_ stt !acc = do
971 res <- step (adaptState gst) stt
972 case res of
973 Yield x s -> do
974 if cmp x prev
975 then do
976 acc' <- fstep acc x
977 go SPEC s acc'
978 else done acc >>= \r -> return $ Yield r (Just s, Just x)
979 Skip s -> go SPEC s acc
980 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
982 stepOuter _ _ (Nothing,_) = return Stop
984{-# INLINE_NORMAL groupsRollingBy #-}
985groupsRollingBy :: Monad m
986 => (a -> a -> Bool)
987 -> Fold m a b
988 -> Stream m a
989 -> Stream m b
990groupsRollingBy cmp f (Stream step state) =
991 Stream (stepOuter f) (Just state, Nothing)
992 where
994 {-# INLINE_LATE stepOuter #-}
995 stepOuter (Fold fstep initial done) gst (Just st, Nothing) = do
996 res <- step (adaptState gst) st
997 case res of
998 Yield x s -> do
999 acc <- initial
1000 acc' <- fstep acc x
1001 go SPEC x s acc'
1003 Skip s -> return $ Skip $ (Just s, Nothing)
1004 Stop -> return Stop
1006 where
1007 go !_ prev stt !acc = do
1008 res <- step (adaptState gst) stt
1009 case res of
1010 Yield x s -> do
1011 if cmp prev x
1012 then do
1013 acc' <- fstep acc x
1014 go SPEC x s acc'
1015 else
1016 done acc >>= \r -> return $ Yield r (Just s, Just x)
1017 Skip s -> go SPEC prev s acc
1018 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1020 stepOuter (Fold fstep initial done) gst (Just st, Just prev') = do
1021 acc <- initial
1022 acc' <- fstep acc prev'
1023 go SPEC prev' st acc'
1025 where
1026 go !_ prevv stt !acc = do
1027 res <- step (adaptState gst) stt
1028 case res of
1029 Yield x s -> do
1030 if cmp prevv x
1031 then do
1032 acc' <- fstep acc x
1033 go SPEC x s acc'
1034 else done acc >>= \r -> return $ Yield r (Just s, Just x)
1035 Skip s -> go SPEC prevv s acc
1036 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1038 stepOuter _ _ (Nothing, _) = return Stop
1040{-# INLINE_NORMAL splitBy #-}
1041splitBy :: Monad m => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1042splitBy predicate f (Stream step state) = Stream (step' f) (Just state)
1044 where
1046 {-# INLINE_LATE step' #-}
1047 step' (Fold fstep initial done) gst (Just st) = initial >>= go SPEC st
1049 where
1051 go !_ stt !acc = do
1052 res <- step (adaptState gst) stt
1053 case res of
1054 Yield x s -> do
1055 if predicate x
1056 then done acc >>= \r -> return $ Yield r (Just s)
1057 else do
1058 acc' <- fstep acc x
1059 go SPEC s acc'
1060 Skip s -> go SPEC s acc
1061 Stop -> done acc >>= \r -> return $ Yield r Nothing
1063 step' _ _ Nothing = return Stop
1065{-# INLINE_NORMAL splitSuffixBy #-}
1066splitSuffixBy :: Monad m
1067 => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1068splitSuffixBy predicate f (Stream step state) = Stream (step' f) (Just state)
1070 where
1072 {-# INLINE_LATE step' #-}
1073 step' (Fold fstep initial done) gst (Just st) = do
1074 res <- step (adaptState gst) st
1075 case res of
1076 Yield x s -> do
1077 acc <- initial
1078 if predicate x
1079 then done acc >>= \val -> return $ Yield val (Just s)
1080 else do
1081 acc' <- fstep acc x
1082 go SPEC s acc'
1084 Skip s -> return $ Skip $ Just s
1085 Stop -> return Stop
1087 where
1089 go !_ stt !acc = do
1090 res <- step (adaptState gst) stt
1091 case res of
1092 Yield x s -> do
1093 if predicate x
1094 then done acc >>= \r -> return $ Yield r (Just s)
1095 else do
1096 acc' <- fstep acc x
1097 go SPEC s acc'
1098 Skip s -> go SPEC s acc
1099 Stop -> done acc >>= \r -> return $ Yield r Nothing
1101 step' _ _ Nothing = return Stop
1103{-# INLINE_NORMAL wordsBy #-}
1104wordsBy :: Monad m => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1105wordsBy predicate f (Stream step state) = Stream (stepOuter f) (Just state)
1107 where
1109 {-# INLINE_LATE stepOuter #-}
1110 stepOuter (Fold fstep initial done) gst (Just st) = do
1111 res <- step (adaptState gst) st
1112 case res of
1113 Yield x s -> do
1114 if predicate x
1115 then return $ Skip (Just s)
1116 else do
1117 acc <- initial
1118 acc' <- fstep acc x
1119 go SPEC s acc'
1121 Skip s -> return $ Skip $ Just s
1122 Stop -> return Stop
1124 where
1126 go !_ stt !acc = do
1127 res <- step (adaptState gst) stt
1128 case res of
1129 Yield x s -> do
1130 if predicate x
1131 then done acc >>= \r -> return $ Yield r (Just s)
1132 else do
1133 acc' <- fstep acc x
1134 go SPEC s acc'
1135 Skip s -> go SPEC s acc
1136 Stop -> done acc >>= \r -> return $ Yield r Nothing
1138 stepOuter _ _ Nothing = return Stop
1140-- String search algorithms:
1141-- http://www-igm.univ-mlv.fr/~lecroq/string/index.html
1144-- TODO can we unify the splitting operations using a splitting configuration
1145-- like in the split package.
1147data SplitStyle = Infix | Suffix | Prefix deriving (Eq, Show)
1149data SplitOptions = SplitOptions
1150 { style :: SplitStyle
1151 , withSep :: Bool -- ^ keep the separators in output
1152 -- , compact :: Bool -- ^ treat multiple consecutive separators as one
1153 -- , trimHead :: Bool -- ^ drop blank at head
1154 -- , trimTail :: Bool -- ^ drop blank at tail
1158data SplitOnState s a =
1159 GO_START
1160 | GO_EMPTY_PAT s
1161 | GO_SINGLE_PAT s a
1162 | GO_SHORT_PAT s
1163 | GO_KARP_RABIN s !(RB.Ring a) !(Ptr a)
1164 | GO_DONE
1166{-# INLINE_NORMAL splitOn #-}
1167splitOn
1168 :: forall m a b. (MonadIO m, Storable a, Enum a, Eq a)
1169 => Array a
1170 -> Fold m a b
1171 -> Stream m a
1172 -> Stream m b
1173splitOn patArr@Array{..} (Fold fstep initial done) (Stream step state) =
1174 Stream stepOuter GO_START
1176 where
1178 patLen = A.length patArr
1179 maxIndex = patLen - 1
1180 elemBits = sizeOf (undefined :: a) * 8
1182 {-# INLINE_LATE stepOuter #-}
1183 stepOuter _ GO_START =
1184 if patLen == 0
1185 then return $ Skip $ GO_EMPTY_PAT state
1186 else if patLen == 1
1187 then do
1188 r <- liftIO $ (A.unsafeIndexIO patArr 0)
1189 return $ Skip $ GO_SINGLE_PAT state r
1190 else if sizeOf (undefined :: a) * patLen
1191 <= sizeOf (undefined :: Word)
1192 then return $ Skip $ GO_SHORT_PAT state
1193 else do
1194 (rb, rhead) <- liftIO $ RB.new patLen
1195 return $ Skip $ GO_KARP_RABIN state rb rhead
1197 stepOuter gst (GO_SINGLE_PAT stt pat) = initial >>= go SPEC stt
1199 where
1201 go !_ st !acc = do
1202 res <- step (adaptState gst) st
1203 case res of
1204 Yield x s -> do
1205 if pat == x
1206 then do
1207 r <- done acc
1208 return $ Yield r (GO_SINGLE_PAT s pat)
1209 else fstep acc x >>= go SPEC s
1210 Skip s -> go SPEC s acc
1211 Stop -> done acc >>= \r -> return $ Yield r GO_DONE
1213 stepOuter gst (GO_SHORT_PAT stt) = initial >>= go0 SPEC 0 (0 :: Word) stt
1215 where
1217 mask :: Word
1218 mask = (1 `shiftL` (elemBits * patLen)) - 1
1220 addToWord wrd a = (wrd `shiftL` elemBits) .|. fromIntegral (fromEnum a)
1222 patWord :: Word
1223 patWord = mask .&. A.foldl' addToWord 0 patArr
1225 go0 !_ !idx wrd st !acc = do
1226 res <- step (adaptState gst) st
1227 case res of
1228 Yield x s -> do
1229 let wrd' = addToWord wrd x
1230 if idx == maxIndex
1231 then do
1232 if wrd' .&. mask == patWord
1233 then do
1234 r <- done acc
1235 return $ Yield r (GO_SHORT_PAT s)
1236 else go1 SPEC wrd' s acc
1237 else go0 SPEC (idx + 1) wrd' s acc
1238 Skip s -> go0 SPEC idx wrd s acc
1239 Stop -> do
1240 acc' <- if idx /= 0
1241 then go2 wrd idx acc
1242 else return acc
1243 done acc' >>= \r -> return $ Yield r GO_DONE
1245 {-# INLINE go1 #-}
1246 go1 !_ wrd st !acc = do
1247 res <- step (adaptState gst) st
1248 case res of
1249 Yield x s -> do
1250 let wrd' = addToWord wrd x
1251 old = (mask .&. wrd) `shiftR` (elemBits * (patLen - 1))
1252 acc' <- fstep acc (toEnum $ fromIntegral old)
1253 if wrd' .&. mask == patWord
1254 then done acc' >>= \r -> return $ Yield r (GO_SHORT_PAT s)
1255 else go1 SPEC wrd' s acc'
1256 Skip s -> go1 SPEC wrd s acc
1257 Stop -> do
1258 acc' <- go2 wrd patLen acc
1259 done acc' >>= \r -> return $ Yield r GO_DONE
1261 go2 !wrd !n !acc | n > 0 = do
1262 let old = (mask .&. wrd) `shiftR` (elemBits * (n - 1))
1263 fstep acc (toEnum $ fromIntegral old) >>= go2 wrd (n - 1)
1264 go2 _ _ acc = return acc
1266 stepOuter gst (GO_KARP_RABIN stt rb rhead) = do
1267 initial >>= go0 SPEC 0 rhead stt
1269 where
1271 k = 2891336453 :: Word32
1272 coeff = k ^ patLen
1273 addCksum cksum a = cksum * k + fromIntegral (fromEnum a)
1274 deltaCksum cksum old new =
1275 addCksum cksum new - coeff * fromIntegral (fromEnum old)
1277 -- XXX shall we use a random starting hash or 1 instead of 0?
1278 patHash = A.foldl' addCksum 0 patArr
1280 -- rh == ringHead
1281 go0 !_ !idx !rh st !acc = do
1282 res <- step (adaptState gst) st
1283 case res of
1284 Yield x s -> do
1285 rh' <- liftIO $ RB.unsafeInsert rb rh x
1286 if idx == maxIndex
1287 then do
1288 let fold = RB.unsafeFoldRing (RB.ringBound rb)
1289 let !ringHash = fold addCksum 0 rb
1290 if ringHash == patHash
1291 then go2 SPEC ringHash rh' s acc
1292 else go1 SPEC ringHash rh' s acc
1293 else go0 SPEC (idx + 1) rh' s acc
1294 Skip s -> go0 SPEC idx rh s acc
1295 Stop -> do
1296 !acc' <- if idx /= 0
1297 then RB.unsafeFoldRingM rh fstep acc rb
1298 else return acc
1299 done acc' >>= \r -> return $ Yield r GO_DONE
1301 -- XXX Theoretically this code can do 4 times faster if GHC generates
1302 -- optimal code. If we use just "(cksum' == patHash)" condition it goes
1303 -- 4x faster, as soon as we add the "RB.unsafeEqArray rb v" condition
1304 -- the generated code changes drastically and becomes 4x slower. Need
1305 -- to investigate what is going on with GHC.
1306 {-# INLINE go1 #-}
1307 go1 !_ !cksum !rh st !acc = do
1308 res <- step (adaptState gst) st
1309 case res of
1310 Yield x s -> do
1311 old <- liftIO $ peek rh
1312 let cksum' = deltaCksum cksum old x
1313 acc' <- fstep acc old
1315 if (cksum' == patHash)
1316 then do
1317 rh' <- liftIO (RB.unsafeInsert rb rh x)
1318 go2 SPEC cksum' rh' s acc'
1319 else do
1320 rh' <- liftIO (RB.unsafeInsert rb rh x)
1321 go1 SPEC cksum' rh' s acc'
1322 Skip s -> go1 SPEC cksum rh s acc
1323 Stop -> do
1324 acc' <- RB.unsafeFoldRingFullM rh fstep acc rb
1325 done acc' >>= \r -> return $ Yield r GO_DONE
1327 go2 !_ !cksum' !rh' s !acc' = do
1328 if RB.unsafeEqArray rb rh' patArr
1329 then do
1330 r <- done acc'
1331 return $ Yield r (GO_KARP_RABIN s rb rhead)
1332 else go1 SPEC cksum' rh' s acc'
1334 stepOuter gst (GO_EMPTY_PAT st) = do
1335 res <- step (adaptState gst) st
1336 case res of
1337 Yield x s -> do
1338 acc <- initial
1339 acc' <- fstep acc x
1340 done acc' >>= \r -> return $ Yield r (GO_EMPTY_PAT s)
1341 Skip s -> return $ Skip (GO_EMPTY_PAT s)
1342 Stop -> return Stop
1344 stepOuter _ GO_DONE = return Stop
1346{-# INLINE_NORMAL splitSuffixOn #-}
1347splitSuffixOn
1348 :: forall m a b. (MonadIO m, Storable a, Enum a, Eq a)
1349 => Bool
1350 -> Array a
1351 -> Fold m a b
1352 -> Stream m a
1353 -> Stream m b
1354splitSuffixOn withSep patArr@Array{..} (Fold fstep initial done)
1355 (Stream step state) =
1356 Stream stepOuter GO_START
1358 where
1360 patLen = A.length patArr
1361 maxIndex = patLen - 1
1362 elemBits = sizeOf (undefined :: a) * 8
1364 {-# INLINE_LATE stepOuter #-}
1365 stepOuter _ GO_START =
1366 if patLen == 0
1367 then return $ Skip $ GO_EMPTY_PAT state
1368 else if patLen == 1
1369 then do
1370 r <- liftIO $ (A.unsafeIndexIO patArr 0)
1371 return $ Skip $ GO_SINGLE_PAT state r
1372 else if sizeOf (undefined :: a) * patLen
1373 <= sizeOf (undefined :: Word)
1374 then return $ Skip $ GO_SHORT_PAT state
1375 else do
1376 (rb, rhead) <- liftIO $ RB.new patLen
1377 return $ Skip $ GO_KARP_RABIN state rb rhead
1379 stepOuter gst (GO_SINGLE_PAT stt pat) = do
1380 -- This first part is the only difference between splitOn and
1381 -- splitSuffixOn.
1382 -- If the last element is a separator do not issue a blank segment.
1383 res <- step (adaptState gst) stt
1384 case res of
1385 Yield x s -> do
1386 acc <- initial
1387 if pat == x
1388 then do
1389 acc' <- if withSep then fstep acc x else return acc
1390 done acc' >>= \r -> return $ Yield r (GO_SINGLE_PAT s pat)
1391 else fstep acc x >>= go SPEC s
1392 Skip s -> return $ Skip $ (GO_SINGLE_PAT s pat)
1393 Stop -> return Stop
1395 where
1397 -- This is identical for splitOn and splitSuffixOn
1398 go !_ st !acc = do
1399 res <- step (adaptState gst) st
1400 case res of
1401 Yield x s -> do
1402 if pat == x
1403 then do
1404 acc' <- if withSep then fstep acc x else return acc
1405 r <- done acc'
1406 return $ Yield r (GO_SINGLE_PAT s pat)
1407 else fstep acc x >>= go SPEC s
1408 Skip s -> go SPEC s acc
1409 Stop -> done acc >>= \r -> return $ Yield r GO_DONE
1411 stepOuter gst (GO_SHORT_PAT stt) = do
1413 -- Call "initial" only if the stream yields an element, otherwise we
1414 -- may call "initial" but never yield anything. initial may produce a
1415 -- side effect, therefore we will end up doing and discard a side
1416 -- effect.
1418 let idx = 0
1419 let wrd = 0
1420 res <- step (adaptState gst) stt
1421 case res of
1422 Yield x s -> do
1423 acc <- initial
1424 let wrd' = addToWord wrd x
1425 acc' <- if withSep then fstep acc x else return acc
1426 if idx == maxIndex
1427 then do
1428 if wrd' .&. mask == patWord
1429 then done acc' >>= \r -> return $ Yield r (GO_SHORT_PAT s)
1430 else go0 SPEC (idx + 1) wrd' s acc'
1431 else go0 SPEC (idx + 1) wrd' s acc'
1432 Skip s -> return $ Skip (GO_SHORT_PAT s)
1433 Stop -> return Stop
1435 where
1437 mask :: Word
1438 mask = (1 `shiftL` (elemBits * patLen)) - 1
1440 addToWord wrd a = (wrd `shiftL` elemBits) .|. fromIntegral (fromEnum a)
1442 patWord :: Word
1443 patWord = mask .&. A.foldl' addToWord 0 patArr
1445 go0 !_ !idx wrd st !acc = do
1446 res <- step (adaptState gst) st
1447 case res of
1448 Yield x s -> do
1449 let wrd' = addToWord wrd x
1450 acc' <- if withSep then fstep acc x else return acc
1451 if idx == maxIndex
1452 then do
1453 if wrd' .&. mask == patWord
1454 then do
1455 r <- done acc'
1456 return $ Yield r (GO_SHORT_PAT s)
1457 else go1 SPEC wrd' s acc'
1458 else go0 SPEC (idx + 1) wrd' s acc'
1459 Skip s -> go0 SPEC idx wrd s acc
1460 Stop -> do
1461 if (idx == maxIndex) && (wrd .&. mask == patWord)
1462 then return Stop
1463 else do
1464 acc' <- if idx /= 0 && not withSep
1465 then go2 wrd idx acc
1466 else return acc
1467 done acc' >>= \r -> return $ Yield r GO_DONE
1469 {-# INLINE go1 #-}
1470 go1 !_ wrd st !acc = do
1471 res <- step (adaptState gst) st
1472 case res of
1473 Yield x s -> do
1474 let wrd' = addToWord wrd x
1475 old = (mask .&. wrd) `shiftR` (elemBits * (patLen - 1))
1476 acc' <- if withSep
1477 then fstep acc x
1478 else fstep acc (toEnum $ fromIntegral old)
1479 if wrd' .&. mask == patWord
1480 then done acc' >>= \r -> return $ Yield r (GO_SHORT_PAT s)
1481 else go1 SPEC wrd' s acc'
1482 Skip s -> go1 SPEC wrd s acc
1483 Stop ->
1484 -- If the last sequence is a separator do not issue a blank
1485 -- segment.
1486 if wrd .&. mask == patWord
1487 then return Stop
1488 else do
1489 acc' <- if withSep
1490 then return acc
1491 else go2 wrd patLen acc
1492 done acc' >>= \r -> return $ Yield r GO_DONE
1494 go2 !wrd !n !acc | n > 0 = do
1495 let old = (mask .&. wrd) `shiftR` (elemBits * (n - 1))
1496 fstep acc (toEnum $ fromIntegral old) >>= go2 wrd (n - 1)
1497 go2 _ _ acc = return acc
1499 stepOuter gst (GO_KARP_RABIN stt rb rhead) = do
1500 let idx = 0
1501 res <- step (adaptState gst) stt
1502 case res of
1503 Yield x s -> do
1504 acc <- initial
1505 acc' <- if withSep then fstep acc x else return acc
1506 rh' <- liftIO (RB.unsafeInsert rb rhead x)
1507 if idx == maxIndex
1508 then do
1509 let fold = RB.unsafeFoldRing (RB.ringBound rb)
1510 let !ringHash = fold addCksum 0 rb
1511 if ringHash == patHash
1512 then go2 SPEC ringHash rh' s acc'
1513 else go0 SPEC (idx + 1) rh' s acc'
1514 else go0 SPEC (idx + 1) rh' s acc'
1515 Skip s -> return $ Skip (GO_KARP_RABIN s rb rhead)
1516 Stop -> return Stop
1518 where
1520 k = 2891336453 :: Word32
1521 coeff = k ^ patLen
1522 addCksum cksum a = cksum * k + fromIntegral (fromEnum a)
1523 deltaCksum cksum old new =
1524 addCksum cksum new - coeff * fromIntegral (fromEnum old)
1526 -- XXX shall we use a random starting hash or 1 instead of 0?
1527 patHash = A.foldl' addCksum 0 patArr
1529 -- rh == ringHead
1530 go0 !_ !idx !rh st !acc = do
1531 res <- step (adaptState gst) st
1532 case res of
1533 Yield x s -> do
1534 acc' <- if withSep then fstep acc x else return acc
1535 rh' <- liftIO (RB.unsafeInsert rb rh x)
1536 if idx == maxIndex
1537 then do
1538 let fold = RB.unsafeFoldRing (RB.ringBound rb)
1539 let !ringHash = fold addCksum 0 rb
1540 if ringHash == patHash
1541 then go2 SPEC ringHash rh' s acc'
1542 else go1 SPEC ringHash rh' s acc'
1543 else go0 SPEC (idx + 1) rh' s acc'
1544 Skip s -> go0 SPEC idx rh s acc
1545 Stop -> do
1546 -- do not issue a blank segment when we end at pattern
1547 if (idx == maxIndex) && RB.unsafeEqArray rb rh patArr
1548 then return Stop
1549 else do
1550 !acc' <- if idx /= 0 && not withSep
1551 then RB.unsafeFoldRingM rh fstep acc rb
1552 else return acc
1553 done acc' >>= \r -> return $ Yield r GO_DONE
1555 -- XXX Theoretically this code can do 4 times faster if GHC generates
1556 -- optimal code. If we use just "(cksum' == patHash)" condition it goes
1557 -- 4x faster, as soon as we add the "RB.unsafeEqArray rb v" condition
1558 -- the generated code changes drastically and becomes 4x slower. Need
1559 -- to investigate what is going on with GHC.
1560 {-# INLINE go1 #-}
1561 go1 !_ !cksum !rh st !acc = do
1562 res <- step (adaptState gst) st
1563 case res of
1564 Yield x s -> do
1565 old <- liftIO $ peek rh
1566 let cksum' = deltaCksum cksum old x
1567 acc' <- if withSep
1568 then fstep acc x
1569 else fstep acc old
1571 if (cksum' == patHash)
1572 then do
1573 rh' <- liftIO (RB.unsafeInsert rb rh x)
1574 go2 SPEC cksum' rh' s acc'
1575 else do
1576 rh' <- liftIO (RB.unsafeInsert rb rh x)
1577 go1 SPEC cksum' rh' s acc'
1578 Skip s -> go1 SPEC cksum rh s acc
1579 Stop -> do
1580 if RB.unsafeEqArray rb rh patArr
1581 then return Stop
1582 else do
1583 acc' <- if withSep
1584 then return acc
1585 else RB.unsafeFoldRingFullM rh fstep acc rb
1586 done acc' >>= \r -> return $ Yield r GO_DONE
1588 go2 !_ !cksum' !rh' s !acc' = do
1589 if RB.unsafeEqArray rb rh' patArr
1590 then do
1591 r <- done acc'
1592 return $ Yield r (GO_KARP_RABIN s rb rhead)
1593 else go1 SPEC cksum' rh' s acc'
1595 stepOuter gst (GO_EMPTY_PAT st) = do
1596 res <- step (adaptState gst) st
1597 case res of
1598 Yield x s -> do
1599 acc <- initial
1600 acc' <- fstep acc x
1601 done acc' >>= \r -> return $ Yield r (GO_EMPTY_PAT s)
1602 Skip s -> return $ Skip (GO_EMPTY_PAT s)
1603 Stop -> return Stop
1605 stepOuter _ GO_DONE = return Stop
1607------------------------------------------------------------------------------
1608-- Substreams
1609------------------------------------------------------------------------------
1611{-# INLINE_NORMAL isPrefixOf #-}
1612isPrefixOf :: (Eq a, Monad m) => Stream m a -> Stream m a -> m Bool
1613isPrefixOf (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
1614 where
1615 go (sa, sb, Nothing) = do
1616 r <- stepa defState sa
1617 case r of
1618 Yield x sa' -> go (sa', sb, Just x)
1619 Skip sa' -> go (sa', sb, Nothing)
1620 Stop -> return True
1622 go (sa, sb, Just x) = do
1623 r <- stepb defState sb
1624 case r of
1625 Yield y sb' ->
1626 if x == y
1627 then go (sa, sb', Nothing)
1628 else return False
1629 Skip sb' -> go (sa, sb', Just x)
1630 Stop -> return False
1632{-# INLINE_NORMAL isSubsequenceOf #-}
1633isSubsequenceOf :: (Eq a, Monad m) => Stream m a -> Stream m a -> m Bool
1634isSubsequenceOf (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
1635 where
1636 go (sa, sb, Nothing) = do
1637 r <- stepa defState sa
1638 case r of
1639 Yield x sa' -> go (sa', sb, Just x)
1640 Skip sa' -> go (sa', sb, Nothing)
1641 Stop -> return True
1643 go (sa, sb, Just x) = do
1644 r <- stepb defState sb
1645 case r of
1646 Yield y sb' ->
1647 if x == y
1648 then go (sa, sb', Nothing)
1649 else go (sa, sb', Just x)
1650 Skip sb' -> go (sa, sb', Just x)
1651 Stop -> return False
1653{-# INLINE_NORMAL stripPrefix #-}
1654stripPrefix
1655 :: (Eq a, Monad m)
1656 => Stream m a -> Stream m a -> m (Maybe (Stream m a))
1657stripPrefix (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
1658 where
1659 go (sa, sb, Nothing) = do
1660 r <- stepa defState sa
1661 case r of
1662 Yield x sa' -> go (sa', sb, Just x)
1663 Skip sa' -> go (sa', sb, Nothing)
1664 Stop -> return $ Just (Stream stepb sb)
1666 go (sa, sb, Just x) = do
1667 r <- stepb defState sb
1668 case r of
1669 Yield y sb' ->
1670 if x == y
1671 then go (sa, sb', Nothing)
1672 else return Nothing
1673 Skip sb' -> go (sa, sb', Just x)
1674 Stop -> return Nothing
1676------------------------------------------------------------------------------
1677-- Map and Fold
1678------------------------------------------------------------------------------
1680-- | Execute a monadic action for each element of the 'Stream'
1681{-# INLINE_NORMAL mapM_ #-}
1682mapM_ :: Monad m => (a -> m b) -> Stream m a -> m ()
1683mapM_ m = drain . mapM m
1685------------------------------------------------------------------------------
1686-- Exceptions
1687------------------------------------------------------------------------------
1689-- | Run a side effect before the stream yields its first element.
1690{-# INLINE_NORMAL before #-}
1691before :: Monad m => m b -> Stream m a -> Stream m a
1692before action (Stream step state) = Stream step' Nothing
1694 where
1696 {-# INLINE_LATE step' #-}
1697 step' _ Nothing = action >> return (Skip (Just state))
1699 step' gst (Just st) = do
1700 res <- step gst st
1701 case res of
1702 Yield x s -> return $ Yield x (Just s)
1703 Skip s -> return $ Skip (Just s)
1704 Stop -> return Stop
1706-- | Run a side effect whenever the stream stops normally.
1707{-# INLINE_NORMAL after #-}
1708after :: Monad m => m b -> Stream m a -> Stream m a
1709after action (Stream step state) = Stream step' state
1711 where
1713 {-# INLINE_LATE step' #-}
1714 step' gst st = do
1715 res <- step gst st
1716 case res of
1717 Yield x s -> return $ Yield x s
1718 Skip s -> return $ Skip s
1719 Stop -> action >> return Stop
1721-- XXX These combinators are expensive due to the call to
1722-- onException/handle/try on each step. Therefore, when possible, they should
1723-- be called in an outer loop where we perform less iterations. For example, we
1724-- cannot call them on each iteration in a char stream, instead we can call
1725-- them when doing an IO on an array.
1727-- XXX For high performance error checks in busy streams we may need another
1728-- Error constructor in step.
1730-- | Run a side effect whenever the stream aborts due to an exception. The
1731-- exception is not caught simply rethrown.
1732{-# INLINE_NORMAL onException #-}
1733onException :: MonadCatch m => m b -> Stream m a -> Stream m a
1734onException action (Stream step state) = Stream step' state
1736 where
1738 {-# INLINE_LATE step' #-}
1739 step' gst st = do
1740 res <- step gst st `MC.onException` action
1741 case res of
1742 Yield x s -> return $ Yield x s
1743 Skip s -> return $ Skip s
1744 Stop -> return Stop
1746-- | Run a side effect whenever the stream stops normally or aborts due to an
1747-- exception.
1748{-# INLINE finally #-}
1749finally :: MonadCatch m => m b -> Stream m a -> Stream m a
1750finally action xs = after action $ onException action xs
1752-- XXX bracket is like concatMap, it generates a stream and then flattens it.
1753-- Like concatMap it has 10x worse performance compared to linear fused
1754-- compositions.
1756-- | Run the first action before the stream starts and remember its output,
1757-- generate a stream using the output, run the second action using the
1758-- remembered value as an argument whenever the stream ends normally or due to
1759-- an exception.
1760{-# INLINE_NORMAL bracket #-}
1761bracket :: MonadCatch m => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
1762bracket bef aft bet = Stream step' Nothing
1764 where
1766 {-# INLINE_LATE step' #-}
1767 step' _ Nothing = bef >>= \x -> return (Skip (Just (bet x, x)))
1769 -- NOTE: It is important to use UnStream instead of the Stream pattern
1770 -- here, otherwise we get huge perf degradation, see note in concatMap.
1771 step' gst (Just (UnStream step state, v)) = do
1772 res <- step gst state `MC.onException` aft v
1773 case res of
1774 Yield x s -> return $ Yield x (Just (Stream step s, v))
1775 Skip s -> return $ Skip (Just (Stream step s, v))
1776 Stop -> aft v >> return Stop
1778-- | When evaluating a stream if an exception occurs, stream evaluation aborts
1779-- and the specified exception handler is run with the exception as argument.
1780{-# INLINE_NORMAL handle #-}
1781handle :: (MonadCatch m, Exception e) => (e -> m a) -> Stream m a -> Stream m a
1782handle f (Stream step state) = Stream step' (Just state)
1784 where
1786 {-# INLINE_LATE step' #-}
1787 step' gst (Just st) = do
1788 res <- MC.try $ step gst st
1789 case res of
1790 Left e -> f e >>= \x -> return (Yield x Nothing)
1791 Right r -> case r of
1792 Yield x s -> return $ Yield x (Just s)
1793 Skip s -> return $ Skip (Just s)
1794 Stop -> return Stop
1796 step' _ Nothing = return Stop
1798-------------------------------------------------------------------------------
1799-- General transformation
1800-------------------------------------------------------------------------------
1802{-# INLINE_NORMAL transform #-}
1803transform :: Monad m => Pipe m a b -> Stream m a -> Stream m b
1804transform (Pipe pstep1 pstep2 pstate) (Stream step state) =
1805 Stream step' (Consume pstate, state)
1807 where
1809 {-# INLINE_LATE step' #-}
1811 step' gst (Consume pst, st) = pst `seq` do
1812 r <- step (adaptState gst) st
1813 case r of
1814 Yield x s -> do
1815 res <- pstep1 pst x
1816 case res of
1817 Pipe.Yield b pst' -> return $ Yield b (pst', s)
1818 Pipe.Continue pst' -> return $ Skip (pst', s)
1819 Skip s -> return $ Skip (Consume pst, s)
1820 Stop -> return Stop
1822 step' _ (Produce pst, st) = pst `seq` do
1823 res <- pstep2 pst
1824 case res of
1825 Pipe.Yield b pst' -> return $ Yield b (pst', st)
1826 Pipe.Continue pst' -> return $ Skip (pst', st)
1828------------------------------------------------------------------------------
1829-- Transformation by Folding (Scans)
1830------------------------------------------------------------------------------
1832------------------------------------------------------------------------------
1833-- Prescans
1834------------------------------------------------------------------------------
1836-- XXX Is a prescan useful, discarding the last step does not sound useful? I
1837-- am not sure about the utility of this function, so this is implemented but
1838-- not exposed. We can expose it if someone provides good reasons why this is
1839-- useful.
1841-- XXX We have to execute the stream one step ahead to know that we are at the
1842-- last step. The vector implementation of prescan executes the last fold step
1843-- but does not yield the result. This means we have executed the effect but
1844-- discarded value. This does not sound right. In this implementation we are
1845-- not executing the last fold step.
1846{-# INLINE_NORMAL prescanlM' #-}
1847prescanlM' :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
1848prescanlM' f mz (Stream step state) = Stream step' (state, mz)
1849 where
1850 {-# INLINE_LATE step' #-}
1851 step' gst (st, prev) = do
1852 r <- step (adaptState gst) st
1853 case r of
1854 Yield x s -> do
1855 acc <- prev
1856 return $ Yield acc (s, f acc x)
1857 Skip s -> return $ Skip (s, prev)
1858 Stop -> return Stop
1860{-# INLINE prescanl' #-}
1861prescanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
1862prescanl' f z = prescanlM' (\a b -> return (f a b)) (return z)
1864------------------------------------------------------------------------------
1865-- Monolithic postscans (postscan followed by a map)
1866------------------------------------------------------------------------------
1868-- The performance of a modular postscan followed by a map seems to be
1869-- equivalent to this monolithic scan followed by map therefore we may not need
1870-- this implementation. We just have it for performance comparison and in case
1871-- modular version does not perform well in some situation.
1873{-# INLINE_NORMAL postscanlMx' #-}
1874postscanlMx' :: Monad m
1875 => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> Stream m b
1876postscanlMx' fstep begin done (Stream step state) = do
1877 Stream step' (state, begin)
1878 where
1879 {-# INLINE_LATE step' #-}
1880 step' gst (st, acc) = do
1881 r <- step (adaptState gst) st
1882 case r of
1883 Yield x s -> do
1884 old <- acc
1885 y <- fstep old x
1886 v <- done y
1887 v `seq` y `seq` return (Yield v (s, return y))
1888 Skip s -> return $ Skip (s, acc)
1889 Stop -> return Stop
1891{-# INLINE_NORMAL postscanlx' #-}
1892postscanlx' :: Monad m
1893 => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> Stream m b
1894postscanlx' fstep begin done s =
1895 postscanlMx' (\b a -> return (fstep b a)) (return begin) (return . done) s
1897-- XXX do we need consM strict to evaluate the begin value?
1898{-# INLINE scanlMx' #-}
1899scanlMx' :: Monad m
1900 => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> Stream m b
1901scanlMx' fstep begin done s =
1902 (begin >>= \x -> x `seq` done x) `consM` postscanlMx' fstep begin done s
1904{-# INLINE scanlx' #-}
1905scanlx' :: Monad m
1906 => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> Stream m b
1907scanlx' fstep begin done s =
1908 scanlMx' (\b a -> return (fstep b a)) (return begin) (return . done) s
1910------------------------------------------------------------------------------
1911-- postscans
1912------------------------------------------------------------------------------
1914-- XXX if we make the initial value of the accumulator monadic then should we
1915-- execute it even if the stream is empty? In that case we would have generated
1916-- the effect but discarded the value, but that is what a fold does when the
1917-- stream is empty. Whatever we decide, need to reconcile this with prescan.
1918-- If we execute the initial value here without even using it then it is ok to
1919-- execute the last step there as well without using the value.
1920-- Looking at the duality with right fold, in case of right fold we always
1921-- perform the action when the construction terminates, so in case of left fold
1922-- we should perform it only when the reduction starts.
1923{-# INLINE_NORMAL postscanlM' #-}
1924postscanlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
1925postscanlM' fstep begin (Stream step state) =
1926 begin `seq` Stream step' (state, begin)
1927 where
1928 {-# INLINE_LATE step' #-}
1929 step' gst (st, acc) = acc `seq` do
1930 r <- step (adaptState gst) st
1931 case r of
1932 Yield x s -> do
1933 y <- fstep acc x
1934 y `seq` return (Yield y (s, y))
1935 Skip s -> return $ Skip (s, acc)
1936 Stop -> return Stop
1938{-# INLINE_NORMAL postscanl' #-}
1939postscanl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
1940postscanl' f = postscanlM' (\a b -> return (f a b))
1942{-# INLINE_NORMAL postscanlM #-}
1943postscanlM :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
1944postscanlM fstep begin (Stream step state) = Stream step' (state, begin)
1945 where
1946 {-# INLINE_LATE step' #-}
1947 step' gst (st, acc) = do
1948 r <- step (adaptState gst) st
1949 case r of
1950 Yield x s -> do
1951 y <- fstep acc x
1952 return (Yield y (s, y))
1953 Skip s -> return $ Skip (s, acc)
1954 Stop -> return Stop
1956{-# INLINE_NORMAL postscanl #-}
1957postscanl :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
1958postscanl f = postscanlM (\a b -> return (f a b))
1960{-# INLINE_NORMAL scanlM' #-}
1961scanlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
1962scanlM' fstep begin s = begin `seq` (begin `cons` postscanlM' fstep begin s)
1964{-# INLINE scanl' #-}
1965scanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
1966scanl' f = scanlM' (\a b -> return (f a b))
1968{-# INLINE_NORMAL scanlM #-}
1969scanlM :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
1970scanlM fstep begin s = begin `cons` postscanlM fstep begin s
1972{-# INLINE scanl #-}
1973scanl :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
1974scanl f = scanlM (\a b -> return (f a b))
1976{-# INLINE_NORMAL scanl1M #-}
1977scanl1M :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
1978scanl1M fstep (Stream step state) = Stream step' (state, Nothing)
1979 where
1980 {-# INLINE_LATE step' #-}
1981 step' gst (st, Nothing) = do
1982 r <- step gst st
1983 case r of
1984 Yield x s -> return $ Yield x (s, Just x)
1985 Skip s -> return $ Skip (s, Nothing)
1986 Stop -> return Stop
1988 step' gst (st, Just acc) = do
1989 r <- step gst st
1990 case r of
1991 Yield y s -> do
1992 z <- fstep acc y
1993 return $ Yield z (s, Just z)
1994 Skip s -> return $ Skip (s, Just acc)
1995 Stop -> return Stop
1997{-# INLINE scanl1 #-}
1998scanl1 :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
1999scanl1 f = scanl1M (\x y -> return (f x y))
2001{-# INLINE_NORMAL scanl1M' #-}
2002scanl1M' :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
2003scanl1M' fstep (Stream step state) = Stream step' (state, Nothing)
2004 where
2005 {-# INLINE_LATE step' #-}
2006 step' gst (st, Nothing) = do
2007 r <- step gst st
2008 case r of
2009 Yield x s -> x `seq` return $ Yield x (s, Just x)
2010 Skip s -> return $ Skip (s, Nothing)
2011 Stop -> return Stop
2013 step' gst (st, Just acc) = acc `seq` do
2014 r <- step gst st
2015 case r of
2016 Yield y s -> do
2017 z <- fstep acc y
2018 z `seq` return $ Yield z (s, Just z)
2019 Skip s -> return $ Skip (s, Just acc)
2020 Stop -> return Stop
2022{-# INLINE scanl1' #-}
2023scanl1' :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
2024scanl1' f = scanl1M' (\x y -> return (f x y))
2026{-# INLINE tap #-}
2027tap :: Monad m => Sink m a -> Stream m a -> Stream m a
2028tap (Sink fstep) (Stream step state) = Stream step' state
2030 where
2032 step' gst st = do
2033 r <- step gst st
2034 case r of
2035 Yield x s -> do
2036 fstep x
2037 return $ Yield x s
2038 Skip s -> return $ Skip s
2039 Stop -> return $ Stop
2041-------------------------------------------------------------------------------
2042-- Filtering
2043-------------------------------------------------------------------------------
2045{-# INLINE_NORMAL takeWhileM #-}
2046takeWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
2047takeWhileM f (Stream step state) = Stream step' state
2048 where
2049 {-# INLINE_LATE step' #-}
2050 step' gst st = do
2051 r <- step gst st
2052 case r of
2053 Yield x s -> do
2054 b <- f x
2055 return $ if b then Yield x s else Stop
2056 Skip s -> return $ Skip s
2057 Stop -> return Stop
2059{-# INLINE takeWhile #-}
2060takeWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
2061takeWhile f = takeWhileM (return . f)
2063{-# INLINE_NORMAL drop #-}
2064drop :: Monad m => Int -> Stream m a -> Stream m a
2065drop n (Stream step state) = Stream step' (state, Just n)
2066 where
2067 {-# INLINE_LATE step' #-}
2068 step' gst (st, Just i)
2069 | i > 0 = do
2070 r <- step gst st
2071 return $
2072 case r of
2073 Yield _ s -> Skip (s, Just (i - 1))
2074 Skip s -> Skip (s, Just i)
2075 Stop -> Stop
2076 | otherwise = return $ Skip (st, Nothing)
2078 step' gst (st, Nothing) = do
2079 r <- step gst st
2080 return $
2081 case r of
2082 Yield x s -> Yield x (s, Nothing)
2083 Skip s -> Skip (s, Nothing)
2084 Stop -> Stop
2086data DropWhileState s a
2087 = DropWhileDrop s
2088 | DropWhileYield a s
2089 | DropWhileNext s
2091{-# INLINE_NORMAL dropWhileM #-}
2092dropWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
2093dropWhileM f (Stream step state) = Stream step' (DropWhileDrop state)
2094 where
2095 {-# INLINE_LATE step' #-}
2096 step' gst (DropWhileDrop st) = do
2097 r <- step gst st
2098 case r of
2099 Yield x s -> do
2100 b <- f x
2101 if b
2102 then return $ Skip (DropWhileDrop s)
2103 else return $ Skip (DropWhileYield x s)
2104 Skip s -> return $ Skip (DropWhileDrop s)
2105 Stop -> return Stop
2107 step' gst (DropWhileNext st) = do
2108 r <- step gst st
2109 case r of
2110 Yield x s -> return $ Skip (DropWhileYield x s)
2111 Skip s -> return $ Skip (DropWhileNext s)
2112 Stop -> return Stop
2114 step' _ (DropWhileYield x st) = return $ Yield x (DropWhileNext st)
2116{-# INLINE dropWhile #-}
2117dropWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
2118dropWhile f = dropWhileM (return . f)
2120{-# INLINE_NORMAL filterM #-}
2121filterM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
2122filterM f (Stream step state) = Stream step' state
2123 where
2124 {-# INLINE_LATE step' #-}
2125 step' gst st = do
2126 r <- step gst st
2127 case r of
2128 Yield x s -> do
2129 b <- f x
2130 return $ if b
2131 then Yield x s
2132 else Skip s
2133 Skip s -> return $ Skip s
2134 Stop -> return Stop
2136{-# INLINE filter #-}
2137filter :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
2138filter f = filterM (return . f)
2140{-# INLINE_NORMAL uniq #-}
2141uniq :: (Eq a, Monad m) => Stream m a -> Stream m a
2142uniq (Stream step state) = Stream step' (Nothing, state)
2143 where
2144 {-# INLINE_LATE step' #-}
2145 step' gst (Nothing, st) = do
2146 r <- step gst st
2147 case r of
2148 Yield x s -> return $ Yield x (Just x, s)
2149 Skip s -> return $ Skip (Nothing, s)
2150 Stop -> return Stop
2151 step' gst (Just x, st) = do
2152 r <- step gst st
2153 case r of
2154 Yield y s | x == y -> return $ Skip (Just x, s)
2155 | otherwise -> return $ Yield y (Just y, s)
2156 Skip s -> return $ Skip (Just x, s)
2157 Stop -> return Stop
2159------------------------------------------------------------------------------
2160-- Transformation by Mapping
2161------------------------------------------------------------------------------
2163{-# INLINE_NORMAL sequence #-}
2164sequence :: Monad m => Stream m (m a) -> Stream m a
2165sequence (Stream step state) = Stream step' state
2166 where
2167 {-# INLINE_LATE step' #-}
2168 step' gst st = do
2169 r <- step (adaptState gst) st
2170 case r of
2171 Yield x s -> x >>= \a -> return (Yield a s)
2172 Skip s -> return $ Skip s
2173 Stop -> return Stop
2175------------------------------------------------------------------------------
2176-- Inserting
2177------------------------------------------------------------------------------
2179data LoopState x s = FirstYield s
2180 | InterspersingYield s
2181 | YieldAndCarry x s
2183{-# INLINE_NORMAL intersperseM #-}
2184intersperseM :: Monad m => m a -> Stream m a -> Stream m a
2185intersperseM m (Stream step state) = Stream step' (FirstYield state)
2186 where
2187 {-# INLINE_LATE step' #-}
2188 step' gst (FirstYield st) = do
2189 r <- step gst st
2190 return $
2191 case r of
2192 Yield x s -> Skip (YieldAndCarry x s)
2193 Skip s -> Skip (FirstYield s)
2194 Stop -> Stop
2196 step' gst (InterspersingYield st) = do
2197 r <- step gst st
2198 case r of
2199 Yield x s -> do
2200 a <- m
2201 return $ Yield a (YieldAndCarry x s)
2202 Skip s -> return $ Skip $ InterspersingYield s
2203 Stop -> return Stop
2205 step' _ (YieldAndCarry x st) = return $ Yield x (InterspersingYield st)
2207{-# INLINE intersperse #-}
2208intersperse :: Monad m => a -> Stream m a -> Stream m a
2209intersperse a = intersperseM (return a)
2211{-# INLINE_NORMAL insertBy #-}
2212insertBy :: Monad m => (a -> a -> Ordering) -> a -> Stream m a -> Stream m a
2213insertBy cmp a (Stream step state) = Stream step' (state, False, Nothing)
2214 where
2215 {-# INLINE_LATE step' #-}
2216 step' gst (st, False, _) = do
2217 r <- step gst st
2218 case r of
2219 Yield x s -> case cmp a x of
2220 GT -> return $ Yield x (s, False, Nothing)
2221 _ -> return $ Yield a (s, True, Just x)
2222 Skip s -> return $ Skip (s, False, Nothing)
2223 Stop -> return $ Yield a (st, True, Nothing)
2225 step' _ (_, True, Nothing) = return Stop
2227 step' gst (st, True, Just prev) = do
2228 r <- step gst st
2229 case r of
2230 Yield x s -> return $ Yield prev (s, True, Just x)
2231 Skip s -> return $ Skip (s, True, Just prev)
2232 Stop -> return $ Yield prev (st, True, Nothing)
2234------------------------------------------------------------------------------
2235-- Deleting
2236------------------------------------------------------------------------------
2238{-# INLINE_NORMAL deleteBy #-}
2239deleteBy :: Monad m => (a -> a -> Bool) -> a -> Stream m a -> Stream m a
2240deleteBy eq x (Stream step state) = Stream step' (state, False)
2241 where
2242 {-# INLINE_LATE step' #-}
2243 step' gst (st, False) = do
2244 r <- step gst st
2245 case r of
2246 Yield y s -> return $
2247 if eq x y then Skip (s, True) else Yield y (s, False)
2248 Skip s -> return $ Skip (s, False)
2249 Stop -> return Stop
2251 step' gst (st, True) = do
2252 r <- step gst st
2253 case r of
2254 Yield y s -> return $ Yield y (s, True)
2255 Skip s -> return $ Skip (s, True)
2256 Stop -> return Stop
2258------------------------------------------------------------------------------
2259-- Transformation by Map and Filter
2260------------------------------------------------------------------------------
2262-- XXX Will this always fuse properly?
2263{-# INLINE_NORMAL mapMaybe #-}
2264mapMaybe :: Monad m => (a -> Maybe b) -> Stream m a -> Stream m b
2265mapMaybe f = fmap fromJust . filter isJust . map f
2267{-# INLINE_NORMAL mapMaybeM #-}
2268mapMaybeM :: Monad m => (a -> m (Maybe b)) -> Stream m a -> Stream m b
2269mapMaybeM f = fmap fromJust . filter isJust . mapM f
2271------------------------------------------------------------------------------
2272-- Zipping
2273------------------------------------------------------------------------------
2275{-# INLINE_NORMAL indexed #-}
2276indexed :: Monad m => Stream m a -> Stream m (Int, a)
2277indexed (Stream step state) = Stream step' (state, 0)
2278 where
2279 {-# INLINE_LATE step' #-}
2280 step' gst (st, i) = i `seq` do
2281 r <- step (adaptState gst) st
2282 case r of
2283 Yield x s -> return $ Yield (i, x) (s, i+1)
2284 Skip s -> return $ Skip (s, i)
2285 Stop -> return Stop
2287{-# INLINE_NORMAL indexedR #-}
2288indexedR :: Monad m => Int -> Stream m a -> Stream m (Int, a)
2289indexedR m (Stream step state) = Stream step' (state, m)
2290 where
2291 {-# INLINE_LATE step' #-}
2292 step' gst (st, i) = i `seq` do
2293 r <- step (adaptState gst) st
2294 case r of
2295 Yield x s -> let i' = i - 1
2296 in return $ Yield (i, x) (s, i')
2297 Skip s -> return $ Skip (s, i)
2298 Stop -> return Stop
2300{-# INLINE_NORMAL zipWithM #-}
2301zipWithM :: Monad m
2302 => (a -> b -> m c) -> Stream m a -> Stream m b -> Stream m c
2303zipWithM f (Stream stepa ta) (Stream stepb tb) = Stream step (ta, tb, Nothing)
2304 where
2305 {-# INLINE_LATE step #-}
2306 step gst (sa, sb, Nothing) = do
2307 r <- stepa (adaptState gst) sa
2308 return $
2309 case r of
2310 Yield x sa' -> Skip (sa', sb, Just x)
2311 Skip sa' -> Skip (sa', sb, Nothing)
2312 Stop -> Stop
2314 step gst (sa, sb, Just x) = do
2315 r <- stepb (adaptState gst) sb
2316 case r of
2317 Yield y sb' -> do
2318 z <- f x y
2319 return $ Yield z (sa, sb', Nothing)
2320 Skip sb' -> return $ Skip (sa, sb', Just x)
2321 Stop -> return Stop
2323{-# RULES "zipWithM xs xs"
2324 forall f xs. zipWithM f xs xs = mapM (\x -> f x x) xs #-}
2326{-# INLINE zipWith #-}
2327zipWith :: Monad m => (a -> b -> c) -> Stream m a -> Stream m b -> Stream m c
2328zipWith f = zipWithM (\a b -> return (f a b))
2330------------------------------------------------------------------------------
2331-- Merging
2332------------------------------------------------------------------------------
2334{-# INLINE_NORMAL mergeByM #-}
2335mergeByM
2336 :: (Monad m)
2337 => (a -> a -> m Ordering) -> Stream m a -> Stream m a -> Stream m a
2338mergeByM cmp (Stream stepa ta) (Stream stepb tb) =
2339 Stream step (Just ta, Just tb, Nothing, Nothing)
2340 where
2341 {-# INLINE_LATE step #-}
2343 -- one of the values is missing, and the corresponding stream is running
2344 step gst (Just sa, sb, Nothing, b) = do
2345 r <- stepa gst sa
2346 return $ case r of
2347 Yield a sa' -> Skip (Just sa', sb, Just a, b)
2348 Skip sa' -> Skip (Just sa', sb, Nothing, b)
2349 Stop -> Skip (Nothing, sb, Nothing, b)
2351 step gst (sa, Just sb, a, Nothing) = do
2352 r <- stepb gst sb
2353 return $ case r of
2354 Yield b sb' -> Skip (sa, Just sb', a, Just b)
2355 Skip sb' -> Skip (sa, Just sb', a, Nothing)
2356 Stop -> Skip (sa, Nothing, a, Nothing)
2358 -- both the values are available
2359 step _ (sa, sb, Just a, Just b) = do
2360 res <- cmp a b
2361 return $ case res of
2362 GT -> Yield b (sa, sb, Just a, Nothing)
2363 _ -> Yield a (sa, sb, Nothing, Just b)
2365 -- one of the values is missing, corresponding stream is done
2366 step _ (Nothing, sb, Nothing, Just b) =
2367 return $ Yield b (Nothing, sb, Nothing, Nothing)
2369 step _ (sa, Nothing, Just a, Nothing) =
2370 return $ Yield a (sa, Nothing, Nothing, Nothing)
2372 step _ (Nothing, Nothing, Nothing, Nothing) = return Stop
2374{-# INLINE mergeBy #-}
2375mergeBy
2376 :: (Monad m)
2377 => (a -> a -> Ordering) -> Stream m a -> Stream m a -> Stream m a
2378mergeBy cmp = mergeByM (\a b -> return $ cmp a b)
2380------------------------------------------------------------------------------
2381-- Transformation comprehensions
2382------------------------------------------------------------------------------
2384{-# INLINE_NORMAL the #-}
2385the :: (Eq a, Monad m) => Stream m a -> m (Maybe a)
2386the (Stream step state) = go state
2387 where
2388 go st = do
2389 r <- step defState st
2390 case r of
2391 Yield x s -> go' x s
2392 Skip s -> go s
2393 Stop -> return Nothing
2394 go' n st = do
2395 r <- step defState st
2396 case r of
2397 Yield x s | x == n -> go' n s
2398 | otherwise -> return Nothing
2399 Skip s -> go' n s
2400 Stop -> return (Just n)
2402--------------------------------------------------------------------------------
2403-- UTF8 Encoding / Decoding
2404--------------------------------------------------------------------------------
2406-- UTF-8 primitives, Lifted from GHC.IO.Encoding.UTF8.
2408ord2 :: Char -> [Word8]
2409ord2 c = assert (n >= 0x80 && n <= 0x07ff) [x1, x2]
2410 where
2411 n = ord c
2412 x1 = fromIntegral $ (n `shiftR` 6) + 0xC0
2413 x2 = fromIntegral $ (n .&. 0x3F) + 0x80
2415ord3 :: Char -> [Word8]
2416ord3 c = assert (n >= 0x0800 && n <= 0xffff) [x1, x2, x3]
2417 where
2418 n = ord c
2419 x1 = fromIntegral $ (n `shiftR` 12) + 0xE0
2420 x2 = fromIntegral $ ((n `shiftR` 6) .&. 0x3F) + 0x80
2421 x3 = fromIntegral $ (n .&. 0x3F) + 0x80
2423ord4 :: Char -> [Word8]
2424ord4 c = assert (n >= 0x10000) [x1, x2, x3, x4]
2425 where
2426 n = ord c
2427 x1 = fromIntegral $ (n `shiftR` 18) + 0xF0
2428 x2 = fromIntegral $ ((n `shiftR` 12) .&. 0x3F) + 0x80
2429 x3 = fromIntegral $ ((n `shiftR` 6) .&. 0x3F) + 0x80
2430 x4 = fromIntegral $ (n .&. 0x3F) + 0x80
2432{-# INLINE chr2 #-}
2433chr2 :: Word8 -> Word8 -> Char
2434chr2 (W8# x1#) (W8# x2#) = C# (chr# (z1# +# z2#))
2435 where
2436 !y1# = word2Int# x1#
2437 !y2# = word2Int# x2#
2438 !z1# = uncheckedIShiftL# (y1# -# 192#) 6#
2439 !z2# = y2# -# 128#
2441{-# INLINE chr3 #-}
2442chr3 :: Word8 -> Word8 -> Word8 -> Char
2443chr3 (W8# x1#) (W8# x2#) (W8# x3#) = C# (chr# (z1# +# z2# +# z3#))
2444 where
2445 !y1# = word2Int# x1#
2446 !y2# = word2Int# x2#
2447 !y3# = word2Int# x3#
2448 !z1# = uncheckedIShiftL# (y1# -# 224#) 12#
2449 !z2# = uncheckedIShiftL# (y2# -# 128#) 6#
2450 !z3# = y3# -# 128#
2452{-# INLINE chr4 #-}
2453chr4 :: Word8 -> Word8 -> Word8 -> Word8 -> Char
2454chr4 (W8# x1#) (W8# x2#) (W8# x3#) (W8# x4#) =
2455 C# (chr# (z1# +# z2# +# z3# +# z4#))
2456 where
2457 !y1# = word2Int# x1#
2458 !y2# = word2Int# x2#
2459 !y3# = word2Int# x3#
2460 !y4# = word2Int# x4#
2461 !z1# = uncheckedIShiftL# (y1# -# 240#) 18#
2462 !z2# = uncheckedIShiftL# (y2# -# 128#) 12#
2463 !z3# = uncheckedIShiftL# (y3# -# 128#) 6#
2464 !z4# = y4# -# 128#
2466{-# INLINE between #-}
2467between ::
2468 Word8 -- ^ byte to check
2469 -> Word8 -- ^ lower bound
2470 -> Word8 -- ^ upper bound
2471 -> Bool
2472between x y z = x >= y && x <= z
2474{-# INLINE validate3 #-}
2475validate3 :: Word8 -> Word8 -> Word8 -> Bool
2476validate3 x1 x2 x3 = validate3_1 ||
2477 validate3_2 ||
2478 validate3_3 ||
2479 validate3_4
2480 where
2481 validate3_1 = (x1 == 0xE0) &&
2482 between x2 0xA0 0xBF &&
2483 between x3 0x80 0xBF
2484 validate3_2 = between x1 0xE1 0xEC &&
2485 between x2 0x80 0xBF &&
2486 between x3 0x80 0xBF
2487 validate3_3 = x1 == 0xED &&
2488 between x2 0x80 0x9F &&
2489 between x3 0x80 0xBF
2490 validate3_4 = between x1 0xEE 0xEF &&
2491 between x2 0x80 0xBF &&
2492 between x3 0x80 0xBF
2494{-# INLINE validate4 #-}
2495validate4 :: Word8 -> Word8 -> Word8 -> Word8 -> Bool
2496validate4 x1 x2 x3 x4 = validate4_1 ||
2497 validate4_2 ||
2498 validate4_3
2499 where
2500 validate4_1 = x1 == 0xF0 &&
2501 between x2 0x90 0xBF &&
2502 between x3 0x80 0xBF &&
2503 between x4 0x80 0xBF
2504 validate4_2 = between x1 0xF1 0xF3 &&
2505 between x2 0x80 0xBF &&
2506 between x3 0x80 0xBF &&
2507 between x4 0x80 0xBF
2508 validate4_3 = x1 == 0xF4 &&
2509 between x2 0x80 0x8F &&
2510 between x3 0x80 0xBF &&
2511 between x4 0x80 0xBF
2513data CodePoint = FreshPoint
2514 | Two !Word8
2515 | GoThree !Word8
2516 | Three !Word8 !Word8
2517 | GoFour !Word8
2518 | Four1 !Word8 !Word8
2519 | Four !Word8 !Word8 !Word8
2520 | YieldPoint !Char
2522{-# INLINE_NORMAL decodeUtf8 #-}
2523decodeUtf8 :: Monad m => Stream m Word8 -> Stream m Char
2524decodeUtf8 (Stream step state) = Stream step' (state, FreshPoint)
2525 where
2526 {-# INLINE_LATE step' #-}
2527 step' gst (st, FreshPoint) = do
2528 r <- step (adaptState gst) st
2529 return $
2530 case r of
2531 Yield c0 s
2532 | c0 <= 0x7f ->
2533 Skip (s, YieldPoint $ unsafeChr (fromIntegral c0))
2534 | c0 >= 0xc0 && c0 <= 0xc1 -> error "Invalid Overlong forms"
2535 | c0 >= 0xc2 && c0 <= 0xdf -> Skip (s, Two c0)
2536 | c0 >= 0xe0 && c0 <= 0xef -> Skip (s, GoThree c0)
2537 | c0 >= 0xf0 -> Skip (s, GoFour c0)
2538 | otherwise -> error "Invalid Sequence"
2539 Skip s -> Skip (s, FreshPoint)
2540 Stop -> Stop
2542 -- Decoding of a UTF8 character that spans two bytes with one byte (c0)
2543 -- known.
2544 step' gst (st, Two c0) = do
2545 r <- step (adaptState gst) st
2546 return $
2547 case r of
2548 Yield c1 s
2549 | c1 < 0x80 || c1 >= 0xc0 -> error "Invalid Sequence"
2550 | otherwise -> Skip (s, YieldPoint $ chr2 c0 c1)
2551 Skip s -> Skip (s, Two c0)
2552 Stop -> error "Input Underflow"
2554 -- Decoding of a UTF8 character that spans three bytes with first byte (c0)
2555 -- known.
2556 step' gst (st, GoThree c0) = do
2557 r <- step (adaptState gst) st
2558 return $
2559 case r of
2560 Yield c1 s
2561 | not (validate3 c0 c1 0x80) -> error "Invalid Sequence"
2562 -- See #3341@ghc
2563 | otherwise -> Skip (s, Three c0 c1)
2564 Skip s -> Skip (s, GoThree c0)
2565 Stop -> error "Input Underflow"
2567 -- Two bytes known of a 3 byte UTF8 encoded unicode character.
2568 step' gst (st, Three c0 c1) = do
2569 r <- step (adaptState gst) st
2570 return $
2571 case r of
2572 Yield c2 s
2573 | not (validate3 c0 c1 c2) -> error "Invalid Sequence"
2574 | otherwise -> Skip (s, YieldPoint $ chr3 c0 c1 c2)
2575 Skip s -> Skip (s, Three c0 c1)
2576 Stop -> error "Input Underflow"
2578 -- One byte known of a four byte UTF8 unicode character.
2579 step' gst (st, GoFour c0) = do
2580 r <- step (adaptState gst) st
2581 return $
2582 case r of
2583 Yield c1 s
2584 | not (validate4 c0 c1 0x80 0x80) -> error "Invalid Sequence"
2585 | otherwise -> Skip (s, Four1 c0 c1)
2586 Skip s -> Skip (s, GoFour c0)
2587 Stop -> error "Input Underflow"
2589 -- Two bytes known of a four byte UTF8 unicode character.
2590 step' gst (st, Four1 c0 c1) = do
2591 r <- step (adaptState gst) st
2592 return $
2593 case r of
2594 Yield c2 s
2595 | not (validate4 c0 c1 c2 0x80) -> error "Invalid Sequence"
2596 | otherwise -> Skip (s, Four c0 c1 c2)
2597 Skip s -> Skip (s, Four1 c0 c1)
2598 Stop -> error "Input Underflow"
2600 -- Three bytes known of a three byte UTF8 unicode character.
2601 step' gst (st, Four c0 c1 c2) = do
2602 r <- step (adaptState gst) st
2603 return $
2604 case r of
2605 Yield c3 s
2606 | not (validate4 c0 c1 c2 c3) -> error "Invalid Sequence"
2607 | otherwise -> Skip (s, YieldPoint $ chr4 c0 c1 c2 c3)
2608 Skip s -> Skip (s, Four c0 c1 c2)
2609 Stop -> error "Input Underflow"
2611 -- This way so that we only need one Yield.
2612 step' _ (s, YieldPoint c) = return $ Yield c (s, FreshPoint)
2614data YieldMany x = None
2615 | YieldMany !x
2617{-# INLINE_NORMAL encodeUtf8 #-}
2618encodeUtf8 :: Monad m => Stream m Char -> Stream m Word8
2619encodeUtf8 (Stream step state) = Stream step' (state, None)
2620 where
2621 {-# INLINE_LATE step' #-}
2622 step' gst (st, None) = do
2623 r <- step (adaptState gst) st
2624 return $
2625 case r of
2626 Yield c s ->
2627 case ord c of
2628 x | x <= 0x7F -> Skip (s, YieldMany [fromIntegral x])
2629 | x <= 0x7FF -> Skip (s, YieldMany (ord2 c))
2630 | x <= 0xFFFF ->
2631 if isSurrogate c
2632 then error "Encountered a surrogate"
2633 else Skip (s, YieldMany (ord3 c))
2634 | otherwise -> Skip (s, YieldMany (ord4 c))
2635 Skip s -> Skip (s, None)
2636 Stop -> Stop
2637 step' _ (s, YieldMany []) = return $ Skip (s, None)
2638 step' _ (s, YieldMany (x:xs)) = return $ Yield x (s, YieldMany xs)