Rearrange code
[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-2009
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 -- * Operations on mutable vectors
19 length, overlaps, new, newWith, read, write, clear, set, copy, grow,
20
21 slice, take, drop, init, tail,
22 unsafeSlice, unsafeInit, unsafeTail,
23
24 -- * Unsafe operations
25 unsafeNew, unsafeNewWith, unsafeRead, unsafeWrite,
26 unsafeCopy, unsafeGrow,
27
28 -- * Internal operations
29 unstream, transform, unstreamR, transformR,
30 unsafeAccum, accum, unsafeUpdate, update, reverse,
31 unstablePartition, unstablePartitionStream
32 ) where
33
34 import qualified Data.Vector.Fusion.Stream as Stream
35 import Data.Vector.Fusion.Stream ( Stream, MStream )
36 import qualified Data.Vector.Fusion.Stream.Monadic as MStream
37 import Data.Vector.Fusion.Stream.Size
38
39 import Control.Monad.Primitive ( PrimMonad, PrimState )
40
41 import GHC.Float (
42 double2Int, int2Double
43 )
44
45 import Prelude hiding ( length, reverse, map, read,
46 take, drop, init, tail )
47
48 #include "vector.h"
49
50 -- | Class of mutable vectors parametrised with a primitive state token.
51 --
52 class MVector v a where
53 -- | Length of the mutable vector. This method should not be
54 -- called directly, use 'length' instead.
55 basicLength :: v s a -> Int
56
57 -- | Yield a part of the mutable vector without copying it. This method
58 -- should not be called directly, use 'unsafeSlice' instead.
59 basicUnsafeSlice :: Int -- ^ starting index
60 -> Int -- ^ length of the slice
61 -> v s a
62 -> v s a
63
64 -- Check whether two vectors overlap. This method should not be
65 -- called directly, use 'overlaps' instead.
66 basicOverlaps :: v s a -> v s a -> Bool
67
68 -- | Create a mutable vector of the given length. This method should not be
69 -- called directly, use 'unsafeNew' instead.
70 basicUnsafeNew :: PrimMonad m => Int -> m (v (PrimState m) a)
71
72 -- | Create a mutable vector of the given length and fill it with an
73 -- initial value. This method should not be called directly, use
74 -- 'unsafeNewWith' instead.
75 basicUnsafeNewWith :: PrimMonad m => Int -> a -> m (v (PrimState m) a)
76
77 -- | Yield the element at the given position. This method should not be
78 -- called directly, use 'unsafeRead' instead.
79 basicUnsafeRead :: PrimMonad m => v (PrimState m) a -> Int -> m a
80
81 -- | Replace the element at the given position. This method should not be
82 -- called directly, use 'unsafeWrite' instead.
83 basicUnsafeWrite :: PrimMonad m => v (PrimState m) a -> Int -> a -> m ()
84
85 -- | Reset all elements of the vector to some undefined value, clearing all
86 -- references to external objects. This is usually a noop for unboxed
87 -- vectors. This method should not be called directly, use 'clear' instead.
88 basicClear :: PrimMonad m => v (PrimState m) a -> m ()
89
90 -- | Set all elements of the vector to the given value. This method should
91 -- not be called directly, use 'set' instead.
92 basicSet :: PrimMonad m => v (PrimState m) a -> a -> m ()
93
94 -- | Copy a vector. The two vectors may not overlap. This method should not
95 -- be called directly, use 'unsafeCopy' instead.
96 basicUnsafeCopy :: PrimMonad m => v (PrimState m) a -- ^ target
97 -> v (PrimState m) a -- ^ source
98 -> m ()
99
100 -- | Grow a vector by the given number of elements. This method should not be
101 -- called directly, use 'unsafeGrow' instead.
102 basicUnsafeGrow :: PrimMonad m => v (PrimState m) a -> Int
103 -> m (v (PrimState m) a)
104
105 {-# INLINE basicUnsafeNewWith #-}
106 basicUnsafeNewWith n x
107 = do
108 v <- basicUnsafeNew n
109 basicSet v x
110 return v
111
112 {-# INLINE basicClear #-}
113 basicClear _ = return ()
114
115 {-# INLINE basicSet #-}
116 basicSet v x = do_set 0
117 where
118 n = basicLength v
119
120 do_set i | i < n = do
121 basicUnsafeWrite v i x
122 do_set (i+1)
123 | otherwise = return ()
124
125 {-# INLINE basicUnsafeCopy #-}
126 basicUnsafeCopy dst src = do_copy 0
127 where
128 n = basicLength src
129
130 do_copy i | i < n = do
131 x <- basicUnsafeRead src i
132 basicUnsafeWrite dst i x
133 do_copy (i+1)
134 | otherwise = return ()
135
136 {-# INLINE basicUnsafeGrow #-}
137 basicUnsafeGrow v by
138 = do
139 v' <- basicUnsafeNew (n+by)
140 basicUnsafeCopy (basicUnsafeSlice 0 n v') v
141 return v'
142 where
143 n = basicLength v
144
145 -- ------------------
146 -- Internal functions
147 -- ------------------
148
149 -- Check whether two vectors overlap.
150 overlaps :: MVector v a => v s a -> v s a -> Bool
151 {-# INLINE overlaps #-}
152 overlaps = basicOverlaps
153
154 unsafeAppend1 :: (PrimMonad m, MVector v a)
155 => v (PrimState m) a -> Int -> a -> m (v (PrimState m) a)
156 {-# INLINE_INNER unsafeAppend1 #-}
157 -- NOTE: The case distinction has to be on the outside because
158 -- GHC creates a join point for the unsafeWrite even when everything
159 -- is inlined. This is bad because with the join point, v isn't getting
160 -- unboxed.
161 unsafeAppend1 v i x
162 | i < length v = do
163 unsafeWrite v i x
164 return v
165 | otherwise = do
166 v' <- enlarge v
167 INTERNAL_CHECK(checkIndex) "unsafeAppend1" i (length v')
168 $ unsafeWrite v' i x
169 return v'
170
171 unsafePrepend1 :: (PrimMonad m, MVector v a)
172 => v (PrimState m) a -> Int -> a -> m (v (PrimState m) a, Int)
173 {-# INLINE_INNER unsafePrepend1 #-}
174 unsafePrepend1 v i x
175 | i /= 0 = do
176 let i' = i-1
177 unsafeWrite v i' x
178 return (v, i')
179 | otherwise = do
180 (v', i) <- enlargeFront v
181 let i' = i-1
182 INTERNAL_CHECK(checkIndex) "unsafePrepend1" i' (length v')
183 $ unsafeWrite v' i' x
184 return (v', i')
185
186 mstream :: (PrimMonad m, MVector v a) => v (PrimState m) a -> MStream m a
187 {-# INLINE mstream #-}
188 mstream v = v `seq` (MStream.unfoldrM get 0 `MStream.sized` Exact n)
189 where
190 n = length v
191
192 {-# INLINE_INNER get #-}
193 get i | i < n = do x <- unsafeRead v i
194 return $ Just (x, i+1)
195 | otherwise = return $ Nothing
196
197 munstream :: (PrimMonad m, MVector v a)
198 => v (PrimState m) a -> MStream m a -> m (v (PrimState m) a)
199 {-# INLINE munstream #-}
200 munstream v s = v `seq` do
201 n' <- MStream.foldM put 0 s
202 return $ unsafeSlice 0 n' v
203 where
204 {-# INLINE_INNER put #-}
205 put i x = do
206 INTERNAL_CHECK(checkIndex) "munstream" i (length v)
207 $ unsafeWrite v i x
208 return (i+1)
209
210 transform :: (PrimMonad m, MVector v a)
211 => (MStream m a -> MStream m a) -> v (PrimState m) a -> m (v (PrimState m) a)
212 {-# INLINE_STREAM transform #-}
213 transform f v = munstream v (f (mstream v))
214
215 mrstream :: (PrimMonad m, MVector v a) => v (PrimState m) a -> MStream m a
216 {-# INLINE mrstream #-}
217 mrstream v = v `seq` (MStream.unfoldrM get n `MStream.sized` Exact n)
218 where
219 n = length v
220
221 {-# INLINE_INNER get #-}
222 get i | j >= 0 = do x <- unsafeRead v j
223 return $ Just (x,j)
224 | otherwise = return Nothing
225 where
226 j = i-1
227
228 munstreamR :: (PrimMonad m, MVector v a)
229 => v (PrimState m) a -> MStream m a -> m (v (PrimState m) a)
230 {-# INLINE munstreamR #-}
231 munstreamR v s = v `seq` do
232 i <- MStream.foldM put n s
233 return $ unsafeSlice i (n-i) v
234 where
235 n = length v
236
237 {-# INLINE_INNER put #-}
238 put i x = do
239 unsafeWrite v j x
240 return j
241 where
242 j = i-1
243
244 transformR :: (PrimMonad m, MVector v a)
245 => (MStream m a -> MStream m a) -> v (PrimState m) a -> m (v (PrimState m) a)
246 {-# INLINE_STREAM transformR #-}
247 transformR f v = munstreamR v (f (mrstream v))
248
249 -- | Create a new mutable vector and fill it with elements from the 'Stream'.
250 -- The vector will grow logarithmically if the 'Size' hint of the 'Stream' is
251 -- inexact.
252 unstream :: (PrimMonad m, MVector v a) => Stream a -> m (v (PrimState m) a)
253 {-# INLINE_STREAM unstream #-}
254 unstream s = case upperBound (Stream.size s) of
255 Just n -> unstreamMax s n
256 Nothing -> unstreamUnknown s
257
258 -- FIXME: I can't think of how to prevent GHC from floating out
259 -- unstreamUnknown. That is bad because SpecConstr then generates two
260 -- specialisations: one for when it is called from unstream (it doesn't know
261 -- the shape of the vector) and one for when the vector has grown. To see the
262 -- problem simply compile this:
263 --
264 -- fromList = Data.Vector.Unboxed.unstream . Stream.fromList
265 --
266
267 unstreamMax
268 :: (PrimMonad m, MVector v a) => Stream a -> Int -> m (v (PrimState m) a)
269 {-# INLINE unstreamMax #-}
270 unstreamMax s n
271 = do
272 v <- INTERNAL_CHECK(checkLength) "unstreamMax" n
273 $ unsafeNew n
274 let put i x = do
275 INTERNAL_CHECK(checkIndex) "unstreamMax" i n
276 $ unsafeWrite v i x
277 return (i+1)
278 n' <- Stream.foldM' put 0 s
279 return $ INTERNAL_CHECK(checkSlice) "unstreamMax" 0 n' n
280 $ unsafeSlice 0 n' v
281
282 unstreamUnknown
283 :: (PrimMonad m, MVector v a) => Stream a -> m (v (PrimState m) a)
284 {-# INLINE unstreamUnknown #-}
285 unstreamUnknown s
286 = do
287 v <- unsafeNew 0
288 (v', n) <- Stream.foldM put (v, 0) s
289 return $ INTERNAL_CHECK(checkSlice) "unstreamUnknown" 0 n (length v')
290 $ unsafeSlice 0 n v'
291 where
292 {-# INLINE_INNER put #-}
293 put (v,i) x = do
294 v' <- unsafeAppend1 v i x
295 return (v',i+1)
296
297 -- | Create a new mutable vector and fill it with elements from the 'Stream'.
298 -- The vector will grow logarithmically if the 'Size' hint of the 'Stream' is
299 -- inexact.
300 unstreamR :: (PrimMonad m, MVector v a) => Stream a -> m (v (PrimState m) a)
301 {-# INLINE_STREAM unstreamR #-}
302 unstreamR s = case upperBound (Stream.size s) of
303 Just n -> unstreamRMax s n
304 Nothing -> unstreamRUnknown s
305
306 unstreamRMax
307 :: (PrimMonad m, MVector v a) => Stream a -> Int -> m (v (PrimState m) a)
308 {-# INLINE unstreamRMax #-}
309 unstreamRMax s n
310 = do
311 v <- INTERNAL_CHECK(checkLength) "unstreamRMax" n
312 $ unsafeNew n
313 let put i x = do
314 let i' = i-1
315 INTERNAL_CHECK(checkIndex) "unstreamRMax" i' n
316 $ unsafeWrite v i' x
317 return i'
318 i <- Stream.foldM' put n s
319 return $ INTERNAL_CHECK(checkSlice) "unstreamRMax" i (n-i) n
320 $ unsafeSlice i (n-i) v
321
322 unstreamRUnknown
323 :: (PrimMonad m, MVector v a) => Stream a -> m (v (PrimState m) a)
324 {-# INLINE unstreamRUnknown #-}
325 unstreamRUnknown s
326 = do
327 v <- unsafeNew 0
328 (v', i) <- Stream.foldM put (v, 0) s
329 let n = length v'
330 return $ INTERNAL_CHECK(checkSlice) "unstreamRUnknown" i (n-i) n
331 $ unsafeSlice i (n-i) v'
332 where
333 {-# INLINE_INNER put #-}
334 put (v,i) x = unsafePrepend1 v i x
335
336 -- Length
337 -- ------
338
339 -- | Length of the mutable vector.
340 length :: MVector v a => v s a -> Int
341 {-# INLINE length #-}
342 length = basicLength
343
344 -- | Check whether the vector is empty
345 null :: MVector v a => v s a -> Bool
346 {-# INLINE null #-}
347 null v = length v == 0
348
349
350 -- Construction
351 -- ------------
352
353 -- | Create a mutable vector of the given length.
354 new :: (PrimMonad m, MVector v a) => Int -> m (v (PrimState m) a)
355 {-# INLINE new #-}
356 new n = BOUNDS_CHECK(checkLength) "new" n
357 $ unsafeNew n
358
359 -- | Create a mutable vector of the given length and fill it with an
360 -- initial value.
361 newWith :: (PrimMonad m, MVector v a) => Int -> a -> m (v (PrimState m) a)
362 {-# INLINE newWith #-}
363 newWith n x = BOUNDS_CHECK(checkLength) "newWith" n
364 $ unsafeNewWith n x
365
366 -- | Create a mutable vector of the given length. The length is not checked.
367 unsafeNew :: (PrimMonad m, MVector v a) => Int -> m (v (PrimState m) a)
368 {-# INLINE unsafeNew #-}
369 unsafeNew n = UNSAFE_CHECK(checkLength) "unsafeNew" n
370 $ basicUnsafeNew n
371
372 -- | Create a mutable vector of the given length and fill it with an
373 -- initial value. The length is not checked.
374 unsafeNewWith :: (PrimMonad m, MVector v a) => Int -> a -> m (v (PrimState m) a)
375 {-# INLINE unsafeNewWith #-}
376 unsafeNewWith n x = UNSAFE_CHECK(checkLength) "unsafeNewWith" n
377 $ basicUnsafeNewWith n x
378
379
380 -- Growing
381 -- -------
382
383 -- | Grow a vector by the given number of elements. The number must be
384 -- positive.
385 grow :: (PrimMonad m, MVector v a)
386 => v (PrimState m) a -> Int -> m (v (PrimState m) a)
387 {-# INLINE grow #-}
388 grow v by = BOUNDS_CHECK(checkLength) "grow" by
389 $ unsafeGrow v by
390
391 growFront :: (PrimMonad m, MVector v a)
392 => v (PrimState m) a -> Int -> m (v (PrimState m) a)
393 {-# INLINE growFront #-}
394 growFront v by = BOUNDS_CHECK(checkLength) "growFront" by
395 $ unsafeGrowFront v by
396
397 enlarge_delta v = max (length v) 1
398
399 -- | Grow a vector logarithmically
400 enlarge :: (PrimMonad m, MVector v a)
401 => v (PrimState m) a -> m (v (PrimState m) a)
402 {-# INLINE enlarge #-}
403 enlarge v = unsafeGrow v (enlarge_delta v)
404
405 enlargeFront :: (PrimMonad m, MVector v a)
406 => v (PrimState m) a -> m (v (PrimState m) a, Int)
407 {-# INLINE enlargeFront #-}
408 enlargeFront v = do
409 v' <- unsafeGrowFront v by
410 return (v', by)
411 where
412 by = enlarge_delta v
413
414 -- | Grow a vector by the given number of elements. The number must be
415 -- positive but this is not checked.
416 unsafeGrow :: (PrimMonad m, MVector v a)
417 => v (PrimState m) a -> Int -> m (v (PrimState m) a)
418 {-# INLINE unsafeGrow #-}
419 unsafeGrow v n = UNSAFE_CHECK(checkLength) "unsafeGrow" n
420 $ basicUnsafeGrow v n
421
422 unsafeGrowFront :: (PrimMonad m, MVector v a)
423 => v (PrimState m) a -> Int -> m (v (PrimState m) a)
424 {-# INLINE unsafeGrowFront #-}
425 unsafeGrowFront v by = UNSAFE_CHECK(checkLength) "unsafeGrowFront" by
426 $ do
427 let n = length v
428 v' <- basicUnsafeNew (by+n)
429 basicUnsafeCopy (basicUnsafeSlice by n v') v
430 return v'
431
432 -- Accessing individual elements
433 -- -----------------------------
434
435 -- | Yield the element at the given position.
436 read :: (PrimMonad m, MVector v a) => v (PrimState m) a -> Int -> m a
437 {-# INLINE read #-}
438 read v i = BOUNDS_CHECK(checkIndex) "read" i (length v)
439 $ unsafeRead v i
440
441 -- | Replace the element at the given position.
442 write :: (PrimMonad m, MVector v a) => v (PrimState m) a -> Int -> a -> m ()
443 {-# INLINE write #-}
444 write v i x = BOUNDS_CHECK(checkIndex) "write" i (length v)
445 $ unsafeWrite v i x
446
447 -- | Swap the elements at the gived positions.
448 swap :: (PrimMonad m, MVector v a) => v (PrimState m) a -> Int -> Int -> m ()
449 {-# INLINE swap #-}
450 swap v i j = BOUNDS_CHECK(checkIndex) "swap" i (length v)
451 $ BOUNDS_CHECK(checkIndex) "swap" j (length v)
452 $ unsafeSwap v i j
453
454 -- | Replace the element at the give position and return the old element.
455 exchange :: (PrimMonad m, MVector v a) => v (PrimState m) a -> Int -> a -> m a
456 {-# INLINE exchange #-}
457 exchange v i x = BOUNDS_CHECK(checkIndex) "exchange" i (length v)
458 $ unsafeExchange v i x
459
460 -- | Yield the element at the given position. No bounds checks are performed.
461 unsafeRead :: (PrimMonad m, MVector v a) => v (PrimState m) a -> Int -> m a
462 {-# INLINE unsafeRead #-}
463 unsafeRead v i = UNSAFE_CHECK(checkIndex) "unsafeRead" i (length v)
464 $ basicUnsafeRead v i
465
466 -- | Replace the element at the given position. No bounds checks are performed.
467 unsafeWrite :: (PrimMonad m, MVector v a)
468 => v (PrimState m) a -> Int -> a -> m ()
469 {-# INLINE unsafeWrite #-}
470 unsafeWrite v i x = UNSAFE_CHECK(checkIndex) "unsafeWrite" i (length v)
471 $ basicUnsafeWrite v i x
472
473 -- | Swap the elements at the given positions. No bounds checks are performed.
474 unsafeSwap :: (PrimMonad m, MVector v a)
475 => v (PrimState m) a -> Int -> Int -> m ()
476 {-# INLINE unsafeSwap #-}
477 unsafeSwap v i j = UNSAFE_CHECK(checkIndex) "unsafeSwap" i (length v)
478 $ UNSAFE_CHECK(checkIndex) "unsafeSwap" j (length v)
479 $ do
480 x <- unsafeRead v i
481 y <- unsafeRead v j
482 unsafeWrite v i y
483 unsafeWrite v j x
484
485 -- | Replace the element at the give position and return the old element. No
486 -- bounds checks are performed.
487 unsafeExchange :: (PrimMonad m, MVector v a)
488 => v (PrimState m) a -> Int -> a -> m a
489 {-# INLINE unsafeExchange #-}
490 unsafeExchange v i x = UNSAFE_CHECK(checkIndex) "unsafeExchange" i (length v)
491 $ do
492 y <- unsafeRead v i
493 unsafeWrite v i x
494 return y
495
496 -- Block operations
497 -- ----------------
498
499 -- | Reset all elements of the vector to some undefined value, clearing all
500 -- references to external objects. This is usually a noop for unboxed vectors.
501 clear :: (PrimMonad m, MVector v a) => v (PrimState m) a -> m ()
502 {-# INLINE clear #-}
503 clear = basicClear
504
505 -- | Set all elements of the vector to the given value.
506 set :: (PrimMonad m, MVector v a) => v (PrimState m) a -> a -> m ()
507 {-# INLINE set #-}
508 set = basicSet
509
510 -- | Copy a vector. The two vectors must have the same length and may not
511 -- overlap.
512 copy :: (PrimMonad m, MVector v a)
513 => v (PrimState m) a -> v (PrimState m) a -> m ()
514 {-# INLINE copy #-}
515 copy dst src = BOUNDS_CHECK(check) "copy" "overlapping vectors"
516 (not (dst `overlaps` src))
517 $ BOUNDS_CHECK(check) "copy" "length mismatch"
518 (length dst == length src)
519 $ unsafeCopy dst src
520
521 -- | Copy a vector. The two vectors must have the same length and may not
522 -- overlap. This is not checked.
523 unsafeCopy :: (PrimMonad m, MVector v a) => v (PrimState m) a -- ^ target
524 -> v (PrimState m) a -- ^ source
525 -> m ()
526 {-# INLINE unsafeCopy #-}
527 unsafeCopy dst src = UNSAFE_CHECK(check) "unsafeCopy" "length mismatch"
528 (length dst == length src)
529 $ UNSAFE_CHECK(check) "unsafeCopy" "overlapping vectors"
530 (not (dst `overlaps` src))
531 $ basicUnsafeCopy dst src
532
533 -- Subvectors
534 -- ----------
535
536 -- | Yield a part of the mutable vector without copying it.
537 slice :: MVector v a => Int -> Int -> v s a -> v s a
538 {-# INLINE slice #-}
539 slice i n v = BOUNDS_CHECK(checkSlice) "slice" i n (length v)
540 $ unsafeSlice i n v
541
542 take :: MVector v a => Int -> v s a -> v s a
543 {-# INLINE take #-}
544 take n v = unsafeSlice 0 (min (max n 0) (length v)) v
545
546 drop :: MVector v a => Int -> v s a -> v s a
547 {-# INLINE drop #-}
548 drop n v = unsafeSlice (min m n') (max 0 (m - n')) v
549 where
550 n' = max n 0
551 m = length v
552
553 init :: MVector v a => v s a -> v s a
554 {-# INLINE init #-}
555 init v = slice 0 (length v - 1) v
556
557 tail :: MVector v a => v s a -> v s a
558 {-# INLINE tail #-}
559 tail v = slice 1 (length v - 1) v
560
561 -- | Yield a part of the mutable vector without copying it. No bounds checks
562 -- are performed.
563 unsafeSlice :: MVector v a => Int -- ^ starting index
564 -> Int -- ^ length of the slice
565 -> v s a
566 -> v s a
567 {-# INLINE unsafeSlice #-}
568 unsafeSlice i n v = UNSAFE_CHECK(checkSlice) "unsafeSlice" i n (length v)
569 $ basicUnsafeSlice i n v
570
571 unsafeInit :: MVector v a => v s a -> v s a
572 {-# INLINE unsafeInit #-}
573 unsafeInit v = unsafeSlice 0 (length v - 1) v
574
575 unsafeTail :: MVector v a => v s a -> v s a
576 {-# INLINE unsafeTail #-}
577 unsafeTail v = unsafeSlice 1 (length v - 1) v
578
579 unsafeTake :: MVector v a => Int -> v s a -> v s a
580 {-# INLINE unsafeTake #-}
581 unsafeTake n v = unsafeSlice 0 n v
582
583 unsafeDrop :: MVector v a => Int -> v s a -> v s a
584 {-# INLINE unsafeDrop #-}
585 unsafeDrop n v = unsafeSlice n (length v - n) v
586
587 -- Permutations
588 -- ------------
589
590 accum :: (PrimMonad m, MVector v a)
591 => (a -> b -> a) -> v (PrimState m) a -> Stream (Int, b) -> m ()
592 {-# INLINE accum #-}
593 accum f !v s = Stream.mapM_ upd s
594 where
595 {-# INLINE_INNER upd #-}
596 upd (i,b) = do
597 a <- BOUNDS_CHECK(checkIndex) "accum" i (length v)
598 $ unsafeRead v i
599 unsafeWrite v i (f a b)
600
601 update :: (PrimMonad m, MVector v a)
602 => v (PrimState m) a -> Stream (Int, a) -> m ()
603 {-# INLINE update #-}
604 update !v s = Stream.mapM_ upd s
605 where
606 {-# INLINE_INNER upd #-}
607 upd (i,b) = BOUNDS_CHECK(checkIndex) "update" i (length v)
608 $ unsafeWrite v i b
609
610 unsafeAccum :: (PrimMonad m, MVector v a)
611 => (a -> b -> a) -> v (PrimState m) a -> Stream (Int, b) -> m ()
612 {-# INLINE unsafeAccum #-}
613 unsafeAccum f !v s = Stream.mapM_ upd s
614 where
615 {-# INLINE_INNER upd #-}
616 upd (i,b) = do
617 a <- UNSAFE_CHECK(checkIndex) "accum" i (length v)
618 $ unsafeRead v i
619 unsafeWrite v i (f a b)
620
621 unsafeUpdate :: (PrimMonad m, MVector v a)
622 => v (PrimState m) a -> Stream (Int, a) -> m ()
623 {-# INLINE unsafeUpdate #-}
624 unsafeUpdate !v s = Stream.mapM_ upd s
625 where
626 {-# INLINE_INNER upd #-}
627 upd (i,b) = UNSAFE_CHECK(checkIndex) "accum" i (length v)
628 $ unsafeWrite v i b
629
630 reverse :: (PrimMonad m, MVector v a) => v (PrimState m) a -> m ()
631 {-# INLINE reverse #-}
632 reverse !v = reverse_loop 0 (length v - 1)
633 where
634 reverse_loop i j | i < j = do
635 unsafeSwap v i j
636 reverse_loop (i + 1) (j - 1)
637 reverse_loop _ _ = return ()
638
639 unstablePartition :: forall m v a. (PrimMonad m, MVector v a)
640 => (a -> Bool) -> v (PrimState m) a -> m Int
641 {-# INLINE unstablePartition #-}
642 unstablePartition f !v = from_left 0 (length v)
643 where
644 -- NOTE: GHC 6.10.4 panics without the signatures on from_left and
645 -- from_right
646 from_left :: Int -> Int -> m Int
647 from_left i j
648 | i == j = return i
649 | otherwise = do
650 x <- unsafeRead v i
651 if f x
652 then from_left (i+1) j
653 else from_right i (j-1)
654
655 from_right :: Int -> Int -> m Int
656 from_right i j
657 | i == j = return i
658 | otherwise = do
659 x <- unsafeRead v j
660 if f x
661 then do
662 y <- unsafeRead v i
663 unsafeWrite v i x
664 unsafeWrite v j y
665 from_left (i+1) j
666 else from_right i (j-1)
667
668 unstablePartitionStream :: (PrimMonad m, MVector v a)
669 => (a -> Bool) -> Stream a -> m (v (PrimState m) a, v (PrimState m) a)
670 {-# INLINE unstablePartitionStream #-}
671 unstablePartitionStream f s
672 = case upperBound (Stream.size s) of
673 Just n -> unstablePartitionMax f s n
674 Nothing -> partitionUnknown f s
675
676 unstablePartitionMax :: (PrimMonad m, MVector v a)
677 => (a -> Bool) -> Stream a -> Int
678 -> m (v (PrimState m) a, v (PrimState m) a)
679 {-# INLINE unstablePartitionMax #-}
680 unstablePartitionMax f s n
681 = do
682 v <- INTERNAL_CHECK(checkLength) "unstablePartitionMax" n
683 $ unsafeNew n
684 let {-# INLINE_INNER put #-}
685 put (i, j) x
686 | f x = do
687 unsafeWrite v i x
688 return (i+1, j)
689 | otherwise = do
690 unsafeWrite v (j-1) x
691 return (i, j-1)
692
693 (i,j) <- Stream.foldM' put (0, n) s
694 return (unsafeSlice 0 i v, unsafeSlice j (n-j) v)
695
696 partitionUnknown :: (PrimMonad m, MVector v a)
697 => (a -> Bool) -> Stream a -> m (v (PrimState m) a, v (PrimState m) a)
698 {-# INLINE partitionUnknown #-}
699 partitionUnknown f s
700 = do
701 v1 <- unsafeNew 0
702 v2 <- unsafeNew 0
703 (v1', n1, v2', n2) <- Stream.foldM' put (v1, 0, v2, 0) s
704 INTERNAL_CHECK(checkSlice) "partitionUnknown" 0 n1 (length v1')
705 $ INTERNAL_CHECK(checkSlice) "partitionUnknown" 0 n2 (length v2')
706 $ return (unsafeSlice 0 n1 v1', unsafeSlice 0 n2 v2')
707 where
708 -- NOTE: The case distinction has to be on the outside because
709 -- GHC creates a join point for the unsafeWrite even when everything
710 -- is inlined. This is bad because with the join point, v isn't getting
711 -- unboxed.
712 {-# INLINE_INNER put #-}
713 put (v1, i1, v2, i2) x
714 | f x = do
715 v1' <- unsafeAppend1 v1 i1 x
716 return (v1', i1+1, v2, i2)
717 | otherwise = do
718 v2' <- unsafeAppend1 v2 i2 x
719 return (v1, i1, v2', i2+1)
720