161098a184e01b72b4ab912fcf4d30223714bd26
[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, 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
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, splitAt, 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 -- | Move the contents of a vector. The two vectors may overlap. This method
126 -- should not be called directly, use 'unsafeMove' instead.
127 basicUnsafeMove :: PrimMonad m => v (PrimState m) a -- ^ target
128 -> v (PrimState m) a -- ^ source
129 -> m ()
130
131 -- | Grow a vector by the given number of elements. This method should not be
132 -- called directly, use 'unsafeGrow' instead.
133 basicUnsafeGrow :: PrimMonad m => v (PrimState m) a -> Int
134 -> m (v (PrimState m) a)
135
136 -- | /DEPRECATED/ in favour of 'basicUnsafeReplicate'
137 basicUnsafeNewWith :: PrimMonad m => Int -> a -> m (v (PrimState m) a)
138
139 {-# INLINE basicUnsafeReplicate #-}
140 basicUnsafeReplicate = basicUnsafeNewWith
141
142 {-# INLINE basicUnsafeNewWith #-}
143 basicUnsafeNewWith n x
144 = do
145 v <- basicUnsafeNew n
146 basicSet v x
147 return v
148
149 {-# INLINE basicClear #-}
150 basicClear _ = return ()
151
152 {-# INLINE basicSet #-}
153 basicSet !v x = do_set 0
154 where
155 !n = basicLength v
156
157 do_set i | i < n = do
158 basicUnsafeWrite v i x
159 do_set (i+1)
160 | otherwise = return ()
161
162 {-# INLINE basicUnsafeCopy #-}
163 basicUnsafeCopy !dst !src = do_copy 0
164 where
165 !n = basicLength src
166
167 do_copy i | i < n = do
168 x <- basicUnsafeRead src i
169 basicUnsafeWrite dst i x
170 do_copy (i+1)
171 | otherwise = return ()
172
173 {-# INLINE basicUnsafeMove #-}
174 basicUnsafeMove !dst !src
175 | basicOverlaps dst src = do
176 srcCopy <- clone src
177 basicUnsafeCopy dst srcCopy
178 | otherwise = basicUnsafeCopy dst src
179
180 {-# INLINE basicUnsafeGrow #-}
181 basicUnsafeGrow v by
182 = do
183 v' <- basicUnsafeNew (n+by)
184 basicUnsafeCopy (basicUnsafeSlice 0 n v') v
185 return v'
186 where
187 n = basicLength v
188
189 {-# DEPRECATED basicUnsafeNewWith "define and use basicUnsafeReplicate instead" #-}
190
191 -- ------------------
192 -- Internal functions
193 -- ------------------
194
195 unsafeAppend1 :: (PrimMonad m, MVector v a)
196 => v (PrimState m) a -> Int -> a -> m (v (PrimState m) a)
197 {-# INLINE_INNER unsafeAppend1 #-}
198 -- NOTE: The case distinction has to be on the outside because
199 -- GHC creates a join point for the unsafeWrite even when everything
200 -- is inlined. This is bad because with the join point, v isn't getting
201 -- unboxed.
202 unsafeAppend1 v i x
203 | i < length v = do
204 unsafeWrite v i x
205 return v
206 | otherwise = do
207 v' <- enlarge v
208 INTERNAL_CHECK(checkIndex) "unsafeAppend1" i (length v')
209 $ unsafeWrite v' i x
210 return v'
211
212 unsafePrepend1 :: (PrimMonad m, MVector v a)
213 => v (PrimState m) a -> Int -> a -> m (v (PrimState m) a, Int)
214 {-# INLINE_INNER unsafePrepend1 #-}
215 unsafePrepend1 v i x
216 | i /= 0 = do
217 let i' = i-1
218 unsafeWrite v i' x
219 return (v, i')
220 | otherwise = do
221 (v', i) <- enlargeFront v
222 let i' = i-1
223 INTERNAL_CHECK(checkIndex) "unsafePrepend1" i' (length v')
224 $ unsafeWrite v' i' x
225 return (v', i')
226
227 mstream :: (PrimMonad m, MVector v a) => v (PrimState m) a -> MStream m a
228 {-# INLINE mstream #-}
229 mstream v = v `seq` n `seq` (MStream.unfoldrM get 0 `MStream.sized` Exact n)
230 where
231 n = length v
232
233 {-# INLINE_INNER get #-}
234 get i | i < n = do x <- unsafeRead v i
235 return $ Just (x, i+1)
236 | otherwise = return $ Nothing
237
238 fill :: (PrimMonad m, MVector v a)
239 => v (PrimState m) a -> MStream m a -> m (v (PrimState m) a)
240 {-# INLINE fill #-}
241 fill v s = v `seq` do
242 n' <- MStream.foldM put 0 s
243 return $ unsafeSlice 0 n' v
244 where
245 {-# INLINE_INNER put #-}
246 put i x = do
247 INTERNAL_CHECK(checkIndex) "fill" i (length v)
248 $ unsafeWrite v i x
249 return (i+1)
250
251 transform :: (PrimMonad m, MVector v a)
252 => (MStream m a -> MStream m a) -> v (PrimState m) a -> m (v (PrimState m) a)
253 {-# INLINE_STREAM transform #-}
254 transform f v = fill v (f (mstream v))
255
256 mstreamR :: (PrimMonad m, MVector v a) => v (PrimState m) a -> MStream m a
257 {-# INLINE mstreamR #-}
258 mstreamR v = v `seq` n `seq` (MStream.unfoldrM get n `MStream.sized` Exact n)
259 where
260 n = length v
261
262 {-# INLINE_INNER get #-}
263 get i | j >= 0 = do x <- unsafeRead v j
264 return $ Just (x,j)
265 | otherwise = return Nothing
266 where
267 j = i-1
268
269 fillR :: (PrimMonad m, MVector v a)
270 => v (PrimState m) a -> MStream m a -> m (v (PrimState m) a)
271 {-# INLINE fillR #-}
272 fillR v s = v `seq` do
273 i <- MStream.foldM put n s
274 return $ unsafeSlice i (n-i) v
275 where
276 n = length v
277
278 {-# INLINE_INNER put #-}
279 put i x = do
280 unsafeWrite v j x
281 return j
282 where
283 j = i-1
284
285 transformR :: (PrimMonad m, MVector v a)
286 => (MStream m a -> MStream m a) -> v (PrimState m) a -> m (v (PrimState m) a)
287 {-# INLINE_STREAM transformR #-}
288 transformR f v = fillR v (f (mstreamR v))
289
290 -- | Create a new mutable vector and fill it with elements from the 'Stream'.
291 -- The vector will grow exponentially if the maximum size of the 'Stream' is
292 -- unknown.
293 unstream :: (PrimMonad m, MVector v a) => Stream a -> m (v (PrimState m) a)
294 -- NOTE: replace INLINE_STREAM by INLINE? (also in unstreamR)
295 {-# INLINE_STREAM unstream #-}
296 unstream s = munstream (Stream.liftStream s)
297
298 -- | Create a new mutable vector and fill it with elements from the monadic
299 -- stream. The vector will grow exponentially if the maximum size of the stream
300 -- is unknown.
301 munstream :: (PrimMonad m, MVector v a) => MStream m a -> m (v (PrimState m) a)
302 {-# INLINE_STREAM munstream #-}
303 munstream s = case upperBound (MStream.size s) of
304 Just n -> munstreamMax s n
305 Nothing -> munstreamUnknown s
306
307 -- FIXME: I can't think of how to prevent GHC from floating out
308 -- unstreamUnknown. That is bad because SpecConstr then generates two
309 -- specialisations: one for when it is called from unstream (it doesn't know
310 -- the shape of the vector) and one for when the vector has grown. To see the
311 -- problem simply compile this:
312 --
313 -- fromList = Data.Vector.Unboxed.unstream . Stream.fromList
314 --
315 -- I'm not sure this still applies (19/04/2010)
316
317 munstreamMax
318 :: (PrimMonad m, MVector v a) => MStream m a -> Int -> m (v (PrimState m) a)
319 {-# INLINE munstreamMax #-}
320 munstreamMax s n
321 = do
322 v <- INTERNAL_CHECK(checkLength) "munstreamMax" n
323 $ unsafeNew n
324 let put i x = do
325 INTERNAL_CHECK(checkIndex) "munstreamMax" i n
326 $ unsafeWrite v i x
327 return (i+1)
328 n' <- MStream.foldM' put 0 s
329 return $ INTERNAL_CHECK(checkSlice) "munstreamMax" 0 n' n
330 $ unsafeSlice 0 n' v
331
332 munstreamUnknown
333 :: (PrimMonad m, MVector v a) => MStream m a -> m (v (PrimState m) a)
334 {-# INLINE munstreamUnknown #-}
335 munstreamUnknown s
336 = do
337 v <- unsafeNew 0
338 (v', n) <- MStream.foldM put (v, 0) s
339 return $ INTERNAL_CHECK(checkSlice) "munstreamUnknown" 0 n (length v')
340 $ unsafeSlice 0 n v'
341 where
342 {-# INLINE_INNER put #-}
343 put (v,i) x = do
344 v' <- unsafeAppend1 v i x
345 return (v',i+1)
346
347 -- | Create a new mutable vector and fill it with elements from the 'Stream'
348 -- from right to left. The vector will grow exponentially if the maximum size
349 -- of the 'Stream' is unknown.
350 unstreamR :: (PrimMonad m, MVector v a) => Stream a -> m (v (PrimState m) a)
351 -- NOTE: replace INLINE_STREAM by INLINE? (also in unstream)
352 {-# INLINE_STREAM unstreamR #-}
353 unstreamR s = munstreamR (Stream.liftStream s)
354
355 -- | Create a new mutable vector and fill it with elements from the monadic
356 -- stream from right to left. The vector will grow exponentially if the maximum
357 -- size of the stream is unknown.
358 munstreamR :: (PrimMonad m, MVector v a) => MStream m a -> m (v (PrimState m) a)
359 {-# INLINE_STREAM munstreamR #-}
360 munstreamR s = case upperBound (MStream.size s) of
361 Just n -> munstreamRMax s n
362 Nothing -> munstreamRUnknown s
363
364 munstreamRMax
365 :: (PrimMonad m, MVector v a) => MStream m a -> Int -> m (v (PrimState m) a)
366 {-# INLINE munstreamRMax #-}
367 munstreamRMax s n
368 = do
369 v <- INTERNAL_CHECK(checkLength) "munstreamRMax" n
370 $ unsafeNew n
371 let put i x = do
372 let i' = i-1
373 INTERNAL_CHECK(checkIndex) "munstreamRMax" i' n
374 $ unsafeWrite v i' x
375 return i'
376 i <- MStream.foldM' put n s
377 return $ INTERNAL_CHECK(checkSlice) "munstreamRMax" i (n-i) n
378 $ unsafeSlice i (n-i) v
379
380 munstreamRUnknown
381 :: (PrimMonad m, MVector v a) => MStream m a -> m (v (PrimState m) a)
382 {-# INLINE munstreamRUnknown #-}
383 munstreamRUnknown s
384 = do
385 v <- unsafeNew 0
386 (v', i) <- MStream.foldM put (v, 0) s
387 let n = length v'
388 return $ INTERNAL_CHECK(checkSlice) "unstreamRUnknown" i (n-i) n
389 $ unsafeSlice i (n-i) v'
390 where
391 {-# INLINE_INNER put #-}
392 put (v,i) x = unsafePrepend1 v i x
393
394 -- Length
395 -- ------
396
397 -- | Length of the mutable vector.
398 length :: MVector v a => v s a -> Int
399 {-# INLINE length #-}
400 length = basicLength
401
402 -- | Check whether the vector is empty
403 null :: MVector v a => v s a -> Bool
404 {-# INLINE null #-}
405 null v = length v == 0
406
407 -- Extracting subvectors
408 -- ---------------------
409
410 -- | Yield a part of the mutable vector without copying it.
411 slice :: MVector v a => Int -> Int -> v s a -> v s a
412 {-# INLINE slice #-}
413 slice i n v = BOUNDS_CHECK(checkSlice) "slice" i n (length v)
414 $ unsafeSlice i n v
415
416 take :: MVector v a => Int -> v s a -> v s a
417 {-# INLINE take #-}
418 take n v = unsafeSlice 0 (min (max n 0) (length v)) v
419
420 drop :: MVector v a => Int -> v s a -> v s a
421 {-# INLINE drop #-}
422 drop n v = unsafeSlice (min m n') (max 0 (m - n')) v
423 where
424 n' = max n 0
425 m = length v
426
427 {-# INLINE splitAt #-}
428 splitAt :: MVector v a => Int -> v s a -> (v s a, v s a)
429 splitAt n v = ( unsafeSlice 0 m v
430 , unsafeSlice m (max 0 (len - n')) v
431 )
432 where
433 m = min n' len
434 n' = max n 0
435 len = length v
436
437 init :: MVector v a => v s a -> v s a
438 {-# INLINE init #-}
439 init v = slice 0 (length v - 1) v
440
441 tail :: MVector v a => v s a -> v s a
442 {-# INLINE tail #-}
443 tail v = slice 1 (length v - 1) v
444
445 -- | Yield a part of the mutable vector without copying it. No bounds checks
446 -- are performed.
447 unsafeSlice :: MVector v a => Int -- ^ starting index
448 -> Int -- ^ length of the slice
449 -> v s a
450 -> v s a
451 {-# INLINE unsafeSlice #-}
452 unsafeSlice i n v = UNSAFE_CHECK(checkSlice) "unsafeSlice" i n (length v)
453 $ basicUnsafeSlice i n v
454
455 unsafeInit :: MVector v a => v s a -> v s a
456 {-# INLINE unsafeInit #-}
457 unsafeInit v = unsafeSlice 0 (length v - 1) v
458
459 unsafeTail :: MVector v a => v s a -> v s a
460 {-# INLINE unsafeTail #-}
461 unsafeTail v = unsafeSlice 1 (length v - 1) v
462
463 unsafeTake :: MVector v a => Int -> v s a -> v s a
464 {-# INLINE unsafeTake #-}
465 unsafeTake n v = unsafeSlice 0 n v
466
467 unsafeDrop :: MVector v a => Int -> v s a -> v s a
468 {-# INLINE unsafeDrop #-}
469 unsafeDrop n v = unsafeSlice n (length v - n) v
470
471 -- Overlapping
472 -- -----------
473
474 -- Check whether two vectors overlap.
475 overlaps :: MVector v a => v s a -> v s a -> Bool
476 {-# INLINE overlaps #-}
477 overlaps = basicOverlaps
478
479 -- Initialisation
480 -- --------------
481
482 -- | Create a mutable vector of the given length.
483 new :: (PrimMonad m, MVector v a) => Int -> m (v (PrimState m) a)
484 {-# INLINE new #-}
485 new n = BOUNDS_CHECK(checkLength) "new" n
486 $ unsafeNew n
487
488 -- | Create a mutable vector of the given length. The length is not checked.
489 unsafeNew :: (PrimMonad m, MVector v a) => Int -> m (v (PrimState m) a)
490 {-# INLINE unsafeNew #-}
491 unsafeNew n = UNSAFE_CHECK(checkLength) "unsafeNew" n
492 $ basicUnsafeNew n
493
494 -- | Create a mutable vector of the given length (0 if the length is negative)
495 -- and fill it with an initial value.
496 replicate :: (PrimMonad m, MVector v a) => Int -> a -> m (v (PrimState m) a)
497 {-# INLINE replicate #-}
498 replicate n x = basicUnsafeReplicate (delay_inline max 0 n) x
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
870 -- Deprecated functions
871 -- --------------------
872
873 -- | /DEPRECATED/ Use 'replicate' instead
874 newWith :: (PrimMonad m, MVector v a) => Int -> a -> m (v (PrimState m) a)
875 {-# INLINE newWith #-}
876 newWith = replicate
877
878 -- | /DEPRECATED/ Use 'replicate' instead
879 unsafeNewWith :: (PrimMonad m, MVector v a) => Int -> a -> m (v (PrimState m) a)
880 {-# INLINE unsafeNewWith #-}
881 unsafeNewWith = replicate
882
883 {-# DEPRECATED newWith, unsafeNewWith "Use replicate instead" #-}
884