1{-# LANGUAGE CPP, FlexibleInstances, Rank2Types, BangPatterns #-} 2 3-- | 4-- Module : Data.Vector.Fusion.Bundle 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-- Bundles for stream fusion 13-- 14 15module Data.Vector.Fusion.Bundle ( 16 -- * Types 17 Step(..), Chunk(..), Bundle, MBundle, 18 19 -- * In-place markers 20 inplace, 21 22 -- * Size hints 23 size, sized, 24 25 -- * Length information 26 length, null, 27 28 -- * Construction 29 empty, singleton, cons, snoc, replicate, generate, (++), 30 31 -- * Accessing individual elements 32 head, last, (!!), (!?), 33 34 -- * Substreams 35 slice, init, tail, take, drop, 36 37 -- * Mapping 38 map, concatMap, flatten, unbox, 39 40 -- * Zipping 41 indexed, indexedR, 42 zipWith, zipWith3, zipWith4, zipWith5, zipWith6, 43 zip, zip3, zip4, zip5, zip6, 44 45 -- * Filtering 46 filter, takeWhile, dropWhile, 47 48 -- * Searching 49 elem, notElem, find, findIndex, 50 51 -- * Folding 52 foldl, foldl1, foldl', foldl1', foldr, foldr1, 53 54 -- * Specialised folds 55 and, or, 56 57 -- * Unfolding 58 unfoldr, unfoldrN, unfoldrExactN, iterateN, 59 60 -- * Scans 61 prescanl, prescanl', 62 postscanl, postscanl', 63 scanl, scanl', 64 scanl1, scanl1', 65 66 -- * Enumerations 67 enumFromStepN, enumFromTo, enumFromThenTo, 68 69 -- * Conversions 70 toList, fromList, fromListN, unsafeFromList, lift, 71 fromVector, reVector, fromVectors, concatVectors, 72 73 -- * Monadic combinators 74 mapM, mapM_, zipWithM, zipWithM_, filterM, mapMaybeM, foldM, fold1M, foldM', fold1M', 75 76 eq, cmp, eqBy, cmpBy 77) where 78 79import Data.Vector.Generic.Base ( Vector ) 80import Data.Vector.Fusion.Bundle.Size 81import Data.Vector.Fusion.Util 82import Data.Vector.Fusion.Stream.Monadic ( Stream(..), Step(..) ) 83import Data.Vector.Fusion.Bundle.Monadic ( Chunk(..), lift ) 84import qualified Data.Vector.Fusion.Bundle.Monadic as M 85import qualified Data.Vector.Fusion.Stream.Monadic as S 86 87import Prelude hiding ( length, null, 88 replicate, (++), 89 head, last, (!!), 90 init, tail, take, drop, 91 map, concatMap, 92 zipWith, zipWith3, zip, zip3, 93 filter, takeWhile, dropWhile, 94 elem, notElem, 95 foldl, foldl1, foldr, foldr1, 96 and, or, 97 scanl, scanl1, 98 enumFromTo, enumFromThenTo, 99 mapM, mapM_ ) 100 101#if MIN_VERSION_base(4,9,0) 102import Data.Functor.Classes (Eq1 (..), Ord1 (..)) 103#endif 104 105import GHC.Base ( build ) 106 107-- Data.Vector.Internal.Check is unused 108#define NOT_VECTOR_MODULE 109#include "vector.h" 110 111-- | The type of pure streams 112type Bundle = M.Bundle Id 113 114-- | Alternative name for monadic streams 115type MBundle = M.Bundle 116 117inplace :: (forall m. Monad m => S.Stream m a -> S.Stream m b) 118 -> (Size -> Size) -> Bundle v a -> Bundle v b 119{-# INLINE_FUSED inplace #-} 120inplace f g b = b `seq` M.fromStream (f (M.elements b)) (g (M.size b)) 121 122{-# RULES 123 124"inplace/inplace [Vector]" 125 forall (f1 :: forall m. Monad m => S.Stream m a -> S.Stream m a) 126 (f2 :: forall m. Monad m => S.Stream m a -> S.Stream m a) 127 g1 g2 s. 128 inplace f1 g1 (inplace f2 g2 s) = inplace (f1 . f2) (g1 . g2) s #-} 129 130 131-- | 'Size' hint of a 'Bundle' 132size :: Bundle v a -> Size 133{-# INLINE size #-} 134size = M.size 135 136-- | Attach a 'Size' hint to a 'Bundle' 137sized :: Bundle v a -> Size -> Bundle v a 138{-# INLINE sized #-} 139sized = M.sized 140 141-- Length 142-- ------ 143 144-- | Length of a 'Bundle' 145length :: Bundle v a -> Int 146{-# INLINE length #-} 147length = unId . M.length 148 149-- | Check if a 'Bundle' is empty 150null :: Bundle v a -> Bool 151{-# INLINE null #-} 152null = unId . M.null 153 154-- Construction 155-- ------------ 156 157-- | Empty 'Bundle' 158empty :: Bundle v a 159{-# INLINE empty #-} 160empty = M.empty 161 162-- | Singleton 'Bundle' 163singleton :: a -> Bundle v a 164{-# INLINE singleton #-} 165singleton = M.singleton 166 167-- | Replicate a value to a given length 168replicate :: Int -> a -> Bundle v a 169{-# INLINE replicate #-} 170replicate = M.replicate 171 172-- | Generate a stream from its indices 173generate :: Int -> (Int -> a) -> Bundle v a 174{-# INLINE generate #-} 175generate = M.generate 176 177-- | Prepend an element 178cons :: a -> Bundle v a -> Bundle v a 179{-# INLINE cons #-} 180cons = M.cons 181 182-- | Append an element 183snoc :: Bundle v a -> a -> Bundle v a 184{-# INLINE snoc #-} 185snoc = M.snoc 186 187infixr 5 ++ 188-- | Concatenate two 'Bundle's 189(++) :: Bundle v a -> Bundle v a -> Bundle v a 190{-# INLINE (++) #-} 191(++) = (M.++) 192 193-- Accessing elements 194-- ------------------ 195 196-- | First element of the 'Bundle' or error if empty 197head :: Bundle v a -> a 198{-# INLINE head #-} 199head = unId . M.head 200 201-- | Last element of the 'Bundle' or error if empty 202last :: Bundle v a -> a 203{-# INLINE last #-} 204last = unId . M.last 205 206infixl 9 !! 207-- | Element at the given position 208(!!) :: Bundle v a -> Int -> a 209{-# INLINE (!!) #-} 210s !! i = unId (s M.!! i) 211 212infixl 9 !? 213-- | Element at the given position or 'Nothing' if out of bounds 214(!?) :: Bundle v a -> Int -> Maybe a 215{-# INLINE (!?) #-} 216s !? i = unId (s M.!? i) 217 218-- Substreams 219-- ---------- 220 221-- | Extract a substream of the given length starting at the given position. 222slice :: Int -- ^ starting index 223 -> Int -- ^ length 224 -> Bundle v a 225 -> Bundle v a 226{-# INLINE slice #-} 227slice = M.slice 228 229-- | All but the last element 230init :: Bundle v a -> Bundle v a 231{-# INLINE init #-} 232init = M.init 233 234-- | All but the first element 235tail :: Bundle v a -> Bundle v a 236{-# INLINE tail #-} 237tail = M.tail 238 239-- | The first @n@ elements 240take :: Int -> Bundle v a -> Bundle v a 241{-# INLINE take #-} 242take = M.take 243 244-- | All but the first @n@ elements 245drop :: Int -> Bundle v a -> Bundle v a 246{-# INLINE drop #-} 247drop = M.drop 248 249-- Mapping 250-- --------------- 251 252-- | Map a function over a 'Bundle' 253map :: (a -> b) -> Bundle v a -> Bundle v b 254{-# INLINE map #-} 255map = M.map 256 257unbox :: Bundle v (Box a) -> Bundle v a 258{-# INLINE unbox #-} 259unbox = M.unbox 260 261concatMap :: (a -> Bundle v b) -> Bundle v a -> Bundle v b 262{-# INLINE concatMap #-} 263concatMap = M.concatMap 264 265-- Zipping 266-- ------- 267 268-- | Pair each element in a 'Bundle' with its index 269indexed :: Bundle v a -> Bundle v (Int,a) 270{-# INLINE indexed #-} 271indexed = M.indexed 272 273-- | Pair each element in a 'Bundle' with its index, starting from the right 274-- and counting down 275indexedR :: Int -> Bundle v a -> Bundle v (Int,a) 276{-# INLINE_FUSED indexedR #-} 277indexedR = M.indexedR 278 279-- | Zip two 'Bundle's with the given function 280zipWith :: (a -> b -> c) -> Bundle v a -> Bundle v b -> Bundle v c 281{-# INLINE zipWith #-} 282zipWith = M.zipWith 283 284-- | Zip three 'Bundle's with the given function 285zipWith3 :: (a -> b -> c -> d) -> Bundle v a -> Bundle v b -> Bundle v c -> Bundle v d 286{-# INLINE zipWith3 #-} 287zipWith3 = M.zipWith3 288 289zipWith4 :: (a -> b -> c -> d -> e) 290 -> Bundle v a -> Bundle v b -> Bundle v c -> Bundle v d 291 -> Bundle v e 292{-# INLINE zipWith4 #-} 293zipWith4 = M.zipWith4 294 295zipWith5 :: (a -> b -> c -> d -> e -> f) 296 -> Bundle v a -> Bundle v b -> Bundle v c -> Bundle v d 297 -> Bundle v e -> Bundle v f 298{-# INLINE zipWith5 #-} 299zipWith5 = M.zipWith5 300 301zipWith6 :: (a -> b -> c -> d -> e -> f -> g) 302 -> Bundle v a -> Bundle v b -> Bundle v c -> Bundle v d 303 -> Bundle v e -> Bundle v f -> Bundle v g 304{-# INLINE zipWith6 #-} 305zipWith6 = M.zipWith6 306 307zip :: Bundle v a -> Bundle v b -> Bundle v (a,b) 308{-# INLINE zip #-} 309zip = M.zip 310 311zip3 :: Bundle v a -> Bundle v b -> Bundle v c -> Bundle v (a,b,c) 312{-# INLINE zip3 #-} 313zip3 = M.zip3 314 315zip4 :: Bundle v a -> Bundle v b -> Bundle v c -> Bundle v d 316 -> Bundle v (a,b,c,d) 317{-# INLINE zip4 #-} 318zip4 = M.zip4 319 320zip5 :: Bundle v a -> Bundle v b -> Bundle v c -> Bundle v d 321 -> Bundle v e -> Bundle v (a,b,c,d,e) 322{-# INLINE zip5 #-} 323zip5 = M.zip5 324 325zip6 :: Bundle v a -> Bundle v b -> Bundle v c -> Bundle v d 326 -> Bundle v e -> Bundle v f -> Bundle v (a,b,c,d,e,f) 327{-# INLINE zip6 #-} 328zip6 = M.zip6 329 330-- Filtering 331-- --------- 332 333-- | Drop elements which do not satisfy the predicate 334filter :: (a -> Bool) -> Bundle v a -> Bundle v a 335{-# INLINE filter #-} 336filter = M.filter 337 338-- | Longest prefix of elements that satisfy the predicate 339takeWhile :: (a -> Bool) -> Bundle v a -> Bundle v a 340{-# INLINE takeWhile #-} 341takeWhile = M.takeWhile 342 343-- | Drop the longest prefix of elements that satisfy the predicate 344dropWhile :: (a -> Bool) -> Bundle v a -> Bundle v a 345{-# INLINE dropWhile #-} 346dropWhile = M.dropWhile 347 348-- Searching 349-- --------- 350 351infix 4 `elem` 352-- | Check whether the 'Bundle' contains an element 353elem :: Eq a => a -> Bundle v a -> Bool 354{-# INLINE elem #-} 355elem x = unId . M.elem x 356 357infix 4 `notElem` 358-- | Inverse of `elem` 359notElem :: Eq a => a -> Bundle v a -> Bool 360{-# INLINE notElem #-} 361notElem x = unId . M.notElem x 362 363-- | Yield 'Just' the first element matching the predicate or 'Nothing' if no 364-- such element exists. 365find :: (a -> Bool) -> Bundle v a -> Maybe a 366{-# INLINE find #-} 367find f = unId . M.find f 368 369-- | Yield 'Just' the index of the first element matching the predicate or 370-- 'Nothing' if no such element exists. 371findIndex :: (a -> Bool) -> Bundle v a -> Maybe Int 372{-# INLINE findIndex #-} 373findIndex f = unId . M.findIndex f 374 375-- Folding 376-- ------- 377 378-- | Left fold 379foldl :: (a -> b -> a) -> a -> Bundle v b -> a 380{-# INLINE foldl #-} 381foldl f z = unId . M.foldl f z 382 383-- | Left fold on non-empty 'Bundle's 384foldl1 :: (a -> a -> a) -> Bundle v a -> a 385{-# INLINE foldl1 #-} 386foldl1 f = unId . M.foldl1 f 387 388-- | Left fold with strict accumulator 389foldl' :: (a -> b -> a) -> a -> Bundle v b -> a 390{-# INLINE foldl' #-} 391foldl' f z = unId . M.foldl' f z 392 393-- | Left fold on non-empty 'Bundle's with strict accumulator 394foldl1' :: (a -> a -> a) -> Bundle v a -> a 395{-# INLINE foldl1' #-} 396foldl1' f = unId . M.foldl1' f 397 398-- | Right fold 399foldr :: (a -> b -> b) -> b -> Bundle v a -> b 400{-# INLINE foldr #-} 401foldr f z = unId . M.foldr f z 402 403-- | Right fold on non-empty 'Bundle's 404foldr1 :: (a -> a -> a) -> Bundle v a -> a 405{-# INLINE foldr1 #-} 406foldr1 f = unId . M.foldr1 f 407 408-- Specialised folds 409-- ----------------- 410 411and :: Bundle v Bool -> Bool 412{-# INLINE and #-} 413and = unId . M.and 414 415or :: Bundle v Bool -> Bool 416{-# INLINE or #-} 417or = unId . M.or 418 419-- Unfolding 420-- --------- 421 422-- | Unfold 423unfoldr :: (s -> Maybe (a, s)) -> s -> Bundle v a 424{-# INLINE unfoldr #-} 425unfoldr = M.unfoldr 426 427-- | Unfold at most @n@ elements 428unfoldrN :: Int -> (s -> Maybe (a, s)) -> s -> Bundle v a 429{-# INLINE unfoldrN #-} 430unfoldrN = M.unfoldrN 431 432-- | Unfold exactly @n@ elements 433-- 434-- @since 0.12.2.0 435unfoldrExactN :: Int -> (s -> (a, s)) -> s -> Bundle v a 436{-# INLINE unfoldrExactN #-} 437unfoldrExactN = M.unfoldrExactN 438 439-- | /O(n)/ Apply function \(\max(n - 1, 0)\) times to an initial value, producing a pure 440-- bundle of exact length \(\max(n, 0)\). Zeroth element will contain the initial value. 441iterateN :: Int -> (a -> a) -> a -> Bundle v a 442{-# INLINE iterateN #-} 443iterateN = M.iterateN 444 445-- Scans 446-- ----- 447 448-- | Prefix scan 449prescanl :: (a -> b -> a) -> a -> Bundle v b -> Bundle v a 450{-# INLINE prescanl #-} 451prescanl = M.prescanl 452 453-- | Prefix scan with strict accumulator 454prescanl' :: (a -> b -> a) -> a -> Bundle v b -> Bundle v a 455{-# INLINE prescanl' #-} 456prescanl' = M.prescanl' 457 458-- | Suffix scan 459postscanl :: (a -> b -> a) -> a -> Bundle v b -> Bundle v a 460{-# INLINE postscanl #-} 461postscanl = M.postscanl 462 463-- | Suffix scan with strict accumulator 464postscanl' :: (a -> b -> a) -> a -> Bundle v b -> Bundle v a 465{-# INLINE postscanl' #-} 466postscanl' = M.postscanl' 467 468-- | Haskell-style scan 469scanl :: (a -> b -> a) -> a -> Bundle v b -> Bundle v a 470{-# INLINE scanl #-} 471scanl = M.scanl 472 473-- | Haskell-style scan with strict accumulator 474scanl' :: (a -> b -> a) -> a -> Bundle v b -> Bundle v a 475{-# INLINE scanl' #-} 476scanl' = M.scanl' 477 478-- | Scan over a non-empty 'Bundle' 479scanl1 :: (a -> a -> a) -> Bundle v a -> Bundle v a 480{-# INLINE scanl1 #-} 481scanl1 = M.scanl1 482 483-- | Scan over a non-empty 'Bundle' with a strict accumulator 484scanl1' :: (a -> a -> a) -> Bundle v a -> Bundle v a 485{-# INLINE scanl1' #-} 486scanl1' = M.scanl1' 487 488 489-- Comparisons 490-- ----------- 491 492-- | Check if two 'Bundle's are equal 493eq :: (Eq a) => Bundle v a -> Bundle v a -> Bool 494{-# INLINE eq #-} 495eq = eqBy (==) 496 497eqBy :: (a -> b -> Bool) -> Bundle v a -> Bundle v b -> Bool 498{-# INLINE eqBy #-} 499eqBy e x y = unId (M.eqBy e x y) 500 501-- | Lexicographically compare two 'Bundle's 502cmp :: (Ord a) => Bundle v a -> Bundle v a -> Ordering 503{-# INLINE cmp #-} 504cmp = cmpBy compare 505 506cmpBy :: (a -> b -> Ordering) -> Bundle v a -> Bundle v b -> Ordering 507{-# INLINE cmpBy #-} 508cmpBy c x y = unId (M.cmpBy c x y) 509 510instance Eq a => Eq (M.Bundle Id v a) where 511 {-# INLINE (==) #-} 512 (==) = eq 513 514instance Ord a => Ord (M.Bundle Id v a) where 515 {-# INLINE compare #-} 516 compare = cmp 517 518#if MIN_VERSION_base(4,9,0) 519instance Eq1 (M.Bundle Id v) where 520 {-# INLINE liftEq #-} 521 liftEq = eqBy 522 523instance Ord1 (M.Bundle Id v) where 524 {-# INLINE liftCompare #-} 525 liftCompare = cmpBy 526#endif 527 528-- Monadic combinators 529-- ------------------- 530 531-- | Apply a monadic action to each element of the stream, producing a monadic 532-- stream of results 533mapM :: Monad m => (a -> m b) -> Bundle v a -> M.Bundle m v b 534{-# INLINE mapM #-} 535mapM f = M.mapM f . lift 536 537-- | Apply a monadic action to each element of the stream 538mapM_ :: Monad m => (a -> m b) -> Bundle v a -> m () 539{-# INLINE mapM_ #-} 540mapM_ f = M.mapM_ f . lift 541 542zipWithM :: Monad m => (a -> b -> m c) -> Bundle v a -> Bundle v b -> M.Bundle m v c 543{-# INLINE zipWithM #-} 544zipWithM f as bs = M.zipWithM f (lift as) (lift bs) 545 546zipWithM_ :: Monad m => (a -> b -> m c) -> Bundle v a -> Bundle v b -> m () 547{-# INLINE zipWithM_ #-} 548zipWithM_ f as bs = M.zipWithM_ f (lift as) (lift bs) 549 550-- | Yield a monadic stream of elements that satisfy the monadic predicate 551filterM :: Monad m => (a -> m Bool) -> Bundle v a -> M.Bundle m v a 552{-# INLINE filterM #-} 553filterM f = M.filterM f . lift 554 555-- | /O(n)/ Apply monadic function to each element of a bundle and 556-- discard elements returning Nothing. 557-- 558-- @since 0.12.2.0 559mapMaybeM :: Monad m => (a -> m (Maybe b)) -> Bundle v a -> M.Bundle m v b 560{-# INLINE mapMaybeM #-} 561mapMaybeM f = M.mapMaybeM f . lift 562 563-- | Monadic fold 564foldM :: Monad m => (a -> b -> m a) -> a -> Bundle v b -> m a 565{-# INLINE foldM #-} 566foldM m z = M.foldM m z . lift 567 568-- | Monadic fold over non-empty stream 569fold1M :: Monad m => (a -> a -> m a) -> Bundle v a -> m a 570{-# INLINE fold1M #-} 571fold1M m = M.fold1M m . lift 572 573-- | Monadic fold with strict accumulator 574foldM' :: Monad m => (a -> b -> m a) -> a -> Bundle v b -> m a 575{-# INLINE foldM' #-} 576foldM' m z = M.foldM' m z . lift 577 578-- | Monad fold over non-empty stream with strict accumulator 579fold1M' :: Monad m => (a -> a -> m a) -> Bundle v a -> m a 580{-# INLINE fold1M' #-} 581fold1M' m = M.fold1M' m . lift 582 583-- Enumerations 584-- ------------ 585 586-- | Yield a 'Bundle' of the given length containing the values @x@, @x+y@, 587-- @x+y+y@ etc. 588enumFromStepN :: Num a => a -> a -> Int -> Bundle v a 589{-# INLINE enumFromStepN #-} 590enumFromStepN = M.enumFromStepN 591 592-- | Enumerate values 593-- 594-- /WARNING:/ This operations can be very inefficient. If at all possible, use 595-- 'enumFromStepN' instead. 596enumFromTo :: Enum a => a -> a -> Bundle v a 597{-# INLINE enumFromTo #-} 598enumFromTo = M.enumFromTo 599 600-- | Enumerate values with a given step. 601-- 602-- /WARNING:/ This operations is very inefficient. If at all possible, use 603-- 'enumFromStepN' instead. 604enumFromThenTo :: Enum a => a -> a -> a -> Bundle v a 605{-# INLINE enumFromThenTo #-} 606enumFromThenTo = M.enumFromThenTo 607 608-- Conversions 609-- ----------- 610 611-- | Convert a 'Bundle' to a list 612toList :: Bundle v a -> [a] 613{-# INLINE toList #-} 614-- toList s = unId (M.toList s) 615toList s = build (\c n -> toListFB c n s) 616 617-- This supports foldr/build list fusion that GHC implements 618toListFB :: (a -> b -> b) -> b -> Bundle v a -> b 619{-# INLINE [0] toListFB #-} 620toListFB c n M.Bundle{M.sElems = Stream step t} = go t 621 where 622 go s = case unId (step s) of 623 Yield x s' -> x `c` go s' 624 Skip s' -> go s' 625 Done -> n 626 627-- | Create a 'Bundle' from a list 628fromList :: [a] -> Bundle v a 629{-# INLINE fromList #-} 630fromList = M.fromList 631 632-- | Create a 'Bundle' from the first @n@ elements of a list 633-- 634-- > fromListN n xs = fromList (take n xs) 635fromListN :: Int -> [a] -> Bundle v a 636{-# INLINE fromListN #-} 637fromListN = M.fromListN 638 639unsafeFromList :: Size -> [a] -> Bundle v a 640{-# INLINE unsafeFromList #-} 641unsafeFromList = M.unsafeFromList 642 643fromVector :: Vector v a => v a -> Bundle v a 644{-# INLINE fromVector #-} 645fromVector = M.fromVector 646 647reVector :: Bundle u a -> Bundle v a 648{-# INLINE reVector #-} 649reVector = M.reVector 650 651fromVectors :: Vector v a => [v a] -> Bundle v a 652{-# INLINE fromVectors #-} 653fromVectors = M.fromVectors 654 655concatVectors :: Vector v a => Bundle u (v a) -> Bundle v a 656{-# INLINE concatVectors #-} 657concatVectors = M.concatVectors 658 659-- | Create a 'Bundle' of values from a 'Bundle' of streamable things 660flatten :: (a -> s) -> (s -> Step s b) -> Size -> Bundle v a -> Bundle v b 661{-# INLINE_FUSED flatten #-} 662flatten mk istep sz = M.flatten (return . mk) (return . istep) sz . lift 663 664