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