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