git.iko.soy/streamly
on 3336bb1 help sign in
streamly/src/Streamly/Streams/StreamD.hs
2704 lines 92 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#if __GLASGOW_HASKELL__ >= 801
16{-# LANGUAGE TypeApplications #-}
17#endif
19#include "inline.hs"
21-- |
22-- Module : Streamly.Streams.StreamD
23-- Copyright : (c) 2018 Harendra Kumar
24-- Copyright : (c) Roman Leshchinskiy 2008-2010
25-- Copyright : (c) The University of Glasgow, 2009
26--
27-- License : BSD3
28-- Maintainer : harendra.kumar@gmail.com
29-- Stability : experimental
30-- Portability : GHC
31--
32-- Direct style re-implementation of CPS style stream in StreamK module. The
33-- symbol or suffix 'D' in this module denotes the "Direct" style. GHC is able
34-- to INLINE and fuse direct style better, providing better performance than
35-- CPS implementation.
36--
37-- @
38-- import qualified Streamly.Streams.StreamD as D
39-- @
41-- Some of the functions in this file have been adapted from the vector
42-- library, https://hackage.haskell.org/package/vector.
44module Streamly.Streams.StreamD
45 (
46 -- * The stream type
47 Step (..)
48 , Stream (..)
50 -- * Construction
51 , nil
52 , cons
54 -- * Deconstruction
55 , uncons
57 -- * Generation
58 -- ** Unfolds
59 , unfoldr
60 , unfoldrM
62 -- ** Specialized Generation
63 -- | Generate a monadic stream from a seed.
64 , repeat
65 , replicate
66 , replicateM
67 , fromIndices
68 , fromIndicesM
69 , generate
70 , generateM
72 -- ** Enumerations
73 , enumerateFromStepIntegral
74 , enumerateFromIntegral
75 , enumerateFromThenIntegral
76 , enumerateFromToIntegral
77 , enumerateFromThenToIntegral
79 , enumerateFromStepNum
80 , numFrom
81 , numFromThen
82 , enumerateFromToFractional
83 , enumerateFromThenToFractional
85 -- ** Conversions
86 -- | Transform an input structure into a stream.
87 -- | Direct style stream does not support @fromFoldable@.
88 , yield
89 , yieldM
90 , fromList
91 , fromListM
92 , fromStreamK
93 , fromStreamD
95 -- * Elimination
96 -- ** General Folds
97 , foldrS
98 , foldrT
99 , foldrM
100 , foldrMx
101 , foldr
102 , foldr1
104 , foldl'
105 , foldlM'
106 , foldlS
107 , foldlT
108 , reverse
109 , reverse'
111 , foldlx'
112 , foldlMx'
114 -- ** Specialized Folds
115 , tap
116 , drain
117 , null
118 , head
119 , tail
120 , last
121 , elem
122 , notElem
123 , all
124 , any
125 , maximum
126 , maximumBy
127 , minimum
128 , minimumBy
129 , findIndices
130 , lookup
131 , findM
132 , find
133 , (!!)
135 -- ** Flattening nested streams
136 , concatMapM
137 , concatMap
138 , concatMapU
140 -- ** Grouping
141 , groupsOf
142 , groupsBy
143 , groupsRollingBy
145 -- ** Splitting
146 , splitBy
147 , splitSuffixBy
148 , wordsBy
149 , splitSuffixBy'
151 , splitOn
152 , splitSuffixOn
154 -- ** Substreams
155 , isPrefixOf
156 , isSubsequenceOf
157 , stripPrefix
159 -- ** Map and Fold
160 , mapM_
162 -- ** Conversions
163 -- | Transform a stream into another type.
164 , toList
165 , toListRev
166 , toStreamK
167 , toStreamD
169 -- * Transformation
170 , transform
172 -- ** By folding (scans)
173 , scanlM'
174 , scanl'
175 , scanlM
176 , scanl
177 , scanl1M'
178 , scanl1'
179 , scanl1M
180 , scanl1
182 , prescanl'
183 , prescanlM'
185 , postscanl
186 , postscanlM
187 , postscanl'
188 , postscanlM'
190 , postscanlx'
191 , postscanlMx'
192 , scanlMx'
193 , scanlx'
195 -- * Filtering
196 , filter
197 , filterM
198 , uniq
199 , take
200 , takeWhile
201 , takeWhileM
202 , drop
203 , dropWhile
204 , dropWhileM
206 -- * Mapping
207 , map
208 , mapM
209 , sequence
211 -- * Inserting
212 , intersperseM
213 , intersperse
214 , insertAfterEach
215 , insertBy
217 -- * Deleting
218 , deleteBy
220 -- ** Map and Filter
221 , mapMaybe
222 , mapMaybeM
224 -- * Zipping
225 , indexed
226 , indexedR
227 , zipWith
228 , zipWithM
230 -- * Comparisons
231 , eqBy
232 , cmpBy
234 -- * Merging
235 , mergeBy
236 , mergeByM
238 -- * Transformation comprehensions
239 , the
241 -- * Exceptions
242 , before
243 , after
244 , bracket
245 , onException
246 , finally
247 , handle
249 -- * UTF8 Encoding / Decoding transformations.
250 , decodeUtf8
251 , encodeUtf8
252 )
253where
255import Control.Exception (Exception)
256import Control.Monad (void)
257import Control.Monad.Catch (MonadCatch)
258import Control.Monad.IO.Class (MonadIO(..))
259import Control.Monad.Trans (MonadTrans(lift))
260import Data.Bits (shiftR, shiftL, (.|.), (.&.))
261#if __GLASGOW_HASKELL__ >= 801
262import Data.Functor.Identity (Identity)
263#endif
264import Data.Maybe (fromJust, isJust)
265import Data.Word (Word32)
266import Foreign.Ptr (Ptr)
267import Foreign.Storable (Storable(..))
268import GHC.Base (assert, Char(..), unsafeChr, (-#), (+#), uncheckedIShiftL#,
269 word2Int#, chr#, ord)
270import GHC.IO.Encoding.Failure (isSurrogate)
271import GHC.Types (SPEC(..))
272import GHC.Word (Word8(..))
273import Prelude
274 hiding (map, mapM, mapM_, repeat, foldr, last, take, filter,
275 takeWhile, drop, dropWhile, all, any, maximum, minimum, elem,
276 notElem, null, head, tail, zipWith, lookup, foldr1, sequence,
277 (!!), scanl, scanl1, concatMap, replicate, enumFromTo, concat,
278 reverse)
280import qualified Control.Monad.Catch as MC
282import Streamly.Memory.Array.Types (Array(..))
283import Streamly.Fold.Types (Fold(..))
284import Streamly.Pipe.Types (Pipe(..), PipeState(..))
285import Streamly.SVar (MonadAsync, defState, adaptState)
286import Streamly.Unfold.Types (Unfold(..))
288import Streamly.Streams.StreamD.Type
290import qualified Streamly.Pipe.Types as Pipe
291import qualified Streamly.Memory.Array.Types as A
292import qualified Streamly.Memory.Ring as RB
293import qualified Streamly.Streams.StreamK as K
295------------------------------------------------------------------------------
296-- Construction
297------------------------------------------------------------------------------
299-- | An empty 'Stream'.
300{-# INLINE_NORMAL nil #-}
301nil :: Monad m => Stream m a
302nil = Stream (\_ _ -> return Stop) ()
304{-# INLINE_NORMAL consM #-}
305consM :: Monad m => m a -> Stream m a -> Stream m a
306consM m (Stream step state) = Stream step1 Nothing
307 where
308 {-# INLINE_LATE step1 #-}
309 step1 _ Nothing = m >>= \x -> return $ Yield x (Just state)
310 step1 gst (Just st) = do
311 r <- step gst st
312 return $
313 case r of
314 Yield a s -> Yield a (Just s)
315 Skip s -> Skip (Just s)
316 Stop -> Stop
318-- XXX implement in terms of consM?
319-- cons x = consM (return x)
320--
321-- | Can fuse but has O(n^2) complexity.
322{-# INLINE_NORMAL cons #-}
323cons :: Monad m => a -> Stream m a -> Stream m a
324cons x (Stream step state) = Stream step1 Nothing
325 where
326 {-# INLINE_LATE step1 #-}
327 step1 _ Nothing = return $ Yield x (Just state)
328 step1 gst (Just st) = do
329 r <- step gst st
330 return $
331 case r of
332 Yield a s -> Yield a (Just s)
333 Skip s -> Skip (Just s)
334 Stop -> Stop
336-------------------------------------------------------------------------------
337-- Deconstruction
338-------------------------------------------------------------------------------
340-- Does not fuse, has the same performance as the StreamK version.
341{-# INLINE_NORMAL uncons #-}
342uncons :: Monad m => Stream m a -> m (Maybe (a, Stream m a))
343uncons (UnStream step state) = go state
344 where
345 go st = do
346 r <- step defState st
347 case r of
348 Yield x s -> return $ Just (x, Stream step s)
349 Skip s -> go s
350 Stop -> return Nothing
352------------------------------------------------------------------------------
353-- Generation by unfold
354------------------------------------------------------------------------------
356{-# INLINE_NORMAL unfoldrM #-}
357unfoldrM :: Monad m => (s -> m (Maybe (a, s))) -> s -> Stream m a
358unfoldrM next state = Stream step state
359 where
360 {-# INLINE_LATE step #-}
361 step _ st = do
362 r <- next st
363 return $ case r of
364 Just (x, s) -> Yield x s
365 Nothing -> Stop
367{-# INLINE_LATE unfoldr #-}
368unfoldr :: Monad m => (s -> Maybe (a, s)) -> s -> Stream m a
369unfoldr f = unfoldrM (return . f)
371------------------------------------------------------------------------------
372-- Specialized Generation
373------------------------------------------------------------------------------
375repeat :: Monad m => a -> Stream m a
376repeat x = Stream (\_ _ -> return $ Yield x ()) ()
378{-# INLINE_NORMAL replicateM #-}
379replicateM :: forall m a. Monad m => Int -> m a -> Stream m a
380replicateM n p = Stream step n
381 where
382 {-# INLINE_LATE step #-}
383 step _ (i :: Int)
384 | i <= 0 = return Stop
385 | otherwise = do
386 x <- p
387 return $ Yield x (i - 1)
389{-# INLINE_NORMAL replicate #-}
390replicate :: Monad m => Int -> a -> Stream m a
391replicate n x = replicateM n (return x)
393-- This would not work properly for floats, therefore we put an Integral
394-- constraint.
395-- | Can be used to enumerate unbounded integrals. This does not check for
396-- overflow or underflow for bounded integrals.
397{-# INLINE_NORMAL enumerateFromStepIntegral #-}
398enumerateFromStepIntegral :: (Integral a, Monad m) => a -> a -> Stream m a
399enumerateFromStepIntegral from stride =
400 from `seq` stride `seq` Stream step from
401 where
402 {-# INLINE_LATE step #-}
403 step _ !x = return $ Yield x $! (x + stride)
405-- We are assuming that "to" is constrained by the type to be within
406-- max/min bounds.
407{-# INLINE enumerateFromToIntegral #-}
408enumerateFromToIntegral :: (Monad m, Integral a) => a -> a -> Stream m a
409enumerateFromToIntegral from to =
410 takeWhile (<= to) $ enumerateFromStepIntegral from 1
412{-# INLINE enumerateFromIntegral #-}
413enumerateFromIntegral :: (Monad m, Integral a, Bounded a) => a -> Stream m a
414enumerateFromIntegral from = enumerateFromToIntegral from maxBound
416data EnumState a = EnumInit | EnumYield a a a | EnumStop
418{-# INLINE_NORMAL enumerateFromThenToIntegralUp #-}
419enumerateFromThenToIntegralUp
420 :: (Monad m, Integral a)
421 => a -> a -> a -> Stream m a
422enumerateFromThenToIntegralUp from next to = Stream step EnumInit
423 where
424 {-# INLINE_LATE step #-}
425 step _ EnumInit =
426 return $
427 if to < next
428 then if to < from
429 then Stop
430 else Yield from EnumStop
431 else -- from <= next <= to
432 let stride = next - from
433 in Skip $ EnumYield from stride (to - stride)
435 step _ (EnumYield x stride toMinus) =
436 return $
437 if x > toMinus
438 then Yield x EnumStop
439 else Yield x $ EnumYield (x + stride) stride toMinus
441 step _ EnumStop = return Stop
443{-# INLINE_NORMAL enumerateFromThenToIntegralDn #-}
444enumerateFromThenToIntegralDn
445 :: (Monad m, Integral a)
446 => a -> a -> a -> Stream m a
447enumerateFromThenToIntegralDn from next to = Stream step EnumInit
448 where
449 {-# INLINE_LATE step #-}
450 step _ EnumInit =
451 return $ if to > next
452 then if to > from
453 then Stop
454 else Yield from EnumStop
455 else -- from >= next >= to
456 let stride = next - from
457 in Skip $ EnumYield from stride (to - stride)
459 step _ (EnumYield x stride toMinus) =
460 return $
461 if x < toMinus
462 then Yield x EnumStop
463 else Yield x $ EnumYield (x + stride) stride toMinus
465 step _ EnumStop = return Stop
467{-# INLINE_NORMAL enumerateFromThenToIntegral #-}
468enumerateFromThenToIntegral
469 :: (Monad m, Integral a)
470 => a -> a -> a -> Stream m a
471enumerateFromThenToIntegral from next to
472 | next >= from = enumerateFromThenToIntegralUp from next to
473 | otherwise = enumerateFromThenToIntegralDn from next to
475{-# INLINE_NORMAL enumerateFromThenIntegral #-}
476enumerateFromThenIntegral
477 :: (Monad m, Integral a, Bounded a)
478 => a -> a -> Stream m a
479enumerateFromThenIntegral from next =
480 if next > from
481 then enumerateFromThenToIntegralUp from next maxBound
482 else enumerateFromThenToIntegralDn from next minBound
484-- For floating point numbers if the increment is less than the precision then
485-- it just gets lost. Therefore we cannot always increment it correctly by just
486-- repeated addition.
487-- 9007199254740992 + 1 + 1 :: Double => 9.007199254740992e15
488-- 9007199254740992 + 2 :: Double => 9.007199254740994e15
490-- Instead we accumulate the increment counter and compute the increment
491-- everytime before adding it to the starting number.
492--
493-- This works for Integrals as well as floating point numbers, but
494-- enumerateFromStepIntegral is faster for integrals.
495{-# INLINE_NORMAL enumerateFromStepNum #-}
496enumerateFromStepNum :: (Monad m, Num a) => a -> a -> Stream m a
497enumerateFromStepNum from stride = Stream step 0
498 where
499 {-# INLINE_LATE step #-}
500 step _ !i = return $ (Yield $! (from + i * stride)) $! (i + 1)
502{-# INLINE_NORMAL numFrom #-}
503numFrom :: (Monad m, Num a) => a -> Stream m a
504numFrom from = enumerateFromStepNum from 1
506{-# INLINE_NORMAL numFromThen #-}
507numFromThen :: (Monad m, Num a) => a -> a -> Stream m a
508numFromThen from next = enumerateFromStepNum from (next - from)
510-- We cannot write a general function for Num. The only way to write code
511-- portable between the two is to use a 'Real' constraint and convert between
512-- Fractional and Integral using fromRational which is horribly slow.
513{-# INLINE_NORMAL enumerateFromToFractional #-}
514enumerateFromToFractional
515 :: (Monad m, Fractional a, Ord a)
516 => a -> a -> Stream m a
517enumerateFromToFractional from to =
518 takeWhile (<= to + 1 / 2) $ enumerateFromStepNum from 1
520{-# INLINE_NORMAL enumerateFromThenToFractional #-}
521enumerateFromThenToFractional
522 :: (Monad m, Fractional a, Ord a)
523 => a -> a -> a -> Stream m a
524enumerateFromThenToFractional from next to =
525 takeWhile predicate $ numFromThen from next
526 where
527 mid = (next - from) / 2
528 predicate | next >= from = (<= to + mid)
529 | otherwise = (>= to + mid)
531-------------------------------------------------------------------------------
532-- Generation by Conversion
533-------------------------------------------------------------------------------
535{-# INLINE_NORMAL fromIndicesM #-}
536fromIndicesM :: Monad m => (Int -> m a) -> Stream m a
537fromIndicesM gen = Stream step 0
538 where
539 {-# INLINE_LATE step #-}
540 step _ i = do
541 x <- gen i
542 return $ Yield x (i + 1)
544{-# INLINE fromIndices #-}
545fromIndices :: Monad m => (Int -> a) -> Stream m a
546fromIndices gen = fromIndicesM (return . gen)
548{-# INLINE_NORMAL generateM #-}
549generateM :: Monad m => Int -> (Int -> m a) -> Stream m a
550generateM n gen = n `seq` Stream step 0
551 where
552 {-# INLINE_LATE step #-}
553 step _ i | i < n = do
554 x <- gen i
555 return $ Yield x (i + 1)
556 | otherwise = return Stop
558{-# INLINE generate #-}
559generate :: Monad m => Int -> (Int -> a) -> Stream m a
560generate n gen = generateM n (return . gen)
562-- XXX we need the MonadAsync constraint because of a rewrite rule.
563-- | Convert a list of monadic actions to a 'Stream'
564{-# INLINE_LATE fromListM #-}
565fromListM :: MonadAsync m => [m a] -> Stream m a
566fromListM = Stream step
567 where
568 {-# INLINE_LATE step #-}
569 step _ (m:ms) = m >>= \x -> return $ Yield x ms
570 step _ [] = return Stop
572{-# INLINE toStreamD #-}
573toStreamD :: (K.IsStream t, Monad m) => t m a -> Stream m a
574toStreamD = fromStreamK . K.toStream
576------------------------------------------------------------------------------
577-- Elimination by Folds
578------------------------------------------------------------------------------
580------------------------------------------------------------------------------
581-- Right Folds
582------------------------------------------------------------------------------
584{-# INLINE_NORMAL foldr1 #-}
585foldr1 :: Monad m => (a -> a -> a) -> Stream m a -> m (Maybe a)
586foldr1 f m = do
587 r <- uncons m
588 case r of
589 Nothing -> return Nothing
590 Just (h, t) -> fmap Just (foldr f h t)
592------------------------------------------------------------------------------
593-- Left Folds
594------------------------------------------------------------------------------
596{-# INLINE_NORMAL foldlT #-}
597foldlT :: (Monad m, Monad (s m), MonadTrans s)
598 => (s m b -> a -> s m b) -> s m b -> Stream m a -> s m b
599foldlT fstep begin (Stream step state) = go SPEC begin state
600 where
601 go !_ acc st = do
602 r <- lift $ step defState st
603 case r of
604 Yield x s -> go SPEC (fstep acc x) s
605 Skip s -> go SPEC acc s
606 Stop -> acc
608-- Note, this is going to have horrible performance, because of the nature of
609-- the stream type (i.e. direct stream vs CPS). Its only for reference, it is
610-- likely be practically unusable.
611{-# INLINE_NORMAL foldlS #-}
612foldlS :: Monad m
613 => (Stream m b -> a -> Stream m b) -> Stream m b -> Stream m a -> Stream m b
614foldlS fstep begin (Stream step state) = Stream step' (Left (state, begin))
615 where
616 step' gst (Left (st, acc)) = do
617 r <- step (adaptState gst) st
618 return $ case r of
619 Yield x s -> Skip (Left (s, fstep acc x))
620 Skip s -> Skip (Left (s, acc))
621 Stop -> Skip (Right acc)
623 step' gst (Right (Stream stp stt)) = do
624 r <- stp (adaptState gst) stt
625 return $ case r of
626 Yield x s -> Yield x (Right (Stream stp s))
627 Skip s -> Skip (Right (Stream stp s))
628 Stop -> Stop
630------------------------------------------------------------------------------
631-- Specialized Folds
632------------------------------------------------------------------------------
634-- | Run a streaming composition, discard the results.
635{-# INLINE_LATE drain #-}
636drain :: Monad m => Stream m a -> m ()
637-- drain = foldrM (\_ xs -> xs) (return ())
638drain (Stream step state) = go SPEC state
639 where
640 go !_ st = do
641 r <- step defState st
642 case r of
643 Yield _ s -> go SPEC s
644 Skip s -> go SPEC s
645 Stop -> return ()
647{-# INLINE_NORMAL null #-}
648null :: Monad m => Stream m a -> m Bool
649null m = foldrM (\_ _ -> return False) (return True) m
651{-# INLINE_NORMAL head #-}
652head :: Monad m => Stream m a -> m (Maybe a)
653head m = foldrM (\x _ -> return (Just x)) (return Nothing) m
655-- Does not fuse, has the same performance as the StreamK version.
656{-# INLINE_NORMAL tail #-}
657tail :: Monad m => Stream m a -> m (Maybe (Stream m a))
658tail (UnStream step state) = go state
659 where
660 go st = do
661 r <- step defState st
662 case r of
663 Yield _ s -> return (Just $ Stream step s)
664 Skip s -> go s
665 Stop -> return Nothing
667-- XXX will it fuse? need custom impl?
668{-# INLINE_NORMAL last #-}
669last :: Monad m => Stream m a -> m (Maybe a)
670last = foldl' (\_ y -> Just y) Nothing
672{-# INLINE_NORMAL elem #-}
673elem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
674-- elem e m = foldrM (\x xs -> if x == e then return True else xs) (return False) m
675elem e (Stream step state) = go state
676 where
677 go st = do
678 r <- step defState st
679 case r of
680 Yield x s
681 | x == e -> return True
682 | otherwise -> go s
683 Skip s -> go s
684 Stop -> return False
686{-# INLINE_NORMAL notElem #-}
687notElem :: (Monad m, Eq a) => a -> Stream m a -> m Bool
688notElem e s = fmap not (elem e s)
690{-# INLINE_NORMAL all #-}
691all :: Monad m => (a -> Bool) -> Stream m a -> m Bool
692-- all p m = foldrM (\x xs -> if p x then xs else return False) (return True) m
693all p (Stream step state) = go state
694 where
695 go st = do
696 r <- step defState st
697 case r of
698 Yield x s
699 | p x -> go s
700 | otherwise -> return False
701 Skip s -> go s
702 Stop -> return True
704{-# INLINE_NORMAL any #-}
705any :: Monad m => (a -> Bool) -> Stream m a -> m Bool
706-- any p m = foldrM (\x xs -> if p x then return True else xs) (return False) m
707any p (Stream step state) = go state
708 where
709 go st = do
710 r <- step defState st
711 case r of
712 Yield x s
713 | p x -> return True
714 | otherwise -> go s
715 Skip s -> go s
716 Stop -> return False
718{-# INLINE_NORMAL maximum #-}
719maximum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
720maximum (Stream step state) = go Nothing state
721 where
722 go Nothing st = do
723 r <- step defState st
724 case r of
725 Yield x s -> go (Just x) s
726 Skip s -> go Nothing s
727 Stop -> return Nothing
728 go (Just acc) st = do
729 r <- step defState st
730 case r of
731 Yield x s
732 | acc <= x -> go (Just x) s
733 | otherwise -> go (Just acc) s
734 Skip s -> go (Just acc) s
735 Stop -> return (Just acc)
737{-# INLINE_NORMAL maximumBy #-}
738maximumBy :: Monad m => (a -> a -> Ordering) -> Stream m a -> m (Maybe a)
739maximumBy cmp (Stream step state) = go Nothing state
740 where
741 go Nothing st = do
742 r <- step defState st
743 case r of
744 Yield x s -> go (Just x) s
745 Skip s -> go Nothing s
746 Stop -> return Nothing
747 go (Just acc) st = do
748 r <- step defState st
749 case r of
750 Yield x s -> case cmp acc x of
751 GT -> go (Just acc) s
752 _ -> go (Just x) s
753 Skip s -> go (Just acc) s
754 Stop -> return (Just acc)
756{-# INLINE_NORMAL minimum #-}
757minimum :: (Monad m, Ord a) => Stream m a -> m (Maybe a)
758minimum (Stream step state) = go Nothing state
759 where
760 go Nothing st = do
761 r <- step defState st
762 case r of
763 Yield x s -> go (Just x) s
764 Skip s -> go Nothing s
765 Stop -> return Nothing
766 go (Just acc) st = do
767 r <- step defState st
768 case r of
769 Yield x s
770 | acc <= x -> go (Just acc) s
771 | otherwise -> go (Just x) s
772 Skip s -> go (Just acc) s
773 Stop -> return (Just acc)
775{-# INLINE_NORMAL minimumBy #-}
776minimumBy :: Monad m => (a -> a -> Ordering) -> Stream m a -> m (Maybe a)
777minimumBy cmp (Stream step state) = go Nothing state
778 where
779 go Nothing st = do
780 r <- step defState st
781 case r of
782 Yield x s -> go (Just x) s
783 Skip s -> go Nothing s
784 Stop -> return Nothing
785 go (Just acc) st = do
786 r <- step defState st
787 case r of
788 Yield x s -> case cmp acc x of
789 GT -> go (Just x) s
790 _ -> go (Just acc) s
791 Skip s -> go (Just acc) s
792 Stop -> return (Just acc)
794{-# INLINE_NORMAL (!!) #-}
795(!!) :: (Monad m) => Stream m a -> Int -> m (Maybe a)
796(Stream step state) !! i = go i state
797 where
798 go n st = do
799 r <- step defState st
800 case r of
801 Yield x s | n < 0 -> return Nothing
802 | n == 0 -> return $ Just x
803 | otherwise -> go (n - 1) s
804 Skip s -> go n s
805 Stop -> return Nothing
807{-# INLINE_NORMAL lookup #-}
808lookup :: (Monad m, Eq a) => a -> Stream m (a, b) -> m (Maybe b)
809lookup e m = foldrM (\(a, b) xs -> if e == a then return (Just b) else xs)
810 (return Nothing) m
812{-# INLINE_NORMAL findM #-}
813findM :: Monad m => (a -> m Bool) -> Stream m a -> m (Maybe a)
814findM p m = foldrM (\x xs -> p x >>= \r -> if r then return (Just x) else xs)
815 (return Nothing) m
817{-# INLINE find #-}
818find :: Monad m => (a -> Bool) -> Stream m a -> m (Maybe a)
819find p = findM (return . p)
821{-# INLINE_NORMAL findIndices #-}
822findIndices :: Monad m => (a -> Bool) -> Stream m a -> Stream m Int
823findIndices p (Stream step state) = Stream step' (state, 0)
824 where
825 {-# INLINE_LATE step' #-}
826 step' gst (st, i) = do
827 r <- step (adaptState gst) st
828 return $ case r of
829 Yield x s -> if p x then Yield i (s, i+1) else Skip (s, i+1)
830 Skip s -> Skip (s, i+1)
831 Stop -> Stop
833{-# INLINE toListRev #-}
834toListRev :: Monad m => Stream m a -> m [a]
835toListRev = foldl' (flip (:)) []
837-- We can implement reverse as:
838--
839-- > reverse = foldlS (flip cons) nil
840--
841-- However, this implementation is unusable because of the horrible performance
842-- of cons. So we just convert it to a list first and then stream from the
843-- list.
844--
845-- XXX Maybe we can use an Array instead of a list here?
846{-# INLINE_NORMAL reverse #-}
847reverse :: Monad m => Stream m a -> Stream m a
848reverse m = Stream step Nothing
849 where
850 {-# INLINE_LATE step #-}
851 step _ Nothing = do
852 xs <- toListRev m
853 return $ Skip (Just xs)
854 step _ (Just (x:xs)) = return $ Yield x (Just xs)
855 step _ (Just []) = return Stop
857-- Much faster reverse for Storables
858{-# INLINE_NORMAL reverse' #-}
859reverse' :: forall m a. (MonadIO m, Storable a) => Stream m a -> Stream m a
860{-
861-- This commented implementation copies the whole stream into one single array
862-- and then streams from that array, this is 3-4x faster than the chunked code
863-- that follows. Though this could be problematic due to unbounded large
864-- allocations. We need to figure out why the chunked code is slower and if we
865-- can optimize the chunked code to work as fast as this one. It may be a
866-- fusion issue?
867import Foreign.ForeignPtr (touchForeignPtr)
868import Foreign.ForeignPtr.Unsafe (unsafeForeignPtrToPtr)
869import Foreign.Ptr (Ptr, plusPtr)
870reverse' m = Stream step Nothing
871 where
872 {-# INLINE_LATE step #-}
873 step _ Nothing = do
874 arr <- A.fromStreamD m
875 let p = aEnd arr `plusPtr` negate (sizeOf (undefined :: a))
876 return $ Skip $ Just (aStart arr, p)
878 step _ (Just (start, p)) | p < unsafeForeignPtrToPtr start = return Stop
880 step _ (Just (start, p)) = do
881 let !x = A.unsafeInlineIO $ do
882 r <- peek p
883 touchForeignPtr start
884 return r
885 next = p `plusPtr` negate (sizeOf (undefined :: a))
886 return $ Yield x (Just (start, next))
887-}
888reverse' m =
889 A.flattenArraysRev
890 $ fromStreamK
891 $ K.reverse
892 $ toStreamK
893 $ A.fromStreamDArraysOf A.defaultChunkSize m
896------------------------------------------------------------------------------
897-- Grouping/Splitting
898------------------------------------------------------------------------------
900{-# INLINE_NORMAL splitSuffixBy' #-}
901splitSuffixBy' :: Monad m
902 => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
903splitSuffixBy' predicate f (Stream step state) =
904 Stream (stepOuter f) (Just state)
906 where
908 {-# INLINE_LATE stepOuter #-}
909 stepOuter (Fold fstep initial done) gst (Just st) = do
910 res <- step (adaptState gst) st
911 case res of
912 Yield x s -> do
913 acc <- initial
914 acc' <- fstep acc x
915 if (predicate x)
916 then done acc' >>= \val -> return $ Yield val (Just s)
917 else go SPEC s acc'
919 Skip s -> return $ Skip $ Just s
920 Stop -> return Stop
922 where
924 go !_ stt !acc = do
925 res <- step (adaptState gst) stt
926 case res of
927 Yield x s -> do
928 acc' <- fstep acc x
929 if (predicate x)
930 then done acc' >>= \val -> return $ Yield val (Just s)
931 else go SPEC s acc'
932 Skip s -> go SPEC s acc
933 Stop -> done acc >>= \val -> return $ Yield val Nothing
935 stepOuter _ _ Nothing = return Stop
937{-# INLINE_NORMAL groupsBy #-}
938groupsBy :: Monad m
939 => (a -> a -> Bool)
940 -> Fold m a b
941 -> Stream m a
942 -> Stream m b
943groupsBy cmp f (Stream step state) = Stream (stepOuter f) (Just state, Nothing)
945 where
947 {-# INLINE_LATE stepOuter #-}
948 stepOuter (Fold fstep initial done) gst (Just st, Nothing) = do
949 res <- step (adaptState gst) st
950 case res of
951 Yield x s -> do
952 acc <- initial
953 acc' <- fstep acc x
954 go SPEC x s acc'
956 Skip s -> return $ Skip $ (Just s, Nothing)
957 Stop -> return Stop
959 where
961 go !_ prev stt !acc = do
962 res <- step (adaptState gst) stt
963 case res of
964 Yield x s -> do
965 if cmp x prev
966 then do
967 acc' <- fstep acc x
968 go SPEC prev s acc'
969 else done acc >>= \r -> return $ Yield r (Just s, Just x)
970 Skip s -> go SPEC prev s acc
971 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
973 stepOuter (Fold fstep initial done) gst (Just st, Just prev) = do
974 acc <- initial
975 acc' <- fstep acc prev
976 go SPEC st acc'
978 where
980 -- XXX code duplicated from the previous equation
981 go !_ stt !acc = do
982 res <- step (adaptState gst) stt
983 case res of
984 Yield x s -> do
985 if cmp x prev
986 then do
987 acc' <- fstep acc x
988 go SPEC s acc'
989 else done acc >>= \r -> return $ Yield r (Just s, Just x)
990 Skip s -> go SPEC s acc
991 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
993 stepOuter _ _ (Nothing,_) = return Stop
995{-# INLINE_NORMAL groupsRollingBy #-}
996groupsRollingBy :: Monad m
997 => (a -> a -> Bool)
998 -> Fold m a b
999 -> Stream m a
1000 -> Stream m b
1001groupsRollingBy cmp f (Stream step state) =
1002 Stream (stepOuter f) (Just state, Nothing)
1003 where
1005 {-# INLINE_LATE stepOuter #-}
1006 stepOuter (Fold fstep initial done) gst (Just st, Nothing) = do
1007 res <- step (adaptState gst) st
1008 case res of
1009 Yield x s -> do
1010 acc <- initial
1011 acc' <- fstep acc x
1012 go SPEC x s acc'
1014 Skip s -> return $ Skip $ (Just s, Nothing)
1015 Stop -> return Stop
1017 where
1018 go !_ prev stt !acc = do
1019 res <- step (adaptState gst) stt
1020 case res of
1021 Yield x s -> do
1022 if cmp prev x
1023 then do
1024 acc' <- fstep acc x
1025 go SPEC x s acc'
1026 else
1027 done acc >>= \r -> return $ Yield r (Just s, Just x)
1028 Skip s -> go SPEC prev s acc
1029 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1031 stepOuter (Fold fstep initial done) gst (Just st, Just prev') = do
1032 acc <- initial
1033 acc' <- fstep acc prev'
1034 go SPEC prev' st acc'
1036 where
1037 go !_ prevv stt !acc = do
1038 res <- step (adaptState gst) stt
1039 case res of
1040 Yield x s -> do
1041 if cmp prevv x
1042 then do
1043 acc' <- fstep acc x
1044 go SPEC x s acc'
1045 else done acc >>= \r -> return $ Yield r (Just s, Just x)
1046 Skip s -> go SPEC prevv s acc
1047 Stop -> done acc >>= \r -> return $ Yield r (Nothing, Nothing)
1049 stepOuter _ _ (Nothing, _) = return Stop
1051{-# INLINE_NORMAL splitBy #-}
1052splitBy :: Monad m => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1053splitBy 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) = initial >>= go SPEC st
1060 where
1062 go !_ stt !acc = do
1063 res <- step (adaptState gst) stt
1064 case res of
1065 Yield x s -> do
1066 if predicate x
1067 then done acc >>= \r -> return $ Yield r (Just s)
1068 else do
1069 acc' <- fstep acc x
1070 go SPEC s acc'
1071 Skip s -> go SPEC s acc
1072 Stop -> done acc >>= \r -> return $ Yield r Nothing
1074 step' _ _ Nothing = return Stop
1076{-# INLINE_NORMAL splitSuffixBy #-}
1077splitSuffixBy :: Monad m
1078 => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1079splitSuffixBy predicate f (Stream step state) = Stream (step' f) (Just state)
1081 where
1083 {-# INLINE_LATE step' #-}
1084 step' (Fold fstep initial done) gst (Just st) = do
1085 res <- step (adaptState gst) st
1086 case res of
1087 Yield x s -> do
1088 acc <- initial
1089 if predicate x
1090 then done acc >>= \val -> return $ Yield val (Just s)
1091 else do
1092 acc' <- fstep acc x
1093 go SPEC s acc'
1095 Skip s -> return $ Skip $ Just s
1096 Stop -> return Stop
1098 where
1100 go !_ stt !acc = do
1101 res <- step (adaptState gst) stt
1102 case res of
1103 Yield x s -> do
1104 if predicate x
1105 then done acc >>= \r -> return $ Yield r (Just s)
1106 else do
1107 acc' <- fstep acc x
1108 go SPEC s acc'
1109 Skip s -> go SPEC s acc
1110 Stop -> done acc >>= \r -> return $ Yield r Nothing
1112 step' _ _ Nothing = return Stop
1114{-# INLINE_NORMAL wordsBy #-}
1115wordsBy :: Monad m => (a -> Bool) -> Fold m a b -> Stream m a -> Stream m b
1116wordsBy predicate f (Stream step state) = Stream (stepOuter f) (Just state)
1118 where
1120 {-# INLINE_LATE stepOuter #-}
1121 stepOuter (Fold fstep initial done) gst (Just st) = do
1122 res <- step (adaptState gst) st
1123 case res of
1124 Yield x s -> do
1125 if predicate x
1126 then return $ Skip (Just s)
1127 else do
1128 acc <- initial
1129 acc' <- fstep acc x
1130 go SPEC s acc'
1132 Skip s -> return $ Skip $ Just s
1133 Stop -> return Stop
1135 where
1137 go !_ stt !acc = do
1138 res <- step (adaptState gst) stt
1139 case res of
1140 Yield x s -> do
1141 if predicate x
1142 then done acc >>= \r -> return $ Yield r (Just s)
1143 else do
1144 acc' <- fstep acc x
1145 go SPEC s acc'
1146 Skip s -> go SPEC s acc
1147 Stop -> done acc >>= \r -> return $ Yield r Nothing
1149 stepOuter _ _ Nothing = return Stop
1151-- String search algorithms:
1152-- http://www-igm.univ-mlv.fr/~lecroq/string/index.html
1155-- TODO can we unify the splitting operations using a splitting configuration
1156-- like in the split package.
1158data SplitStyle = Infix | Suffix | Prefix deriving (Eq, Show)
1160data SplitOptions = SplitOptions
1161 { style :: SplitStyle
1162 , withSep :: Bool -- ^ keep the separators in output
1163 -- , compact :: Bool -- ^ treat multiple consecutive separators as one
1164 -- , trimHead :: Bool -- ^ drop blank at head
1165 -- , trimTail :: Bool -- ^ drop blank at tail
1169data SplitOnState s a =
1170 GO_START
1171 | GO_EMPTY_PAT s
1172 | GO_SINGLE_PAT s a
1173 | GO_SHORT_PAT s
1174 | GO_KARP_RABIN s !(RB.Ring a) !(Ptr a)
1175 | GO_DONE
1177{-# INLINE_NORMAL splitOn #-}
1178splitOn
1179 :: forall m a b. (MonadIO m, Storable a, Enum a, Eq a)
1180 => Array a
1181 -> Fold m a b
1182 -> Stream m a
1183 -> Stream m b
1184splitOn patArr@Array{..} (Fold fstep initial done) (Stream step state) =
1185 Stream stepOuter GO_START
1187 where
1189 patLen = A.length patArr
1190 maxIndex = patLen - 1
1191 elemBits = sizeOf (undefined :: a) * 8
1193 {-# INLINE_LATE stepOuter #-}
1194 stepOuter _ GO_START =
1195 if patLen == 0
1196 then return $ Skip $ GO_EMPTY_PAT state
1197 else if patLen == 1
1198 then do
1199 r <- liftIO $ (A.unsafeIndexIO patArr 0)
1200 return $ Skip $ GO_SINGLE_PAT state r
1201 else if sizeOf (undefined :: a) * patLen
1202 <= sizeOf (undefined :: Word)
1203 then return $ Skip $ GO_SHORT_PAT state
1204 else do
1205 (rb, rhead) <- liftIO $ RB.new patLen
1206 return $ Skip $ GO_KARP_RABIN state rb rhead
1208 stepOuter gst (GO_SINGLE_PAT stt pat) = initial >>= go SPEC stt
1210 where
1212 go !_ st !acc = do
1213 res <- step (adaptState gst) st
1214 case res of
1215 Yield x s -> do
1216 if pat == x
1217 then do
1218 r <- done acc
1219 return $ Yield r (GO_SINGLE_PAT s pat)
1220 else fstep acc x >>= go SPEC s
1221 Skip s -> go SPEC s acc
1222 Stop -> done acc >>= \r -> return $ Yield r GO_DONE
1224 stepOuter gst (GO_SHORT_PAT stt) = initial >>= go0 SPEC 0 (0 :: Word) stt
1226 where
1228 mask :: Word
1229 mask = (1 `shiftL` (elemBits * patLen)) - 1
1231 addToWord wrd a = (wrd `shiftL` elemBits) .|. fromIntegral (fromEnum a)
1233 patWord :: Word
1234 patWord = mask .&. A.foldl' addToWord 0 patArr
1236 go0 !_ !idx wrd st !acc = do
1237 res <- step (adaptState gst) st
1238 case res of
1239 Yield x s -> do
1240 let wrd' = addToWord wrd x
1241 if idx == maxIndex
1242 then do
1243 if wrd' .&. mask == patWord
1244 then do
1245 r <- done acc
1246 return $ Yield r (GO_SHORT_PAT s)
1247 else go1 SPEC wrd' s acc
1248 else go0 SPEC (idx + 1) wrd' s acc
1249 Skip s -> go0 SPEC idx wrd s acc
1250 Stop -> do
1251 acc' <- if idx /= 0
1252 then go2 wrd idx acc
1253 else return acc
1254 done acc' >>= \r -> return $ Yield r GO_DONE
1256 {-# INLINE go1 #-}
1257 go1 !_ wrd st !acc = do
1258 res <- step (adaptState gst) st
1259 case res of
1260 Yield x s -> do
1261 let wrd' = addToWord wrd x
1262 old = (mask .&. wrd) `shiftR` (elemBits * (patLen - 1))
1263 acc' <- fstep acc (toEnum $ fromIntegral old)
1264 if wrd' .&. mask == patWord
1265 then done acc' >>= \r -> return $ Yield r (GO_SHORT_PAT s)
1266 else go1 SPEC wrd' s acc'
1267 Skip s -> go1 SPEC wrd s acc
1268 Stop -> do
1269 acc' <- go2 wrd patLen acc
1270 done acc' >>= \r -> return $ Yield r GO_DONE
1272 go2 !wrd !n !acc | n > 0 = do
1273 let old = (mask .&. wrd) `shiftR` (elemBits * (n - 1))
1274 fstep acc (toEnum $ fromIntegral old) >>= go2 wrd (n - 1)
1275 go2 _ _ acc = return acc
1277 stepOuter gst (GO_KARP_RABIN stt rb rhead) = do
1278 initial >>= go0 SPEC 0 rhead stt
1280 where
1282 k = 2891336453 :: Word32
1283 coeff = k ^ patLen
1284 addCksum cksum a = cksum * k + fromIntegral (fromEnum a)
1285 deltaCksum cksum old new =
1286 addCksum cksum new - coeff * fromIntegral (fromEnum old)
1288 -- XXX shall we use a random starting hash or 1 instead of 0?
1289 patHash = A.foldl' addCksum 0 patArr
1291 -- rh == ringHead
1292 go0 !_ !idx !rh st !acc = do
1293 res <- step (adaptState gst) st
1294 case res of
1295 Yield x s -> do
1296 rh' <- liftIO $ RB.unsafeInsert rb rh x
1297 if idx == maxIndex
1298 then do
1299 let fold = RB.unsafeFoldRing (RB.ringBound rb)
1300 let !ringHash = fold addCksum 0 rb
1301 if ringHash == patHash
1302 then go2 SPEC ringHash rh' s acc
1303 else go1 SPEC ringHash rh' s acc
1304 else go0 SPEC (idx + 1) rh' s acc
1305 Skip s -> go0 SPEC idx rh s acc
1306 Stop -> do
1307 !acc' <- if idx /= 0
1308 then RB.unsafeFoldRingM rh fstep acc rb
1309 else return acc
1310 done acc' >>= \r -> return $ Yield r GO_DONE
1312 -- XXX Theoretically this code can do 4 times faster if GHC generates
1313 -- optimal code. If we use just "(cksum' == patHash)" condition it goes
1314 -- 4x faster, as soon as we add the "RB.unsafeEqArray rb v" condition
1315 -- the generated code changes drastically and becomes 4x slower. Need
1316 -- to investigate what is going on with GHC.
1317 {-# INLINE go1 #-}
1318 go1 !_ !cksum !rh st !acc = do
1319 res <- step (adaptState gst) st
1320 case res of
1321 Yield x s -> do
1322 old <- liftIO $ peek rh
1323 let cksum' = deltaCksum cksum old x
1324 acc' <- fstep acc old
1326 if (cksum' == patHash)
1327 then do
1328 rh' <- liftIO (RB.unsafeInsert rb rh x)
1329 go2 SPEC cksum' rh' s acc'
1330 else do
1331 rh' <- liftIO (RB.unsafeInsert rb rh x)
1332 go1 SPEC cksum' rh' s acc'
1333 Skip s -> go1 SPEC cksum rh s acc
1334 Stop -> do
1335 acc' <- RB.unsafeFoldRingFullM rh fstep acc rb
1336 done acc' >>= \r -> return $ Yield r GO_DONE
1338 go2 !_ !cksum' !rh' s !acc' = do
1339 if RB.unsafeEqArray rb rh' patArr
1340 then do
1341 r <- done acc'
1342 return $ Yield r (GO_KARP_RABIN s rb rhead)
1343 else go1 SPEC cksum' rh' s acc'
1345 stepOuter gst (GO_EMPTY_PAT st) = do
1346 res <- step (adaptState gst) st
1347 case res of
1348 Yield x s -> do
1349 acc <- initial
1350 acc' <- fstep acc x
1351 done acc' >>= \r -> return $ Yield r (GO_EMPTY_PAT s)
1352 Skip s -> return $ Skip (GO_EMPTY_PAT s)
1353 Stop -> return Stop
1355 stepOuter _ GO_DONE = return Stop
1357{-# INLINE_NORMAL splitSuffixOn #-}
1358splitSuffixOn
1359 :: forall m a b. (MonadIO m, Storable a, Enum a, Eq a)
1360 => Bool
1361 -> Array a
1362 -> Fold m a b
1363 -> Stream m a
1364 -> Stream m b
1365splitSuffixOn withSep patArr@Array{..} (Fold fstep initial done)
1366 (Stream step state) =
1367 Stream stepOuter GO_START
1369 where
1371 patLen = A.length patArr
1372 maxIndex = patLen - 1
1373 elemBits = sizeOf (undefined :: a) * 8
1375 {-# INLINE_LATE stepOuter #-}
1376 stepOuter _ GO_START =
1377 if patLen == 0
1378 then return $ Skip $ GO_EMPTY_PAT state
1379 else if patLen == 1
1380 then do
1381 r <- liftIO $ (A.unsafeIndexIO patArr 0)
1382 return $ Skip $ GO_SINGLE_PAT state r
1383 else if sizeOf (undefined :: a) * patLen
1384 <= sizeOf (undefined :: Word)
1385 then return $ Skip $ GO_SHORT_PAT state
1386 else do
1387 (rb, rhead) <- liftIO $ RB.new patLen
1388 return $ Skip $ GO_KARP_RABIN state rb rhead
1390 stepOuter gst (GO_SINGLE_PAT stt pat) = do
1391 -- This first part is the only difference between splitOn and
1392 -- splitSuffixOn.
1393 -- If the last element is a separator do not issue a blank segment.
1394 res <- step (adaptState gst) stt
1395 case res of
1396 Yield x s -> do
1397 acc <- initial
1398 if pat == x
1399 then do
1400 acc' <- if withSep then fstep acc x else return acc
1401 done acc' >>= \r -> return $ Yield r (GO_SINGLE_PAT s pat)
1402 else fstep acc x >>= go SPEC s
1403 Skip s -> return $ Skip $ (GO_SINGLE_PAT s pat)
1404 Stop -> return Stop
1406 where
1408 -- This is identical for splitOn and splitSuffixOn
1409 go !_ st !acc = do
1410 res <- step (adaptState gst) st
1411 case res of
1412 Yield x s -> do
1413 if pat == x
1414 then do
1415 acc' <- if withSep then fstep acc x else return acc
1416 r <- done acc'
1417 return $ Yield r (GO_SINGLE_PAT s pat)
1418 else fstep acc x >>= go SPEC s
1419 Skip s -> go SPEC s acc
1420 Stop -> done acc >>= \r -> return $ Yield r GO_DONE
1422 stepOuter gst (GO_SHORT_PAT stt) = do
1424 -- Call "initial" only if the stream yields an element, otherwise we
1425 -- may call "initial" but never yield anything. initial may produce a
1426 -- side effect, therefore we will end up doing and discard a side
1427 -- effect.
1429 let idx = 0
1430 let wrd = 0
1431 res <- step (adaptState gst) stt
1432 case res of
1433 Yield x s -> do
1434 acc <- initial
1435 let wrd' = addToWord wrd x
1436 acc' <- if withSep then fstep acc x else return acc
1437 if idx == maxIndex
1438 then do
1439 if wrd' .&. mask == patWord
1440 then done acc' >>= \r -> return $ Yield r (GO_SHORT_PAT s)
1441 else go0 SPEC (idx + 1) wrd' s acc'
1442 else go0 SPEC (idx + 1) wrd' s acc'
1443 Skip s -> return $ Skip (GO_SHORT_PAT s)
1444 Stop -> return Stop
1446 where
1448 mask :: Word
1449 mask = (1 `shiftL` (elemBits * patLen)) - 1
1451 addToWord wrd a = (wrd `shiftL` elemBits) .|. fromIntegral (fromEnum a)
1453 patWord :: Word
1454 patWord = mask .&. A.foldl' addToWord 0 patArr
1456 go0 !_ !idx wrd st !acc = do
1457 res <- step (adaptState gst) st
1458 case res of
1459 Yield x s -> do
1460 let wrd' = addToWord wrd x
1461 acc' <- if withSep then fstep acc x else return acc
1462 if idx == maxIndex
1463 then do
1464 if wrd' .&. mask == patWord
1465 then do
1466 r <- done acc'
1467 return $ Yield r (GO_SHORT_PAT s)
1468 else go1 SPEC wrd' s acc'
1469 else go0 SPEC (idx + 1) wrd' s acc'
1470 Skip s -> go0 SPEC idx wrd s acc
1471 Stop -> do
1472 if (idx == maxIndex) && (wrd .&. mask == patWord)
1473 then return Stop
1474 else do
1475 acc' <- if idx /= 0 && not withSep
1476 then go2 wrd idx acc
1477 else return acc
1478 done acc' >>= \r -> return $ Yield r GO_DONE
1480 {-# INLINE go1 #-}
1481 go1 !_ wrd st !acc = do
1482 res <- step (adaptState gst) st
1483 case res of
1484 Yield x s -> do
1485 let wrd' = addToWord wrd x
1486 old = (mask .&. wrd) `shiftR` (elemBits * (patLen - 1))
1487 acc' <- if withSep
1488 then fstep acc x
1489 else fstep acc (toEnum $ fromIntegral old)
1490 if wrd' .&. mask == patWord
1491 then done acc' >>= \r -> return $ Yield r (GO_SHORT_PAT s)
1492 else go1 SPEC wrd' s acc'
1493 Skip s -> go1 SPEC wrd s acc
1494 Stop ->
1495 -- If the last sequence is a separator do not issue a blank
1496 -- segment.
1497 if wrd .&. mask == patWord
1498 then return Stop
1499 else do
1500 acc' <- if withSep
1501 then return acc
1502 else go2 wrd patLen acc
1503 done acc' >>= \r -> return $ Yield r GO_DONE
1505 go2 !wrd !n !acc | n > 0 = do
1506 let old = (mask .&. wrd) `shiftR` (elemBits * (n - 1))
1507 fstep acc (toEnum $ fromIntegral old) >>= go2 wrd (n - 1)
1508 go2 _ _ acc = return acc
1510 stepOuter gst (GO_KARP_RABIN stt rb rhead) = do
1511 let idx = 0
1512 res <- step (adaptState gst) stt
1513 case res of
1514 Yield x s -> do
1515 acc <- initial
1516 acc' <- if withSep then fstep acc x else return acc
1517 rh' <- liftIO (RB.unsafeInsert rb rhead x)
1518 if idx == maxIndex
1519 then do
1520 let fold = RB.unsafeFoldRing (RB.ringBound rb)
1521 let !ringHash = fold addCksum 0 rb
1522 if ringHash == patHash
1523 then go2 SPEC ringHash rh' s acc'
1524 else go0 SPEC (idx + 1) rh' s acc'
1525 else go0 SPEC (idx + 1) rh' s acc'
1526 Skip s -> return $ Skip (GO_KARP_RABIN s rb rhead)
1527 Stop -> return Stop
1529 where
1531 k = 2891336453 :: Word32
1532 coeff = k ^ patLen
1533 addCksum cksum a = cksum * k + fromIntegral (fromEnum a)
1534 deltaCksum cksum old new =
1535 addCksum cksum new - coeff * fromIntegral (fromEnum old)
1537 -- XXX shall we use a random starting hash or 1 instead of 0?
1538 patHash = A.foldl' addCksum 0 patArr
1540 -- rh == ringHead
1541 go0 !_ !idx !rh st !acc = do
1542 res <- step (adaptState gst) st
1543 case res of
1544 Yield x s -> do
1545 acc' <- if withSep then fstep acc x else return acc
1546 rh' <- liftIO (RB.unsafeInsert rb rh x)
1547 if idx == maxIndex
1548 then do
1549 let fold = RB.unsafeFoldRing (RB.ringBound rb)
1550 let !ringHash = fold addCksum 0 rb
1551 if ringHash == patHash
1552 then go2 SPEC ringHash rh' s acc'
1553 else go1 SPEC ringHash rh' s acc'
1554 else go0 SPEC (idx + 1) rh' s acc'
1555 Skip s -> go0 SPEC idx rh s acc
1556 Stop -> do
1557 -- do not issue a blank segment when we end at pattern
1558 if (idx == maxIndex) && RB.unsafeEqArray rb rh patArr
1559 then return Stop
1560 else do
1561 !acc' <- if idx /= 0 && not withSep
1562 then RB.unsafeFoldRingM rh fstep acc rb
1563 else return acc
1564 done acc' >>= \r -> return $ Yield r GO_DONE
1566 -- XXX Theoretically this code can do 4 times faster if GHC generates
1567 -- optimal code. If we use just "(cksum' == patHash)" condition it goes
1568 -- 4x faster, as soon as we add the "RB.unsafeEqArray rb v" condition
1569 -- the generated code changes drastically and becomes 4x slower. Need
1570 -- to investigate what is going on with GHC.
1571 {-# INLINE go1 #-}
1572 go1 !_ !cksum !rh st !acc = do
1573 res <- step (adaptState gst) st
1574 case res of
1575 Yield x s -> do
1576 old <- liftIO $ peek rh
1577 let cksum' = deltaCksum cksum old x
1578 acc' <- if withSep
1579 then fstep acc x
1580 else fstep acc old
1582 if (cksum' == patHash)
1583 then do
1584 rh' <- liftIO (RB.unsafeInsert rb rh x)
1585 go2 SPEC cksum' rh' s acc'
1586 else do
1587 rh' <- liftIO (RB.unsafeInsert rb rh x)
1588 go1 SPEC cksum' rh' s acc'
1589 Skip s -> go1 SPEC cksum rh s acc
1590 Stop -> do
1591 if RB.unsafeEqArray rb rh patArr
1592 then return Stop
1593 else do
1594 acc' <- if withSep
1595 then return acc
1596 else RB.unsafeFoldRingFullM rh fstep acc rb
1597 done acc' >>= \r -> return $ Yield r GO_DONE
1599 go2 !_ !cksum' !rh' s !acc' = do
1600 if RB.unsafeEqArray rb rh' patArr
1601 then do
1602 r <- done acc'
1603 return $ Yield r (GO_KARP_RABIN s rb rhead)
1604 else go1 SPEC cksum' rh' s acc'
1606 stepOuter gst (GO_EMPTY_PAT st) = do
1607 res <- step (adaptState gst) st
1608 case res of
1609 Yield x s -> do
1610 acc <- initial
1611 acc' <- fstep acc x
1612 done acc' >>= \r -> return $ Yield r (GO_EMPTY_PAT s)
1613 Skip s -> return $ Skip (GO_EMPTY_PAT s)
1614 Stop -> return Stop
1616 stepOuter _ GO_DONE = return Stop
1618------------------------------------------------------------------------------
1619-- Substreams
1620------------------------------------------------------------------------------
1622{-# INLINE_NORMAL isPrefixOf #-}
1623isPrefixOf :: (Eq a, Monad m) => Stream m a -> Stream m a -> m Bool
1624isPrefixOf (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
1625 where
1626 go (sa, sb, Nothing) = do
1627 r <- stepa defState sa
1628 case r of
1629 Yield x sa' -> go (sa', sb, Just x)
1630 Skip sa' -> go (sa', sb, Nothing)
1631 Stop -> return True
1633 go (sa, sb, Just x) = do
1634 r <- stepb defState sb
1635 case r of
1636 Yield y sb' ->
1637 if x == y
1638 then go (sa, sb', Nothing)
1639 else return False
1640 Skip sb' -> go (sa, sb', Just x)
1641 Stop -> return False
1643{-# INLINE_NORMAL isSubsequenceOf #-}
1644isSubsequenceOf :: (Eq a, Monad m) => Stream m a -> Stream m a -> m Bool
1645isSubsequenceOf (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
1646 where
1647 go (sa, sb, Nothing) = do
1648 r <- stepa defState sa
1649 case r of
1650 Yield x sa' -> go (sa', sb, Just x)
1651 Skip sa' -> go (sa', sb, Nothing)
1652 Stop -> return True
1654 go (sa, sb, Just x) = do
1655 r <- stepb defState sb
1656 case r of
1657 Yield y sb' ->
1658 if x == y
1659 then go (sa, sb', Nothing)
1660 else go (sa, sb', Just x)
1661 Skip sb' -> go (sa, sb', Just x)
1662 Stop -> return False
1664{-# INLINE_NORMAL stripPrefix #-}
1665stripPrefix
1666 :: (Eq a, Monad m)
1667 => Stream m a -> Stream m a -> m (Maybe (Stream m a))
1668stripPrefix (Stream stepa ta) (Stream stepb tb) = go (ta, tb, Nothing)
1669 where
1670 go (sa, sb, Nothing) = do
1671 r <- stepa defState sa
1672 case r of
1673 Yield x sa' -> go (sa', sb, Just x)
1674 Skip sa' -> go (sa', sb, Nothing)
1675 Stop -> return $ Just (Stream stepb sb)
1677 go (sa, sb, Just x) = do
1678 r <- stepb defState sb
1679 case r of
1680 Yield y sb' ->
1681 if x == y
1682 then go (sa, sb', Nothing)
1683 else return Nothing
1684 Skip sb' -> go (sa, sb', Just x)
1685 Stop -> return Nothing
1687------------------------------------------------------------------------------
1688-- Map and Fold
1689------------------------------------------------------------------------------
1691-- | Execute a monadic action for each element of the 'Stream'
1692{-# INLINE_NORMAL mapM_ #-}
1693mapM_ :: Monad m => (a -> m b) -> Stream m a -> m ()
1694mapM_ m = drain . mapM m
1696-------------------------------------------------------------------------------
1697-- Stream transformations using Unfolds
1698-------------------------------------------------------------------------------
1700-- | @concatMapU unfold stream@ uses @unfold@ to map the input stream elements
1701-- to streams and then flattens the generated streams into a single output
1702-- stream.
1704-- This is like 'concatMap' but uses an unfold with an explicit state to
1705-- generate the stream instead of a 'Stream' type generator. This allows better
1706-- optimization via fusion. This can be many times more efficient than
1707-- 'concatMap'.
1709{-# INLINE_NORMAL concatMapU #-}
1710concatMapU :: Monad m => Unfold m a b -> Stream m a -> Stream m b
1711concatMapU (Unfold istep inject) (Stream ostep u) = Stream step (Left u)
1712 where
1713 {-# INLINE_LATE step #-}
1714 step gst (Left t) = do
1715 r <- ostep (adaptState gst) t
1716 case r of
1717 Yield a t' -> do
1718 s <- inject a
1719 s `seq` return (Skip (Right (s, t')))
1720 Skip t' -> return $ Skip (Left t')
1721 Stop -> return $ Stop
1723 step _ (Right (s, t)) = do
1724 r <- istep s
1725 return $ case r of
1726 Yield x s' -> Yield x (Right (s', t))
1727 Skip s' -> Skip (Right (s', t))
1728 Stop -> Skip (Left t)
1730------------------------------------------------------------------------------
1731-- Exceptions
1732------------------------------------------------------------------------------
1734-- | Run a side effect before the stream yields its first element.
1735{-# INLINE_NORMAL before #-}
1736before :: Monad m => m b -> Stream m a -> Stream m a
1737before action (Stream step state) = Stream step' Nothing
1739 where
1741 {-# INLINE_LATE step' #-}
1742 step' _ Nothing = action >> return (Skip (Just state))
1744 step' gst (Just st) = do
1745 res <- step gst st
1746 case res of
1747 Yield x s -> return $ Yield x (Just s)
1748 Skip s -> return $ Skip (Just s)
1749 Stop -> return Stop
1751-- | Run a side effect whenever the stream stops normally.
1752{-# INLINE_NORMAL after #-}
1753after :: Monad m => m b -> Stream m a -> Stream m a
1754after action (Stream step state) = Stream step' state
1756 where
1758 {-# INLINE_LATE step' #-}
1759 step' gst st = do
1760 res <- step gst st
1761 case res of
1762 Yield x s -> return $ Yield x s
1763 Skip s -> return $ Skip s
1764 Stop -> action >> return Stop
1766-- XXX These combinators are expensive due to the call to
1767-- onException/handle/try on each step. Therefore, when possible, they should
1768-- be called in an outer loop where we perform less iterations. For example, we
1769-- cannot call them on each iteration in a char stream, instead we can call
1770-- them when doing an IO on an array.
1772-- XXX For high performance error checks in busy streams we may need another
1773-- Error constructor in step.
1775-- | Run a side effect whenever the stream aborts due to an exception. The
1776-- exception is not caught simply rethrown.
1777{-# INLINE_NORMAL onException #-}
1778onException :: MonadCatch m => m b -> Stream m a -> Stream m a
1779onException action (Stream step state) = Stream step' state
1781 where
1783 {-# INLINE_LATE step' #-}
1784 step' gst st = do
1785 res <- step gst st `MC.onException` action
1786 case res of
1787 Yield x s -> return $ Yield x s
1788 Skip s -> return $ Skip s
1789 Stop -> return Stop
1791-- | Run a side effect whenever the stream stops normally or aborts due to an
1792-- exception.
1793{-# INLINE finally #-}
1794finally :: MonadCatch m => m b -> Stream m a -> Stream m a
1795finally action xs = after action $ onException action xs
1797-- XXX bracket is like concatMap, it generates a stream and then flattens it.
1798-- Like concatMap it has 10x worse performance compared to linear fused
1799-- compositions.
1801-- | Run the first action before the stream starts and remember its output,
1802-- generate a stream using the output, run the second action using the
1803-- remembered value as an argument whenever the stream ends normally or due to
1804-- an exception.
1805{-# INLINE_NORMAL bracket #-}
1806bracket :: MonadCatch m => m b -> (b -> m c) -> (b -> Stream m a) -> Stream m a
1807bracket bef aft bet = Stream step' Nothing
1809 where
1811 {-# INLINE_LATE step' #-}
1812 step' _ Nothing = bef >>= \x -> return (Skip (Just (bet x, x)))
1814 -- NOTE: It is important to use UnStream instead of the Stream pattern
1815 -- here, otherwise we get huge perf degradation, see note in concatMap.
1816 step' gst (Just (UnStream step state, v)) = do
1817 res <- step gst state `MC.onException` aft v
1818 case res of
1819 Yield x s -> return $ Yield x (Just (Stream step s, v))
1820 Skip s -> return $ Skip (Just (Stream step s, v))
1821 Stop -> aft v >> return Stop
1823-- | When evaluating a stream if an exception occurs, stream evaluation aborts
1824-- and the specified exception handler is run with the exception as argument.
1825{-# INLINE_NORMAL handle #-}
1826handle :: (MonadCatch m, Exception e) => (e -> m a) -> Stream m a -> Stream m a
1827handle f (Stream step state) = Stream step' (Just state)
1829 where
1831 {-# INLINE_LATE step' #-}
1832 step' gst (Just st) = do
1833 res <- MC.try $ step gst st
1834 case res of
1835 Left e -> f e >>= \x -> return (Yield x Nothing)
1836 Right r -> case r of
1837 Yield x s -> return $ Yield x (Just s)
1838 Skip s -> return $ Skip (Just s)
1839 Stop -> return Stop
1841 step' _ Nothing = return Stop
1843-------------------------------------------------------------------------------
1844-- General transformation
1845-------------------------------------------------------------------------------
1847{-# INLINE_NORMAL transform #-}
1848transform :: Monad m => Pipe m a b -> Stream m a -> Stream m b
1849transform (Pipe pstep1 pstep2 pstate) (Stream step state) =
1850 Stream step' (Consume pstate, state)
1852 where
1854 {-# INLINE_LATE step' #-}
1856 step' gst (Consume pst, st) = pst `seq` do
1857 r <- step (adaptState gst) st
1858 case r of
1859 Yield x s -> do
1860 res <- pstep1 pst x
1861 case res of
1862 Pipe.Yield b pst' -> return $ Yield b (pst', s)
1863 Pipe.Continue pst' -> return $ Skip (pst', s)
1864 Skip s -> return $ Skip (Consume pst, s)
1865 Stop -> return Stop
1867 step' _ (Produce pst, st) = pst `seq` do
1868 res <- pstep2 pst
1869 case res of
1870 Pipe.Yield b pst' -> return $ Yield b (pst', st)
1871 Pipe.Continue pst' -> return $ Skip (pst', st)
1873------------------------------------------------------------------------------
1874-- Transformation by Folding (Scans)
1875------------------------------------------------------------------------------
1877------------------------------------------------------------------------------
1878-- Prescans
1879------------------------------------------------------------------------------
1881-- XXX Is a prescan useful, discarding the last step does not sound useful? I
1882-- am not sure about the utility of this function, so this is implemented but
1883-- not exposed. We can expose it if someone provides good reasons why this is
1884-- useful.
1886-- XXX We have to execute the stream one step ahead to know that we are at the
1887-- last step. The vector implementation of prescan executes the last fold step
1888-- but does not yield the result. This means we have executed the effect but
1889-- discarded value. This does not sound right. In this implementation we are
1890-- not executing the last fold step.
1891{-# INLINE_NORMAL prescanlM' #-}
1892prescanlM' :: Monad m => (b -> a -> m b) -> m b -> Stream m a -> Stream m b
1893prescanlM' f mz (Stream step state) = Stream step' (state, mz)
1894 where
1895 {-# INLINE_LATE step' #-}
1896 step' gst (st, prev) = do
1897 r <- step (adaptState gst) st
1898 case r of
1899 Yield x s -> do
1900 acc <- prev
1901 return $ Yield acc (s, f acc x)
1902 Skip s -> return $ Skip (s, prev)
1903 Stop -> return Stop
1905{-# INLINE prescanl' #-}
1906prescanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
1907prescanl' f z = prescanlM' (\a b -> return (f a b)) (return z)
1909------------------------------------------------------------------------------
1910-- Monolithic postscans (postscan followed by a map)
1911------------------------------------------------------------------------------
1913-- The performance of a modular postscan followed by a map seems to be
1914-- equivalent to this monolithic scan followed by map therefore we may not need
1915-- this implementation. We just have it for performance comparison and in case
1916-- modular version does not perform well in some situation.
1918{-# INLINE_NORMAL postscanlMx' #-}
1919postscanlMx' :: Monad m
1920 => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> Stream m b
1921postscanlMx' fstep begin done (Stream step state) = do
1922 Stream step' (state, begin)
1923 where
1924 {-# INLINE_LATE step' #-}
1925 step' gst (st, acc) = do
1926 r <- step (adaptState gst) st
1927 case r of
1928 Yield x s -> do
1929 old <- acc
1930 y <- fstep old x
1931 v <- done y
1932 v `seq` y `seq` return (Yield v (s, return y))
1933 Skip s -> return $ Skip (s, acc)
1934 Stop -> return Stop
1936{-# INLINE_NORMAL postscanlx' #-}
1937postscanlx' :: Monad m
1938 => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> Stream m b
1939postscanlx' fstep begin done s =
1940 postscanlMx' (\b a -> return (fstep b a)) (return begin) (return . done) s
1942-- XXX do we need consM strict to evaluate the begin value?
1943{-# INLINE scanlMx' #-}
1944scanlMx' :: Monad m
1945 => (x -> a -> m x) -> m x -> (x -> m b) -> Stream m a -> Stream m b
1946scanlMx' fstep begin done s =
1947 (begin >>= \x -> x `seq` done x) `consM` postscanlMx' fstep begin done s
1949{-# INLINE scanlx' #-}
1950scanlx' :: Monad m
1951 => (x -> a -> x) -> x -> (x -> b) -> Stream m a -> Stream m b
1952scanlx' fstep begin done s =
1953 scanlMx' (\b a -> return (fstep b a)) (return begin) (return . done) s
1955------------------------------------------------------------------------------
1956-- postscans
1957------------------------------------------------------------------------------
1959{-# INLINE_NORMAL postscanlM' #-}
1960postscanlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
1961postscanlM' fstep begin (Stream step state) =
1962 begin `seq` Stream step' (state, begin)
1963 where
1964 {-# INLINE_LATE step' #-}
1965 step' gst (st, acc) = acc `seq` do
1966 r <- step (adaptState gst) st
1967 case r of
1968 Yield x s -> do
1969 y <- fstep acc x
1970 y `seq` return (Yield y (s, y))
1971 Skip s -> return $ Skip (s, acc)
1972 Stop -> return Stop
1974{-# INLINE_NORMAL postscanl' #-}
1975postscanl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
1976postscanl' f = postscanlM' (\a b -> return (f a b))
1978{-# INLINE_NORMAL postscanlM #-}
1979postscanlM :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
1980postscanlM fstep begin (Stream step state) = Stream step' (state, begin)
1981 where
1982 {-# INLINE_LATE step' #-}
1983 step' gst (st, acc) = do
1984 r <- step (adaptState gst) st
1985 case r of
1986 Yield x s -> do
1987 y <- fstep acc x
1988 return (Yield y (s, y))
1989 Skip s -> return $ Skip (s, acc)
1990 Stop -> return Stop
1992{-# INLINE_NORMAL postscanl #-}
1993postscanl :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a
1994postscanl f = postscanlM (\a b -> return (f a b))
1996{-# INLINE_NORMAL scanlM' #-}
1997scanlM' :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
1998scanlM' fstep begin s = begin `seq` (begin `cons` postscanlM' fstep begin s)
2000{-# INLINE scanl' #-}
2001scanl' :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
2002scanl' f = scanlM' (\a b -> return (f a b))
2004{-# INLINE_NORMAL scanlM #-}
2005scanlM :: Monad m => (b -> a -> m b) -> b -> Stream m a -> Stream m b
2006scanlM fstep begin s = begin `cons` postscanlM fstep begin s
2008{-# INLINE scanl #-}
2009scanl :: Monad m => (b -> a -> b) -> b -> Stream m a -> Stream m b
2010scanl f = scanlM (\a b -> return (f a b))
2012{-# INLINE_NORMAL scanl1M #-}
2013scanl1M :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
2014scanl1M fstep (Stream step state) = Stream step' (state, Nothing)
2015 where
2016 {-# INLINE_LATE step' #-}
2017 step' gst (st, Nothing) = do
2018 r <- step gst st
2019 case r of
2020 Yield x s -> return $ Yield x (s, Just x)
2021 Skip s -> return $ Skip (s, Nothing)
2022 Stop -> return Stop
2024 step' gst (st, Just acc) = do
2025 r <- step gst st
2026 case r of
2027 Yield y s -> do
2028 z <- fstep acc y
2029 return $ Yield z (s, Just z)
2030 Skip s -> return $ Skip (s, Just acc)
2031 Stop -> return Stop
2033{-# INLINE scanl1 #-}
2034scanl1 :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
2035scanl1 f = scanl1M (\x y -> return (f x y))
2037{-# INLINE_NORMAL scanl1M' #-}
2038scanl1M' :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a
2039scanl1M' fstep (Stream step state) = Stream step' (state, Nothing)
2040 where
2041 {-# INLINE_LATE step' #-}
2042 step' gst (st, Nothing) = do
2043 r <- step gst st
2044 case r of
2045 Yield x s -> x `seq` return $ Yield x (s, Just x)
2046 Skip s -> return $ Skip (s, Nothing)
2047 Stop -> return Stop
2049 step' gst (st, Just acc) = acc `seq` do
2050 r <- step gst st
2051 case r of
2052 Yield y s -> do
2053 z <- fstep acc y
2054 z `seq` return $ Yield z (s, Just z)
2055 Skip s -> return $ Skip (s, Just acc)
2056 Stop -> return Stop
2058{-# INLINE scanl1' #-}
2059scanl1' :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a
2060scanl1' f = scanl1M' (\x y -> return (f x y))
2062{-# INLINE tap #-}
2063tap :: Monad m => Fold m a b -> Stream m a -> Stream m a
2064tap (Fold fstep initial extract) (Stream step state) = Stream step' Nothing
2066 where
2068 step' _ Nothing = do
2069 r <- initial
2070 return $ Skip (Just (r, state))
2072 step' gst (Just (acc, st)) = do
2073 r <- step gst st
2074 case r of
2075 Yield x s -> do
2076 acc' <- fstep acc x
2077 return $ Yield x (Just (acc', s))
2078 Skip s -> return $ Skip (Just (acc, s))
2079 Stop -> do
2080 void $ extract acc
2081 return $ Stop
2083-------------------------------------------------------------------------------
2084-- Filtering
2085-------------------------------------------------------------------------------
2087{-# INLINE_NORMAL takeWhileM #-}
2088takeWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
2089takeWhileM f (Stream step state) = Stream step' state
2090 where
2091 {-# INLINE_LATE step' #-}
2092 step' gst st = do
2093 r <- step gst st
2094 case r of
2095 Yield x s -> do
2096 b <- f x
2097 return $ if b then Yield x s else Stop
2098 Skip s -> return $ Skip s
2099 Stop -> return Stop
2101{-# INLINE takeWhile #-}
2102takeWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
2103takeWhile f = takeWhileM (return . f)
2105{-# INLINE_NORMAL drop #-}
2106drop :: Monad m => Int -> Stream m a -> Stream m a
2107drop n (Stream step state) = Stream step' (state, Just n)
2108 where
2109 {-# INLINE_LATE step' #-}
2110 step' gst (st, Just i)
2111 | i > 0 = do
2112 r <- step gst st
2113 return $
2114 case r of
2115 Yield _ s -> Skip (s, Just (i - 1))
2116 Skip s -> Skip (s, Just i)
2117 Stop -> Stop
2118 | otherwise = return $ Skip (st, Nothing)
2120 step' gst (st, Nothing) = do
2121 r <- step gst st
2122 return $
2123 case r of
2124 Yield x s -> Yield x (s, Nothing)
2125 Skip s -> Skip (s, Nothing)
2126 Stop -> Stop
2128data DropWhileState s a
2129 = DropWhileDrop s
2130 | DropWhileYield a s
2131 | DropWhileNext s
2133{-# INLINE_NORMAL dropWhileM #-}
2134dropWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
2135dropWhileM f (Stream step state) = Stream step' (DropWhileDrop state)
2136 where
2137 {-# INLINE_LATE step' #-}
2138 step' gst (DropWhileDrop st) = do
2139 r <- step gst st
2140 case r of
2141 Yield x s -> do
2142 b <- f x
2143 if b
2144 then return $ Skip (DropWhileDrop s)
2145 else return $ Skip (DropWhileYield x s)
2146 Skip s -> return $ Skip (DropWhileDrop s)
2147 Stop -> return Stop
2149 step' gst (DropWhileNext st) = do
2150 r <- step gst st
2151 case r of
2152 Yield x s -> return $ Skip (DropWhileYield x s)
2153 Skip s -> return $ Skip (DropWhileNext s)
2154 Stop -> return Stop
2156 step' _ (DropWhileYield x st) = return $ Yield x (DropWhileNext st)
2158{-# INLINE dropWhile #-}
2159dropWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
2160dropWhile f = dropWhileM (return . f)
2162{-# INLINE_NORMAL filterM #-}
2163filterM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a
2164filterM f (Stream step state) = Stream step' state
2165 where
2166 {-# INLINE_LATE step' #-}
2167 step' gst st = do
2168 r <- step gst st
2169 case r of
2170 Yield x s -> do
2171 b <- f x
2172 return $ if b
2173 then Yield x s
2174 else Skip s
2175 Skip s -> return $ Skip s
2176 Stop -> return Stop
2178{-# INLINE filter #-}
2179filter :: Monad m => (a -> Bool) -> Stream m a -> Stream m a
2180filter f = filterM (return . f)
2182{-# INLINE_NORMAL uniq #-}
2183uniq :: (Eq a, Monad m) => Stream m a -> Stream m a
2184uniq (Stream step state) = Stream step' (Nothing, state)
2185 where
2186 {-# INLINE_LATE step' #-}
2187 step' gst (Nothing, st) = do
2188 r <- step gst st
2189 case r of
2190 Yield x s -> return $ Yield x (Just x, s)
2191 Skip s -> return $ Skip (Nothing, s)
2192 Stop -> return Stop
2193 step' gst (Just x, st) = do
2194 r <- step gst st
2195 case r of
2196 Yield y s | x == y -> return $ Skip (Just x, s)
2197 | otherwise -> return $ Yield y (Just y, s)
2198 Skip s -> return $ Skip (Just x, s)
2199 Stop -> return Stop
2201------------------------------------------------------------------------------
2202-- Transformation by Mapping
2203------------------------------------------------------------------------------
2205{-# INLINE_NORMAL sequence #-}
2206sequence :: Monad m => Stream m (m a) -> Stream m a
2207sequence (Stream step state) = Stream step' state
2208 where
2209 {-# INLINE_LATE step' #-}
2210 step' gst st = do
2211 r <- step (adaptState gst) st
2212 case r of
2213 Yield x s -> x >>= \a -> return (Yield a s)
2214 Skip s -> return $ Skip s
2215 Stop -> return Stop
2217------------------------------------------------------------------------------
2218-- Inserting
2219------------------------------------------------------------------------------
2221data LoopState x s = FirstYield s
2222 | InterspersingYield s
2223 | YieldAndCarry x s
2225{-# INLINE_NORMAL intersperseM #-}
2226intersperseM :: Monad m => m a -> Stream m a -> Stream m a
2227intersperseM m (Stream step state) = Stream step' (FirstYield state)
2228 where
2229 {-# INLINE_LATE step' #-}
2230 step' gst (FirstYield st) = do
2231 r <- step gst st
2232 return $
2233 case r of
2234 Yield x s -> Skip (YieldAndCarry x s)
2235 Skip s -> Skip (FirstYield s)
2236 Stop -> Stop
2238 step' gst (InterspersingYield st) = do
2239 r <- step gst st
2240 case r of
2241 Yield x s -> do
2242 a <- m
2243 return $ Yield a (YieldAndCarry x s)
2244 Skip s -> return $ Skip $ InterspersingYield s
2245 Stop -> return Stop
2247 step' _ (YieldAndCarry x st) = return $ Yield x (InterspersingYield st)
2249data SuffixState s a
2250 = SuffixElem s
2251 | SuffixSuffix s
2252 | SuffixYield a (SuffixState s a)
2254{-# INLINE_NORMAL insertAfterEach #-}
2255insertAfterEach :: forall m a. Monad m => m a -> Stream m a -> Stream m a
2256insertAfterEach action (Stream step state) = Stream step' (SuffixElem state)
2257 where
2258 {-# INLINE_LATE step' #-}
2259 step' gst (SuffixElem st) = do
2260 r <- step gst st
2261 return $ case r of
2262 Yield x s -> Skip (SuffixYield x (SuffixSuffix s))
2263 Skip s -> Skip (SuffixElem s)
2264 Stop -> Stop
2266 step' _ (SuffixSuffix st) = do
2267 action >>= \r -> return $ Skip (SuffixYield r (SuffixElem st))
2269 step' _ (SuffixYield x next) = return $ Yield x next
2271{-# INLINE intersperse #-}
2272intersperse :: Monad m => a -> Stream m a -> Stream m a
2273intersperse a = intersperseM (return a)
2275{-# INLINE_NORMAL insertBy #-}
2276insertBy :: Monad m => (a -> a -> Ordering) -> a -> Stream m a -> Stream m a
2277insertBy cmp a (Stream step state) = Stream step' (state, False, Nothing)
2278 where
2279 {-# INLINE_LATE step' #-}
2280 step' gst (st, False, _) = do
2281 r <- step gst st
2282 case r of
2283 Yield x s -> case cmp a x of
2284 GT -> return $ Yield x (s, False, Nothing)
2285 _ -> return $ Yield a (s, True, Just x)
2286 Skip s -> return $ Skip (s, False, Nothing)
2287 Stop -> return $ Yield a (st, True, Nothing)
2289 step' _ (_, True, Nothing) = return Stop
2291 step' gst (st, True, Just prev) = do
2292 r <- step gst st
2293 case r of
2294 Yield x s -> return $ Yield prev (s, True, Just x)
2295 Skip s -> return $ Skip (s, True, Just prev)
2296 Stop -> return $ Yield prev (st, True, Nothing)
2298------------------------------------------------------------------------------
2299-- Deleting
2300------------------------------------------------------------------------------
2302{-# INLINE_NORMAL deleteBy #-}
2303deleteBy :: Monad m => (a -> a -> Bool) -> a -> Stream m a -> Stream m a
2304deleteBy eq x (Stream step state) = Stream step' (state, False)
2305 where
2306 {-# INLINE_LATE step' #-}
2307 step' gst (st, False) = do
2308 r <- step gst st
2309 case r of
2310 Yield y s -> return $
2311 if eq x y then Skip (s, True) else Yield y (s, False)
2312 Skip s -> return $ Skip (s, False)
2313 Stop -> return Stop
2315 step' gst (st, True) = do
2316 r <- step gst st
2317 case r of
2318 Yield y s -> return $ Yield y (s, True)
2319 Skip s -> return $ Skip (s, True)
2320 Stop -> return Stop
2322------------------------------------------------------------------------------
2323-- Transformation by Map and Filter
2324------------------------------------------------------------------------------
2326-- XXX Will this always fuse properly?
2327{-# INLINE_NORMAL mapMaybe #-}
2328mapMaybe :: Monad m => (a -> Maybe b) -> Stream m a -> Stream m b
2329mapMaybe f = fmap fromJust . filter isJust . map f
2331{-# INLINE_NORMAL mapMaybeM #-}
2332mapMaybeM :: Monad m => (a -> m (Maybe b)) -> Stream m a -> Stream m b
2333mapMaybeM f = fmap fromJust . filter isJust . mapM f
2335------------------------------------------------------------------------------
2336-- Zipping
2337------------------------------------------------------------------------------
2339{-# INLINE_NORMAL indexed #-}
2340indexed :: Monad m => Stream m a -> Stream m (Int, a)
2341indexed (Stream step state) = Stream step' (state, 0)
2342 where
2343 {-# INLINE_LATE step' #-}
2344 step' gst (st, i) = i `seq` do
2345 r <- step (adaptState gst) st
2346 case r of
2347 Yield x s -> return $ Yield (i, x) (s, i+1)
2348 Skip s -> return $ Skip (s, i)
2349 Stop -> return Stop
2351{-# INLINE_NORMAL indexedR #-}
2352indexedR :: Monad m => Int -> Stream m a -> Stream m (Int, a)
2353indexedR m (Stream step state) = Stream step' (state, m)
2354 where
2355 {-# INLINE_LATE step' #-}
2356 step' gst (st, i) = i `seq` do
2357 r <- step (adaptState gst) st
2358 case r of
2359 Yield x s -> let i' = i - 1
2360 in return $ Yield (i, x) (s, i')
2361 Skip s -> return $ Skip (s, i)
2362 Stop -> return Stop
2364{-# INLINE_NORMAL zipWithM #-}
2365zipWithM :: Monad m
2366 => (a -> b -> m c) -> Stream m a -> Stream m b -> Stream m c
2367zipWithM f (Stream stepa ta) (Stream stepb tb) = Stream step (ta, tb, Nothing)
2368 where
2369 {-# INLINE_LATE step #-}
2370 step gst (sa, sb, Nothing) = do
2371 r <- stepa (adaptState gst) sa
2372 return $
2373 case r of
2374 Yield x sa' -> Skip (sa', sb, Just x)
2375 Skip sa' -> Skip (sa', sb, Nothing)
2376 Stop -> Stop
2378 step gst (sa, sb, Just x) = do
2379 r <- stepb (adaptState gst) sb
2380 case r of
2381 Yield y sb' -> do
2382 z <- f x y
2383 return $ Yield z (sa, sb', Nothing)
2384 Skip sb' -> return $ Skip (sa, sb', Just x)
2385 Stop -> return Stop
2387#if __GLASGOW_HASKELL__ >= 801
2388{-# RULES "zipWithM xs xs"
2389 forall f xs. zipWithM @Identity f xs xs = mapM (\x -> f x x) xs #-}
2390#endif
2392{-# INLINE zipWith #-}
2393zipWith :: Monad m => (a -> b -> c) -> Stream m a -> Stream m b -> Stream m c
2394zipWith f = zipWithM (\a b -> return (f a b))
2396------------------------------------------------------------------------------
2397-- Merging
2398------------------------------------------------------------------------------
2400{-# INLINE_NORMAL mergeByM #-}
2401mergeByM
2402 :: (Monad m)
2403 => (a -> a -> m Ordering) -> Stream m a -> Stream m a -> Stream m a
2404mergeByM cmp (Stream stepa ta) (Stream stepb tb) =
2405 Stream step (Just ta, Just tb, Nothing, Nothing)
2406 where
2407 {-# INLINE_LATE step #-}
2409 -- one of the values is missing, and the corresponding stream is running
2410 step gst (Just sa, sb, Nothing, b) = do
2411 r <- stepa gst sa
2412 return $ case r of
2413 Yield a sa' -> Skip (Just sa', sb, Just a, b)
2414 Skip sa' -> Skip (Just sa', sb, Nothing, b)
2415 Stop -> Skip (Nothing, sb, Nothing, b)
2417 step gst (sa, Just sb, a, Nothing) = do
2418 r <- stepb gst sb
2419 return $ case r of
2420 Yield b sb' -> Skip (sa, Just sb', a, Just b)
2421 Skip sb' -> Skip (sa, Just sb', a, Nothing)
2422 Stop -> Skip (sa, Nothing, a, Nothing)
2424 -- both the values are available
2425 step _ (sa, sb, Just a, Just b) = do
2426 res <- cmp a b
2427 return $ case res of
2428 GT -> Yield b (sa, sb, Just a, Nothing)
2429 _ -> Yield a (sa, sb, Nothing, Just b)
2431 -- one of the values is missing, corresponding stream is done
2432 step _ (Nothing, sb, Nothing, Just b) =
2433 return $ Yield b (Nothing, sb, Nothing, Nothing)
2435 step _ (sa, Nothing, Just a, Nothing) =
2436 return $ Yield a (sa, Nothing, Nothing, Nothing)
2438 step _ (Nothing, Nothing, Nothing, Nothing) = return Stop
2440{-# INLINE mergeBy #-}
2441mergeBy
2442 :: (Monad m)
2443 => (a -> a -> Ordering) -> Stream m a -> Stream m a -> Stream m a
2444mergeBy cmp = mergeByM (\a b -> return $ cmp a b)
2446------------------------------------------------------------------------------
2447-- Transformation comprehensions
2448------------------------------------------------------------------------------
2450{-# INLINE_NORMAL the #-}
2451the :: (Eq a, Monad m) => Stream m a -> m (Maybe a)
2452the (Stream step state) = go state
2453 where
2454 go st = do
2455 r <- step defState st
2456 case r of
2457 Yield x s -> go' x s
2458 Skip s -> go s
2459 Stop -> return Nothing
2460 go' n st = do
2461 r <- step defState st
2462 case r of
2463 Yield x s | x == n -> go' n s
2464 | otherwise -> return Nothing
2465 Skip s -> go' n s
2466 Stop -> return (Just n)
2468--------------------------------------------------------------------------------
2469-- UTF8 Encoding / Decoding
2470--------------------------------------------------------------------------------
2472-- UTF-8 primitives, Lifted from GHC.IO.Encoding.UTF8.
2474ord2 :: Char -> [Word8]
2475ord2 c = assert (n >= 0x80 && n <= 0x07ff) [x1, x2]
2476 where
2477 n = ord c
2478 x1 = fromIntegral $ (n `shiftR` 6) + 0xC0
2479 x2 = fromIntegral $ (n .&. 0x3F) + 0x80
2481ord3 :: Char -> [Word8]
2482ord3 c = assert (n >= 0x0800 && n <= 0xffff) [x1, x2, x3]
2483 where
2484 n = ord c
2485 x1 = fromIntegral $ (n `shiftR` 12) + 0xE0
2486 x2 = fromIntegral $ ((n `shiftR` 6) .&. 0x3F) + 0x80
2487 x3 = fromIntegral $ (n .&. 0x3F) + 0x80
2489ord4 :: Char -> [Word8]
2490ord4 c = assert (n >= 0x10000) [x1, x2, x3, x4]
2491 where
2492 n = ord c
2493 x1 = fromIntegral $ (n `shiftR` 18) + 0xF0
2494 x2 = fromIntegral $ ((n `shiftR` 12) .&. 0x3F) + 0x80
2495 x3 = fromIntegral $ ((n `shiftR` 6) .&. 0x3F) + 0x80
2496 x4 = fromIntegral $ (n .&. 0x3F) + 0x80
2498{-# INLINE chr2 #-}
2499chr2 :: Word8 -> Word8 -> Char
2500chr2 (W8# x1#) (W8# x2#) = C# (chr# (z1# +# z2#))
2501 where
2502 !y1# = word2Int# x1#
2503 !y2# = word2Int# x2#
2504 !z1# = uncheckedIShiftL# (y1# -# 192#) 6#
2505 !z2# = y2# -# 128#
2507{-# INLINE chr3 #-}
2508chr3 :: Word8 -> Word8 -> Word8 -> Char
2509chr3 (W8# x1#) (W8# x2#) (W8# x3#) = C# (chr# (z1# +# z2# +# z3#))
2510 where
2511 !y1# = word2Int# x1#
2512 !y2# = word2Int# x2#
2513 !y3# = word2Int# x3#
2514 !z1# = uncheckedIShiftL# (y1# -# 224#) 12#
2515 !z2# = uncheckedIShiftL# (y2# -# 128#) 6#
2516 !z3# = y3# -# 128#
2518{-# INLINE chr4 #-}
2519chr4 :: Word8 -> Word8 -> Word8 -> Word8 -> Char
2520chr4 (W8# x1#) (W8# x2#) (W8# x3#) (W8# x4#) =
2521 C# (chr# (z1# +# z2# +# z3# +# z4#))
2522 where
2523 !y1# = word2Int# x1#
2524 !y2# = word2Int# x2#
2525 !y3# = word2Int# x3#
2526 !y4# = word2Int# x4#
2527 !z1# = uncheckedIShiftL# (y1# -# 240#) 18#
2528 !z2# = uncheckedIShiftL# (y2# -# 128#) 12#
2529 !z3# = uncheckedIShiftL# (y3# -# 128#) 6#
2530 !z4# = y4# -# 128#
2532{-# INLINE between #-}
2533between ::
2534 Word8 -- ^ byte to check
2535 -> Word8 -- ^ lower bound
2536 -> Word8 -- ^ upper bound
2537 -> Bool
2538between x y z = x >= y && x <= z
2540{-# INLINE validate3 #-}
2541validate3 :: Word8 -> Word8 -> Word8 -> Bool
2542validate3 x1 x2 x3 = validate3_1 ||
2543 validate3_2 ||
2544 validate3_3 ||
2545 validate3_4
2546 where
2547 validate3_1 = (x1 == 0xE0) &&
2548 between x2 0xA0 0xBF &&
2549 between x3 0x80 0xBF
2550 validate3_2 = between x1 0xE1 0xEC &&
2551 between x2 0x80 0xBF &&
2552 between x3 0x80 0xBF
2553 validate3_3 = x1 == 0xED &&
2554 between x2 0x80 0x9F &&
2555 between x3 0x80 0xBF
2556 validate3_4 = between x1 0xEE 0xEF &&
2557 between x2 0x80 0xBF &&
2558 between x3 0x80 0xBF
2560{-# INLINE validate4 #-}
2561validate4 :: Word8 -> Word8 -> Word8 -> Word8 -> Bool
2562validate4 x1 x2 x3 x4 = validate4_1 ||
2563 validate4_2 ||
2564 validate4_3
2565 where
2566 validate4_1 = x1 == 0xF0 &&
2567 between x2 0x90 0xBF &&
2568 between x3 0x80 0xBF &&
2569 between x4 0x80 0xBF
2570 validate4_2 = between x1 0xF1 0xF3 &&
2571 between x2 0x80 0xBF &&
2572 between x3 0x80 0xBF &&
2573 between x4 0x80 0xBF
2574 validate4_3 = x1 == 0xF4 &&
2575 between x2 0x80 0x8F &&
2576 between x3 0x80 0xBF &&
2577 between x4 0x80 0xBF
2579data CodePoint = FreshPoint
2580 | Two !Word8
2581 | GoThree !Word8
2582 | Three !Word8 !Word8
2583 | GoFour !Word8
2584 | Four1 !Word8 !Word8
2585 | Four !Word8 !Word8 !Word8
2586 | YieldPoint !Char
2588{-# INLINE_NORMAL decodeUtf8 #-}
2589decodeUtf8 :: Monad m => Stream m Word8 -> Stream m Char
2590decodeUtf8 (Stream step state) = Stream step' (state, FreshPoint)
2591 where
2592 {-# INLINE_LATE step' #-}
2593 step' gst (st, FreshPoint) = do
2594 r <- step (adaptState gst) st
2595 return $
2596 case r of
2597 Yield c0 s
2598 | c0 <= 0x7f ->
2599 Skip (s, YieldPoint $ unsafeChr (fromIntegral c0))
2600 | c0 >= 0xc0 && c0 <= 0xc1 -> error "Invalid Overlong forms"
2601 | c0 >= 0xc2 && c0 <= 0xdf -> Skip (s, Two c0)
2602 | c0 >= 0xe0 && c0 <= 0xef -> Skip (s, GoThree c0)
2603 | c0 >= 0xf0 -> Skip (s, GoFour c0)
2604 | otherwise -> error "Invalid Sequence"
2605 Skip s -> Skip (s, FreshPoint)
2606 Stop -> Stop
2608 -- Decoding of a UTF8 character that spans two bytes with one byte (c0)
2609 -- known.
2610 step' gst (st, Two c0) = do
2611 r <- step (adaptState gst) st
2612 return $
2613 case r of
2614 Yield c1 s
2615 | c1 < 0x80 || c1 >= 0xc0 -> error "Invalid Sequence"
2616 | otherwise -> Skip (s, YieldPoint $ chr2 c0 c1)
2617 Skip s -> Skip (s, Two c0)
2618 Stop -> error "Input Underflow"
2620 -- Decoding of a UTF8 character that spans three bytes with first byte (c0)
2621 -- known.
2622 step' gst (st, GoThree c0) = do
2623 r <- step (adaptState gst) st
2624 return $
2625 case r of
2626 Yield c1 s
2627 | not (validate3 c0 c1 0x80) -> error "Invalid Sequence"
2628 -- See #3341@ghc
2629 | otherwise -> Skip (s, Three c0 c1)
2630 Skip s -> Skip (s, GoThree c0)
2631 Stop -> error "Input Underflow"
2633 -- Two bytes known of a 3 byte UTF8 encoded unicode character.
2634 step' gst (st, Three c0 c1) = do
2635 r <- step (adaptState gst) st
2636 return $
2637 case r of
2638 Yield c2 s
2639 | not (validate3 c0 c1 c2) -> error "Invalid Sequence"
2640 | otherwise -> Skip (s, YieldPoint $ chr3 c0 c1 c2)
2641 Skip s -> Skip (s, Three c0 c1)
2642 Stop -> error "Input Underflow"
2644 -- One byte known of a four byte UTF8 unicode character.
2645 step' gst (st, GoFour c0) = do
2646 r <- step (adaptState gst) st
2647 return $
2648 case r of
2649 Yield c1 s
2650 | not (validate4 c0 c1 0x80 0x80) -> error "Invalid Sequence"
2651 | otherwise -> Skip (s, Four1 c0 c1)
2652 Skip s -> Skip (s, GoFour c0)
2653 Stop -> error "Input Underflow"
2655 -- Two bytes known of a four byte UTF8 unicode character.
2656 step' gst (st, Four1 c0 c1) = do
2657 r <- step (adaptState gst) st
2658 return $
2659 case r of
2660 Yield c2 s
2661 | not (validate4 c0 c1 c2 0x80) -> error "Invalid Sequence"
2662 | otherwise -> Skip (s, Four c0 c1 c2)
2663 Skip s -> Skip (s, Four1 c0 c1)
2664 Stop -> error "Input Underflow"
2666 -- Three bytes known of a three byte UTF8 unicode character.
2667 step' gst (st, Four c0 c1 c2) = do
2668 r <- step (adaptState gst) st
2669 return $
2670 case r of
2671 Yield c3 s
2672 | not (validate4 c0 c1 c2 c3) -> error "Invalid Sequence"
2673 | otherwise -> Skip (s, YieldPoint $ chr4 c0 c1 c2 c3)
2674 Skip s -> Skip (s, Four c0 c1 c2)
2675 Stop -> error "Input Underflow"
2677 -- This way so that we only need one Yield.
2678 step' _ (s, YieldPoint c) = return $ Yield c (s, FreshPoint)
2680data YieldMany x = None
2681 | YieldMany !x
2683{-# INLINE_NORMAL encodeUtf8 #-}
2684encodeUtf8 :: Monad m => Stream m Char -> Stream m Word8
2685encodeUtf8 (Stream step state) = Stream step' (state, None)
2686 where
2687 {-# INLINE_LATE step' #-}
2688 step' gst (st, None) = do
2689 r <- step (adaptState gst) st
2690 return $
2691 case r of
2692 Yield c s ->
2693 case ord c of
2694 x | x <= 0x7F -> Skip (s, YieldMany [fromIntegral x])
2695 | x <= 0x7FF -> Skip (s, YieldMany (ord2 c))
2696 | x <= 0xFFFF ->
2697 if isSurrogate c
2698 then error "Encountered a surrogate"
2699 else Skip (s, YieldMany (ord3 c))
2700 | otherwise -> Skip (s, YieldMany (ord4 c))
2701 Skip s -> Skip (s, None)
2702 Stop -> Stop
2703 step' _ (s, YieldMany []) = return $ Skip (s, None)
2704 step' _ (s, YieldMany (x:xs)) = return $ Yield x (s, YieldMany xs)