7837140e673be31bb258d9ac62fb68a0a2474fd5
[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, clone,
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 -- | Create a copy of a mutable vector.
478 clone :: (PrimMonad m, MVector v a) => v (PrimState m) a -> m (v (PrimState m) a)
479 {-# INLINE clone #-}
480 clone v = do
481 v' <- unsafeNew (length v)
482 unsafeCopy v' v
483 return v'
484
485 -- Growing
486 -- -------
487
488 -- | Grow a vector by the given number of elements. The number must be
489 -- positive.
490 grow :: (PrimMonad m, MVector v a)
491 => v (PrimState m) a -> Int -> m (v (PrimState m) a)
492 {-# INLINE grow #-}
493 grow v by = BOUNDS_CHECK(checkLength) "grow" by
494 $ unsafeGrow v by
495
496 growFront :: (PrimMonad m, MVector v a)
497 => v (PrimState m) a -> Int -> m (v (PrimState m) a)
498 {-# INLINE growFront #-}
499 growFront v by = BOUNDS_CHECK(checkLength) "growFront" by
500 $ unsafeGrowFront v by
501
502 enlarge_delta v = max (length v) 1
503
504 -- | Grow a vector logarithmically
505 enlarge :: (PrimMonad m, MVector v a)
506 => v (PrimState m) a -> m (v (PrimState m) a)
507 {-# INLINE enlarge #-}
508 enlarge v = unsafeGrow v (enlarge_delta v)
509
510 enlargeFront :: (PrimMonad m, MVector v a)
511 => v (PrimState m) a -> m (v (PrimState m) a, Int)
512 {-# INLINE enlargeFront #-}
513 enlargeFront v = do
514 v' <- unsafeGrowFront v by
515 return (v', by)
516 where
517 by = enlarge_delta v
518
519 -- | Grow a vector by the given number of elements. The number must be
520 -- positive but this is not checked.
521 unsafeGrow :: (PrimMonad m, MVector v a)
522 => v (PrimState m) a -> Int -> m (v (PrimState m) a)
523 {-# INLINE unsafeGrow #-}
524 unsafeGrow v n = UNSAFE_CHECK(checkLength) "unsafeGrow" n
525 $ basicUnsafeGrow v n
526
527 unsafeGrowFront :: (PrimMonad m, MVector v a)
528 => v (PrimState m) a -> Int -> m (v (PrimState m) a)
529 {-# INLINE unsafeGrowFront #-}
530 unsafeGrowFront v by = UNSAFE_CHECK(checkLength) "unsafeGrowFront" by
531 $ do
532 let n = length v
533 v' <- basicUnsafeNew (by+n)
534 basicUnsafeCopy (basicUnsafeSlice by n v') v
535 return v'
536
537 -- Restricting memory usage
538 -- ------------------------
539
540 -- | Reset all elements of the vector to some undefined value, clearing all
541 -- references to external objects. This is usually a noop for unboxed vectors.
542 clear :: (PrimMonad m, MVector v a) => v (PrimState m) a -> m ()
543 {-# INLINE clear #-}
544 clear = basicClear
545
546 -- Accessing individual elements
547 -- -----------------------------
548
549 -- | Yield the element at the given position.
550 read :: (PrimMonad m, MVector v a) => v (PrimState m) a -> Int -> m a
551 {-# INLINE read #-}
552 read v i = BOUNDS_CHECK(checkIndex) "read" i (length v)
553 $ unsafeRead v i
554
555 -- | Replace the element at the given position.
556 write :: (PrimMonad m, MVector v a) => v (PrimState m) a -> Int -> a -> m ()
557 {-# INLINE write #-}
558 write v i x = BOUNDS_CHECK(checkIndex) "write" i (length v)
559 $ unsafeWrite v i x
560
561 -- | Swap the elements at the given positions.
562 swap :: (PrimMonad m, MVector v a) => v (PrimState m) a -> Int -> Int -> m ()
563 {-# INLINE swap #-}
564 swap v i j = BOUNDS_CHECK(checkIndex) "swap" i (length v)
565 $ BOUNDS_CHECK(checkIndex) "swap" j (length v)
566 $ unsafeSwap v i j
567
568 -- | Replace the element at the give position and return the old element.
569 exchange :: (PrimMonad m, MVector v a) => v (PrimState m) a -> Int -> a -> m a
570 {-# INLINE exchange #-}
571 exchange v i x = BOUNDS_CHECK(checkIndex) "exchange" i (length v)
572 $ unsafeExchange v i x
573
574 -- | Yield the element at the given position. No bounds checks are performed.
575 unsafeRead :: (PrimMonad m, MVector v a) => v (PrimState m) a -> Int -> m a
576 {-# INLINE unsafeRead #-}
577 unsafeRead v i = UNSAFE_CHECK(checkIndex) "unsafeRead" i (length v)
578 $ basicUnsafeRead v i
579
580 -- | Replace the element at the given position. No bounds checks are performed.
581 unsafeWrite :: (PrimMonad m, MVector v a)
582 => v (PrimState m) a -> Int -> a -> m ()
583 {-# INLINE unsafeWrite #-}
584 unsafeWrite v i x = UNSAFE_CHECK(checkIndex) "unsafeWrite" i (length v)
585 $ basicUnsafeWrite v i x
586
587 -- | Swap the elements at the given positions. No bounds checks are performed.
588 unsafeSwap :: (PrimMonad m, MVector v a)
589 => v (PrimState m) a -> Int -> Int -> m ()
590 {-# INLINE unsafeSwap #-}
591 unsafeSwap v i j = UNSAFE_CHECK(checkIndex) "unsafeSwap" i (length v)
592 $ UNSAFE_CHECK(checkIndex) "unsafeSwap" j (length v)
593 $ do
594 x <- unsafeRead v i
595 y <- unsafeRead v j
596 unsafeWrite v i y
597 unsafeWrite v j x
598
599 -- | Replace the element at the give position and return the old element. No
600 -- bounds checks are performed.
601 unsafeExchange :: (PrimMonad m, MVector v a)
602 => v (PrimState m) a -> Int -> a -> m a
603 {-# INLINE unsafeExchange #-}
604 unsafeExchange v i x = UNSAFE_CHECK(checkIndex) "unsafeExchange" i (length v)
605 $ do
606 y <- unsafeRead v i
607 unsafeWrite v i x
608 return y
609
610 -- Filling and copying
611 -- -------------------
612
613 -- | Set all elements of the vector to the given value.
614 set :: (PrimMonad m, MVector v a) => v (PrimState m) a -> a -> m ()
615 {-# INLINE set #-}
616 set = basicSet
617
618 -- | Copy a vector. The two vectors must have the same length and may not
619 -- overlap.
620 copy :: (PrimMonad m, MVector v a)
621 => v (PrimState m) a -> v (PrimState m) a -> m ()
622 {-# INLINE copy #-}
623 copy dst src = BOUNDS_CHECK(check) "copy" "overlapping vectors"
624 (not (dst `overlaps` src))
625 $ BOUNDS_CHECK(check) "copy" "length mismatch"
626 (length dst == length src)
627 $ unsafeCopy dst src
628
629 -- | Copy a vector. The two vectors must have the same length and may not
630 -- overlap. This is not checked.
631 unsafeCopy :: (PrimMonad m, MVector v a) => v (PrimState m) a -- ^ target
632 -> v (PrimState m) a -- ^ source
633 -> m ()
634 {-# INLINE unsafeCopy #-}
635 unsafeCopy dst src = UNSAFE_CHECK(check) "unsafeCopy" "length mismatch"
636 (length dst == length src)
637 $ UNSAFE_CHECK(check) "unsafeCopy" "overlapping vectors"
638 (not (dst `overlaps` src))
639 $ (dst `seq` src `seq` basicUnsafeCopy dst src)
640
641
642 -- Permutations
643 -- ------------
644
645 accum :: (PrimMonad m, MVector v a)
646 => (a -> b -> a) -> v (PrimState m) a -> Stream (Int, b) -> m ()
647 {-# INLINE accum #-}
648 accum f !v s = Stream.mapM_ upd s
649 where
650 {-# INLINE_INNER upd #-}
651 upd (i,b) = do
652 a <- BOUNDS_CHECK(checkIndex) "accum" i n
653 $ unsafeRead v i
654 unsafeWrite v i (f a b)
655
656 !n = length v
657
658 update :: (PrimMonad m, MVector v a)
659 => v (PrimState m) a -> Stream (Int, a) -> m ()
660 {-# INLINE update #-}
661 update !v s = Stream.mapM_ upd s
662 where
663 {-# INLINE_INNER upd #-}
664 upd (i,b) = BOUNDS_CHECK(checkIndex) "update" i n
665 $ unsafeWrite v i b
666
667 !n = length v
668
669 unsafeAccum :: (PrimMonad m, MVector v a)
670 => (a -> b -> a) -> v (PrimState m) a -> Stream (Int, b) -> m ()
671 {-# INLINE unsafeAccum #-}
672 unsafeAccum f !v s = Stream.mapM_ upd s
673 where
674 {-# INLINE_INNER upd #-}
675 upd (i,b) = do
676 a <- UNSAFE_CHECK(checkIndex) "accum" i n
677 $ unsafeRead v i
678 unsafeWrite v i (f a b)
679
680 !n = length v
681
682 unsafeUpdate :: (PrimMonad m, MVector v a)
683 => v (PrimState m) a -> Stream (Int, a) -> m ()
684 {-# INLINE unsafeUpdate #-}
685 unsafeUpdate !v s = Stream.mapM_ upd s
686 where
687 {-# INLINE_INNER upd #-}
688 upd (i,b) = UNSAFE_CHECK(checkIndex) "accum" i n
689 $ unsafeWrite v i b
690
691 !n = length v
692
693 reverse :: (PrimMonad m, MVector v a) => v (PrimState m) a -> m ()
694 {-# INLINE reverse #-}
695 reverse !v = reverse_loop 0 (length v - 1)
696 where
697 reverse_loop i j | i < j = do
698 unsafeSwap v i j
699 reverse_loop (i + 1) (j - 1)
700 reverse_loop _ _ = return ()
701
702 unstablePartition :: forall m v a. (PrimMonad m, MVector v a)
703 => (a -> Bool) -> v (PrimState m) a -> m Int
704 {-# INLINE unstablePartition #-}
705 unstablePartition f !v = from_left 0 (length v)
706 where
707 -- NOTE: GHC 6.10.4 panics without the signatures on from_left and
708 -- from_right
709 from_left :: Int -> Int -> m Int
710 from_left i j
711 | i == j = return i
712 | otherwise = do
713 x <- unsafeRead v i
714 if f x
715 then from_left (i+1) j
716 else from_right i (j-1)
717
718 from_right :: Int -> Int -> m Int
719 from_right i j
720 | i == j = return i
721 | otherwise = do
722 x <- unsafeRead v j
723 if f x
724 then do
725 y <- unsafeRead v i
726 unsafeWrite v i x
727 unsafeWrite v j y
728 from_left (i+1) j
729 else from_right i (j-1)
730
731 unstablePartitionStream :: (PrimMonad m, MVector v a)
732 => (a -> Bool) -> Stream a -> m (v (PrimState m) a, v (PrimState m) a)
733 {-# INLINE unstablePartitionStream #-}
734 unstablePartitionStream f s
735 = case upperBound (Stream.size s) of
736 Just n -> unstablePartitionMax f s n
737 Nothing -> partitionUnknown f s
738
739 unstablePartitionMax :: (PrimMonad m, MVector v a)
740 => (a -> Bool) -> Stream a -> Int
741 -> m (v (PrimState m) a, v (PrimState m) a)
742 {-# INLINE unstablePartitionMax #-}
743 unstablePartitionMax f s n
744 = do
745 v <- INTERNAL_CHECK(checkLength) "unstablePartitionMax" n
746 $ unsafeNew n
747 let {-# INLINE_INNER put #-}
748 put (i, j) x
749 | f x = do
750 unsafeWrite v i x
751 return (i+1, j)
752 | otherwise = do
753 unsafeWrite v (j-1) x
754 return (i, j-1)
755
756 (i,j) <- Stream.foldM' put (0, n) s
757 return (unsafeSlice 0 i v, unsafeSlice j (n-j) v)
758
759 partitionStream :: (PrimMonad m, MVector v a)
760 => (a -> Bool) -> Stream a -> m (v (PrimState m) a, v (PrimState m) a)
761 {-# INLINE partitionStream #-}
762 partitionStream f s
763 = case upperBound (Stream.size s) of
764 Just n -> partitionMax f s n
765 Nothing -> partitionUnknown f s
766
767 partitionMax :: (PrimMonad m, MVector v a)
768 => (a -> Bool) -> Stream a -> Int -> m (v (PrimState m) a, v (PrimState m) a)
769 {-# INLINE partitionMax #-}
770 partitionMax f s n
771 = do
772 v <- INTERNAL_CHECK(checkLength) "unstablePartitionMax" n
773 $ unsafeNew n
774
775 let {-# INLINE_INNER put #-}
776 put (i,j) x
777 | f x = do
778 unsafeWrite v i x
779 return (i+1,j)
780
781 | otherwise = let j' = j-1 in
782 do
783 unsafeWrite v j' x
784 return (i,j')
785
786 (i,j) <- Stream.foldM' put (0,n) s
787 INTERNAL_CHECK(check) "partitionMax" "invalid indices" (i <= j)
788 $ return ()
789 let l = unsafeSlice 0 i v
790 r = unsafeSlice j (n-j) v
791 reverse r
792 return (l,r)
793
794 partitionUnknown :: (PrimMonad m, MVector v a)
795 => (a -> Bool) -> Stream a -> m (v (PrimState m) a, v (PrimState m) a)
796 {-# INLINE partitionUnknown #-}
797 partitionUnknown f s
798 = do
799 v1 <- unsafeNew 0
800 v2 <- unsafeNew 0
801 (v1', n1, v2', n2) <- Stream.foldM' put (v1, 0, v2, 0) s
802 INTERNAL_CHECK(checkSlice) "partitionUnknown" 0 n1 (length v1')
803 $ INTERNAL_CHECK(checkSlice) "partitionUnknown" 0 n2 (length v2')
804 $ return (unsafeSlice 0 n1 v1', unsafeSlice 0 n2 v2')
805 where
806 -- NOTE: The case distinction has to be on the outside because
807 -- GHC creates a join point for the unsafeWrite even when everything
808 -- is inlined. This is bad because with the join point, v isn't getting
809 -- unboxed.
810 {-# INLINE_INNER put #-}
811 put (v1, i1, v2, i2) x
812 | f x = do
813 v1' <- unsafeAppend1 v1 i1 x
814 return (v1', i1+1, v2, i2)
815 | otherwise = do
816 v2' <- unsafeAppend1 v2 i2 x
817 return (v1, i1, v2', i2+1)
818
819 -- Deprecated functions
820 -- --------------------
821
822 -- | /DEPRECATED/ Use 'replicate' instead
823 newWith :: (PrimMonad m, MVector v a) => Int -> a -> m (v (PrimState m) a)
824 {-# INLINE newWith #-}
825 newWith = replicate
826
827 -- | /DEPRECATED/ Use 'replicate' instead
828 unsafeNewWith :: (PrimMonad m, MVector v a) => Int -> a -> m (v (PrimState m) a)
829 {-# INLINE unsafeNewWith #-}
830 unsafeNewWith = replicate
831
832 {-# DEPRECATED newWith, unsafeNewWith "Use replicate instead" #-}
833