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