on c06d2e2 help sign in
text/Data/Text/Fusion.hs
886 lines 37 kB history raw
1{-# LANGUAGE ExistentialQuantification, BangPatterns, MagicHash #-}
2--TODO: eliminate the need to supress these warnings:
3{-# OPTIONS_GHC -fno-warn-incomplete-patterns -fno-warn-unused-matches #-}
5module Data.Text.Fusion where
7import Prelude hiding (map, tail, head, foldr, filter,concat)
9import Data.Char (ord)
10import Control.Exception(assert)
11import Control.Monad(liftM2)
12import Control.Monad.ST(runST,ST)
13import Data.Array.Base
14import Data.Bits
15import qualified Data.ByteString as B
16import qualified Data.ByteString.Unsafe as B
17import Data.ByteString.Internal(ByteString(..),mallocByteString,memcpy)
18import qualified Data.List as L
19import Data.Word(Word8, Word16, Word32)
20import Foreign.ForeignPtr(withForeignPtr,ForeignPtr)
21import Foreign.Storable(pokeByteOff)
22import GHC.Exts (Int(..), (+#))
23import System.IO.Unsafe(unsafePerformIO)
25import Data.Text.Internal(Text(..),empty)
26import Data.Text.UnsafeChar(unsafeChr,unsafeChr8,unsafeChr32)
27import qualified Data.Text.Utf8 as U8
28import qualified Data.Text.Utf16 as U16
29import qualified Data.Text.Utf32 as U32
31default(Int)
33infixl 2 :!:
34data PairS a b = !a :!: !b
36data Switch = S1 | S2
38data EitherS a b = LeftS !a | RightS !b
40data Stream a = forall s. Stream (s -> Step s a) !s {-# UNPACK #-}!Int
42data Step s a = Done
43 | Skip !s
44 | Yield !a !s
46data Encoding = ASCII | Utf8 | Utf16BE | Utf16LE | Utf32BE | Utf32LE
47 deriving (Read, Show, Eq)
49-- | /O(n)/ Convert a Text into a Stream Char.
50stream :: Text -> Stream Char
51stream (Text arr off len) = Stream next off len
52 where
53 end = off+len
54 {-# INLINE next #-}
55 next !i
56 | i >= end = Done
57 | n >= 0xD800 && n <= 0xDBFF = Yield (U16.chr2 n n2) (i + 2)
58 | otherwise = Yield (unsafeChr n) (i + 1)
59 where
60 n = unsafeAt arr i
61 n2 = unsafeAt arr (i + 1)
62{-# INLINE [0] stream #-}
64-- | /O(n)/ Convert a Stream Char into a Text.
65unstream :: Stream Char -> Text
66unstream (Stream next0 s0 len) = x `seq` Text (fst x) 0 (snd x)
67 where
68 x :: ((UArray Int Word16),Int)
69 x = runST ((unsafeNewArray_ (0,len+1) :: ST s (STUArray s Int Word16))
70 >>= (\arr -> loop arr 0 (len+1) s0))
71 loop arr !i !max !s
72 | i + 1 > max = do arr' <- unsafeNewArray_ (0,max*2)
73 case next0 s of
74 Done -> liftM2 (,) (unsafeFreezeSTUArray arr) (return i)
75 _ -> copy arr arr' >> loop arr' i (max*2) s
76 | otherwise = case next0 s of
77 Done -> liftM2 (,) (unsafeFreezeSTUArray arr) (return i)
78 Skip s' -> loop arr i max s'
79 Yield x s'
80 | n < 0x10000 -> do
81 unsafeWrite arr i (fromIntegral n :: Word16)
82 loop arr (i+1) max s'
83 | otherwise -> do
84 unsafeWrite arr i l
85 unsafeWrite arr (i + 1) r
86 loop arr (i+2) max s'
87 where
88 n :: Int
89 n = ord x
90 m :: Int
91 m = n - 0x10000
92 l :: Word16
93 l = fromIntegral $ (shiftR m 10) + (0xD800 :: Int)
94 r :: Word16
95 r = fromIntegral $ (m .&. (0x3FF :: Int)) + (0xDC00 :: Int)
96{-# INLINE [0] unstream #-}
99copy src dest = (do
100 (_,max) <- getBounds src
101 copy_loop 0 max)
102 where
103 copy_loop i max
104 | i > max = return ()
105 | otherwise = do v <- unsafeRead src i
106 unsafeWrite dest i v
107 copy_loop (i+1) max
109-- | /O(n)/ Determines if two streams are equal.
110eq :: Ord a => Stream a -> Stream a -> Bool
111eq (Stream next1 s1 _) (Stream next2 s2 _) = compare (next1 s1) (next2 s2)
112 where
113 compare Done Done = True
114 compare Done _ = False
115 compare _ Done = False
116 compare (Skip s1') (Skip s2') = compare (next1 s1') (next2 s2')
117 compare (Skip s1') x2 = compare (next1 s1') x2
118 compare x1 (Skip s2') = compare x1 (next2 s2')
119 compare (Yield x1 s1') (Yield x2 s2') = x1 == x2 &&
120 compare (next1 s1') (next2 s2')
121{-# SPECIALISE eq :: Stream Char -> Stream Char -> Bool #-}
124-- | /O(n) Convert a ByteString into a Stream Char, using the specified encoding standard.
125stream_bs :: Encoding -> ByteString -> Stream Char
126stream_bs ASCII bs = Stream next 0 l
127 where
128 l = B.length bs
129 {-# INLINE next #-}
130 next i
131 | i >= l = Done
132 | otherwise = Yield (unsafeChr8 x1) (i+1)
133 where
134 x1 = B.unsafeIndex bs i
135stream_bs Utf8 bs = Stream next 0 l
136 where
137 l = B.length bs
138 {-# INLINE next #-}
139 next i
140 | i >= l = Done
141 | U8.validate1 x1 = Yield (unsafeChr8 x1) (i+1)
142 | i+1 < l && U8.validate2 x1 x2 = Yield (U8.chr2 x1 x2) (i+2)
143 | i+2 < l && U8.validate3 x1 x2 x3 = Yield (U8.chr3 x1 x2 x3) (i+3)
144 | i+3 < l && U8.validate4 x1 x2 x3 x4 = Yield (U8.chr4 x1 x2 x3 x4) (i+4)
145 | otherwise = error "bsStream: bad UTF-8 stream"
146 where
147 x1 = index i
148 x2 = index (i + 1)
149 x3 = index (i + 2)
150 x4 = index (i + 3)
151 index = B.unsafeIndex bs
152stream_bs Utf16LE bs = Stream next 0 l
153 where
154 l = B.length bs
155 {-# INLINE next #-}
156 next i
157 | i >= l = Done
158 | i+1 < l && U16.validate1 x1 = Yield (unsafeChr x1) (i+2)
159 | i+3 < l && U16.validate2 x1 x2 = Yield (U16.chr2 x1 x2) (i+4)
160 | otherwise = error $ "bsStream: bad UTF-16LE stream"
161 where
162 x1 = (shiftL (index (i + 1)) 8) + (index i)
163 x2 = (shiftL (index (i + 3)) 8) + (index (i + 2))
164 index = fromIntegral . B.unsafeIndex bs :: Int -> Word16
165stream_bs Utf16BE bs = Stream next 0 l
166 where
167 l = B.length bs
168 {-# INLINE next #-}
169 next i
170 | i >= l = Done
171 | i+1 < l && U16.validate1 x1 = Yield (unsafeChr x1) (i+2)
172 | i+3 < l && U16.validate2 x1 x2 = Yield (U16.chr2 x1 x2) (i+4)
173 | otherwise = error $ "bsStream: bad UTF16-BE stream "
174 where
175 x1 = (shiftL (index i) 8) + (index (i + 1))
176 x2 = (shiftL (index (i + 2)) 8) + (index (i + 3))
177 index = fromIntegral . B.unsafeIndex bs :: Int -> Word16
178stream_bs Utf32BE bs = Stream next 0 l
179 where
180 l = B.length bs
181 {-# INLINE next #-}
182 next i
183 | i >= l = Done
184 | i+3 < l && U32.validate x = Yield (unsafeChr32 x) (i+4)
185 | otherwise = error "bsStream: bad UTF-32BE stream"
186 where
187 x = shiftL x1 24 + shiftL x2 16 + shiftL x3 8 + x4
188 x1 = index i
189 x2 = index (i+1)
190 x3 = index (i+2)
191 x4 = index (i+3)
192 index = fromIntegral . B.unsafeIndex bs :: Int -> Word32
193stream_bs Utf32LE bs = Stream next 0 l
194 where
195 l = B.length bs
196 {-# INLINE next #-}
197 next i
198 | i >= l = Done
199 | i+3 < l && U32.validate x = Yield (unsafeChr32 x) (i+4)
200 | otherwise = error "bsStream: bad UTF-32LE stream"
201 where
202 x = shiftL x4 24 + shiftL x3 16 + shiftL x2 8 + x1
203 x1 = index i
204 x2 = index $ i+1
205 x3 = index $ i+2
206 x4 = index $ i+3
207 index = fromIntegral . B.unsafeIndex bs :: Int -> Word32
208{-# INLINE [0] stream_bs #-}
210-- | /O(n)/ Convert a Stream Char into a Stream Word8 using the specified encoding standard.
211restream :: Encoding -> Stream Char -> Stream Word8
212restream ASCII (Stream next0 s0 len) = Stream next s0 (len*2)
213 where
214 next !s = case next0 s of
215 Done -> Done
216 Skip s' -> Skip s'
217 Yield x xs -> Yield x' xs
218 where x' = fromIntegral (ord x) :: Word8
219restream Utf8 (Stream next0 s0 len) =
220 Stream next ((Just s0) :!: Nothing :!: Nothing :!: Nothing) (len*2)
221 where
222 {-# INLINE next #-}
223 next ((Just s) :!: Nothing :!: Nothing :!: Nothing) = case next0 s of
224 Done -> Done
225 Skip s' -> Skip ((Just s') :!: Nothing :!: Nothing :!: Nothing)
226 Yield x xs
227 | n <= 0x7F -> Yield c ((Just xs) :!: Nothing :!: Nothing :!: Nothing)
228 | n <= 0x07FF -> Yield (fst c2) ((Just xs) :!: (Just $ snd c2) :!: Nothing :!: Nothing)
229 | n <= 0xFFFF -> Yield (fst3 c3) ((Just xs) :!: (Just $ snd3 c3) :!: (Just $ trd3 c3) :!: Nothing)
230 | otherwise -> Yield (fst4 c4) ((Just xs) :!: (Just $ snd4 c4) :!: (Just $ trd4 c4) :!: (Just $ fth4 c4))
231 where
232 n = ord x
233 c = fromIntegral n
234 c2 = U8.ord2 x
235 c3 = U8.ord3 x
236 c4 = U8.ord4 x
237 next ((Just s) :!: (Just x2) :!: Nothing :!: Nothing) = Yield x2 ((Just s) :!: Nothing :!: Nothing :!: Nothing)
238 next ((Just s) :!: (Just x2) :!: x3 :!: Nothing) = Yield x2 ((Just s) :!: x3 :!: Nothing :!: Nothing)
239 next ((Just s) :!: (Just x2) :!: x3 :!: x4) = Yield x2 ((Just s) :!: x3 :!: x4 :!: Nothing)
240restream Utf16BE (Stream next0 s0 len) =
241 Stream next (Just s0 :!: Nothing :!: Nothing :!: Nothing) (len*2)
242 where
243 {-# INLINE next #-}
244 next (Just s :!: Nothing :!: Nothing :!: Nothing) = case next0 s of
245 Done -> Done
246 Skip s' -> Skip (Just s' :!: Nothing :!: Nothing :!: Nothing)
247 Yield x xs
248 | n < 0x10000 -> Yield (fromIntegral $ shiftR n 8) (Just xs :!: Just (fromIntegral n) :!: Nothing :!: Nothing)
249 | otherwise -> Yield c1 (Just xs :!: Just c2 :!: Just c3 :!: Just c4)
250 where
251 n = ord x
252 n1 = n - 0x10000
253 c1 = fromIntegral (shiftR n1 18 + 0xD8)
254 c2 = fromIntegral (shiftR n1 10)
255 n2 = n1 .&. 0x3FF
256 c3 = fromIntegral (shiftR n2 8 + 0xDC)
257 c4 = fromIntegral n2
258 next ((Just s) :!: (Just x2) :!: Nothing :!: Nothing) = Yield x2 ((Just s) :!: Nothing :!: Nothing :!: Nothing)
259 next ((Just s) :!: (Just x2) :!: x3 :!: Nothing) = Yield x2 ((Just s) :!: x3 :!: Nothing :!: Nothing)
260 next ((Just s) :!: (Just x2) :!: x3 :!: x4) = Yield x2 ((Just s) :!: x3 :!: x4 :!: Nothing)
261restream Utf16LE (Stream next0 s0 len) =
262 Stream next (Just s0 :!: Nothing :!: Nothing :!: Nothing) (len*2)
263 where
264 {-# INLINE next #-}
265 next (Just s :!: Nothing :!: Nothing :!: Nothing) = case next0 s of
266 Done -> Done
267 Skip s' -> Skip (Just s' :!: Nothing :!: Nothing :!: Nothing)
268 Yield x xs
269 | n < 0x10000 -> Yield (fromIntegral n) (Just xs :!: Just (fromIntegral $ shiftR n 8) :!: Nothing :!: Nothing)
270 | otherwise -> Yield c1 (Just xs :!: Just c2 :!: Just c3 :!: Just c4)
271 where
272 n = ord x
273 n1 = n - 0x10000
274 c2 = fromIntegral (shiftR n1 18 + 0xD8)
275 c1 = fromIntegral (shiftR n1 10)
276 n2 = n1 .&. 0x3FF
277 c4 = fromIntegral (shiftR n2 8 + 0xDC)
278 c3 = fromIntegral n2
279 next ((Just s) :!: (Just x2) :!: Nothing :!: Nothing) = Yield x2 ((Just s) :!: Nothing :!: Nothing :!: Nothing)
280 next ((Just s) :!: (Just x2) :!: x3 :!: Nothing) = Yield x2 ((Just s) :!: x3 :!: Nothing :!: Nothing)
281 next ((Just s) :!: (Just x2) :!: x3 :!: x4) = Yield x2 ((Just s) :!: x3 :!: x4 :!: Nothing)
282restream Utf32BE (Stream next0 s0 len) =
283 Stream next (Just s0 :!: Nothing :!: Nothing :!: Nothing) (len*2)
284 where
285 {-# INLINE next #-}
286 next (Just s :!: Nothing :!: Nothing :!: Nothing) = case next0 s of
287 Done -> Done
288 Skip s' -> Skip (Just s' :!: Nothing :!: Nothing :!: Nothing)
289 Yield x xs -> Yield c1 (Just xs :!: Just c2 :!: Just c3 :!: Just c4)
290 where
291 n = ord x
292 c1 = fromIntegral $ shiftR n 24
293 c2 = fromIntegral $ shiftR n 16
294 c3 = fromIntegral $ shiftR n 8
295 c4 = fromIntegral n
296 next ((Just s) :!: (Just x2) :!: Nothing :!: Nothing) = Yield x2 ((Just s) :!: Nothing :!: Nothing :!: Nothing)
297 next ((Just s) :!: (Just x2) :!: x3 :!: Nothing) = Yield x2 ((Just s) :!: x3 :!: Nothing :!: Nothing)
298 next ((Just s) :!: (Just x2) :!: x3 :!: x4) = Yield x2 ((Just s) :!: x3 :!: x4 :!: Nothing)
299restream Utf32LE (Stream next0 s0 len) =
300 Stream next (Just s0 :!: Nothing :!: Nothing :!: Nothing) (len*2)
301 where
302 {-# INLINE next #-}
303 next (Just s :!: Nothing :!: Nothing :!: Nothing) = case next0 s of
304 Done -> Done
305 Skip s' -> Skip (Just s' :!: Nothing :!: Nothing :!: Nothing)
306 Yield x xs -> Yield c1 (Just xs :!: Just c2 :!: Just c3 :!: Just c4)
307 where
308 n = ord x
309 c4 = fromIntegral $ shiftR n 24
310 c3 = fromIntegral $ shiftR n 16
311 c2 = fromIntegral $ shiftR n 8
312 c1 = fromIntegral n
313 next ((Just s) :!: (Just x2) :!: Nothing :!: Nothing) = Yield x2 ((Just s) :!: Nothing :!: Nothing :!: Nothing)
314 next ((Just s) :!: (Just x2) :!: x3 :!: Nothing) = Yield x2 ((Just s) :!: x3 :!: Nothing :!: Nothing)
315 next ((Just s) :!: (Just x2) :!: x3 :!: x4) = Yield x2 ((Just s) :!: x3 :!: x4 :!: Nothing)
316{-# INLINE restream #-}
319fst3 (x1,_,_) = x1
320snd3 (_,x2,_) = x2
321trd3 (_,_,x3) = x3
322fst4 (x1,_,_,_) = x1
323snd4 (_,x2,_,_) = x2
324trd4 (_,_,x3,_) = x3
325fth4 (_,_,_,x4) = x4
327-- | /O(n)/ Convert a Stream Word8 to a ByteString
328unstream_bs :: Stream Word8 -> ByteString
329unstream_bs (Stream next s0 len) = unsafePerformIO $ do
330 fp0 <- mallocByteString len
331 loop fp0 len 0 s0
332 where
333 loop !fp !n !off !s = case next s of
334 Done -> trimUp fp n off
335 Skip s' -> loop fp n off s'
336 Yield x s'
337 | n == off -> realloc fp n off s' x
338 | otherwise -> do
339 withForeignPtr fp $ \p -> pokeByteOff p off x
340 loop fp n (off+1) s'
341 {-# NOINLINE realloc #-}
342 realloc fp n off s x = do
343 let n' = n+n
344 fp' <- copy0 fp n n'
345 withForeignPtr fp' $ \p -> pokeByteOff p off x
346 loop fp' n' (off+1) s
347 {-# NOINLINE trimUp #-}
348 trimUp fp _ off = return $! PS fp 0 off
349 copy0 :: ForeignPtr Word8 -> Int -> Int -> IO (ForeignPtr Word8)
350 copy0 !src !srcLen !destLen = assert (srcLen <= destLen) $ do
351 dest <- mallocByteString destLen
352 withForeignPtr src $ \src' ->
353 withForeignPtr dest $ \dest' ->
354 memcpy dest' src' (fromIntegral destLen)
355 return dest
356{-# RULES "STREAM stream/unstream fusion" forall s. stream (unstream s) = s #-}
358-- ----------------------------------------------------------------------------
359-- * Basic stream functions
361-- | /O(n)/ Adds a character to the front of a Stream Char.
362cons :: Char -> Stream Char -> Stream Char
363cons w (Stream next0 s0 len) = Stream next (S2 :!: s0) (len+2)
364 where
365 {-# INLINE next #-}
366 next (S2 :!: s) = Yield w (S1 :!: s)
367 next (S1 :!: s) = case next0 s of
368 Done -> Done
369 Skip s' -> Skip (S1 :!: s')
370 Yield x s' -> Yield x (S1 :!: s')
371{-# INLINE [0] cons #-}
373-- | /O(n)/ Adds a character to the end of a stream.
374snoc :: Stream Char -> Char -> Stream Char
375snoc (Stream next0 xs0 len) w = Stream next (Just xs0) (len+2)
376 where
377 {-# INLINE next #-}
378 next (Just xs) = case next0 xs of
379 Done -> Yield w Nothing
380 Skip xs' -> Skip (Just xs')
381 Yield x xs' -> Yield x (Just xs')
382 next Nothing = Done
383{-# INLINE [0] snoc #-}
385-- | /O(n)/ Appends one Stream to the other.
386append :: Stream Char -> Stream Char -> Stream Char
387append (Stream next0 s01 len1) (Stream next1 s02 len2) =
388 Stream next (Left s01) (len1 + len2)
389 where
390 {-# INLINE next #-}
391 next (Left s1) = case next0 s1 of
392 Done -> Skip (Right s02)
393 Skip s1' -> Skip (Left s1')
394 Yield x s1' -> Yield x (Left s1')
395 next (Right s2) = case next1 s2 of
396 Done -> Done
397 Skip s2' -> Skip (Right s2')
398 Yield x s2' -> Yield x (Right s2')
399{-# INLINE [0] append #-}
401-- | /O(1)/ Returns the first character of a Text, which must be non-empty.
402-- Subject to array fusion.
403head :: Stream Char -> Char
404head (Stream next s0 len) = loop_head s0
405 where
406 loop_head !s = case next s of
407 Yield x _ -> x
408 Skip s' -> loop_head s'
409 Done -> error "head: Empty list"
410{-# INLINE [0] head #-}
412-- | /O(n)/ Returns the last character of a Stream Char, which must be non-empty.
413last :: Stream Char -> Char
414last (Stream next s0 len) = loop0_last s0
415 where
416 loop0_last !s = case next s of
417 Done -> error "last: Empty list"
418 Skip s' -> seq s' $ loop0_last s'
419 Yield x s' -> seq s' $ loop_last x s'
420 loop_last !x !s = case next s of
421 Done -> x
422 Skip s' -> seq s' $ loop_last x s'
423 Yield x' s' -> seq s' $ loop_last x' s'
424{-# INLINE[0] last #-}
426-- | /O(1)/ Returns all characters after the head of a Stream Char, which must
427-- be non-empty.
428tail :: Stream Char -> Stream Char
429tail (Stream next0 s0 len) = Stream next (False :!: s0) (len-1)
430 where
431 {-# INLINE next #-}
432 next (False :!: s) = case next0 s of
433 Done -> error "tail"
434 Skip s' -> Skip (False :!: s')
435 Yield _ s' -> Skip (True :!: s')
436 next (True :!: s) = case next0 s of
437 Done -> Done
438 Skip s' -> Skip (True :!: s')
439 Yield x s' -> Yield x (True :!: s')
440{-# INLINE [0] tail #-}
443-- | /O(1)/ Returns all but the last character of a Stream Char, which
444-- must be non-empty.
445init :: Stream Char -> Stream Char
446init (Stream next0 s0 len) = Stream next (Nothing :!: s0) (len-1)
447 where
448 {-# INLINE next #-}
449 next (Nothing :!: s) = case next0 s of
450 Done -> errorEmptyList "init"
451 Skip s' -> seq s' $ Skip (Nothing :!: s')
452 Yield x s' -> seq s' $ Skip (Just x :!: s')
453 next (Just x :!: s) = case next0 s of
454 Done -> Done
455 Skip s' -> seq s' $ Skip (Just x :!: s')
456 Yield x' s' -> seq s' $ Yield x (Just x' :!: s')
457{-# INLINE [0] init #-}
459-- | /O(1)/ Tests whether a Stream Char is empty or not.
460null :: Stream Char -> Bool
461null (Stream next s0 len) = loop_null s0
462 where
463 loop_null !s = case next s of
464 Done -> True
465 Yield _ _ -> False
466 Skip s' -> loop_null s'
467{-# INLINE[0] null #-}
469-- | /O(n)/ Returns the number of characters in a text.
470length :: Stream Char -> Int
471length (Stream next s0 len) = loop_length 0# s0
472 where
474 loop_length z# s = case next s of
475 Done -> (I# z#)
476 Skip s' -> loop_length z# s'
477 Yield _ s' -> loop_length (z# +# 1#) s'
478{-# INLINE[0] length #-}
480-- ----------------------------------------------------------------------------
481-- * Stream transformations
483-- | /O(n)/ 'map' @f @xs is the Stream Char obtained by applying @f@ to each element of
484-- @xs@.
485map :: (Char -> Char) -> Stream Char -> Stream Char
486map f (Stream next0 s0 len) = Stream next s0 len
487 where
488 {-# INLINE next #-}
489 next !s = case next0 s of
490 Done -> Done
491 Skip s' -> Skip s'
492 Yield x s' -> Yield (f x) s'
493{-# INLINE [0] map #-}
495{-#
496 RULES "STREAM map/map fusion" forall f g s.
497 map f (map g s) = map (\x -> f (g x)) s
498 #-}
500-- | /O(n)/ The 'intersperse' function takes a character and places it between each of
501-- the characters of a Stream.
502intersperse :: Char -> Stream Char -> Stream Char
503intersperse c (Stream next0 s0 len) = Stream next (s0 :!: Nothing :!: S1) len
504 where
505 {-# INLINE next #-}
506 next (s :!: Nothing :!: S1) = case next0 s of
507 Done -> Done
508 Skip s' -> Skip (s' :!: Nothing :!: S1)
509 Yield x s' -> Skip (s' :!: Just x :!: S1)
510 next (s :!: Just x :!: S1) = Yield x (s :!: Nothing :!: S2)
511 next (s :!: Nothing :!: S2) = case next0 s of
512 Done -> Done
513 Skip s' -> Skip (s' :!: Nothing :!: S2)
514 Yield x s' -> Yield c (s' :!: Just x :!: S1)
516-- ----------------------------------------------------------------------------
517-- * Reducing Streams (folds)
519-- | foldl, applied to a binary operator, a starting value (typically the
520-- left-identity of the operator), and a Stream, reduces the Stream using the
521-- binary operator, from left to right.
522foldl :: (b -> Char -> b) -> b -> Stream Char -> b
523foldl f z0 (Stream next s0 len) = loop_foldl z0 s0
524 where
525 loop_foldl z !s = case next s of
526 Done -> z
527 Skip s' -> loop_foldl z s'
528 Yield x s' -> loop_foldl (f z x) s'
529{-# INLINE [0] foldl #-}
531-- | A strict version of foldl.
532foldl' :: (b -> Char -> b) -> b -> Stream Char -> b
533foldl' f z0 (Stream next s0 len) = loop_foldl' z0 s0
534 where
535 loop_foldl' !z !s = case next s of
536 Done -> z
537 Skip s' -> loop_foldl' z s'
538 Yield x s' -> loop_foldl' (f z x) s'
539{-# INLINE [0] foldl' #-}
541-- | foldl1 is a variant of foldl that has no starting value argument,
542-- and thus must be applied to non-empty Streams.
543foldl1 :: (Char -> Char -> Char) -> Stream Char -> Char
544foldl1 f (Stream next s0 len) = loop0_foldl1 s0
545 where
546 loop0_foldl1 !s = case next s of
547 Skip s' -> loop0_foldl1 s'
548 Yield x s' -> loop_foldl1 x s'
549 Done -> errorEmptyList "foldl1"
550 loop_foldl1 z !s = case next s of
551 Done -> z
552 Skip s' -> loop_foldl1 z s'
553 Yield x s' -> loop_foldl1 (f z x) s'
554{-# INLINE [0] foldl1 #-}
556-- | A strict version of foldl1.
557foldl1' :: (Char -> Char -> Char) -> Stream Char -> Char
558foldl1' f (Stream next s0 len) = loop0_foldl1' s0
559 where
560 loop0_foldl1' !s = case next s of
561 Skip s' -> loop0_foldl1' s'
562 Yield x s' -> loop_foldl1' x s'
563 Done -> errorEmptyList "foldl1"
564 loop_foldl1' !z !s = case next s of
565 Done -> z
566 Skip s' -> loop_foldl1' z s'
567 Yield x s' -> loop_foldl1' (f z x) s'
568{-# INLINE [0] foldl1' #-}
570-- | 'foldr', applied to a binary operator, a starting value (typically the
571-- right-identity of the operator), and a stream, reduces the stream using the
572-- binary operator, from right to left.
573foldr :: (Char -> b -> b) -> b -> Stream Char -> b
574foldr f z (Stream next s0 len) = loop_foldr s0
575 where
576 loop_foldr !s = case next s of
577 Done -> z
578 Skip s' -> loop_foldr s'
579 Yield x s' -> f x (loop_foldr s')
580{-# INLINE [0] foldr #-}
582-- | foldr1 is a variant of 'foldr' that has no starting value argument,
583-- and thust must be applied to non-empty streams.
584-- Subject to array fusion.
585foldr1 :: (Char -> Char -> Char) -> Stream Char -> Char
586foldr1 f (Stream next s0 len) = loop0_foldr1 s0
587 where
588 loop0_foldr1 !s = case next s of
589 Done -> error "foldr1"
590 Skip s' -> loop0_foldr1 s'
591 Yield x s' -> loop_foldr1 x s'
593 loop_foldr1 x !s = case next s of
594 Done -> x
595 Skip s' -> loop_foldr1 x s'
596 Yield x' s' -> f x (loop_foldr1 x' s')
597{-# INLINE [0] foldr1 #-}
599-- ----------------------------------------------------------------------------
600-- ** Special folds
602-- | /O(n)/ Concatenate a list of streams. Subject to array fusion.
603concat :: [Stream Char] -> Stream Char
604concat = L.foldr append (Stream next Done 0)
605 where
606 next Done = Done
608-- | Map a function over a stream that results in a steram and concatenate the
609-- results.
610concatMap :: (Char -> Stream Char) -> Stream Char -> Stream Char
611concatMap f = foldr (append . f) (stream empty)
613-- | /O(n)/ any @p @xs determines if any character in the stream
614-- @xs@ satisifes the predicate @p@.
615any :: (Char -> Bool) -> Stream Char -> Bool
616any p (Stream next0 s0 len) = loop_any s0
617 where
618 loop_any !s = case next0 s of
619 Done -> False
620 Skip s' -> seq s' $ loop_any s'
621 Yield x s' | p x -> True
622 | otherwise -> seq s' $ loop_any s'
624-- | /O(n)/ all @p @xs determines if all characters in the 'Text'
625-- @xs@ satisify the predicate @p@.
626all :: (Char -> Bool) -> Stream Char -> Bool
627all p (Stream next0 s0 len) = loop_all s0
628 where
629 loop_all !s = case next0 s of
630 Done -> True
631 Skip s' -> seq s' $ loop_all s'
632 Yield x s' | p x -> seq s' $ loop_all s'
633 | otherwise -> False
635-- | /O(n)/ maximum returns the maximum value from a stream, which must be
636-- non-empty.
637maximum :: Stream Char -> Char
638maximum (Stream next0 s0 len) = loop0_maximum s0
639 where
640 loop0_maximum !s = case next0 s of
641 Done -> errorEmptyList "maximum"
642 Skip s' -> seq s' $ loop0_maximum s'
643 Yield x s' -> seq s' $ loop_maximum x s'
644 loop_maximum !z !s = case next0 s of
645 Done -> z
646 Skip s' -> seq s' $ loop_maximum z s'
647 Yield x s'
648 | x > z -> seq s' $ loop_maximum x s'
649 | otherwise -> seq s' $ loop_maximum z s'
651-- | /O(n)/ minimum returns the minimum value from a 'Text', which must be
652-- non-empty.
653minimum :: Stream Char -> Char
654minimum (Stream next0 s0 len) = loop0_minimum s0
655 where
656 loop0_minimum !s = case next0 s of
657 Done -> errorEmptyList "minimum"
658 Skip s' -> seq s' $ loop0_minimum s'
659 Yield x s' -> seq s' $ loop_minimum x s'
660 loop_minimum !z !s = case next0 s of
661 Done -> z
662 Skip s' -> seq s' $ loop_minimum z s'
663 Yield x s'
664 | x < z -> seq s' $ loop_minimum x s'
665 | otherwise -> seq s' $ loop_minimum z s'
670-- -----------------------------------------------------------------------------
671-- * Building streams
673-- -----------------------------------------------------------------------------
674-- ** Generating and unfolding streams
676-- | /O(n)/, where @n@ is the length of the result. The unfoldr function
677-- is analogous to the List 'unfoldr'. unfoldr builds a stream
678-- from a seed value. The function takes the element and returns
679-- Nothing if it is done producing the stream or returns Just
680-- (a,b), in which case, a is the next Char in the string, and b is
681-- the seed value for further production.
682unfoldr :: (a -> Maybe (Char,a)) -> a -> Stream Char
683unfoldr f s0 = Stream next s0 1
684 where
685 {-# INLINE next #-}
686 next !s = case f s of
687 Nothing -> Done
688 Just (w, s') -> Yield w s'
689{-# INLINE [0] unfoldr #-}
691-- | O(n) Like unfoldr, unfoldrN builds a stream from a seed
692-- value. However, the length of the result should be limited by the
693-- first argument to unfoldrN. This function is more efficient than
694-- unfoldr when the maximum length of the result and correct,
695-- otherwise its complexity performance is similar to 'unfoldr'
696unfoldrN :: Int -> (a -> Maybe (Char,a)) -> a -> Stream Char
697unfoldrN n f s0 = Stream next (0 :!: s0) (n*2)
698 where
699 {-# INLINE next #-}
700 next (z :!: s) = case f s of
701 Nothing -> Done
702 Just (w, s') | z >= n -> Done
703 | otherwise -> Yield w ((z + 1) :!: s')
704-------------------------------------------------------------------------------
705-- * Substreams
707-- | /O(n)/ take n, applied to a stream, returns the prefix of the
708-- stream of length @n@, or the stream itself if @n@ is greater than the
709-- length of the stream.
710take :: Int -> Stream Char -> Stream Char
711take n0 (Stream next0 s0 len) = Stream next (n0 :!: s0) len
712 where
713 {-# INLINE next #-}
714 next (n :!: s) | n <= 0 = Done
715 | otherwise = case next0 s of
716 Done -> Done
717 Skip s' -> Skip (n :!: s')
718 Yield x s' -> Yield x ((n-1) :!: s')
719{-# INLINE [0] take #-}
721-- | /O(n)/ drop n, applied to a stream, returns the suffix of the
722-- stream of length @n@, or the empty stream if @n@ is greater than the
723-- length of the stream.
724drop :: Int -> Stream Char -> Stream Char
725drop n0 (Stream next0 s0 len) = Stream next (Just ((max 0 n0)) :!: s0) (len - n0)
726 where
727 {-# INLINE next #-}
728 next (Just !n :!: s)
729 | n == 0 = Skip (Nothing :!: s)
730 | otherwise = case next0 s of
731 Done -> Done
732 Skip s' -> Skip (Just n :!: s')
733 Yield _ s' -> Skip (Just (n-1) :!: s')
734 next (Nothing :!: s) = case next0 s of
735 Done -> Done
736 Skip s' -> Skip (Nothing :!: s')
737 Yield x s' -> Yield x (Nothing :!: s')
738{-# INLINE [0] drop #-}
740-- | takeWhile, applied to a predicate @p@ and a stream, returns the
741-- longest prefix (possibly empty) of elements that satisfy p.
742takeWhile :: (Char -> Bool) -> Stream Char -> Stream Char
743takeWhile p (Stream next0 s0 len) = Stream next s0 len
744 where
745 {-# INLINE next #-}
746 next !s = case next0 s of
747 Done -> Done
748 Skip s' -> Skip s'
749 Yield x s' | p x -> Yield x s'
750 | otherwise -> Done
751{-# INLINE [0] takeWhile #-}
753-- | dropWhile @p @xs returns the suffix remaining after takeWhile @p @xs.
754dropWhile :: (Char -> Bool) -> Stream Char -> Stream Char
755dropWhile p (Stream next0 s0 len) = Stream next (S1 :!: s0) len
756 where
757 {-# INLINE next #-}
758 next (S1 :!: s) = case next0 s of
759 Done -> Done
760 Skip s' -> Skip (S1 :!: s')
761 Yield x s' | p x -> Skip (S1 :!: s')
762 | otherwise -> Yield x (S2 :!: s')
763 next (S2 :!: s) = case next0 s of
764 Done -> Done
765 Skip s' -> Skip (S2 :!: s')
766 Yield x s' -> Yield x (S2 :!: s')
767{-# INLINE [0] dropWhile #-}
769-- ----------------------------------------------------------------------------
770-- * Searching
772-------------------------------------------------------------------------------
773-- ** Searching by equality
775-- | /O(n)/ elem is the stream membership predicate.
776elem :: Char -> Stream Char -> Bool
777elem w (Stream next s0 len) = loop_elem s0
778 where
779 loop_elem !s = case next s of
780 Done -> False
781 Skip s' -> loop_elem s'
782 Yield x s' | x == w -> True
783 | otherwise -> loop_elem s'
784{-# INLINE [0] elem #-}
786-------------------------------------------------------------------------------
787-- ** Searching with a predicate
789-- | /O(n)/ The 'find' function takes a predicate and a stream,
790-- and returns the first element in matching the predicate, or 'Nothing'
791-- if there is no such element.
793find :: (Char -> Bool) -> Stream Char -> Maybe Char
794find p (Stream next s0 len) = loop_find s0
795 where
796 loop_find !s = case next s of
797 Done -> Nothing
798 Skip s' -> loop_find s'
799 Yield x s' | p x -> Just x
800 | otherwise -> loop_find s'
801{-# INLINE [0] find #-}
803-- | /O(n)/ 'filter', applied to a predicate and a stream,
804-- returns a stream containing those characters that satisfy the
805-- predicate.
806filter :: (Char -> Bool) -> Stream Char -> Stream Char
807filter p (Stream next0 s0 len) = Stream next s0 len
808 where
809 {-# INLINE next #-}
810 next !s = case next0 s of
811 Done -> Done
812 Skip s' -> Skip s'
813 Yield x s' | p x -> Yield x s'
814 | otherwise -> Skip s'
815{-# INLINE [0] filter #-}
817{-# RULES
818 "Stream filter/filter fusion" forall p q s.
819 filter p (filter q s) = filter (\x -> q x && p x) s
820 #-}
822-------------------------------------------------------------------------------
823-- ** Indexing streams
825-- | /O(1)/ stream index (subscript) operator, starting from 0.
826index :: Stream Char -> Int -> Char
827index (Stream next s0 len) n0
828 | n0 < 0 = error "Stream.(!!): negative index"
829 | otherwise = loop_index n0 s0
830 where
831 loop_index !n !s = case next s of
832 Done -> error "Stream.(!!): index too large"
833 Skip s' -> loop_index n s'
834 Yield x s' | n == 0 -> x
835 | otherwise -> loop_index (n-1) s'
836{-# INLINE [0] index #-}
838-- | The 'findIndex' function takes a predicate and a stream and
839-- returns the index of the first element in the stream
840-- satisfying the predicate.
841findIndex :: (Char -> Bool) -> Stream Char -> Maybe Int
842findIndex p (Stream next s0 len) = loop_findIndex 0 s0
843 where
844 loop_findIndex !i !s = case next s of
845 Done -> Nothing
846 Skip s' -> loop_findIndex i s' -- hmm. not caught by QC
847 Yield x s' | p x -> Just i
848 | otherwise -> loop_findIndex (i+1) s'
849{-# INLINE [0] findIndex #-}
851-- | /O(n)/ The 'elemIndex' function returns the index of the first
852-- element in the given stream which is equal to the query
853-- element, or 'Nothing' if there is no such element.
854elemIndex :: Char -> Stream Char -> Maybe Int
855elemIndex a (Stream next s0 len) = loop_elemIndex 0 s0
856 where
857 loop_elemIndex !i !s = case next s of
858 Done -> Nothing
859 Skip s' -> loop_elemIndex i s'
860 Yield x s' | a == x -> Just i
861 | otherwise -> loop_elemIndex (i+1) s'
862{-# INLINE [0] elemIndex #-}
864-------------------------------------------------------------------------------
865-- * Zipping
867-- | zipWith generalises 'zip' by zipping with the function given as
868-- the first argument, instead of a tupling function.
869zipWith :: (Char -> Char -> Char) -> Stream Char -> Stream Char -> Stream Char
870zipWith f (Stream next0 sa0 len1) (Stream next1 sb0 len2) = Stream next (sa0 :!: sb0 :!: Nothing) (min len1 len2)
871 where
872 {-# INLINE next #-}
873 next (sa :!: sb :!: Nothing) = case next0 sa of
874 Done -> Done
875 Skip sa' -> Skip (sa' :!: sb :!: Nothing)
876 Yield a sa' -> Skip (sa' :!: sb :!: Just a)
878 next (sa' :!: sb :!: Just a) = case next1 sb of
879 Done -> Done
880 Skip sb' -> Skip (sa' :!: sb' :!: Just a)
881 Yield b sb' -> Yield (f a b) (sa' :!: sb' :!: Nothing)
882{-# INLINE [0] zipWith #-}
884errorEmptyList :: String -> a
885errorEmptyList fun =
886 error ("Prelude." ++ fun ++ ": empty list")