1{-# LANGUAGE CPP, ExistentialQuantification, MultiParamTypeClasses, FlexibleInstances, Rank2Types, BangPatterns, KindSignatures, GADTs, ScopedTypeVariables #-} 2 3-- | 4-- Module : Data.Vector.Fusion.Stream.Monadic 5-- Copyright : (c) Roman Leshchinskiy 2008-2010 6-- License : BSD-style 7-- 8-- Maintainer : Roman Leshchinskiy <rl@cse.unsw.edu.au> 9-- Stability : experimental 10-- Portability : non-portable 11-- 12-- Monadic stream combinators. 13-- 14 15module Data.Vector.Fusion.Stream.Monadic ( 16 Stream(..), Step(..), SPEC(..), 17 18 -- * Length 19 length, null, 20 21 -- * Construction 22 empty, singleton, cons, snoc, replicate, replicateM, generate, generateM, (++), 23 24 -- * Accessing elements 25 head, last, (!!), (!?), 26 27 -- * Substreams 28 slice, init, tail, take, drop, 29 30 -- * Mapping 31 map, mapM, mapM_, trans, unbox, concatMap, flatten, 32 33 -- * Zipping 34 indexed, indexedR, zipWithM_, 35 zipWithM, zipWith3M, zipWith4M, zipWith5M, zipWith6M, 36 zipWith, zipWith3, zipWith4, zipWith5, zipWith6, 37 zip, zip3, zip4, zip5, zip6, 38 39 -- * Comparisons 40 eqBy, cmpBy, 41 42 -- * Filtering 43 filter, filterM, uniq, mapMaybe, mapMaybeM, catMaybes, takeWhile, takeWhileM, dropWhile, dropWhileM, 44 45 -- * Searching 46 elem, notElem, find, findM, findIndex, findIndexM, 47 48 -- * Folding 49 foldl, foldlM, foldl1, foldl1M, foldM, fold1M, 50 foldl', foldlM', foldl1', foldl1M', foldM', fold1M', 51 foldr, foldrM, foldr1, foldr1M, 52 53 -- * Specialised folds 54 and, or, concatMapM, 55 56 -- * Unfolding 57 unfoldr, unfoldrM, 58 unfoldrN, unfoldrNM, 59 unfoldrExactN, unfoldrExactNM, 60 iterateN, iterateNM, 61 62 -- * Scans 63 prescanl, prescanlM, prescanl', prescanlM', 64 postscanl, postscanlM, postscanl', postscanlM', 65 scanl, scanlM, scanl', scanlM', 66 scanl1, scanl1M, scanl1', scanl1M', 67 68 -- * Enumerations 69 enumFromStepN, enumFromTo, enumFromThenTo, 70 71 -- * Conversions 72 toList, fromList, fromListN 73) where 74 75import Data.Vector.Fusion.Util ( Box(..) ) 76 77import Data.Char ( ord ) 78import GHC.Base ( unsafeChr ) 79import Control.Monad ( liftM ) 80import Prelude hiding ( length, null, 81 replicate, (++), 82 head, last, (!!), 83 init, tail, take, drop, 84 map, mapM, mapM_, concatMap, 85 zipWith, zipWith3, zip, zip3, 86 filter, takeWhile, dropWhile, 87 elem, notElem, 88 foldl, foldl1, foldr, foldr1, 89 and, or, 90 scanl, scanl1, 91 enumFromTo, enumFromThenTo ) 92 93import Data.Int ( Int8, Int16, Int32 ) 94import Data.Word ( Word8, Word16, Word32, Word64 ) 95 96#if !MIN_VERSION_base(4,8,0) 97import Data.Word ( Word8, Word16, Word32, Word, Word64 ) 98#endif 99 100#if __GLASGOW_HASKELL__ >= 708 101import GHC.Types ( SPEC(..) ) 102#elif __GLASGOW_HASKELL__ >= 700 103import GHC.Exts ( SpecConstrAnnotation(..) ) 104#endif 105 106#include "vector.h" 107#include "MachDeps.h" 108 109#if WORD_SIZE_IN_BITS > 32 110import Data.Int ( Int64 ) 111#endif 112 113#if __GLASGOW_HASKELL__ < 708 114data SPEC = SPEC | SPEC2 115#if __GLASGOW_HASKELL__ >= 700 116{-# ANN type SPEC ForceSpecConstr #-} 117#endif 118#endif 119 120emptyStream :: String 121{-# NOINLINE emptyStream #-} 122emptyStream = "empty stream" 123 124#define EMPTY_STREAM (\state -> ERROR state emptyStream) 125 126-- | Result of taking a single step in a stream 127data Step s a where 128 Yield :: a -> s -> Step s a 129 Skip :: s -> Step s a 130 Done :: Step s a 131 132instance Functor (Step s) where 133 {-# INLINE fmap #-} 134 fmap f (Yield x s) = Yield (f x) s 135 fmap _ (Skip s) = Skip s 136 fmap _ Done = Done 137#if MIN_VERSION_base(4,8,0) 138 {-# INLINE (<$) #-} 139 (<$) = fmap . const 140#endif 141 142-- | Monadic streams 143data Stream m a = forall s. Stream (s -> m (Step s a)) s 144 145-- Length 146-- ------ 147 148-- | Length of a 'Stream' 149length :: Monad m => Stream m a -> m Int 150{-# INLINE_FUSED length #-} 151length = foldl' (\n _ -> n+1) 0 152 153-- | Check if a 'Stream' is empty 154null :: Monad m => Stream m a -> m Bool 155{-# INLINE_FUSED null #-} 156null (Stream step t) = null_loop t 157 where 158 null_loop s = do 159 r <- step s 160 case r of 161 Yield _ _ -> return False 162 Skip s' -> null_loop s' 163 Done -> return True 164 165-- Construction 166-- ------------ 167 168-- | Empty 'Stream' 169empty :: Monad m => Stream m a 170{-# INLINE_FUSED empty #-} 171empty = Stream (const (return Done)) () 172 173-- | Singleton 'Stream' 174singleton :: Monad m => a -> Stream m a 175{-# INLINE_FUSED singleton #-} 176singleton x = Stream (return . step) True 177 where 178 {-# INLINE_INNER step #-} 179 step True = Yield x False 180 step False = Done 181 182-- | Replicate a value to a given length 183replicate :: Monad m => Int -> a -> Stream m a 184{-# INLINE_FUSED replicate #-} 185replicate n x = replicateM n (return x) 186 187-- | Yield a 'Stream' of values obtained by performing the monadic action the 188-- given number of times 189replicateM :: Monad m => Int -> m a -> Stream m a 190{-# INLINE_FUSED replicateM #-} 191replicateM n p = Stream step n 192 where 193 {-# INLINE_INNER step #-} 194 step i | i <= 0 = return Done 195 | otherwise = do { x <- p; return $ Yield x (i-1) } 196 197generate :: Monad m => Int -> (Int -> a) -> Stream m a 198{-# INLINE generate #-} 199generate n f = generateM n (return . f) 200 201-- | Generate a stream from its indices 202generateM :: Monad m => Int -> (Int -> m a) -> Stream m a 203{-# INLINE_FUSED generateM #-} 204generateM n f = n `seq` Stream step 0 205 where 206 {-# INLINE_INNER step #-} 207 step i | i < n = do 208 x <- f i 209 return $ Yield x (i+1) 210 | otherwise = return Done 211 212-- | Prepend an element 213cons :: Monad m => a -> Stream m a -> Stream m a 214{-# INLINE cons #-} 215cons x s = singleton x ++ s 216 217-- | Append an element 218snoc :: Monad m => Stream m a -> a -> Stream m a 219{-# INLINE snoc #-} 220snoc s x = s ++ singleton x 221 222infixr 5 ++ 223-- | Concatenate two 'Stream's 224(++) :: Monad m => Stream m a -> Stream m a -> Stream m a 225{-# INLINE_FUSED (++) #-} 226Stream stepa ta ++ Stream stepb tb = Stream step (Left ta) 227 where 228 {-# INLINE_INNER step #-} 229 step (Left sa) = do 230 r <- stepa sa 231 case r of 232 Yield x sa' -> return $ Yield x (Left sa') 233 Skip sa' -> return $ Skip (Left sa') 234 Done -> return $ Skip (Right tb) 235 step (Right sb) = do 236 r <- stepb sb 237 case r of 238 Yield x sb' -> return $ Yield x (Right sb') 239 Skip sb' -> return $ Skip (Right sb') 240 Done -> return $ Done 241 242-- Accessing elements 243-- ------------------ 244 245-- | First element of the 'Stream' or error if empty 246head :: Monad m => Stream m a -> m a 247{-# INLINE_FUSED head #-} 248head (Stream step t) = head_loop SPEC t 249 where 250 head_loop !_ s 251 = do 252 r <- step s 253 case r of 254 Yield x _ -> return x 255 Skip s' -> head_loop SPEC s' 256 Done -> EMPTY_STREAM "head" 257 258 259 260-- | Last element of the 'Stream' or error if empty 261last :: Monad m => Stream m a -> m a 262{-# INLINE_FUSED last #-} 263last (Stream step t) = last_loop0 SPEC t 264 where 265 last_loop0 !_ s 266 = do 267 r <- step s 268 case r of 269 Yield x s' -> last_loop1 SPEC x s' 270 Skip s' -> last_loop0 SPEC s' 271 Done -> EMPTY_STREAM "last" 272 273 last_loop1 !_ x s 274 = do 275 r <- step s 276 case r of 277 Yield y s' -> last_loop1 SPEC y s' 278 Skip s' -> last_loop1 SPEC x s' 279 Done -> return x 280 281infixl 9 !! 282-- | Element at the given position 283(!!) :: Monad m => Stream m a -> Int -> m a 284{-# INLINE (!!) #-} 285Stream step t !! j | j < 0 = ERROR "!!" "negative index" 286 | otherwise = index_loop SPEC t j 287 where 288 index_loop !_ s i 289 = i `seq` 290 do 291 r <- step s 292 case r of 293 Yield x s' | i == 0 -> return x 294 | otherwise -> index_loop SPEC s' (i-1) 295 Skip s' -> index_loop SPEC s' i 296 Done -> EMPTY_STREAM "!!" 297 298infixl 9 !? 299-- | Element at the given position or 'Nothing' if out of bounds 300(!?) :: Monad m => Stream m a -> Int -> m (Maybe a) 301{-# INLINE (!?) #-} 302Stream step t !? j = index_loop SPEC t j 303 where 304 index_loop !_ s i 305 = i `seq` 306 do 307 r <- step s 308 case r of 309 Yield x s' | i == 0 -> return (Just x) 310 | otherwise -> index_loop SPEC s' (i-1) 311 Skip s' -> index_loop SPEC s' i 312 Done -> return Nothing 313 314-- Substreams 315-- ---------- 316 317-- | Extract a substream of the given length starting at the given position. 318slice :: Monad m => Int -- ^ starting index 319 -> Int -- ^ length 320 -> Stream m a 321 -> Stream m a 322{-# INLINE slice #-} 323slice i n s = take n (drop i s) 324 325-- | All but the last element 326init :: Monad m => Stream m a -> Stream m a 327{-# INLINE_FUSED init #-} 328init (Stream step t) = Stream step' (Nothing, t) 329 where 330 {-# INLINE_INNER step' #-} 331 step' (Nothing, s) = liftM (\r -> 332 case r of 333 Yield x s' -> Skip (Just x, s') 334 Skip s' -> Skip (Nothing, s') 335 Done -> EMPTY_STREAM "init" 336 ) (step s) 337 338 step' (Just x, s) = liftM (\r -> 339 case r of 340 Yield y s' -> Yield x (Just y, s') 341 Skip s' -> Skip (Just x, s') 342 Done -> Done 343 ) (step s) 344 345-- | All but the first element 346tail :: Monad m => Stream m a -> Stream m a 347{-# INLINE_FUSED tail #-} 348tail (Stream step t) = Stream step' (Left t) 349 where 350 {-# INLINE_INNER step' #-} 351 step' (Left s) = liftM (\r -> 352 case r of 353 Yield _ s' -> Skip (Right s') 354 Skip s' -> Skip (Left s') 355 Done -> EMPTY_STREAM "tail" 356 ) (step s) 357 358 step' (Right s) = liftM (\r -> 359 case r of 360 Yield x s' -> Yield x (Right s') 361 Skip s' -> Skip (Right s') 362 Done -> Done 363 ) (step s) 364 365-- | The first @n@ elements 366take :: Monad m => Int -> Stream m a -> Stream m a 367{-# INLINE_FUSED take #-} 368take n (Stream step t) = n `seq` Stream step' (t, 0) 369 where 370 {-# INLINE_INNER step' #-} 371 step' (s, i) | i < n = liftM (\r -> 372 case r of 373 Yield x s' -> Yield x (s', i+1) 374 Skip s' -> Skip (s', i) 375 Done -> Done 376 ) (step s) 377 step' (_, _) = return Done 378 379-- | All but the first @n@ elements 380drop :: Monad m => Int -> Stream m a -> Stream m a 381{-# INLINE_FUSED drop #-} 382drop n (Stream step t) = Stream step' (t, Just n) 383 where 384 {-# INLINE_INNER step' #-} 385 step' (s, Just i) | i > 0 = liftM (\r -> 386 case r of 387 Yield _ s' -> Skip (s', Just (i-1)) 388 Skip s' -> Skip (s', Just i) 389 Done -> Done 390 ) (step s) 391 | otherwise = return $ Skip (s, Nothing) 392 393 step' (s, Nothing) = liftM (\r -> 394 case r of 395 Yield x s' -> Yield x (s', Nothing) 396 Skip s' -> Skip (s', Nothing) 397 Done -> Done 398 ) (step s) 399 400-- Mapping 401-- ------- 402 403instance Monad m => Functor (Stream m) where 404 {-# INLINE fmap #-} 405 fmap = map 406 407-- | Map a function over a 'Stream' 408map :: Monad m => (a -> b) -> Stream m a -> Stream m b 409{-# INLINE map #-} 410map f = mapM (return . f) 411 412 413-- | Map a monadic function over a 'Stream' 414mapM :: Monad m => (a -> m b) -> Stream m a -> Stream m b 415{-# INLINE_FUSED mapM #-} 416mapM f (Stream step t) = Stream step' t 417 where 418 {-# INLINE_INNER step' #-} 419 step' s = do 420 r <- step s 421 case r of 422 Yield x s' -> liftM (`Yield` s') (f x) 423 Skip s' -> return (Skip s') 424 Done -> return Done 425 426consume :: Monad m => Stream m a -> m () 427{-# INLINE_FUSED consume #-} 428consume (Stream step t) = consume_loop SPEC t 429 where 430 consume_loop !_ s 431 = do 432 r <- step s 433 case r of 434 Yield _ s' -> consume_loop SPEC s' 435 Skip s' -> consume_loop SPEC s' 436 Done -> return () 437 438-- | Execute a monadic action for each element of the 'Stream' 439mapM_ :: Monad m => (a -> m b) -> Stream m a -> m () 440{-# INLINE_FUSED mapM_ #-} 441mapM_ m = consume . mapM m 442 443-- | Transform a 'Stream' to use a different monad 444trans :: (Monad m, Monad m') 445 => (forall z. m z -> m' z) -> Stream m a -> Stream m' a 446{-# INLINE_FUSED trans #-} 447trans f (Stream step s) = Stream (f . step) s 448 449unbox :: Monad m => Stream m (Box a) -> Stream m a 450{-# INLINE_FUSED unbox #-} 451unbox (Stream step t) = Stream step' t 452 where 453 {-# INLINE_INNER step' #-} 454 step' s = do 455 r <- step s 456 case r of 457 Yield (Box x) s' -> return $ Yield x s' 458 Skip s' -> return $ Skip s' 459 Done -> return $ Done 460 461-- Zipping 462-- ------- 463 464-- | Pair each element in a 'Stream' with its index 465indexed :: Monad m => Stream m a -> Stream m (Int,a) 466{-# INLINE_FUSED indexed #-} 467indexed (Stream step t) = Stream step' (t,0) 468 where 469 {-# INLINE_INNER step' #-} 470 step' (s,i) = i `seq` 471 do 472 r <- step s 473 case r of 474 Yield x s' -> return $ Yield (i,x) (s', i+1) 475 Skip s' -> return $ Skip (s', i) 476 Done -> return Done 477 478-- | Pair each element in a 'Stream' with its index, starting from the right 479-- and counting down 480indexedR :: Monad m => Int -> Stream m a -> Stream m (Int,a) 481{-# INLINE_FUSED indexedR #-} 482indexedR m (Stream step t) = Stream step' (t,m) 483 where 484 {-# INLINE_INNER step' #-} 485 step' (s,i) = i `seq` 486 do 487 r <- step s 488 case r of 489 Yield x s' -> let i' = i-1 490 in 491 return $ Yield (i',x) (s', i') 492 Skip s' -> return $ Skip (s', i) 493 Done -> return Done 494 495-- | Zip two 'Stream's with the given monadic function 496zipWithM :: Monad m => (a -> b -> m c) -> Stream m a -> Stream m b -> Stream m c 497{-# INLINE_FUSED zipWithM #-} 498zipWithM f (Stream stepa ta) (Stream stepb tb) = Stream step (ta, tb, Nothing) 499 where 500 {-# INLINE_INNER step #-} 501 step (sa, sb, Nothing) = liftM (\r -> 502 case r of 503 Yield x sa' -> Skip (sa', sb, Just x) 504 Skip sa' -> Skip (sa', sb, Nothing) 505 Done -> Done 506 ) (stepa sa) 507 508 step (sa, sb, Just x) = do 509 r <- stepb sb 510 case r of 511 Yield y sb' -> 512 do 513 z <- f x y 514 return $ Yield z (sa, sb', Nothing) 515 Skip sb' -> return $ Skip (sa, sb', Just x) 516 Done -> return $ Done 517 518zipWithM_ :: Monad m => (a -> b -> m c) -> Stream m a -> Stream m b -> m () 519{-# INLINE zipWithM_ #-} 520zipWithM_ f sa sb = consume (zipWithM f sa sb) 521 522zipWith3M :: Monad m => (a -> b -> c -> m d) -> Stream m a -> Stream m b -> Stream m c -> Stream m d 523{-# INLINE_FUSED zipWith3M #-} 524zipWith3M f (Stream stepa ta) 525 (Stream stepb tb) 526 (Stream stepc tc) = Stream step (ta, tb, tc, Nothing) 527 where 528 {-# INLINE_INNER step #-} 529 step (sa, sb, sc, Nothing) = do 530 r <- stepa sa 531 return $ case r of 532 Yield x sa' -> Skip (sa', sb, sc, Just (x, Nothing)) 533 Skip sa' -> Skip (sa', sb, sc, Nothing) 534 Done -> Done 535 536 step (sa, sb, sc, Just (x, Nothing)) = do 537 r <- stepb sb 538 return $ case r of 539 Yield y sb' -> Skip (sa, sb', sc, Just (x, Just y)) 540 Skip sb' -> Skip (sa, sb', sc, Just (x, Nothing)) 541 Done -> Done 542 543 step (sa, sb, sc, Just (x, Just y)) = do 544 r <- stepc sc 545 case r of 546 Yield z sc' -> f x y z >>= (\res -> return $ Yield res (sa, sb, sc', Nothing)) 547 Skip sc' -> return $ Skip (sa, sb, sc', Just (x, Just y)) 548 Done -> return $ Done 549 550zipWith4M :: Monad m => (a -> b -> c -> d -> m e) 551 -> Stream m a -> Stream m b -> Stream m c -> Stream m d 552 -> Stream m e 553{-# INLINE zipWith4M #-} 554zipWith4M f sa sb sc sd 555 = zipWithM (\(a,b) (c,d) -> f a b c d) (zip sa sb) (zip sc sd) 556 557zipWith5M :: Monad m => (a -> b -> c -> d -> e -> m f) 558 -> Stream m a -> Stream m b -> Stream m c -> Stream m d 559 -> Stream m e -> Stream m f 560{-# INLINE zipWith5M #-} 561zipWith5M f sa sb sc sd se 562 = zipWithM (\(a,b,c) (d,e) -> f a b c d e) (zip3 sa sb sc) (zip sd se) 563 564zipWith6M :: Monad m => (a -> b -> c -> d -> e -> f -> m g) 565 -> Stream m a -> Stream m b -> Stream m c -> Stream m d 566 -> Stream m e -> Stream m f -> Stream m g 567{-# INLINE zipWith6M #-} 568zipWith6M fn sa sb sc sd se sf 569 = zipWithM (\(a,b,c) (d,e,f) -> fn a b c d e f) (zip3 sa sb sc) 570 (zip3 sd se sf) 571 572zipWith :: Monad m => (a -> b -> c) -> Stream m a -> Stream m b -> Stream m c 573{-# INLINE zipWith #-} 574zipWith f = zipWithM (\a b -> return (f a b)) 575 576zipWith3 :: Monad m => (a -> b -> c -> d) 577 -> Stream m a -> Stream m b -> Stream m c -> Stream m d 578{-# INLINE zipWith3 #-} 579zipWith3 f = zipWith3M (\a b c -> return (f a b c)) 580 581zipWith4 :: Monad m => (a -> b -> c -> d -> e) 582 -> Stream m a -> Stream m b -> Stream m c -> Stream m d 583 -> Stream m e 584{-# INLINE zipWith4 #-} 585zipWith4 f = zipWith4M (\a b c d -> return (f a b c d)) 586 587zipWith5 :: Monad m => (a -> b -> c -> d -> e -> f) 588 -> Stream m a -> Stream m b -> Stream m c -> Stream m d 589 -> Stream m e -> Stream m f 590{-# INLINE zipWith5 #-} 591zipWith5 f = zipWith5M (\a b c d e -> return (f a b c d e)) 592 593zipWith6 :: Monad m => (a -> b -> c -> d -> e -> f -> g) 594 -> Stream m a -> Stream m b -> Stream m c -> Stream m d 595 -> Stream m e -> Stream m f -> Stream m g 596{-# INLINE zipWith6 #-} 597zipWith6 fn = zipWith6M (\a b c d e f -> return (fn a b c d e f)) 598 599zip :: Monad m => Stream m a -> Stream m b -> Stream m (a,b) 600{-# INLINE zip #-} 601zip = zipWith (,) 602 603zip3 :: Monad m => Stream m a -> Stream m b -> Stream m c -> Stream m (a,b,c) 604{-# INLINE zip3 #-} 605zip3 = zipWith3 (,,) 606 607zip4 :: Monad m => Stream m a -> Stream m b -> Stream m c -> Stream m d 608 -> Stream m (a,b,c,d) 609{-# INLINE zip4 #-} 610zip4 = zipWith4 (,,,) 611 612zip5 :: Monad m => Stream m a -> Stream m b -> Stream m c -> Stream m d 613 -> Stream m e -> Stream m (a,b,c,d,e) 614{-# INLINE zip5 #-} 615zip5 = zipWith5 (,,,,) 616 617zip6 :: Monad m => Stream m a -> Stream m b -> Stream m c -> Stream m d 618 -> Stream m e -> Stream m f -> Stream m (a,b,c,d,e,f) 619{-# INLINE zip6 #-} 620zip6 = zipWith6 (,,,,,) 621 622-- Comparisons 623-- ----------- 624 625-- | Check if two 'Stream's are equal 626eqBy :: (Monad m) => (a -> b -> Bool) -> Stream m a -> Stream m b -> m Bool 627{-# INLINE_FUSED eqBy #-} 628eqBy eq (Stream step1 t1) (Stream step2 t2) = eq_loop0 SPEC t1 t2 629 where 630 eq_loop0 !_ s1 s2 = do 631 r <- step1 s1 632 case r of 633 Yield x s1' -> eq_loop1 SPEC x s1' s2 634 Skip s1' -> eq_loop0 SPEC s1' s2 635 Done -> eq_null s2 636 637 eq_loop1 !_ x s1 s2 = do 638 r <- step2 s2 639 case r of 640 Yield y s2' 641 | eq x y -> eq_loop0 SPEC s1 s2' 642 | otherwise -> return False 643 Skip s2' -> eq_loop1 SPEC x s1 s2' 644 Done -> return False 645 646 eq_null s2 = do 647 r <- step2 s2 648 case r of 649 Yield _ _ -> return False 650 Skip s2' -> eq_null s2' 651 Done -> return True 652 653-- | Lexicographically compare two 'Stream's 654cmpBy :: (Monad m) => (a -> b -> Ordering) -> Stream m a -> Stream m b -> m Ordering 655{-# INLINE_FUSED cmpBy #-} 656cmpBy cmp (Stream step1 t1) (Stream step2 t2) = cmp_loop0 SPEC t1 t2 657 where 658 cmp_loop0 !_ s1 s2 = do 659 r <- step1 s1 660 case r of 661 Yield x s1' -> cmp_loop1 SPEC x s1' s2 662 Skip s1' -> cmp_loop0 SPEC s1' s2 663 Done -> cmp_null s2 664 665 cmp_loop1 !_ x s1 s2 = do 666 r <- step2 s2 667 case r of 668 Yield y s2' -> case x `cmp` y of 669 EQ -> cmp_loop0 SPEC s1 s2' 670 c -> return c 671 Skip s2' -> cmp_loop1 SPEC x s1 s2' 672 Done -> return GT 673 674 cmp_null s2 = do 675 r <- step2 s2 676 case r of 677 Yield _ _ -> return LT 678 Skip s2' -> cmp_null s2' 679 Done -> return EQ 680 681-- Filtering 682-- --------- 683 684-- | Drop elements which do not satisfy the predicate 685filter :: Monad m => (a -> Bool) -> Stream m a -> Stream m a 686{-# INLINE filter #-} 687filter f = filterM (return . f) 688 689mapMaybe :: Monad m => (a -> Maybe b) -> Stream m a -> Stream m b 690{-# INLINE_FUSED mapMaybe #-} 691mapMaybe f (Stream step t) = Stream step' t 692 where 693 {-# INLINE_INNER step' #-} 694 step' s = do 695 r <- step s 696 case r of 697 Yield x s' -> do 698 return $ case f x of 699 Nothing -> Skip s' 700 Just b' -> Yield b' s' 701 Skip s' -> return $ Skip s' 702 Done -> return $ Done 703 704catMaybes :: Monad m => Stream m (Maybe a) -> Stream m a 705catMaybes = mapMaybe id 706 707-- | Drop elements which do not satisfy the monadic predicate 708filterM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a 709{-# INLINE_FUSED filterM #-} 710filterM f (Stream step t) = Stream step' t 711 where 712 {-# INLINE_INNER step' #-} 713 step' s = do 714 r <- step s 715 case r of 716 Yield x s' -> do 717 b <- f x 718 return $ if b then Yield x s' 719 else Skip s' 720 Skip s' -> return $ Skip s' 721 Done -> return $ Done 722 723-- | Apply monadic function to each element and drop all Nothings 724-- 725-- @since 0.12.2.0 726mapMaybeM :: Monad m => (a -> m (Maybe b)) -> Stream m a -> Stream m b 727{-# INLINE_FUSED mapMaybeM #-} 728mapMaybeM f (Stream step t) = Stream step' t 729 where 730 {-# INLINE_INNER step' #-} 731 step' s = do 732 r <- step s 733 case r of 734 Yield x s' -> do 735 fx <- f x 736 return $ case fx of 737 Nothing -> Skip s' 738 Just b -> Yield b s' 739 Skip s' -> return $ Skip s' 740 Done -> return $ Done 741 742-- | Drop repeated adjacent elements. 743uniq :: (Eq a, Monad m) => Stream m a -> Stream m a 744{-# INLINE_FUSED uniq #-} 745uniq (Stream step st) = Stream step' (Nothing,st) 746 where 747 {-# INLINE_INNER step' #-} 748 step' (Nothing, s) = do r <- step s 749 case r of 750 Yield x s' -> return $ Yield x (Just x , s') 751 Skip s' -> return $ Skip (Nothing, s') 752 Done -> return Done 753 step' (Just x0, s) = do r <- step s 754 case r of 755 Yield x s' | x == x0 -> return $ Skip (Just x0, s') 756 | otherwise -> return $ Yield x (Just x , s') 757 Skip s' -> return $ Skip (Just x0, s') 758 Done -> return Done 759 760-- | Longest prefix of elements that satisfy the predicate 761takeWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a 762{-# INLINE takeWhile #-} 763takeWhile f = takeWhileM (return . f) 764 765-- | Longest prefix of elements that satisfy the monadic predicate 766takeWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a 767{-# INLINE_FUSED takeWhileM #-} 768takeWhileM f (Stream step t) = Stream step' t 769 where 770 {-# INLINE_INNER step' #-} 771 step' s = do 772 r <- step s 773 case r of 774 Yield x s' -> do 775 b <- f x 776 return $ if b then Yield x s' else Done 777 Skip s' -> return $ Skip s' 778 Done -> return $ Done 779 780-- | Drop the longest prefix of elements that satisfy the predicate 781dropWhile :: Monad m => (a -> Bool) -> Stream m a -> Stream m a 782{-# INLINE dropWhile #-} 783dropWhile f = dropWhileM (return . f) 784 785data DropWhile s a = DropWhile_Drop s | DropWhile_Yield a s | DropWhile_Next s 786 787-- | Drop the longest prefix of elements that satisfy the monadic predicate 788dropWhileM :: Monad m => (a -> m Bool) -> Stream m a -> Stream m a 789{-# INLINE_FUSED dropWhileM #-} 790dropWhileM f (Stream step t) = Stream step' (DropWhile_Drop t) 791 where 792 -- NOTE: we jump through hoops here to have only one Yield; local data 793 -- declarations would be nice! 794 795 {-# INLINE_INNER step' #-} 796 step' (DropWhile_Drop s) 797 = do 798 r <- step s 799 case r of 800 Yield x s' -> do 801 b <- f x 802 return $ if b then Skip (DropWhile_Drop s') 803 else Skip (DropWhile_Yield x s') 804 Skip s' -> return $ Skip (DropWhile_Drop s') 805 Done -> return $ Done 806 807 step' (DropWhile_Yield x s) = return $ Yield x (DropWhile_Next s) 808 809 step' (DropWhile_Next s) 810 = liftM (\r -> 811 case r of 812 Yield x s' -> Skip (DropWhile_Yield x s') 813 Skip s' -> Skip (DropWhile_Next s') 814 Done -> Done 815 ) (step s) 816 817-- Searching 818-- --------- 819 820infix 4 `elem` 821-- | Check whether the 'Stream' contains an element 822elem :: (Monad m, Eq a) => a -> Stream m a -> m Bool 823{-# INLINE_FUSED elem #-} 824elem x (Stream step t) = elem_loop SPEC t 825 where 826 elem_loop !_ s 827 = do 828 r <- step s 829 case r of 830 Yield y s' | x == y -> return True 831 | otherwise -> elem_loop SPEC s' 832 Skip s' -> elem_loop SPEC s' 833 Done -> return False 834 835infix 4 `notElem` 836-- | Inverse of `elem` 837notElem :: (Monad m, Eq a) => a -> Stream m a -> m Bool 838{-# INLINE notElem #-} 839notElem x s = liftM not (elem x s) 840 841-- | Yield 'Just' the first element that satisfies the predicate or 'Nothing' 842-- if no such element exists. 843find :: Monad m => (a -> Bool) -> Stream m a -> m (Maybe a) 844{-# INLINE find #-} 845find f = findM (return . f) 846 847-- | Yield 'Just' the first element that satisfies the monadic predicate or 848-- 'Nothing' if no such element exists. 849findM :: Monad m => (a -> m Bool) -> Stream m a -> m (Maybe a) 850{-# INLINE_FUSED findM #-} 851findM f (Stream step t) = find_loop SPEC t 852 where 853 find_loop !_ s 854 = do 855 r <- step s 856 case r of 857 Yield x s' -> do 858 b <- f x 859 if b then return $ Just x 860 else find_loop SPEC s' 861 Skip s' -> find_loop SPEC s' 862 Done -> return Nothing 863 864-- | Yield 'Just' the index of the first element that satisfies the predicate 865-- or 'Nothing' if no such element exists. 866findIndex :: Monad m => (a -> Bool) -> Stream m a -> m (Maybe Int) 867{-# INLINE_FUSED findIndex #-} 868findIndex f = findIndexM (return . f) 869 870-- | Yield 'Just' the index of the first element that satisfies the monadic 871-- predicate or 'Nothing' if no such element exists. 872findIndexM :: Monad m => (a -> m Bool) -> Stream m a -> m (Maybe Int) 873{-# INLINE_FUSED findIndexM #-} 874findIndexM f (Stream step t) = findIndex_loop SPEC t 0 875 where 876 findIndex_loop !_ s i 877 = do 878 r <- step s 879 case r of 880 Yield x s' -> do 881 b <- f x 882 if b then return $ Just i 883 else findIndex_loop SPEC s' (i+1) 884 Skip s' -> findIndex_loop SPEC s' i 885 Done -> return Nothing 886 887-- Folding 888-- ------- 889 890-- | Left fold 891foldl :: Monad m => (a -> b -> a) -> a -> Stream m b -> m a 892{-# INLINE foldl #-} 893foldl f = foldlM (\a b -> return (f a b)) 894 895-- | Left fold with a monadic operator 896foldlM :: Monad m => (a -> b -> m a) -> a -> Stream m b -> m a 897{-# INLINE_FUSED foldlM #-} 898foldlM m w (Stream step t) = foldlM_loop SPEC w t 899 where 900 foldlM_loop !_ z s 901 = do 902 r <- step s 903 case r of 904 Yield x s' -> do { z' <- m z x; foldlM_loop SPEC z' s' } 905 Skip s' -> foldlM_loop SPEC z s' 906 Done -> return z 907 908-- | Same as 'foldlM' 909foldM :: Monad m => (a -> b -> m a) -> a -> Stream m b -> m a 910{-# INLINE foldM #-} 911foldM = foldlM 912 913-- | Left fold over a non-empty 'Stream' 914foldl1 :: Monad m => (a -> a -> a) -> Stream m a -> m a 915{-# INLINE foldl1 #-} 916foldl1 f = foldl1M (\a b -> return (f a b)) 917 918-- | Left fold over a non-empty 'Stream' with a monadic operator 919foldl1M :: Monad m => (a -> a -> m a) -> Stream m a -> m a 920{-# INLINE_FUSED foldl1M #-} 921foldl1M f (Stream step t) = foldl1M_loop SPEC t 922 where 923 foldl1M_loop !_ s 924 = do 925 r <- step s 926 case r of 927 Yield x s' -> foldlM f x (Stream step s') 928 Skip s' -> foldl1M_loop SPEC s' 929 Done -> EMPTY_STREAM "foldl1M" 930 931-- | Same as 'foldl1M' 932fold1M :: Monad m => (a -> a -> m a) -> Stream m a -> m a 933{-# INLINE fold1M #-} 934fold1M = foldl1M 935 936-- | Left fold with a strict accumulator 937foldl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> m a 938{-# INLINE foldl' #-} 939foldl' f = foldlM' (\a b -> return (f a b)) 940 941-- | Left fold with a strict accumulator and a monadic operator 942foldlM' :: Monad m => (a -> b -> m a) -> a -> Stream m b -> m a 943{-# INLINE_FUSED foldlM' #-} 944foldlM' m w (Stream step t) = foldlM'_loop SPEC w t 945 where 946 foldlM'_loop !_ z s 947 = z `seq` 948 do 949 r <- step s 950 case r of 951 Yield x s' -> do { z' <- m z x; foldlM'_loop SPEC z' s' } 952 Skip s' -> foldlM'_loop SPEC z s' 953 Done -> return z 954 955-- | Same as 'foldlM'' 956foldM' :: Monad m => (a -> b -> m a) -> a -> Stream m b -> m a 957{-# INLINE foldM' #-} 958foldM' = foldlM' 959 960-- | Left fold over a non-empty 'Stream' with a strict accumulator 961foldl1' :: Monad m => (a -> a -> a) -> Stream m a -> m a 962{-# INLINE foldl1' #-} 963foldl1' f = foldl1M' (\a b -> return (f a b)) 964 965-- | Left fold over a non-empty 'Stream' with a strict accumulator and a 966-- monadic operator 967foldl1M' :: Monad m => (a -> a -> m a) -> Stream m a -> m a 968{-# INLINE_FUSED foldl1M' #-} 969foldl1M' f (Stream step t) = foldl1M'_loop SPEC t 970 where 971 foldl1M'_loop !_ s 972 = do 973 r <- step s 974 case r of 975 Yield x s' -> foldlM' f x (Stream step s') 976 Skip s' -> foldl1M'_loop SPEC s' 977 Done -> EMPTY_STREAM "foldl1M'" 978 979-- | Same as 'foldl1M'' 980fold1M' :: Monad m => (a -> a -> m a) -> Stream m a -> m a 981{-# INLINE fold1M' #-} 982fold1M' = foldl1M' 983 984-- | Right fold 985foldr :: Monad m => (a -> b -> b) -> b -> Stream m a -> m b 986{-# INLINE foldr #-} 987foldr f = foldrM (\a b -> return (f a b)) 988 989-- | Right fold with a monadic operator 990foldrM :: Monad m => (a -> b -> m b) -> b -> Stream m a -> m b 991{-# INLINE_FUSED foldrM #-} 992foldrM f z (Stream step t) = foldrM_loop SPEC t 993 where 994 foldrM_loop !_ s 995 = do 996 r <- step s 997 case r of 998 Yield x s' -> f x =<< foldrM_loop SPEC s' 999 Skip s' -> foldrM_loop SPEC s' 1000 Done -> return z 1001 1002-- | Right fold over a non-empty stream 1003foldr1 :: Monad m => (a -> a -> a) -> Stream m a -> m a 1004{-# INLINE foldr1 #-} 1005foldr1 f = foldr1M (\a b -> return (f a b)) 1006 1007-- | Right fold over a non-empty stream with a monadic operator 1008foldr1M :: Monad m => (a -> a -> m a) -> Stream m a -> m a 1009{-# INLINE_FUSED foldr1M #-} 1010foldr1M f (Stream step t) = foldr1M_loop0 SPEC t 1011 where 1012 foldr1M_loop0 !_ s 1013 = do 1014 r <- step s 1015 case r of 1016 Yield x s' -> foldr1M_loop1 SPEC x s' 1017 Skip s' -> foldr1M_loop0 SPEC s' 1018 Done -> EMPTY_STREAM "foldr1M" 1019 1020 foldr1M_loop1 !_ x s 1021 = do 1022 r <- step s 1023 case r of 1024 Yield y s' -> f x =<< foldr1M_loop1 SPEC y s' 1025 Skip s' -> foldr1M_loop1 SPEC x s' 1026 Done -> return x 1027 1028-- Specialised folds 1029-- ----------------- 1030 1031and :: Monad m => Stream m Bool -> m Bool 1032{-# INLINE_FUSED and #-} 1033and (Stream step t) = and_loop SPEC t 1034 where 1035 and_loop !_ s 1036 = do 1037 r <- step s 1038 case r of 1039 Yield False _ -> return False 1040 Yield True s' -> and_loop SPEC s' 1041 Skip s' -> and_loop SPEC s' 1042 Done -> return True 1043 1044or :: Monad m => Stream m Bool -> m Bool 1045{-# INLINE_FUSED or #-} 1046or (Stream step t) = or_loop SPEC t 1047 where 1048 or_loop !_ s 1049 = do 1050 r <- step s 1051 case r of 1052 Yield False s' -> or_loop SPEC s' 1053 Yield True _ -> return True 1054 Skip s' -> or_loop SPEC s' 1055 Done -> return False 1056 1057concatMap :: Monad m => (a -> Stream m b) -> Stream m a -> Stream m b 1058{-# INLINE concatMap #-} 1059concatMap f = concatMapM (return . f) 1060 1061concatMapM :: Monad m => (a -> m (Stream m b)) -> Stream m a -> Stream m b 1062{-# INLINE_FUSED concatMapM #-} 1063concatMapM f (Stream step t) = Stream concatMap_go (Left t) 1064 where 1065 concatMap_go (Left s) = do 1066 r <- step s 1067 case r of 1068 Yield a s' -> do 1069 b_stream <- f a 1070 return $ Skip (Right (b_stream, s')) 1071 Skip s' -> return $ Skip (Left s') 1072 Done -> return Done 1073 concatMap_go (Right (Stream inner_step inner_s, s)) = do 1074 r <- inner_step inner_s 1075 case r of 1076 Yield b inner_s' -> return $ Yield b (Right (Stream inner_step inner_s', s)) 1077 Skip inner_s' -> return $ Skip (Right (Stream inner_step inner_s', s)) 1078 Done -> return $ Skip (Left s) 1079 1080-- | Create a 'Stream' of values from a 'Stream' of streamable things 1081flatten :: Monad m => (a -> m s) -> (s -> m (Step s b)) -> Stream m a -> Stream m b 1082{-# INLINE_FUSED flatten #-} 1083flatten mk istep (Stream ostep u) = Stream step (Left u) 1084 where 1085 {-# INLINE_INNER step #-} 1086 step (Left t) = do 1087 r <- ostep t 1088 case r of 1089 Yield a t' -> do 1090 s <- mk a 1091 s `seq` return (Skip (Right (s,t'))) 1092 Skip t' -> return $ Skip (Left t') 1093 Done -> return $ Done 1094 1095 1096 step (Right (s,t)) = do 1097 r <- istep s 1098 case r of 1099 Yield x s' -> return $ Yield x (Right (s',t)) 1100 Skip s' -> return $ Skip (Right (s',t)) 1101 Done -> return $ Skip (Left t) 1102 1103-- Unfolding 1104-- --------- 1105 1106-- | Unfold 1107unfoldr :: Monad m => (s -> Maybe (a, s)) -> s -> Stream m a 1108{-# INLINE_FUSED unfoldr #-} 1109unfoldr f = unfoldrM (return . f) 1110 1111-- | Unfold with a monadic function 1112unfoldrM :: Monad m => (s -> m (Maybe (a, s))) -> s -> Stream m a 1113{-# INLINE_FUSED unfoldrM #-} 1114unfoldrM f t = Stream step t 1115 where 1116 {-# INLINE_INNER step #-} 1117 step s = liftM (\r -> 1118 case r of 1119 Just (x, s') -> Yield x s' 1120 Nothing -> Done 1121 ) (f s) 1122 1123unfoldrN :: Monad m => Int -> (s -> Maybe (a, s)) -> s -> Stream m a 1124{-# INLINE_FUSED unfoldrN #-} 1125unfoldrN n f = unfoldrNM n (return . f) 1126 1127-- | Unfold at most @n@ elements with a monadic function. 1128unfoldrNM :: Monad m => Int -> (s -> m (Maybe (a, s))) -> s -> Stream m a 1129{-# INLINE_FUSED unfoldrNM #-} 1130unfoldrNM m f t = Stream step (t,m) 1131 where 1132 {-# INLINE_INNER step #-} 1133 step (s,n) | n <= 0 = return Done 1134 | otherwise = liftM (\r -> 1135 case r of 1136 Just (x,s') -> Yield x (s',n-1) 1137 Nothing -> Done 1138 ) (f s) 1139 1140-- | Unfold exactly @n@ elements 1141-- 1142-- @since 0.12.2.0 1143unfoldrExactN :: Monad m => Int -> (s -> (a, s)) -> s -> Stream m a 1144{-# INLINE_FUSED unfoldrExactN #-} 1145unfoldrExactN n f = unfoldrExactNM n (return . f) 1146 1147-- | Unfold exactly @n@ elements with a monadic function. 1148-- 1149-- @since 0.12.2.0 1150unfoldrExactNM :: Monad m => Int -> (s -> m (a, s)) -> s -> Stream m a 1151{-# INLINE_FUSED unfoldrExactNM #-} 1152unfoldrExactNM m f t = Stream step (t,m) 1153 where 1154 {-# INLINE_INNER step #-} 1155 step (s,n) | n <= 0 = return Done 1156 | otherwise = do (x,s') <- f s 1157 return $ Yield x (s',n-1) 1158 1159-- | /O(n)/ Apply monadic function \(\max(n - 1, 0)\) times to an initial value, 1160-- producing a stream of \(\max(n, 0)\) values. 1161iterateNM :: Monad m => Int -> (a -> m a) -> a -> Stream m a 1162{-# INLINE_FUSED iterateNM #-} 1163iterateNM n f x0 = Stream step (x0,n) 1164 where 1165 {-# INLINE_INNER step #-} 1166 step (x,i) | i <= 0 = return Done 1167 | i == n = return $ Yield x (x,i-1) 1168 | otherwise = do a <- f x 1169 return $ Yield a (a,i-1) 1170 1171-- | /O(n)/ Apply function \(\max(n - 1, 0)\) times to an initial value, 1172-- producing a stream of \(\max(n, 0)\) values. 1173iterateN :: Monad m => Int -> (a -> a) -> a -> Stream m a 1174{-# INLINE_FUSED iterateN #-} 1175iterateN n f x0 = iterateNM n (return . f) x0 1176 1177-- Scans 1178-- ----- 1179 1180-- | Prefix scan 1181prescanl :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a 1182{-# INLINE prescanl #-} 1183prescanl f = prescanlM (\a b -> return (f a b)) 1184 1185-- | Prefix scan with a monadic operator 1186prescanlM :: Monad m => (a -> b -> m a) -> a -> Stream m b -> Stream m a 1187{-# INLINE_FUSED prescanlM #-} 1188prescanlM f w (Stream step t) = Stream step' (t,w) 1189 where 1190 {-# INLINE_INNER step' #-} 1191 step' (s,x) = do 1192 r <- step s 1193 case r of 1194 Yield y s' -> do 1195 z <- f x y 1196 return $ Yield x (s', z) 1197 Skip s' -> return $ Skip (s', x) 1198 Done -> return Done 1199 1200-- | Prefix scan with strict accumulator 1201prescanl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a 1202{-# INLINE prescanl' #-} 1203prescanl' f = prescanlM' (\a b -> return (f a b)) 1204 1205-- | Prefix scan with strict accumulator and a monadic operator 1206prescanlM' :: Monad m => (a -> b -> m a) -> a -> Stream m b -> Stream m a 1207{-# INLINE_FUSED prescanlM' #-} 1208prescanlM' f w (Stream step t) = Stream step' (t,w) 1209 where 1210 {-# INLINE_INNER step' #-} 1211 step' (s,x) = x `seq` 1212 do 1213 r <- step s 1214 case r of 1215 Yield y s' -> do 1216 z <- f x y 1217 return $ Yield x (s', z) 1218 Skip s' -> return $ Skip (s', x) 1219 Done -> return Done 1220 1221-- | Suffix scan 1222postscanl :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a 1223{-# INLINE postscanl #-} 1224postscanl f = postscanlM (\a b -> return (f a b)) 1225 1226-- | Suffix scan with a monadic operator 1227postscanlM :: Monad m => (a -> b -> m a) -> a -> Stream m b -> Stream m a 1228{-# INLINE_FUSED postscanlM #-} 1229postscanlM f w (Stream step t) = Stream step' (t,w) 1230 where 1231 {-# INLINE_INNER step' #-} 1232 step' (s,x) = do 1233 r <- step s 1234 case r of 1235 Yield y s' -> do 1236 z <- f x y 1237 return $ Yield z (s',z) 1238 Skip s' -> return $ Skip (s',x) 1239 Done -> return Done 1240 1241-- | Suffix scan with strict accumulator 1242postscanl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a 1243{-# INLINE postscanl' #-} 1244postscanl' f = postscanlM' (\a b -> return (f a b)) 1245 1246-- | Suffix scan with strict acccumulator and a monadic operator 1247postscanlM' :: Monad m => (a -> b -> m a) -> a -> Stream m b -> Stream m a 1248{-# INLINE_FUSED postscanlM' #-} 1249postscanlM' f w (Stream step t) = w `seq` Stream step' (t,w) 1250 where 1251 {-# INLINE_INNER step' #-} 1252 step' (s,x) = x `seq` 1253 do 1254 r <- step s 1255 case r of 1256 Yield y s' -> do 1257 z <- f x y 1258 z `seq` return (Yield z (s',z)) 1259 Skip s' -> return $ Skip (s',x) 1260 Done -> return Done 1261 1262-- | Haskell-style scan 1263scanl :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a 1264{-# INLINE scanl #-} 1265scanl f = scanlM (\a b -> return (f a b)) 1266 1267-- | Haskell-style scan with a monadic operator 1268scanlM :: Monad m => (a -> b -> m a) -> a -> Stream m b -> Stream m a 1269{-# INLINE scanlM #-} 1270scanlM f z s = z `cons` postscanlM f z s 1271 1272-- | Haskell-style scan with strict accumulator 1273scanl' :: Monad m => (a -> b -> a) -> a -> Stream m b -> Stream m a 1274{-# INLINE scanl' #-} 1275scanl' f = scanlM' (\a b -> return (f a b)) 1276 1277-- | Haskell-style scan with strict accumulator and a monadic operator 1278scanlM' :: Monad m => (a -> b -> m a) -> a -> Stream m b -> Stream m a 1279{-# INLINE scanlM' #-} 1280scanlM' f z s = z `seq` (z `cons` postscanlM f z s) 1281 1282-- | Scan over a non-empty 'Stream' 1283scanl1 :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a 1284{-# INLINE scanl1 #-} 1285scanl1 f = scanl1M (\x y -> return (f x y)) 1286 1287-- | Scan over a non-empty 'Stream' with a monadic operator 1288scanl1M :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a 1289{-# INLINE_FUSED scanl1M #-} 1290scanl1M f (Stream step t) = Stream step' (t, Nothing) 1291 where 1292 {-# INLINE_INNER step' #-} 1293 step' (s, Nothing) = do 1294 r <- step s 1295 case r of 1296 Yield x s' -> return $ Yield x (s', Just x) 1297 Skip s' -> return $ Skip (s', Nothing) 1298 Done -> EMPTY_STREAM "scanl1M" 1299 1300 step' (s, Just x) = do 1301 r <- step s 1302 case r of 1303 Yield y s' -> do 1304 z <- f x y 1305 return $ Yield z (s', Just z) 1306 Skip s' -> return $ Skip (s', Just x) 1307 Done -> return Done 1308 1309-- | Scan over a non-empty 'Stream' with a strict accumulator 1310scanl1' :: Monad m => (a -> a -> a) -> Stream m a -> Stream m a 1311{-# INLINE scanl1' #-} 1312scanl1' f = scanl1M' (\x y -> return (f x y)) 1313 1314-- | Scan over a non-empty 'Stream' with a strict accumulator and a monadic 1315-- operator 1316scanl1M' :: Monad m => (a -> a -> m a) -> Stream m a -> Stream m a 1317{-# INLINE_FUSED scanl1M' #-} 1318scanl1M' f (Stream step t) = Stream step' (t, Nothing) 1319 where 1320 {-# INLINE_INNER step' #-} 1321 step' (s, Nothing) = do 1322 r <- step s 1323 case r of 1324 Yield x s' -> x `seq` return (Yield x (s', Just x)) 1325 Skip s' -> return $ Skip (s', Nothing) 1326 Done -> EMPTY_STREAM "scanl1M" 1327 1328 step' (s, Just x) = x `seq` 1329 do 1330 r <- step s 1331 case r of 1332 Yield y s' -> do 1333 z <- f x y 1334 z `seq` return (Yield z (s', Just z)) 1335 Skip s' -> return $ Skip (s', Just x) 1336 Done -> return Done 1337 1338-- Enumerations 1339-- ------------ 1340 1341-- The Enum class is broken for this, there just doesn't seem to be a 1342-- way to implement this generically. We have to specialise for as many types 1343-- as we can but this doesn't help in polymorphic loops. 1344 1345-- | Yield a 'Stream' of the given length containing the values @x@, @x+y@, 1346-- @x+y+y@ etc. 1347enumFromStepN :: (Num a, Monad m) => a -> a -> Int -> Stream m a 1348{-# INLINE_FUSED enumFromStepN #-} 1349enumFromStepN x y n = x `seq` y `seq` n `seq` Stream step (x,n) 1350 where 1351 {-# INLINE_INNER step #-} 1352 step (w,m) | m > 0 = return $ Yield w (w+y,m-1) 1353 | otherwise = return $ Done 1354 1355-- | Enumerate values 1356-- 1357-- /WARNING:/ This operation can be very inefficient. If at all possible, use 1358-- 'enumFromStepN' instead. 1359enumFromTo :: (Enum a, Monad m) => a -> a -> Stream m a 1360{-# INLINE_FUSED enumFromTo #-} 1361enumFromTo x y = fromList [x .. y] 1362 1363-- NOTE: We use (x+1) instead of (succ x) below because the latter checks for 1364-- overflow which can't happen here. 1365 1366-- FIXME: add "too large" test for Int 1367enumFromTo_small :: (Integral a, Monad m) => a -> a -> Stream m a 1368{-# INLINE_FUSED enumFromTo_small #-} 1369enumFromTo_small x y = x `seq` y `seq` Stream step (Just x) 1370 where 1371 {-# INLINE_INNER step #-} 1372 step Nothing = return $ Done 1373 step (Just z) | z == y = return $ Yield z Nothing 1374 | z < y = return $ Yield z (Just (z+1)) 1375 | otherwise = return $ Done 1376 1377{-# RULES 1378 1379"enumFromTo<Int8> [Stream]" 1380 enumFromTo = enumFromTo_small :: Monad m => Int8 -> Int8 -> Stream m Int8 1381 1382"enumFromTo<Int16> [Stream]" 1383 enumFromTo = enumFromTo_small :: Monad m => Int16 -> Int16 -> Stream m Int16 1384 1385"enumFromTo<Word8> [Stream]" 1386 enumFromTo = enumFromTo_small :: Monad m => Word8 -> Word8 -> Stream m Word8 1387 1388"enumFromTo<Word16> [Stream]" 1389 enumFromTo = enumFromTo_small :: Monad m => Word16 -> Word16 -> Stream m Word16 #-} 1390 1391 1392#if WORD_SIZE_IN_BITS > 32 1393 1394{-# RULES 1395 1396"enumFromTo<Int32> [Stream]" 1397 enumFromTo = enumFromTo_small :: Monad m => Int32 -> Int32 -> Stream m Int32 1398 1399"enumFromTo<Word32> [Stream]" 1400 enumFromTo = enumFromTo_small :: Monad m => Word32 -> Word32 -> Stream m Word32 #-} 1401 1402 1403#endif 1404 1405-- NOTE: We could implement a generic "too large" test: 1406-- 1407-- len x y | x > y = 0 1408-- | n > 0 && n <= fromIntegral (maxBound :: Int) = fromIntegral n 1409-- | otherwise = error 1410-- where 1411-- n = y-x+1 1412-- 1413-- Alas, GHC won't eliminate unnecessary comparisons (such as n >= 0 for 1414-- unsigned types). See http://hackage.haskell.org/trac/ghc/ticket/3744 1415-- 1416 1417enumFromTo_int :: forall m. Monad m => Int -> Int -> Stream m Int 1418{-# INLINE_FUSED enumFromTo_int #-} 1419enumFromTo_int x y = x `seq` y `seq` Stream step (Just x) 1420 where 1421 -- {-# INLINE [0] len #-} 1422 -- len :: Int -> Int -> Int 1423 -- len u v | u > v = 0 1424 -- | otherwise = BOUNDS_CHECK(check) "enumFromTo" "vector too large" 1425 -- (n > 0) 1426 -- $ n 1427 -- where 1428 -- n = v-u+1 1429 1430 {-# INLINE_INNER step #-} 1431 step Nothing = return $ Done 1432 step (Just z) | z == y = return $ Yield z Nothing 1433 | z < y = return $ Yield z (Just (z+1)) 1434 | otherwise = return $ Done 1435 1436 1437enumFromTo_intlike :: (Integral a, Monad m) => a -> a -> Stream m a 1438{-# INLINE_FUSED enumFromTo_intlike #-} 1439enumFromTo_intlike x y = x `seq` y `seq` Stream step (Just x) 1440 where 1441 {-# INLINE_INNER step #-} 1442 step Nothing = return $ Done 1443 step (Just z) | z == y = return $ Yield z Nothing 1444 | z < y = return $ Yield z (Just (z+1)) 1445 | otherwise = return $ Done 1446 1447{-# RULES 1448 1449"enumFromTo<Int> [Stream]" 1450 enumFromTo = enumFromTo_int :: Monad m => Int -> Int -> Stream m Int 1451 1452#if WORD_SIZE_IN_BITS > 32 1453 1454"enumFromTo<Int64> [Stream]" 1455 enumFromTo = enumFromTo_intlike :: Monad m => Int64 -> Int64 -> Stream m Int64 #-} 1456 1457#else 1458 1459"enumFromTo<Int32> [Stream]" 1460 enumFromTo = enumFromTo_intlike :: Monad m => Int32 -> Int32 -> Stream m Int32 #-} 1461 1462#endif 1463 1464enumFromTo_big_word :: (Integral a, Monad m) => a -> a -> Stream m a 1465{-# INLINE_FUSED enumFromTo_big_word #-} 1466enumFromTo_big_word x y = x `seq` y `seq` Stream step (Just x) 1467 where 1468 {-# INLINE_INNER step #-} 1469 step Nothing = return $ Done 1470 step (Just z) | z == y = return $ Yield z Nothing 1471 | z < y = return $ Yield z (Just (z+1)) 1472 | otherwise = return $ Done 1473 1474{-# RULES 1475 1476"enumFromTo<Word> [Stream]" 1477 enumFromTo = enumFromTo_big_word :: Monad m => Word -> Word -> Stream m Word 1478 1479"enumFromTo<Word64> [Stream]" 1480 enumFromTo = enumFromTo_big_word 1481 :: Monad m => Word64 -> Word64 -> Stream m Word64 1482 1483#if WORD_SIZE_IN_BITS == 32 1484 1485"enumFromTo<Word32> [Stream]" 1486 enumFromTo = enumFromTo_big_word 1487 :: Monad m => Word32 -> Word32 -> Stream m Word32 1488 1489#endif 1490 1491"enumFromTo<Integer> [Stream]" 1492 enumFromTo = enumFromTo_big_word 1493 :: Monad m => Integer -> Integer -> Stream m Integer #-} 1494 1495 1496 1497#if WORD_SIZE_IN_BITS > 32 1498 1499-- FIXME: the "too large" test is totally wrong 1500enumFromTo_big_int :: (Integral a, Monad m) => a -> a -> Stream m a 1501{-# INLINE_FUSED enumFromTo_big_int #-} 1502enumFromTo_big_int x y = x `seq` y `seq` Stream step (Just x) 1503 where 1504 {-# INLINE_INNER step #-} 1505 step Nothing = return $ Done 1506 step (Just z) | z == y = return $ Yield z Nothing 1507 | z < y = return $ Yield z (Just (z+1)) 1508 | otherwise = return $ Done 1509 1510{-# RULES 1511 1512"enumFromTo<Int64> [Stream]" 1513 enumFromTo = enumFromTo_big_int :: Monad m => Int64 -> Int64 -> Stream m Int64 #-} 1514 1515 1516 1517#endif 1518 1519enumFromTo_char :: Monad m => Char -> Char -> Stream m Char 1520{-# INLINE_FUSED enumFromTo_char #-} 1521enumFromTo_char x y = x `seq` y `seq` Stream step xn 1522 where 1523 xn = ord x 1524 yn = ord y 1525 1526 {-# INLINE_INNER step #-} 1527 step zn | zn <= yn = return $ Yield (unsafeChr zn) (zn+1) 1528 | otherwise = return $ Done 1529 1530{-# RULES 1531 1532"enumFromTo<Char> [Stream]" 1533 enumFromTo = enumFromTo_char #-} 1534 1535 1536 1537------------------------------------------------------------------------ 1538 1539-- Specialise enumFromTo for Float and Double. 1540-- Also, try to do something about pairs? 1541 1542enumFromTo_double :: (Monad m, Ord a, RealFrac a) => a -> a -> Stream m a 1543{-# INLINE_FUSED enumFromTo_double #-} 1544enumFromTo_double n m = n `seq` m `seq` Stream step ini 1545 where 1546 lim = m + 1/2 -- important to float out 1547 1548-- GHC changed definition of Enum for Double in GHC8.6 so we have to 1549-- accomodate both definitions in order to preserve validity of 1550-- rewrite rule 1551-- 1552-- ISSUE: https://gitlab.haskell.org/ghc/ghc/issues/15081 1553-- COMMIT: https://gitlab.haskell.org/ghc/ghc/commit/4ffaf4b67773af4c72d92bb8b6c87b1a7d34ac0f 1554#if MIN_VERSION_base(4,12,0) 1555 ini = 0 1556 step x | x' <= lim = return $ Yield x' (x+1) 1557 | otherwise = return $ Done 1558 where 1559 x' = x + n 1560#else 1561 ini = n 1562 step x | x <= lim = return $ Yield x (x+1) 1563 | otherwise = return $ Done 1564#endif 1565 1566{-# RULES 1567 1568"enumFromTo<Double> [Stream]" 1569 enumFromTo = enumFromTo_double :: Monad m => Double -> Double -> Stream m Double 1570 1571"enumFromTo<Float> [Stream]" 1572 enumFromTo = enumFromTo_double :: Monad m => Float -> Float -> Stream m Float #-} 1573 1574 1575 1576------------------------------------------------------------------------ 1577 1578-- | Enumerate values with a given step. 1579-- 1580-- /WARNING:/ This operation is very inefficient. If at all possible, use 1581-- 'enumFromStepN' instead. 1582enumFromThenTo :: (Enum a, Monad m) => a -> a -> a -> Stream m a 1583{-# INLINE_FUSED enumFromThenTo #-} 1584enumFromThenTo x y z = fromList [x, y .. z] 1585 1586-- FIXME: Specialise enumFromThenTo. 1587 1588-- Conversions 1589-- ----------- 1590 1591-- | Convert a 'Stream' to a list 1592toList :: Monad m => Stream m a -> m [a] 1593{-# INLINE toList #-} 1594toList = foldr (:) [] 1595 1596-- | Convert a list to a 'Stream' 1597fromList :: Monad m => [a] -> Stream m a 1598{-# INLINE fromList #-} 1599fromList zs = Stream step zs 1600 where 1601 step (x:xs) = return (Yield x xs) 1602 step [] = return Done 1603 1604-- | Convert the first @n@ elements of a list to a 'Bundle' 1605fromListN :: Monad m => Int -> [a] -> Stream m a 1606{-# INLINE_FUSED fromListN #-} 1607fromListN m zs = Stream step (zs,m) 1608 where 1609 {-# INLINE_INNER step #-} 1610 step (_, n) | n <= 0 = return Done 1611 step (x:xs,n) = return (Yield x (xs,n-1)) 1612 step ([],_) = return Done 1613 1614{- 1615fromVector :: (Monad m, Vector v a) => v a -> Stream m a 1616{-# INLINE_FUSED fromVector #-} 1617fromVector v = v `seq` n `seq` Stream (Unf step 0) 1618 (Unf vstep True) 1619 (Just v) 1620 (Exact n) 1621 where 1622 n = basicLength v 1623 1624 {-# INLINE step #-} 1625 step i | i >= n = return Done 1626 | otherwise = case basicUnsafeIndexM v i of 1627 Box x -> return $ Yield x (i+1) 1628 1629 1630 {-# INLINE vstep #-} 1631 vstep True = return (Yield (Chunk (basicLength v) (\mv -> basicUnsafeCopy mv v)) False) 1632 vstep False = return Done 1633 1634fromVectors :: forall m a. (Monad m, Vector v a) => [v a] -> Stream m a 1635{-# INLINE_FUSED fromVectors #-} 1636fromVectors vs = Stream (Unf pstep (Left vs)) 1637 (Unf vstep vs) 1638 Nothing 1639 (Exact n) 1640 where 1641 n = List.foldl' (\k v -> k + basicLength v) 0 vs 1642 1643 pstep (Left []) = return Done 1644 pstep (Left (v:vs)) = basicLength v `seq` return (Skip (Right (v,0,vs))) 1645 1646 pstep (Right (v,i,vs)) 1647 | i >= basicLength v = return $ Skip (Left vs) 1648 | otherwise = case basicUnsafeIndexM v i of 1649 Box x -> return $ Yield x (Right (v,i+1,vs)) 1650 1651 -- FIXME: work around bug in GHC 7.6.1 1652 vstep :: [v a] -> m (Step [v a] (Chunk v a)) 1653 vstep [] = return Done 1654 vstep (v:vs) = return $ Yield (Chunk (basicLength v) 1655 (\mv -> INTERNAL_CHECK(check) "concatVectors" "length mismatch" 1656 (M.basicLength mv == basicLength v) 1657 $ basicUnsafeCopy mv v)) vs 1658 1659 1660concatVectors :: (Monad m, Vector v a) => Stream m (v a) -> Stream m a 1661{-# INLINE_FUSED concatVectors #-} 1662concatVectors (Stream step s} 1663 = Stream (Unf pstep (Left s)) 1664 (Unf vstep s) 1665 Nothing 1666 Unknown 1667 where 1668 pstep (Left s) = do 1669 r <- step s 1670 case r of 1671 Yield v s' -> basicLength v `seq` return (Skip (Right (v,0,s'))) 1672 Skip s' -> return (Skip (Left s')) 1673 Done -> return Done 1674 1675 pstep (Right (v,i,s)) 1676 | i >= basicLength v = return (Skip (Left s)) 1677 | otherwise = case basicUnsafeIndexM v i of 1678 Box x -> return (Yield x (Right (v,i+1,s))) 1679 1680 1681 vstep s = do 1682 r <- step s 1683 case r of 1684 Yield v s' -> return (Yield (Chunk (basicLength v) 1685 (\mv -> INTERNAL_CHECK(check) "concatVectors" "length mismatch" 1686 (M.basicLength mv == basicLength v) 1687 $ basicUnsafeCopy mv v)) s') 1688 Skip s' -> return (Skip s') 1689 Done -> return Done 1690 1691reVector :: Monad m => Stream m a -> Stream m a 1692{-# INLINE_FUSED reVector #-} 1693reVector (Stream step s, sSize = n} = Stream step s n 1694 1695{-# RULES 1696 1697"reVector [Vector]" 1698 reVector = id 1699 1700"reVector/reVector [Vector]" forall s. 1701 reVector (reVector s) = s #-} 1702 1703 1704-} 1705 1706