Allow streams to produce entire vectors as well as individual elements
[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, unstablePartitionStream, partitionStream
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.Stream as Stream
64 import Data.Vector.Fusion.Stream ( Stream, MStream )
65 import qualified Data.Vector.Fusion.Stream.Monadic as MStream
66 import Data.Vector.Fusion.Stream.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 MStream.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) <- MStream.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 -> MStream m u a
242 {-# INLINE mstream #-}
243 mstream v = v `seq` n `seq` (MStream.unfoldrM get 0 `MStream.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 -> MStream m u a
255 -> m (v (PrimState m) a)
256 {-# INLINE fill #-}
257 fill v s = v `seq` do
258 n' <- MStream.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 => (MStream m u a -> MStream m u a)
269 -> v (PrimState m) a
270 -> m (v (PrimState m) a)
271 {-# INLINE_STREAM transform #-}
272 transform f v = fill v (f (mstream v))
273
274 mstreamR :: (PrimMonad m, MVector v a)
275 => v (PrimState m) a
276 -> MStream m u a
277 {-# INLINE mstreamR #-}
278 mstreamR v = v `seq` n `seq` (MStream.unfoldrM get n `MStream.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 -> MStream m u a
292 -> m (v (PrimState m) a)
293 {-# INLINE fillR #-}
294 fillR v s = v `seq` do
295 i <- MStream.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 => (MStream m u a -> MStream m u a)
309 -> v (PrimState m) a
310 -> m (v (PrimState m) a)
311 {-# INLINE_STREAM transformR #-}
312 transformR f v = fillR v (f (mstreamR v))
313
314 -- | Create a new mutable vector and fill it with elements from the 'Stream'.
315 -- The vector will grow exponentially if the maximum size of the 'Stream' is
316 -- unknown.
317 unstream :: (PrimMonad m, MVector v a)
318 => Stream u a -> m (v (PrimState m) a)
319 -- NOTE: replace INLINE_STREAM by INLINE? (also in unstreamR)
320 {-# INLINE_STREAM unstream #-}
321 unstream s = munstream (Stream.liftStream 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 => MStream m u a -> m (v (PrimState m) a)
328 {-# INLINE_STREAM munstream #-}
329 munstream s = case upperBound (MStream.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 . Stream.fromList
340 --
341 -- I'm not sure this still applies (19/04/2010)
342
343 munstreamMax :: (PrimMonad m, MVector v a)
344 => MStream 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' <- MStream.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 => MStream m u a -> m (v (PrimState m) a)
360 {-# INLINE munstreamUnknown #-}
361 munstreamUnknown s
362 = do
363 v <- unsafeNew 0
364 (v', n) <- MStream.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 'Stream'.
380 -- The vector will grow exponentially if the maximum size of the 'Stream' is
381 -- unknown.
382 vunstream :: (PrimMonad m, V.Vector v a)
383 => Stream v a -> m (V.Mutable v (PrimState m) a)
384 -- NOTE: replace INLINE_STREAM by INLINE? (also in unstreamR)
385 {-# INLINE_STREAM vunstream #-}
386 vunstream s = vmunstream (Stream.liftStream 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 => MStream m v a -> m (V.Mutable v (PrimState m) a)
393 {-# INLINE_STREAM vmunstream #-}
394 vmunstream s = case upperBound (MStream.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 . Stream.fromList
405 --
406 -- I'm not sure this still applies (19/04/2010)
407
408 vmunstreamMax :: (PrimMonad m, V.Vector v a)
409 => MStream 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 put i x = do
416 INTERNAL_CHECK(checkIndex) "munstreamMax" i n
417 $ unsafeWrite v i x
418 return (i+1)
419
420 {-# INLINE_INNER copy #-}
421 copy i u = do
422 let n = V.basicLength u
423 V.basicUnsafeCopy (basicUnsafeSlice i n v) u
424 return (i+n)
425
426 n' <- MStream.vfoldlM' put copy 0 s
427 return $ INTERNAL_CHECK(checkSlice) "munstreamMax" 0 n' n
428 $ unsafeSlice 0 n' v
429
430 vmunstreamUnknown :: (PrimMonad m, V.Vector v a)
431 => MStream m v a -> m (V.Mutable v (PrimState m) a)
432 {-# INLINE vmunstreamUnknown #-}
433 vmunstreamUnknown s
434 = do
435 v <- unsafeNew 0
436 (v', n) <- MStream.vfoldlM put copy (v, 0) s
437 return $ INTERNAL_CHECK(checkSlice) "munstreamUnknown" 0 n (length v')
438 $ unsafeSlice 0 n v'
439 where
440 {-# INLINE_INNER put #-}
441 put (v,i) x = do
442 v' <- unsafeAppend1 v i x
443 return (v',i+1)
444
445 {-# INLINE_INNER copy #-}
446 copy (v,i) u = do
447 let n = V.basicLength u
448 j = i+n
449 v' <- if basicLength v < j
450 then unsafeGrow v (delay_inline max (enlarge_delta v) (j - basicLength v))
451 else return v
452 V.basicUnsafeCopy (basicUnsafeSlice i n v') u
453 return (v',j)
454
455
456
457
458 -- | Create a new mutable vector and fill it with elements from the 'Stream'
459 -- from right to left. The vector will grow exponentially if the maximum size
460 -- of the 'Stream' is unknown.
461 unstreamR :: (PrimMonad m, MVector v a)
462 => Stream u a -> m (v (PrimState m) a)
463 -- NOTE: replace INLINE_STREAM by INLINE? (also in unstream)
464 {-# INLINE_STREAM unstreamR #-}
465 unstreamR s = munstreamR (Stream.liftStream s)
466
467 -- | Create a new mutable vector and fill it with elements from the monadic
468 -- stream from right to left. The vector will grow exponentially if the maximum
469 -- size of the stream is unknown.
470 munstreamR :: (PrimMonad m, MVector v a)
471 => MStream m u a -> m (v (PrimState m) a)
472 {-# INLINE_STREAM munstreamR #-}
473 munstreamR s = case upperBound (MStream.size s) of
474 Just n -> munstreamRMax s n
475 Nothing -> munstreamRUnknown s
476
477 munstreamRMax :: (PrimMonad m, MVector v a)
478 => MStream m u a -> Int -> m (v (PrimState m) a)
479 {-# INLINE munstreamRMax #-}
480 munstreamRMax s n
481 = do
482 v <- INTERNAL_CHECK(checkLength) "munstreamRMax" n
483 $ unsafeNew n
484 let put i x = do
485 let i' = i-1
486 INTERNAL_CHECK(checkIndex) "munstreamRMax" i' n
487 $ unsafeWrite v i' x
488 return i'
489 i <- MStream.foldM' put n s
490 return $ INTERNAL_CHECK(checkSlice) "munstreamRMax" i (n-i) n
491 $ unsafeSlice i (n-i) v
492
493 munstreamRUnknown :: (PrimMonad m, MVector v a)
494 => MStream m u a -> m (v (PrimState m) a)
495 {-# INLINE munstreamRUnknown #-}
496 munstreamRUnknown s
497 = do
498 v <- unsafeNew 0
499 (v', i) <- MStream.foldM put (v, 0) s
500 let n = length v'
501 return $ INTERNAL_CHECK(checkSlice) "unstreamRUnknown" i (n-i) n
502 $ unsafeSlice i (n-i) v'
503 where
504 {-# INLINE_INNER put #-}
505 put (v,i) x = unsafePrepend1 v i x
506
507 -- Length
508 -- ------
509
510 -- | Length of the mutable vector.
511 length :: MVector v a => v s a -> Int
512 {-# INLINE length #-}
513 length = basicLength
514
515 -- | Check whether the vector is empty
516 null :: MVector v a => v s a -> Bool
517 {-# INLINE null #-}
518 null v = length v == 0
519
520 -- Extracting subvectors
521 -- ---------------------
522
523 -- | Yield a part of the mutable vector without copying it.
524 slice :: MVector v a => Int -> Int -> v s a -> v s a
525 {-# INLINE slice #-}
526 slice i n v = BOUNDS_CHECK(checkSlice) "slice" i n (length v)
527 $ unsafeSlice i n v
528
529 take :: MVector v a => Int -> v s a -> v s a
530 {-# INLINE take #-}
531 take n v = unsafeSlice 0 (min (max n 0) (length v)) v
532
533 drop :: MVector v a => Int -> v s a -> v s a
534 {-# INLINE drop #-}
535 drop n v = unsafeSlice (min m n') (max 0 (m - n')) v
536 where
537 n' = max n 0
538 m = length v
539
540 {-# INLINE splitAt #-}
541 splitAt :: MVector v a => Int -> v s a -> (v s a, v s a)
542 splitAt n v = ( unsafeSlice 0 m v
543 , unsafeSlice m (max 0 (len - n')) v
544 )
545 where
546 m = min n' len
547 n' = max n 0
548 len = length v
549
550 init :: MVector v a => v s a -> v s a
551 {-# INLINE init #-}
552 init v = slice 0 (length v - 1) v
553
554 tail :: MVector v a => v s a -> v s a
555 {-# INLINE tail #-}
556 tail v = slice 1 (length v - 1) v
557
558 -- | Yield a part of the mutable vector without copying it. No bounds checks
559 -- are performed.
560 unsafeSlice :: MVector v a => Int -- ^ starting index
561 -> Int -- ^ length of the slice
562 -> v s a
563 -> v s a
564 {-# INLINE unsafeSlice #-}
565 unsafeSlice i n v = UNSAFE_CHECK(checkSlice) "unsafeSlice" i n (length v)
566 $ basicUnsafeSlice i n v
567
568 unsafeInit :: MVector v a => v s a -> v s a
569 {-# INLINE unsafeInit #-}
570 unsafeInit v = unsafeSlice 0 (length v - 1) v
571
572 unsafeTail :: MVector v a => v s a -> v s a
573 {-# INLINE unsafeTail #-}
574 unsafeTail v = unsafeSlice 1 (length v - 1) v
575
576 unsafeTake :: MVector v a => Int -> v s a -> v s a
577 {-# INLINE unsafeTake #-}
578 unsafeTake n v = unsafeSlice 0 n v
579
580 unsafeDrop :: MVector v a => Int -> v s a -> v s a
581 {-# INLINE unsafeDrop #-}
582 unsafeDrop n v = unsafeSlice n (length v - n) v
583
584 -- Overlapping
585 -- -----------
586
587 -- Check whether two vectors overlap.
588 overlaps :: MVector v a => v s a -> v s a -> Bool
589 {-# INLINE overlaps #-}
590 overlaps = basicOverlaps
591
592 -- Initialisation
593 -- --------------
594
595 -- | Create a mutable vector of the given length.
596 new :: (PrimMonad m, MVector v a) => Int -> m (v (PrimState m) a)
597 {-# INLINE new #-}
598 new n = BOUNDS_CHECK(checkLength) "new" n
599 $ unsafeNew n
600
601 -- | Create a mutable vector of the given length. The length is not checked.
602 unsafeNew :: (PrimMonad m, MVector v a) => Int -> m (v (PrimState m) a)
603 {-# INLINE unsafeNew #-}
604 unsafeNew n = UNSAFE_CHECK(checkLength) "unsafeNew" n
605 $ basicUnsafeNew n
606
607 -- | Create a mutable vector of the given length (0 if the length is negative)
608 -- and fill it with an initial value.
609 replicate :: (PrimMonad m, MVector v a) => Int -> a -> m (v (PrimState m) a)
610 {-# INLINE replicate #-}
611 replicate n x = basicUnsafeReplicate (delay_inline max 0 n) x
612
613 -- | Create a mutable vector of the given length (0 if the length is negative)
614 -- and fill it with values produced by repeatedly executing the monadic action.
615 replicateM :: (PrimMonad m, MVector v a) => Int -> m a -> m (v (PrimState m) a)
616 {-# INLINE replicateM #-}
617 replicateM n m = munstream (MStream.replicateM n m)
618
619 -- | Create a copy of a mutable vector.
620 clone :: (PrimMonad m, MVector v a) => v (PrimState m) a -> m (v (PrimState m) a)
621 {-# INLINE clone #-}
622 clone v = do
623 v' <- unsafeNew (length v)
624 unsafeCopy v' v
625 return v'
626
627 -- Growing
628 -- -------
629
630 -- | Grow a vector by the given number of elements. The number must be
631 -- positive.
632 grow :: (PrimMonad m, MVector v a)
633 => v (PrimState m) a -> Int -> m (v (PrimState m) a)
634 {-# INLINE grow #-}
635 grow v by = BOUNDS_CHECK(checkLength) "grow" by
636 $ unsafeGrow v by
637
638 growFront :: (PrimMonad m, MVector v a)
639 => v (PrimState m) a -> Int -> m (v (PrimState m) a)
640 {-# INLINE growFront #-}
641 growFront v by = BOUNDS_CHECK(checkLength) "growFront" by
642 $ unsafeGrowFront v by
643
644 enlarge_delta v = max (length v) 1
645
646 -- | Grow a vector logarithmically
647 enlarge :: (PrimMonad m, MVector v a)
648 => v (PrimState m) a -> m (v (PrimState m) a)
649 {-# INLINE enlarge #-}
650 enlarge v = unsafeGrow v (enlarge_delta v)
651
652 enlargeFront :: (PrimMonad m, MVector v a)
653 => v (PrimState m) a -> m (v (PrimState m) a, Int)
654 {-# INLINE enlargeFront #-}
655 enlargeFront v = do
656 v' <- unsafeGrowFront v by
657 return (v', by)
658 where
659 by = enlarge_delta v
660
661 -- | Grow a vector by the given number of elements. The number must be
662 -- positive but this is not checked.
663 unsafeGrow :: (PrimMonad m, MVector v a)
664 => v (PrimState m) a -> Int -> m (v (PrimState m) a)
665 {-# INLINE unsafeGrow #-}
666 unsafeGrow v n = UNSAFE_CHECK(checkLength) "unsafeGrow" n
667 $ basicUnsafeGrow v n
668
669 unsafeGrowFront :: (PrimMonad m, MVector v a)
670 => v (PrimState m) a -> Int -> m (v (PrimState m) a)
671 {-# INLINE unsafeGrowFront #-}
672 unsafeGrowFront v by = UNSAFE_CHECK(checkLength) "unsafeGrowFront" by
673 $ do
674 let n = length v
675 v' <- basicUnsafeNew (by+n)
676 basicUnsafeCopy (basicUnsafeSlice by n v') v
677 return v'
678
679 -- Restricting memory usage
680 -- ------------------------
681
682 -- | Reset all elements of the vector to some undefined value, clearing all
683 -- references to external objects. This is usually a noop for unboxed vectors.
684 clear :: (PrimMonad m, MVector v a) => v (PrimState m) a -> m ()
685 {-# INLINE clear #-}
686 clear = basicClear
687
688 -- Accessing individual elements
689 -- -----------------------------
690
691 -- | Yield the element at the given position.
692 read :: (PrimMonad m, MVector v a) => v (PrimState m) a -> Int -> m a
693 {-# INLINE read #-}
694 read v i = BOUNDS_CHECK(checkIndex) "read" i (length v)
695 $ unsafeRead v i
696
697 -- | Replace the element at the given position.
698 write :: (PrimMonad m, MVector v a) => v (PrimState m) a -> Int -> a -> m ()
699 {-# INLINE write #-}
700 write v i x = BOUNDS_CHECK(checkIndex) "write" i (length v)
701 $ unsafeWrite v i x
702
703 -- | Swap the elements at the given positions.
704 swap :: (PrimMonad m, MVector v a) => v (PrimState m) a -> Int -> Int -> m ()
705 {-# INLINE swap #-}
706 swap v i j = BOUNDS_CHECK(checkIndex) "swap" i (length v)
707 $ BOUNDS_CHECK(checkIndex) "swap" j (length v)
708 $ unsafeSwap v i j
709
710 -- | Replace the element at the give position and return the old element.
711 exchange :: (PrimMonad m, MVector v a) => v (PrimState m) a -> Int -> a -> m a
712 {-# INLINE exchange #-}
713 exchange v i x = BOUNDS_CHECK(checkIndex) "exchange" i (length v)
714 $ unsafeExchange v i x
715
716 -- | Yield the element at the given position. No bounds checks are performed.
717 unsafeRead :: (PrimMonad m, MVector v a) => v (PrimState m) a -> Int -> m a
718 {-# INLINE unsafeRead #-}
719 unsafeRead v i = UNSAFE_CHECK(checkIndex) "unsafeRead" i (length v)
720 $ basicUnsafeRead v i
721
722 -- | Replace the element at the given position. No bounds checks are performed.
723 unsafeWrite :: (PrimMonad m, MVector v a)
724 => v (PrimState m) a -> Int -> a -> m ()
725 {-# INLINE unsafeWrite #-}
726 unsafeWrite v i x = UNSAFE_CHECK(checkIndex) "unsafeWrite" i (length v)
727 $ basicUnsafeWrite v i x
728
729 -- | Swap the elements at the given positions. No bounds checks are performed.
730 unsafeSwap :: (PrimMonad m, MVector v a)
731 => v (PrimState m) a -> Int -> Int -> m ()
732 {-# INLINE unsafeSwap #-}
733 unsafeSwap v i j = UNSAFE_CHECK(checkIndex) "unsafeSwap" i (length v)
734 $ UNSAFE_CHECK(checkIndex) "unsafeSwap" j (length v)
735 $ do
736 x <- unsafeRead v i
737 y <- unsafeRead v j
738 unsafeWrite v i y
739 unsafeWrite v j x
740
741 -- | Replace the element at the give position and return the old element. No
742 -- bounds checks are performed.
743 unsafeExchange :: (PrimMonad m, MVector v a)
744 => v (PrimState m) a -> Int -> a -> m a
745 {-# INLINE unsafeExchange #-}
746 unsafeExchange v i x = UNSAFE_CHECK(checkIndex) "unsafeExchange" i (length v)
747 $ do
748 y <- unsafeRead v i
749 unsafeWrite v i x
750 return y
751
752 -- Filling and copying
753 -- -------------------
754
755 -- | Set all elements of the vector to the given value.
756 set :: (PrimMonad m, MVector v a) => v (PrimState m) a -> a -> m ()
757 {-# INLINE set #-}
758 set = basicSet
759
760 -- | Copy a vector. The two vectors must have the same length and may not
761 -- overlap.
762 copy :: (PrimMonad m, MVector v a)
763 => v (PrimState m) a -> v (PrimState m) a -> m ()
764 {-# INLINE copy #-}
765 copy dst src = BOUNDS_CHECK(check) "copy" "overlapping vectors"
766 (not (dst `overlaps` src))
767 $ BOUNDS_CHECK(check) "copy" "length mismatch"
768 (length dst == length src)
769 $ unsafeCopy dst src
770
771 -- | Move the contents of a vector. The two vectors must have the same
772 -- length.
773 --
774 -- If the vectors do not overlap, then this is equivalent to 'copy'.
775 -- Otherwise, the copying is performed as if the source vector were
776 -- copied to a temporary vector and then the temporary vector was copied
777 -- to the target vector.
778 move :: (PrimMonad m, MVector v a)
779 => v (PrimState m) a -> v (PrimState m) a -> m ()
780 {-# INLINE move #-}
781 move dst src = BOUNDS_CHECK(check) "move" "length mismatch"
782 (length dst == length src)
783 $ unsafeMove dst src
784
785 -- | Copy a vector. The two vectors must have the same length and may not
786 -- overlap. This is not checked.
787 unsafeCopy :: (PrimMonad m, MVector v a) => v (PrimState m) a -- ^ target
788 -> v (PrimState m) a -- ^ source
789 -> m ()
790 {-# INLINE unsafeCopy #-}
791 unsafeCopy dst src = UNSAFE_CHECK(check) "unsafeCopy" "length mismatch"
792 (length dst == length src)
793 $ UNSAFE_CHECK(check) "unsafeCopy" "overlapping vectors"
794 (not (dst `overlaps` src))
795 $ (dst `seq` src `seq` basicUnsafeCopy dst src)
796
797 -- | Move the contents of a vector. The two vectors must have the same
798 -- length, but this is not checked.
799 --
800 -- If the vectors do not overlap, then this is equivalent to 'unsafeCopy'.
801 -- Otherwise, the copying is performed as if the source vector were
802 -- copied to a temporary vector and then the temporary vector was copied
803 -- to the target vector.
804 unsafeMove :: (PrimMonad m, MVector v a) => v (PrimState m) a -- ^ target
805 -> v (PrimState m) a -- ^ source
806 -> m ()
807 {-# INLINE unsafeMove #-}
808 unsafeMove dst src = UNSAFE_CHECK(check) "unsafeMove" "length mismatch"
809 (length dst == length src)
810 $ (dst `seq` src `seq` basicUnsafeMove dst src)
811
812 -- Permutations
813 -- ------------
814
815 accum :: (PrimMonad m, MVector v a)
816 => (a -> b -> a) -> v (PrimState m) a -> Stream u (Int, b) -> m ()
817 {-# INLINE accum #-}
818 accum f !v s = Stream.mapM_ upd s
819 where
820 {-# INLINE_INNER upd #-}
821 upd (i,b) = do
822 a <- BOUNDS_CHECK(checkIndex) "accum" i n
823 $ unsafeRead v i
824 unsafeWrite v i (f a b)
825
826 !n = length v
827
828 update :: (PrimMonad m, MVector v a)
829 => v (PrimState m) a -> Stream u (Int, a) -> m ()
830 {-# INLINE update #-}
831 update !v s = Stream.mapM_ upd s
832 where
833 {-# INLINE_INNER upd #-}
834 upd (i,b) = BOUNDS_CHECK(checkIndex) "update" i n
835 $ unsafeWrite v i b
836
837 !n = length v
838
839 unsafeAccum :: (PrimMonad m, MVector v a)
840 => (a -> b -> a) -> v (PrimState m) a -> Stream u (Int, b) -> m ()
841 {-# INLINE unsafeAccum #-}
842 unsafeAccum f !v s = Stream.mapM_ upd s
843 where
844 {-# INLINE_INNER upd #-}
845 upd (i,b) = do
846 a <- UNSAFE_CHECK(checkIndex) "accum" i n
847 $ unsafeRead v i
848 unsafeWrite v i (f a b)
849
850 !n = length v
851
852 unsafeUpdate :: (PrimMonad m, MVector v a)
853 => v (PrimState m) a -> Stream u (Int, a) -> m ()
854 {-# INLINE unsafeUpdate #-}
855 unsafeUpdate !v s = Stream.mapM_ upd s
856 where
857 {-# INLINE_INNER upd #-}
858 upd (i,b) = UNSAFE_CHECK(checkIndex) "accum" i n
859 $ unsafeWrite v i b
860
861 !n = length v
862
863 reverse :: (PrimMonad m, MVector v a) => v (PrimState m) a -> m ()
864 {-# INLINE reverse #-}
865 reverse !v = reverse_loop 0 (length v - 1)
866 where
867 reverse_loop i j | i < j = do
868 unsafeSwap v i j
869 reverse_loop (i + 1) (j - 1)
870 reverse_loop _ _ = return ()
871
872 unstablePartition :: forall m v a. (PrimMonad m, MVector v a)
873 => (a -> Bool) -> v (PrimState m) a -> m Int
874 {-# INLINE unstablePartition #-}
875 unstablePartition f !v = from_left 0 (length v)
876 where
877 -- NOTE: GHC 6.10.4 panics without the signatures on from_left and
878 -- from_right
879 from_left :: Int -> Int -> m Int
880 from_left i j
881 | i == j = return i
882 | otherwise = do
883 x <- unsafeRead v i
884 if f x
885 then from_left (i+1) j
886 else from_right i (j-1)
887
888 from_right :: Int -> Int -> m Int
889 from_right i j
890 | i == j = return i
891 | otherwise = do
892 x <- unsafeRead v j
893 if f x
894 then do
895 y <- unsafeRead v i
896 unsafeWrite v i x
897 unsafeWrite v j y
898 from_left (i+1) j
899 else from_right i (j-1)
900
901 unstablePartitionStream :: (PrimMonad m, MVector v a)
902 => (a -> Bool) -> Stream u a -> m (v (PrimState m) a, v (PrimState m) a)
903 {-# INLINE unstablePartitionStream #-}
904 unstablePartitionStream f s
905 = case upperBound (Stream.size s) of
906 Just n -> unstablePartitionMax f s n
907 Nothing -> partitionUnknown f s
908
909 unstablePartitionMax :: (PrimMonad m, MVector v a)
910 => (a -> Bool) -> Stream u a -> Int
911 -> m (v (PrimState m) a, v (PrimState m) a)
912 {-# INLINE unstablePartitionMax #-}
913 unstablePartitionMax f s n
914 = do
915 v <- INTERNAL_CHECK(checkLength) "unstablePartitionMax" n
916 $ unsafeNew n
917 let {-# INLINE_INNER put #-}
918 put (i, j) x
919 | f x = do
920 unsafeWrite v i x
921 return (i+1, j)
922 | otherwise = do
923 unsafeWrite v (j-1) x
924 return (i, j-1)
925
926 (i,j) <- Stream.foldM' put (0, n) s
927 return (unsafeSlice 0 i v, unsafeSlice j (n-j) v)
928
929 partitionStream :: (PrimMonad m, MVector v a)
930 => (a -> Bool) -> Stream u a -> m (v (PrimState m) a, v (PrimState m) a)
931 {-# INLINE partitionStream #-}
932 partitionStream f s
933 = case upperBound (Stream.size s) of
934 Just n -> partitionMax f s n
935 Nothing -> partitionUnknown f s
936
937 partitionMax :: (PrimMonad m, MVector v a)
938 => (a -> Bool) -> Stream u a -> Int -> m (v (PrimState m) a, v (PrimState m) a)
939 {-# INLINE partitionMax #-}
940 partitionMax f s n
941 = do
942 v <- INTERNAL_CHECK(checkLength) "unstablePartitionMax" n
943 $ unsafeNew n
944
945 let {-# INLINE_INNER put #-}
946 put (i,j) x
947 | f x = do
948 unsafeWrite v i x
949 return (i+1,j)
950
951 | otherwise = let j' = j-1 in
952 do
953 unsafeWrite v j' x
954 return (i,j')
955
956 (i,j) <- Stream.foldM' put (0,n) s
957 INTERNAL_CHECK(check) "partitionMax" "invalid indices" (i <= j)
958 $ return ()
959 let l = unsafeSlice 0 i v
960 r = unsafeSlice j (n-j) v
961 reverse r
962 return (l,r)
963
964 partitionUnknown :: (PrimMonad m, MVector v a)
965 => (a -> Bool) -> Stream u a -> m (v (PrimState m) a, v (PrimState m) a)
966 {-# INLINE partitionUnknown #-}
967 partitionUnknown f s
968 = do
969 v1 <- unsafeNew 0
970 v2 <- unsafeNew 0
971 (v1', n1, v2', n2) <- Stream.foldM' put (v1, 0, v2, 0) s
972 INTERNAL_CHECK(checkSlice) "partitionUnknown" 0 n1 (length v1')
973 $ INTERNAL_CHECK(checkSlice) "partitionUnknown" 0 n2 (length v2')
974 $ return (unsafeSlice 0 n1 v1', unsafeSlice 0 n2 v2')
975 where
976 -- NOTE: The case distinction has to be on the outside because
977 -- GHC creates a join point for the unsafeWrite even when everything
978 -- is inlined. This is bad because with the join point, v isn't getting
979 -- unboxed.
980 {-# INLINE_INNER put #-}
981 put (v1, i1, v2, i2) x
982 | f x = do
983 v1' <- unsafeAppend1 v1 i1 x
984 return (v1', i1+1, v2, i2)
985 | otherwise = do
986 v2' <- unsafeAppend1 v2 i2 x
987 return (v1, i1, v2', i2+1)
988