{take|drop}WhileSlice
[darcs-mirrors/vector.git] / Data / Vector / IVector.hs
1 {-# LANGUAGE Rank2Types, MultiParamTypeClasses #-}
2 -- |
3 -- Module : Data.Vector.IVector
4 -- Copyright : (c) Roman Leshchinskiy 2008
5 -- License : BSD-style
6 --
7 -- Maintainer : rl@cse.unsw.edu.au
8 -- Stability : experimental
9 -- Portability : non-portable
10 --
11 -- Generic interface to pure vectors
12 --
13
14 #include "phases.h"
15
16 module Data.Vector.IVector (
17 -- * Immutable vectors
18 IVector,
19
20 -- * Length information
21 length,
22
23 -- * Construction
24 empty, singleton, cons, snoc, replicate, (++),
25
26 -- * Subvectors
27 slice, subvector, takeSlice, take, dropSlice, drop,
28
29 -- * Mapping and zipping
30 map, zipWith,
31
32 -- * Filtering
33 filter, takeWhileSlice, takeWhile, dropWhileSlice, dropWhile,
34
35 -- * Searching
36 elem, notElem, find, findIndex,
37
38 -- * Folding
39 foldl, foldl1, foldl', foldl1', foldr, foldr1,
40
41 -- * Conversion to/from lists
42 toList, fromList,
43
44 -- * Conversion to/from Streams
45 stream, unstream,
46
47 -- * MVector-based initialisation
48 create,
49
50 -- * Unsafe functions
51 unsafeSlice, unsafeIndex
52 ) where
53
54 import qualified Data.Vector.MVector as MVector
55 import Data.Vector.MVector ( MVector )
56
57 import qualified Data.Vector.Stream as Stream
58 import Data.Vector.Stream ( Stream )
59 import Data.Vector.Stream.Size
60
61 import Control.Exception ( assert )
62
63 import Prelude hiding ( length,
64 replicate, (++),
65 head, last,
66 init, tail, take, drop,
67 map, zipWith,
68 filter, takeWhile, dropWhile,
69 elem, notElem,
70 foldl, foldl1, foldr, foldr1 )
71
72 -- | Class of immutable vectors. Just like with 'MVector', the type of the
73 -- elements can be restricted by using GADTs.
74 --
75 class IVector v a where
76 -- | Construct a pure vector from a monadic initialiser.
77 create :: (forall mv m. MVector mv m a => m (mv m a)) -> v a
78
79 -- | Length of the vector
80 length :: v a -> Int
81
82 -- | Yield a part of the vector without copying it. No range checks!
83 unsafeSlice :: v a -> Int -> Int -> v a
84
85 -- | Apply the given function to the element at the given position. This
86 -- interface prevents us from being too lazy. Suppose we had
87 --
88 -- > unsafeIndex' :: v a -> Int -> a
89 --
90 -- instead. Now, if we wanted to copy a vector, we'd do something like
91 --
92 -- > copy mv v ... = ... unsafeWrite mv i (unsafeIndex' v i) ...
93 --
94 -- For lazy vectors, the indexing would not be evaluated which means that we
95 -- would retain a reference to the original vector in each element we write.
96 -- This would be bad!
97 --
98 -- With 'unsafeIndex', we can do
99 --
100 -- > copy mv v ... = ... unsafeIndex v i (unsafeWrite mv i) ...
101 --
102 -- which does not have this problem.
103 --
104 unsafeIndex :: v a -> Int -> (a -> b) -> b
105
106 -- | Convert a vector to a 'Stream'
107 stream :: IVector v a => v a -> Stream a
108 {-# INLINE_STREAM stream #-}
109 stream v = v `seq` (Stream.unfold get 0 `Stream.sized` Exact n)
110 where
111 n = length v
112
113 {-# INLINE get #-}
114 get i | i < n = unsafeIndex v i $ \x -> Just (x, i+1)
115 | otherwise = Nothing
116
117 -- | Create a vector from a 'Stream'
118 unstream :: IVector v a => Stream a -> v a
119 {-# INLINE_STREAM unstream #-}
120 unstream s = create (MVector.unstream s)
121
122 {-# RULES
123
124 "stream/unstream [Vector.IVector]" forall s.
125 stream (unstream s) = s
126
127 #-}
128
129 -- Construction
130 -- ------------
131
132 -- | Empty vector
133 empty :: IVector v a => v a
134 {-# INLINE empty #-}
135 empty = unstream Stream.empty
136
137 -- | Vector with exaclty one element
138 singleton :: IVector v a => a -> v a
139 {-# INLINE singleton #-}
140 singleton x = unstream (Stream.singleton x)
141
142 -- | Vector of the given length with the given value in each position
143 replicate :: IVector v a => Int -> a -> v a
144 {-# INLINE replicate #-}
145 replicate n = unstream . Stream.replicate n
146
147 -- | Prepend an element
148 cons :: IVector v a => a -> v a -> v a
149 {-# INLINE cons #-}
150 cons x = unstream . Stream.cons x . stream
151
152 -- | Append an element
153 snoc :: IVector v a => v a -> a -> v a
154 {-# INLINE snoc #-}
155 snoc v = unstream . Stream.snoc (stream v)
156
157 infixr 5 ++
158 -- | Concatenate two vectors
159 (++) :: IVector v a => v a -> v a -> v a
160 {-# INLINE (++) #-}
161 v ++ w = unstream (stream v Stream.++ stream w)
162
163 -- Subarrays
164 -- ---------
165
166 -- | Yield a part of the vector without copying it. Safer version of
167 -- 'unsafeSlice'.
168 slice :: IVector v a => v a -> Int -- ^ starting index
169 -> Int -- ^ length
170 -> v a
171 {-# INLINE slice #-}
172 slice v i n = assert (i >= 0 && n >= 0 && i+n <= length v)
173 $ unsafeSlice v i n
174
175 -- | Yield the first @n@ elements without copying.
176 takeSlice :: IVector v a => Int -> v a -> v a
177 {-# INLINE takeSlice #-}
178 takeSlice n v = slice v 0 n
179
180 -- | Copy the first @n@ elements to a new vector.
181 take :: IVector v a => Int -> v a -> v a
182 {-# INLINE take #-}
183 take n = unstream . Stream.take n . stream
184
185 -- | Yield all but the first @n@ elements without copying.
186 dropSlice :: IVector v a => Int -> v a -> v a
187 {-# INLINE dropSlice #-}
188 dropSlice n v = slice v n (length v - n)
189
190 -- | Copy all but the first @n@ elements to a new vectors.
191 drop :: IVector v a => Int -> v a -> v a
192 {-# INLINE drop #-}
193 drop n = unstream . Stream.drop n . stream
194
195 -- Mapping/zipping
196 -- ---------------
197
198 -- | Map a function over a vector
199 map :: (IVector v a, IVector v b) => (a -> b) -> v a -> v b
200 {-# INLINE map #-}
201 map f = unstream . Stream.map f . stream
202
203 -- | Zip two vectors with the given function.
204 zipWith :: (IVector v a, IVector v b, IVector v c) => (a -> b -> c) -> v a -> v b -> v c
205 {-# INLINE zipWith #-}
206 zipWith f xs ys = unstream (Stream.zipWith f (stream xs) (stream ys))
207
208 -- Filtering
209 -- ---------
210
211 -- | Drop elements which do not satisfy the predicate
212 filter :: IVector v a => (a -> Bool) -> v a -> v a
213 {-# INLINE filter #-}
214 filter f = unstream . Stream.filter f . stream
215
216 -- | Yield the longest prefix of elements satisfying the predicate without
217 -- copying.
218 takeWhileSlice :: IVector v a => (a -> Bool) -> v a -> v a
219 {-# INLINE takeWhileSlice #-}
220 takeWhileSlice f v = case findIndex (not . f) v of
221 Just n -> takeSlice n v
222 Nothing -> v
223
224 -- | Copy the longest prefix of elements satisfying the predicate to a new
225 -- vector
226 takeWhile :: IVector v a => (a -> Bool) -> v a -> v a
227 {-# INLINE takeWhile #-}
228 takeWhile f = unstream . Stream.takeWhile f . stream
229
230 -- | Drop the longest prefix of elements that satisfy the predicate without
231 -- copying
232 dropWhileSlice :: IVector v a => (a -> Bool) -> v a -> v a
233 {-# INLINE dropWhileSlice #-}
234 dropWhileSlice f v = case findIndex (not . f) v of
235 Just n -> dropSlice n v
236 Nothing -> v
237
238 -- | Drop the longest prefix of elements that satisfy the predicate and copy
239 -- the rest to a new vector.
240 dropWhile :: IVector v a => (a -> Bool) -> v a -> v a
241 {-# INLINE dropWhile #-}
242 dropWhile f = unstream . Stream.dropWhile f . stream
243
244 -- Searching
245 -- ---------
246
247 infix 4 `elem`
248 -- | Check whether the vector contains an element
249 elem :: (IVector v a, Eq a) => a -> v a -> Bool
250 {-# INLINE elem #-}
251 elem x = Stream.elem x . stream
252
253 infix 4 `notElem`
254 -- | Inverse of `elem`
255 notElem :: (IVector v a, Eq a) => a -> v a -> Bool
256 {-# INLINE notElem #-}
257 notElem x = Stream.notElem x . stream
258
259 -- | Yield 'Just' the first element matching the predicate or 'Nothing' if no
260 -- such element exists.
261 find :: IVector v a => (a -> Bool) -> v a -> Maybe a
262 {-# INLINE find #-}
263 find f = Stream.find f . stream
264
265 -- | Yield 'Just' the index of the first element matching the predicate or
266 -- 'Nothing' if no such element exists.
267 findIndex :: IVector v a => (a -> Bool) -> v a -> Maybe Int
268 {-# INLINE findIndex #-}
269 findIndex f = Stream.findIndex f . stream
270
271 -- Folding
272 -- -------
273
274 -- | Left fold
275 foldl :: IVector v b => (a -> b -> a) -> a -> v b -> a
276 {-# INLINE foldl #-}
277 foldl f z = Stream.foldl f z . stream
278
279 -- | Lefgt fold on non-empty vectors
280 foldl1 :: IVector v a => (a -> a -> a) -> v a -> a
281 {-# INLINE foldl1 #-}
282 foldl1 f = Stream.foldl1 f . stream
283
284 -- | Left fold with strict accumulator
285 foldl' :: IVector v b => (a -> b -> a) -> a -> v b -> a
286 {-# INLINE foldl' #-}
287 foldl' f z = Stream.foldl' f z . stream
288
289 -- | Left fold on non-empty vectors with strict accumulator
290 foldl1' :: IVector v a => (a -> a -> a) -> v a -> a
291 {-# INLINE foldl1' #-}
292 foldl1' f = Stream.foldl1' f . stream
293
294 -- | Right fold
295 foldr :: IVector v a => (a -> b -> b) -> b -> v a -> b
296 {-# INLINE foldr #-}
297 foldr f z = Stream.foldr f z . stream
298
299 -- | Right fold on non-empty vectors
300 foldr1 :: IVector v a => (a -> a -> a) -> v a -> a
301 {-# INLINE foldr1 #-}
302 foldr1 f = Stream.foldr1 f . stream
303
304 -- | Convert a vector to a list
305 toList :: IVector v a => v a -> [a]
306 {-# INLINE toList #-}
307 toList = Stream.toList . stream
308
309 -- | Convert a list to a vector
310 fromList :: IVector v a => [a] -> v a
311 {-# INLINE fromList #-}
312 fromList = unstream . Stream.fromList
313