Be explicitly strict in vector lengths more often to avoid triggering LiberateCase
[darcs-mirrors/vector.git] / Data / Vector / Generic / Mutable.hs
1 {-# LANGUAGE MultiParamTypeClasses, BangPatterns, ScopedTypeVariables #-}
2 -- |
3 -- Module : Data.Vector.Generic.Mutable
4 -- Copyright : (c) Roman Leshchinskiy 2008-2010
5 -- License : BSD-style
6 --
7 -- Maintainer : Roman Leshchinskiy <rl@cse.unsw.edu.au>
8 -- Stability : experimental
9 -- Portability : non-portable
10 --
11 -- Generic interface to mutable vectors
12 --
13
14 module Data.Vector.Generic.Mutable (
15 -- * Class of mutable vector types
16 MVector(..),
17
18 -- * Accessors
19
20 -- ** Length information
21 length, null,
22
23 -- ** Extracting subvectors
24 slice, init, tail, take, drop,
25 unsafeSlice, unsafeInit, unsafeTail, unsafeTake, unsafeDrop,
26
27 -- ** Overlapping
28 overlaps,
29
30 -- * Construction
31
32 -- ** Initialisation
33 new, unsafeNew, replicate,
34
35 -- ** Growing
36 grow, unsafeGrow,
37
38 -- ** Restricting memory usage
39 clear,
40
41 -- * Accessing individual elements
42 read, write, swap,
43 unsafeRead, unsafeWrite, unsafeSwap,
44
45 -- * Modifying vectors
46
47 -- ** Filling and copying
48 set, copy, unsafeCopy,
49
50 -- * Internal operations
51 unstream, unstreamR,
52 munstream, munstreamR,
53 transform, transformR,
54 fill, fillR,
55 unsafeAccum, accum, unsafeUpdate, update, reverse,
56 unstablePartition, unstablePartitionStream, partitionStream,
57
58 -- * Deprecated operations
59 newWith, unsafeNewWith
60 ) where
61
62 import qualified Data.Vector.Fusion.Stream as Stream
63 import Data.Vector.Fusion.Stream ( Stream, MStream )
64 import qualified Data.Vector.Fusion.Stream.Monadic as MStream
65 import Data.Vector.Fusion.Stream.Size
66 import Data.Vector.Fusion.Util ( delay_inline )
67
68 import Control.Monad.Primitive ( PrimMonad, PrimState )
69
70 import Prelude hiding ( length, null, replicate, reverse, map, read,
71 take, drop, init, tail )
72
73 #include "vector.h"
74
75 -- | Class of mutable vectors parametrised with a primitive state token.
76 --
77 class MVector v a where
78 -- | Length of the mutable vector. This method should not be
79 -- called directly, use 'length' instead.
80 basicLength :: v s a -> Int
81
82 -- | Yield a part of the mutable vector without copying it. This method
83 -- should not be called directly, use 'unsafeSlice' instead.
84 basicUnsafeSlice :: Int -- ^ starting index
85 -> Int -- ^ length of the slice
86 -> v s a
87 -> v s a
88
89 -- Check whether two vectors overlap. This method should not be
90 -- called directly, use 'overlaps' instead.
91 basicOverlaps :: v s a -> v s a -> Bool
92
93 -- | Create a mutable vector of the given length. This method should not be
94 -- called directly, use 'unsafeNew' instead.
95 basicUnsafeNew :: PrimMonad m => Int -> m (v (PrimState m) a)
96
97 -- | Create a mutable vector of the given length and fill it with an
98 -- initial value. This method should not be called directly, use
99 -- 'replicate' instead.
100 basicUnsafeReplicate :: PrimMonad m => Int -> a -> m (v (PrimState m) a)
101
102 -- | Yield the element at the given position. This method should not be
103 -- called directly, use 'unsafeRead' instead.
104 basicUnsafeRead :: PrimMonad m => v (PrimState m) a -> Int -> m a
105
106 -- | Replace the element at the given position. This method should not be
107 -- called directly, use 'unsafeWrite' instead.
108 basicUnsafeWrite :: PrimMonad m => v (PrimState m) a -> Int -> a -> m ()
109
110 -- | Reset all elements of the vector to some undefined value, clearing all
111 -- references to external objects. This is usually a noop for unboxed
112 -- vectors. This method should not be called directly, use 'clear' instead.
113 basicClear :: PrimMonad m => v (PrimState m) a -> m ()
114
115 -- | Set all elements of the vector to the given value. This method should
116 -- not be called directly, use 'set' instead.
117 basicSet :: PrimMonad m => v (PrimState m) a -> a -> m ()
118
119 -- | Copy a vector. The two vectors may not overlap. This method should not
120 -- be called directly, use 'unsafeCopy' instead.
121 basicUnsafeCopy :: PrimMonad m => v (PrimState m) a -- ^ target
122 -> v (PrimState m) a -- ^ source
123 -> m ()
124
125 -- | Grow a vector by the given number of elements. This method should not be
126 -- called directly, use 'unsafeGrow' instead.
127 basicUnsafeGrow :: PrimMonad m => v (PrimState m) a -> Int
128 -> m (v (PrimState m) a)
129
130 -- | /DEPRECATED/ in favour of 'basicUnsafeReplicate'
131 basicUnsafeNewWith :: PrimMonad m => Int -> a -> m (v (PrimState m) a)
132
133 {-# INLINE basicUnsafeReplicate #-}
134 basicUnsafeReplicate = basicUnsafeNewWith
135
136 {-# INLINE basicUnsafeNewWith #-}
137 basicUnsafeNewWith n x
138 = do
139 v <- basicUnsafeNew n
140 basicSet v x
141 return v
142
143 {-# INLINE basicClear #-}
144 basicClear _ = return ()
145
146 {-# INLINE basicSet #-}
147 basicSet v x = do_set 0
148 where
149 n = basicLength v
150
151 do_set i | i < n = do
152 basicUnsafeWrite v i x
153 do_set (i+1)
154 | otherwise = return ()
155
156 {-# INLINE basicUnsafeCopy #-}
157 basicUnsafeCopy dst src = do_copy 0
158 where
159 n = basicLength src
160
161 do_copy i | i < n = do
162 x <- basicUnsafeRead src i
163 basicUnsafeWrite dst i x
164 do_copy (i+1)
165 | otherwise = return ()
166
167 {-# INLINE basicUnsafeGrow #-}
168 basicUnsafeGrow v by
169 = do
170 v' <- basicUnsafeNew (n+by)
171 basicUnsafeCopy (basicUnsafeSlice 0 n v') v
172 return v'
173 where
174 n = basicLength v
175
176 {-# DEPRECATED basicUnsafeNewWith "define and use basicUnsafeReplicate instead" #-}
177
178 -- ------------------
179 -- Internal functions
180 -- ------------------
181
182 unsafeAppend1 :: (PrimMonad m, MVector v a)
183 => v (PrimState m) a -> Int -> a -> m (v (PrimState m) a)
184 {-# INLINE_INNER unsafeAppend1 #-}
185 -- NOTE: The case distinction has to be on the outside because
186 -- GHC creates a join point for the unsafeWrite even when everything
187 -- is inlined. This is bad because with the join point, v isn't getting
188 -- unboxed.
189 unsafeAppend1 v i x
190 | i < length v = do
191 unsafeWrite v i x
192 return v
193 | otherwise = do
194 v' <- enlarge v
195 INTERNAL_CHECK(checkIndex) "unsafeAppend1" i (length v')
196 $ unsafeWrite v' i x
197 return v'
198
199 unsafePrepend1 :: (PrimMonad m, MVector v a)
200 => v (PrimState m) a -> Int -> a -> m (v (PrimState m) a, Int)
201 {-# INLINE_INNER unsafePrepend1 #-}
202 unsafePrepend1 v i x
203 | i /= 0 = do
204 let i' = i-1
205 unsafeWrite v i' x
206 return (v, i')
207 | otherwise = do
208 (v', i) <- enlargeFront v
209 let i' = i-1
210 INTERNAL_CHECK(checkIndex) "unsafePrepend1" i' (length v')
211 $ unsafeWrite v' i' x
212 return (v', i')
213
214 mstream :: (PrimMonad m, MVector v a) => v (PrimState m) a -> MStream m a
215 {-# INLINE mstream #-}
216 mstream v = v `seq` n `seq` (MStream.unfoldrM get 0 `MStream.sized` Exact n)
217 where
218 n = length v
219
220 {-# INLINE_INNER get #-}
221 get i | i < n = do x <- unsafeRead v i
222 return $ Just (x, i+1)
223 | otherwise = return $ Nothing
224
225 fill :: (PrimMonad m, MVector v a)
226 => v (PrimState m) a -> MStream m a -> m (v (PrimState m) a)
227 {-# INLINE fill #-}
228 fill v s = v `seq` do
229 n' <- MStream.foldM put 0 s
230 return $ unsafeSlice 0 n' v
231 where
232 {-# INLINE_INNER put #-}
233 put i x = do
234 INTERNAL_CHECK(checkIndex) "fill" i (length v)
235 $ unsafeWrite v i x
236 return (i+1)
237
238 transform :: (PrimMonad m, MVector v a)
239 => (MStream m a -> MStream m a) -> v (PrimState m) a -> m (v (PrimState m) a)
240 {-# INLINE_STREAM transform #-}
241 transform f v = fill v (f (mstream v))
242
243 mstreamR :: (PrimMonad m, MVector v a) => v (PrimState m) a -> MStream m a
244 {-# INLINE mstreamR #-}
245 mstreamR v = v `seq` n `seq` (MStream.unfoldrM get n `MStream.sized` Exact n)
246 where
247 n = length v
248
249 {-# INLINE_INNER get #-}
250 get i | j >= 0 = do x <- unsafeRead v j
251 return $ Just (x,j)
252 | otherwise = return Nothing
253 where
254 j = i-1
255
256 fillR :: (PrimMonad m, MVector v a)
257 => v (PrimState m) a -> MStream m a -> m (v (PrimState m) a)
258 {-# INLINE fillR #-}
259 fillR v s = v `seq` do
260 i <- MStream.foldM put n s
261 return $ unsafeSlice i (n-i) v
262 where
263 n = length v
264
265 {-# INLINE_INNER put #-}
266 put i x = do
267 unsafeWrite v j x
268 return j
269 where
270 j = i-1
271
272 transformR :: (PrimMonad m, MVector v a)
273 => (MStream m a -> MStream m a) -> v (PrimState m) a -> m (v (PrimState m) a)
274 {-# INLINE_STREAM transformR #-}
275 transformR f v = fillR v (f (mstreamR v))
276
277 -- | Create a new mutable vector and fill it with elements from the 'Stream'.
278 -- The vector will grow exponentially if the maximum size of the 'Stream' is
279 -- unknown.
280 unstream :: (PrimMonad m, MVector v a) => Stream a -> m (v (PrimState m) a)
281 -- NOTE: replace INLINE_STREAM by INLINE? (also in unstreamR)
282 {-# INLINE_STREAM unstream #-}
283 unstream s = munstream (Stream.liftStream s)
284
285 -- | Create a new mutable vector and fill it with elements from the monadic
286 -- stream. The vector will grow exponentially if the maximum size of the stream
287 -- is unknown.
288 munstream :: (PrimMonad m, MVector v a) => MStream m a -> m (v (PrimState m) a)
289 {-# INLINE_STREAM munstream #-}
290 munstream s = case upperBound (MStream.size s) of
291 Just n -> munstreamMax s n
292 Nothing -> munstreamUnknown s
293
294 -- FIXME: I can't think of how to prevent GHC from floating out
295 -- unstreamUnknown. That is bad because SpecConstr then generates two
296 -- specialisations: one for when it is called from unstream (it doesn't know
297 -- the shape of the vector) and one for when the vector has grown. To see the
298 -- problem simply compile this:
299 --
300 -- fromList = Data.Vector.Unboxed.unstream . Stream.fromList
301 --
302 -- I'm not sure this still applies (19/04/2010)
303
304 munstreamMax
305 :: (PrimMonad m, MVector v a) => MStream m a -> Int -> m (v (PrimState m) a)
306 {-# INLINE munstreamMax #-}
307 munstreamMax s n
308 = do
309 v <- INTERNAL_CHECK(checkLength) "munstreamMax" n
310 $ unsafeNew n
311 let put i x = do
312 INTERNAL_CHECK(checkIndex) "munstreamMax" i n
313 $ unsafeWrite v i x
314 return (i+1)
315 n' <- MStream.foldM' put 0 s
316 return $ INTERNAL_CHECK(checkSlice) "munstreamMax" 0 n' n
317 $ unsafeSlice 0 n' v
318
319 munstreamUnknown
320 :: (PrimMonad m, MVector v a) => MStream m a -> m (v (PrimState m) a)
321 {-# INLINE munstreamUnknown #-}
322 munstreamUnknown s
323 = do
324 v <- unsafeNew 0
325 (v', n) <- MStream.foldM put (v, 0) s
326 return $ INTERNAL_CHECK(checkSlice) "munstreamUnknown" 0 n (length v')
327 $ unsafeSlice 0 n v'
328 where
329 {-# INLINE_INNER put #-}
330 put (v,i) x = do
331 v' <- unsafeAppend1 v i x
332 return (v',i+1)
333
334 -- | Create a new mutable vector and fill it with elements from the 'Stream'
335 -- from right to left. The vector will grow exponentially if the maximum size
336 -- of the 'Stream' is unknown.
337 unstreamR :: (PrimMonad m, MVector v a) => Stream a -> m (v (PrimState m) a)
338 -- NOTE: replace INLINE_STREAM by INLINE? (also in unstream)
339 {-# INLINE_STREAM unstreamR #-}
340 unstreamR s = munstreamR (Stream.liftStream s)
341
342 -- | Create a new mutable vector and fill it with elements from the monadic
343 -- stream from right to left. The vector will grow exponentially if the maximum
344 -- size of the stream is unknown.
345 munstreamR :: (PrimMonad m, MVector v a) => MStream m a -> m (v (PrimState m) a)
346 {-# INLINE_STREAM munstreamR #-}
347 munstreamR s = case upperBound (MStream.size s) of
348 Just n -> munstreamRMax s n
349 Nothing -> munstreamRUnknown s
350
351 munstreamRMax
352 :: (PrimMonad m, MVector v a) => MStream m a -> Int -> m (v (PrimState m) a)
353 {-# INLINE munstreamRMax #-}
354 munstreamRMax s n
355 = do
356 v <- INTERNAL_CHECK(checkLength) "munstreamRMax" n
357 $ unsafeNew n
358 let put i x = do
359 let i' = i-1
360 INTERNAL_CHECK(checkIndex) "munstreamRMax" i' n
361 $ unsafeWrite v i' x
362 return i'
363 i <- MStream.foldM' put n s
364 return $ INTERNAL_CHECK(checkSlice) "munstreamRMax" i (n-i) n
365 $ unsafeSlice i (n-i) v
366
367 munstreamRUnknown
368 :: (PrimMonad m, MVector v a) => MStream m a -> m (v (PrimState m) a)
369 {-# INLINE munstreamRUnknown #-}
370 munstreamRUnknown s
371 = do
372 v <- unsafeNew 0
373 (v', i) <- MStream.foldM put (v, 0) s
374 let n = length v'
375 return $ INTERNAL_CHECK(checkSlice) "unstreamRUnknown" i (n-i) n
376 $ unsafeSlice i (n-i) v'
377 where
378 {-# INLINE_INNER put #-}
379 put (v,i) x = unsafePrepend1 v i x
380
381 -- Length
382 -- ------
383
384 -- | Length of the mutable vector.
385 length :: MVector v a => v s a -> Int
386 {-# INLINE length #-}
387 length = basicLength
388
389 -- | Check whether the vector is empty
390 null :: MVector v a => v s a -> Bool
391 {-# INLINE null #-}
392 null v = length v == 0
393
394 -- Extracting subvectors
395 -- ---------------------
396
397 -- | Yield a part of the mutable vector without copying it.
398 slice :: MVector v a => Int -> Int -> v s a -> v s a
399 {-# INLINE slice #-}
400 slice i n v = BOUNDS_CHECK(checkSlice) "slice" i n (length v)
401 $ unsafeSlice i n v
402
403 take :: MVector v a => Int -> v s a -> v s a
404 {-# INLINE take #-}
405 take n v = unsafeSlice 0 (min (max n 0) (length v)) v
406
407 drop :: MVector v a => Int -> v s a -> v s a
408 {-# INLINE drop #-}
409 drop n v = unsafeSlice (min m n') (max 0 (m - n')) v
410 where
411 n' = max n 0
412 m = length v
413
414 init :: MVector v a => v s a -> v s a
415 {-# INLINE init #-}
416 init v = slice 0 (length v - 1) v
417
418 tail :: MVector v a => v s a -> v s a
419 {-# INLINE tail #-}
420 tail v = slice 1 (length v - 1) v
421
422 -- | Yield a part of the mutable vector without copying it. No bounds checks
423 -- are performed.
424 unsafeSlice :: MVector v a => Int -- ^ starting index
425 -> Int -- ^ length of the slice
426 -> v s a
427 -> v s a
428 {-# INLINE unsafeSlice #-}
429 unsafeSlice i n v = UNSAFE_CHECK(checkSlice) "unsafeSlice" i n (length v)
430 $ basicUnsafeSlice i n v
431
432 unsafeInit :: MVector v a => v s a -> v s a
433 {-# INLINE unsafeInit #-}
434 unsafeInit v = unsafeSlice 0 (length v - 1) v
435
436 unsafeTail :: MVector v a => v s a -> v s a
437 {-# INLINE unsafeTail #-}
438 unsafeTail v = unsafeSlice 1 (length v - 1) v
439
440 unsafeTake :: MVector v a => Int -> v s a -> v s a
441 {-# INLINE unsafeTake #-}
442 unsafeTake n v = unsafeSlice 0 n v
443
444 unsafeDrop :: MVector v a => Int -> v s a -> v s a
445 {-# INLINE unsafeDrop #-}
446 unsafeDrop n v = unsafeSlice n (length v - n) v
447
448 -- Overlapping
449 -- -----------
450
451 -- Check whether two vectors overlap.
452 overlaps :: MVector v a => v s a -> v s a -> Bool
453 {-# INLINE overlaps #-}
454 overlaps = basicOverlaps
455
456 -- Initialisation
457 -- --------------
458
459 -- | Create a mutable vector of the given length.
460 new :: (PrimMonad m, MVector v a) => Int -> m (v (PrimState m) a)
461 {-# INLINE new #-}
462 new n = BOUNDS_CHECK(checkLength) "new" n
463 $ unsafeNew n
464
465 -- | Create a mutable vector of the given length. The length is not checked.
466 unsafeNew :: (PrimMonad m, MVector v a) => Int -> m (v (PrimState m) a)
467 {-# INLINE unsafeNew #-}
468 unsafeNew n = UNSAFE_CHECK(checkLength) "unsafeNew" n
469 $ basicUnsafeNew n
470
471 -- | Create a mutable vector of the given length (0 if the length is negative)
472 -- and fill it with an initial value.
473 replicate :: (PrimMonad m, MVector v a) => Int -> a -> m (v (PrimState m) a)
474 {-# INLINE replicate #-}
475 replicate n x = basicUnsafeReplicate (delay_inline max 0 n) x
476
477 -- Growing
478 -- -------
479
480 -- | Grow a vector by the given number of elements. The number must be
481 -- positive.
482 grow :: (PrimMonad m, MVector v a)
483 => v (PrimState m) a -> Int -> m (v (PrimState m) a)
484 {-# INLINE grow #-}
485 grow v by = BOUNDS_CHECK(checkLength) "grow" by
486 $ unsafeGrow v by
487
488 growFront :: (PrimMonad m, MVector v a)
489 => v (PrimState m) a -> Int -> m (v (PrimState m) a)
490 {-# INLINE growFront #-}
491 growFront v by = BOUNDS_CHECK(checkLength) "growFront" by
492 $ unsafeGrowFront v by
493
494 enlarge_delta v = max (length v) 1
495
496 -- | Grow a vector logarithmically
497 enlarge :: (PrimMonad m, MVector v a)
498 => v (PrimState m) a -> m (v (PrimState m) a)
499 {-# INLINE enlarge #-}
500 enlarge v = unsafeGrow v (enlarge_delta v)
501
502 enlargeFront :: (PrimMonad m, MVector v a)
503 => v (PrimState m) a -> m (v (PrimState m) a, Int)
504 {-# INLINE enlargeFront #-}
505 enlargeFront v = do
506 v' <- unsafeGrowFront v by
507 return (v', by)
508 where
509 by = enlarge_delta v
510
511 -- | Grow a vector by the given number of elements. The number must be
512 -- positive but this is not checked.
513 unsafeGrow :: (PrimMonad m, MVector v a)
514 => v (PrimState m) a -> Int -> m (v (PrimState m) a)
515 {-# INLINE unsafeGrow #-}
516 unsafeGrow v n = UNSAFE_CHECK(checkLength) "unsafeGrow" n
517 $ basicUnsafeGrow v n
518
519 unsafeGrowFront :: (PrimMonad m, MVector v a)
520 => v (PrimState m) a -> Int -> m (v (PrimState m) a)
521 {-# INLINE unsafeGrowFront #-}
522 unsafeGrowFront v by = UNSAFE_CHECK(checkLength) "unsafeGrowFront" by
523 $ do
524 let n = length v
525 v' <- basicUnsafeNew (by+n)
526 basicUnsafeCopy (basicUnsafeSlice by n v') v
527 return v'
528
529 -- Restricting memory usage
530 -- ------------------------
531
532 -- | Reset all elements of the vector to some undefined value, clearing all
533 -- references to external objects. This is usually a noop for unboxed vectors.
534 clear :: (PrimMonad m, MVector v a) => v (PrimState m) a -> m ()
535 {-# INLINE clear #-}
536 clear = basicClear
537
538 -- Accessing individual elements
539 -- -----------------------------
540
541 -- | Yield the element at the given position.
542 read :: (PrimMonad m, MVector v a) => v (PrimState m) a -> Int -> m a
543 {-# INLINE read #-}
544 read v i = BOUNDS_CHECK(checkIndex) "read" i (length v)
545 $ unsafeRead v i
546
547 -- | Replace the element at the given position.
548 write :: (PrimMonad m, MVector v a) => v (PrimState m) a -> Int -> a -> m ()
549 {-# INLINE write #-}
550 write v i x = BOUNDS_CHECK(checkIndex) "write" i (length v)
551 $ unsafeWrite v i x
552
553 -- | Swap the elements at the given positions.
554 swap :: (PrimMonad m, MVector v a) => v (PrimState m) a -> Int -> Int -> m ()
555 {-# INLINE swap #-}
556 swap v i j = BOUNDS_CHECK(checkIndex) "swap" i (length v)
557 $ BOUNDS_CHECK(checkIndex) "swap" j (length v)
558 $ unsafeSwap v i j
559
560 -- | Replace the element at the give position and return the old element.
561 exchange :: (PrimMonad m, MVector v a) => v (PrimState m) a -> Int -> a -> m a
562 {-# INLINE exchange #-}
563 exchange v i x = BOUNDS_CHECK(checkIndex) "exchange" i (length v)
564 $ unsafeExchange v i x
565
566 -- | Yield the element at the given position. No bounds checks are performed.
567 unsafeRead :: (PrimMonad m, MVector v a) => v (PrimState m) a -> Int -> m a
568 {-# INLINE unsafeRead #-}
569 unsafeRead v i = UNSAFE_CHECK(checkIndex) "unsafeRead" i (length v)
570 $ basicUnsafeRead v i
571
572 -- | Replace the element at the given position. No bounds checks are performed.
573 unsafeWrite :: (PrimMonad m, MVector v a)
574 => v (PrimState m) a -> Int -> a -> m ()
575 {-# INLINE unsafeWrite #-}
576 unsafeWrite v i x = UNSAFE_CHECK(checkIndex) "unsafeWrite" i (length v)
577 $ basicUnsafeWrite v i x
578
579 -- | Swap the elements at the given positions. No bounds checks are performed.
580 unsafeSwap :: (PrimMonad m, MVector v a)
581 => v (PrimState m) a -> Int -> Int -> m ()
582 {-# INLINE unsafeSwap #-}
583 unsafeSwap v i j = UNSAFE_CHECK(checkIndex) "unsafeSwap" i (length v)
584 $ UNSAFE_CHECK(checkIndex) "unsafeSwap" j (length v)
585 $ do
586 x <- unsafeRead v i
587 y <- unsafeRead v j
588 unsafeWrite v i y
589 unsafeWrite v j x
590
591 -- | Replace the element at the give position and return the old element. No
592 -- bounds checks are performed.
593 unsafeExchange :: (PrimMonad m, MVector v a)
594 => v (PrimState m) a -> Int -> a -> m a
595 {-# INLINE unsafeExchange #-}
596 unsafeExchange v i x = UNSAFE_CHECK(checkIndex) "unsafeExchange" i (length v)
597 $ do
598 y <- unsafeRead v i
599 unsafeWrite v i x
600 return y
601
602 -- Filling and copying
603 -- -------------------
604
605 -- | Set all elements of the vector to the given value.
606 set :: (PrimMonad m, MVector v a) => v (PrimState m) a -> a -> m ()
607 {-# INLINE set #-}
608 set = basicSet
609
610 -- | Copy a vector. The two vectors must have the same length and may not
611 -- overlap.
612 copy :: (PrimMonad m, MVector v a)
613 => v (PrimState m) a -> v (PrimState m) a -> m ()
614 {-# INLINE copy #-}
615 copy dst src = BOUNDS_CHECK(check) "copy" "overlapping vectors"
616 (not (dst `overlaps` src))
617 $ BOUNDS_CHECK(check) "copy" "length mismatch"
618 (length dst == length src)
619 $ unsafeCopy dst src
620
621 -- | Copy a vector. The two vectors must have the same length and may not
622 -- overlap. This is not checked.
623 unsafeCopy :: (PrimMonad m, MVector v a) => v (PrimState m) a -- ^ target
624 -> v (PrimState m) a -- ^ source
625 -> m ()
626 {-# INLINE unsafeCopy #-}
627 unsafeCopy dst src = UNSAFE_CHECK(check) "unsafeCopy" "length mismatch"
628 (length dst == length src)
629 $ UNSAFE_CHECK(check) "unsafeCopy" "overlapping vectors"
630 (not (dst `overlaps` src))
631 $ (dst `seq` src `seq` basicUnsafeCopy dst src)
632
633
634 -- Permutations
635 -- ------------
636
637 accum :: (PrimMonad m, MVector v a)
638 => (a -> b -> a) -> v (PrimState m) a -> Stream (Int, b) -> m ()
639 {-# INLINE accum #-}
640 accum f !v s = Stream.mapM_ upd s
641 where
642 {-# INLINE_INNER upd #-}
643 upd (i,b) = do
644 a <- BOUNDS_CHECK(checkIndex) "accum" i n
645 $ unsafeRead v i
646 unsafeWrite v i (f a b)
647
648 !n = length v
649
650 update :: (PrimMonad m, MVector v a)
651 => v (PrimState m) a -> Stream (Int, a) -> m ()
652 {-# INLINE update #-}
653 update !v s = Stream.mapM_ upd s
654 where
655 {-# INLINE_INNER upd #-}
656 upd (i,b) = BOUNDS_CHECK(checkIndex) "update" i n
657 $ unsafeWrite v i b
658
659 !n = length v
660
661 unsafeAccum :: (PrimMonad m, MVector v a)
662 => (a -> b -> a) -> v (PrimState m) a -> Stream (Int, b) -> m ()
663 {-# INLINE unsafeAccum #-}
664 unsafeAccum f !v s = Stream.mapM_ upd s
665 where
666 {-# INLINE_INNER upd #-}
667 upd (i,b) = do
668 a <- UNSAFE_CHECK(checkIndex) "accum" i n
669 $ unsafeRead v i
670 unsafeWrite v i (f a b)
671
672 !n = length v
673
674 unsafeUpdate :: (PrimMonad m, MVector v a)
675 => v (PrimState m) a -> Stream (Int, a) -> m ()
676 {-# INLINE unsafeUpdate #-}
677 unsafeUpdate !v s = Stream.mapM_ upd s
678 where
679 {-# INLINE_INNER upd #-}
680 upd (i,b) = UNSAFE_CHECK(checkIndex) "accum" i n
681 $ unsafeWrite v i b
682
683 !n = length v
684
685 reverse :: (PrimMonad m, MVector v a) => v (PrimState m) a -> m ()
686 {-# INLINE reverse #-}
687 reverse !v = reverse_loop 0 (length v - 1)
688 where
689 reverse_loop i j | i < j = do
690 unsafeSwap v i j
691 reverse_loop (i + 1) (j - 1)
692 reverse_loop _ _ = return ()
693
694 unstablePartition :: forall m v a. (PrimMonad m, MVector v a)
695 => (a -> Bool) -> v (PrimState m) a -> m Int
696 {-# INLINE unstablePartition #-}
697 unstablePartition f !v = from_left 0 (length v)
698 where
699 -- NOTE: GHC 6.10.4 panics without the signatures on from_left and
700 -- from_right
701 from_left :: Int -> Int -> m Int
702 from_left i j
703 | i == j = return i
704 | otherwise = do
705 x <- unsafeRead v i
706 if f x
707 then from_left (i+1) j
708 else from_right i (j-1)
709
710 from_right :: Int -> Int -> m Int
711 from_right i j
712 | i == j = return i
713 | otherwise = do
714 x <- unsafeRead v j
715 if f x
716 then do
717 y <- unsafeRead v i
718 unsafeWrite v i x
719 unsafeWrite v j y
720 from_left (i+1) j
721 else from_right i (j-1)
722
723 unstablePartitionStream :: (PrimMonad m, MVector v a)
724 => (a -> Bool) -> Stream a -> m (v (PrimState m) a, v (PrimState m) a)
725 {-# INLINE unstablePartitionStream #-}
726 unstablePartitionStream f s
727 = case upperBound (Stream.size s) of
728 Just n -> unstablePartitionMax f s n
729 Nothing -> partitionUnknown f s
730
731 unstablePartitionMax :: (PrimMonad m, MVector v a)
732 => (a -> Bool) -> Stream a -> Int
733 -> m (v (PrimState m) a, v (PrimState m) a)
734 {-# INLINE unstablePartitionMax #-}
735 unstablePartitionMax f s n
736 = do
737 v <- INTERNAL_CHECK(checkLength) "unstablePartitionMax" n
738 $ unsafeNew n
739 let {-# INLINE_INNER put #-}
740 put (i, j) x
741 | f x = do
742 unsafeWrite v i x
743 return (i+1, j)
744 | otherwise = do
745 unsafeWrite v (j-1) x
746 return (i, j-1)
747
748 (i,j) <- Stream.foldM' put (0, n) s
749 return (unsafeSlice 0 i v, unsafeSlice j (n-j) v)
750
751 partitionStream :: (PrimMonad m, MVector v a)
752 => (a -> Bool) -> Stream a -> m (v (PrimState m) a, v (PrimState m) a)
753 {-# INLINE partitionStream #-}
754 partitionStream f s
755 = case upperBound (Stream.size s) of
756 Just n -> partitionMax f s n
757 Nothing -> partitionUnknown f s
758
759 partitionMax :: (PrimMonad m, MVector v a)
760 => (a -> Bool) -> Stream a -> Int -> m (v (PrimState m) a, v (PrimState m) a)
761 {-# INLINE partitionMax #-}
762 partitionMax f s n
763 = do
764 v <- INTERNAL_CHECK(checkLength) "unstablePartitionMax" n
765 $ unsafeNew n
766
767 let {-# INLINE_INNER put #-}
768 put (i,j) x
769 | f x = do
770 unsafeWrite v i x
771 return (i+1,j)
772
773 | otherwise = let j' = j-1 in
774 do
775 unsafeWrite v j' x
776 return (i,j')
777
778 (i,j) <- Stream.foldM' put (0,n) s
779 INTERNAL_CHECK(check) "partitionMax" "invalid indices" (i <= j)
780 $ return ()
781 let l = unsafeSlice 0 i v
782 r = unsafeSlice j (n-j) v
783 reverse r
784 return (l,r)
785
786 partitionUnknown :: (PrimMonad m, MVector v a)
787 => (a -> Bool) -> Stream a -> m (v (PrimState m) a, v (PrimState m) a)
788 {-# INLINE partitionUnknown #-}
789 partitionUnknown f s
790 = do
791 v1 <- unsafeNew 0
792 v2 <- unsafeNew 0
793 (v1', n1, v2', n2) <- Stream.foldM' put (v1, 0, v2, 0) s
794 INTERNAL_CHECK(checkSlice) "partitionUnknown" 0 n1 (length v1')
795 $ INTERNAL_CHECK(checkSlice) "partitionUnknown" 0 n2 (length v2')
796 $ return (unsafeSlice 0 n1 v1', unsafeSlice 0 n2 v2')
797 where
798 -- NOTE: The case distinction has to be on the outside because
799 -- GHC creates a join point for the unsafeWrite even when everything
800 -- is inlined. This is bad because with the join point, v isn't getting
801 -- unboxed.
802 {-# INLINE_INNER put #-}
803 put (v1, i1, v2, i2) x
804 | f x = do
805 v1' <- unsafeAppend1 v1 i1 x
806 return (v1', i1+1, v2, i2)
807 | otherwise = do
808 v2' <- unsafeAppend1 v2 i2 x
809 return (v1, i1, v2', i2+1)
810
811 -- Deprecated functions
812 -- --------------------
813
814 -- | /DEPRECATED/ Use 'replicate' instead
815 newWith :: (PrimMonad m, MVector v a) => Int -> a -> m (v (PrimState m) a)
816 {-# INLINE newWith #-}
817 newWith = replicate
818
819 -- | /DEPRECATED/ Use 'replicate' instead
820 unsafeNewWith :: (PrimMonad m, MVector v a) => Int -> a -> m (v (PrimState m) a)
821 {-# INLINE unsafeNewWith #-}
822 unsafeNewWith = replicate
823
824 {-# DEPRECATED newWith, unsafeNewWith "Use replicate instead" #-}
825