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