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