1{-# LANGUAGE CPP, 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 14module 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 growFront, unsafeGrowFront, 38 39 -- ** Restricting memory usage 40 clear, 41 42 -- * Accessing individual elements 43 read, write, modify, swap, exchange, 44 unsafeRead, unsafeWrite, unsafeModify, unsafeSwap, unsafeExchange, 45 46 -- * Modifying vectors 47 nextPermutation, 48 49 -- ** Filling and copying 50 set, copy, move, unsafeCopy, unsafeMove, 51 52 -- * Internal operations 53 mstream, mstreamR, 54 unstream, unstreamR, vunstream, 55 munstream, munstreamR, 56 transform, transformR, 57 fill, fillR, 58 unsafeAccum, accum, unsafeUpdate, update, reverse, 59 unstablePartition, unstablePartitionBundle, partitionBundle, 60 partitionWithBundle 61) where 62 63import Data.Vector.Generic.Mutable.Base 64import qualified Data.Vector.Generic.Base as V 65 66import qualified Data.Vector.Fusion.Bundle as Bundle 67import Data.Vector.Fusion.Bundle ( Bundle, MBundle, Chunk(..) ) 68import qualified Data.Vector.Fusion.Bundle.Monadic as MBundle 69import Data.Vector.Fusion.Stream.Monadic ( Stream ) 70import qualified Data.Vector.Fusion.Stream.Monadic as Stream 71import Data.Vector.Fusion.Bundle.Size 72import Data.Vector.Fusion.Util ( delay_inline ) 73 74import Control.Monad.Primitive ( PrimMonad, PrimState ) 75 76import Prelude hiding ( length, null, replicate, reverse, map, read, 77 take, drop, splitAt, init, tail ) 78 79#include "vector.h" 80 81{- 82type family Immutable (v :: * -> * -> *) :: * -> * 83 84-- | Class of mutable vectors parametrised with a primitive state token. 85-- 86class MBundle.Pointer u a => MVector v a where 87 -- | Length of the mutable vector. This method should not be 88 -- called directly, use 'length' instead. 89 basicLength :: v s a -> Int 90 91 -- | Yield a part of the mutable vector without copying it. This method 92 -- should not be called directly, use 'unsafeSlice' instead. 93 basicUnsafeSlice :: Int -- ^ starting index 94 -> Int -- ^ length of the slice 95 -> v s a 96 -> v s a 97 98 -- Check whether two vectors overlap. This method should not be 99 -- called directly, use 'overlaps' instead. 100 basicOverlaps :: v s a -> v s a -> Bool 101 102 -- | Create a mutable vector of the given length. This method should not be 103 -- called directly, use 'unsafeNew' instead. 104 basicUnsafeNew :: PrimMonad m => Int -> m (v (PrimState m) a) 105 106 -- | Create a mutable vector of the given length and fill it with an 107 -- initial value. This method should not be called directly, use 108 -- 'replicate' instead. 109 basicUnsafeReplicate :: PrimMonad m => Int -> a -> m (v (PrimState m) a) 110 111 -- | Yield the element at the given position. This method should not be 112 -- called directly, use 'unsafeRead' instead. 113 basicUnsafeRead :: PrimMonad m => v (PrimState m) a -> Int -> m a 114 115 -- | Replace the element at the given position. This method should not be 116 -- called directly, use 'unsafeWrite' instead. 117 basicUnsafeWrite :: PrimMonad m => v (PrimState m) a -> Int -> a -> m () 118 119 -- | Reset all elements of the vector to some undefined value, clearing all 120 -- references to external objects. This is usually a noop for unboxed 121 -- vectors. This method should not be called directly, use 'clear' instead. 122 basicClear :: PrimMonad m => v (PrimState m) a -> m () 123 124 -- | Set all elements of the vector to the given value. This method should 125 -- not be called directly, use 'set' instead. 126 basicSet :: PrimMonad m => v (PrimState m) a -> a -> m () 127 128 basicUnsafeCopyPointer :: PrimMonad m => v (PrimState m) a 129 -> Immutable v a 130 -> m () 131 132 -- | Copy a vector. The two vectors may not overlap. This method should not 133 -- be called directly, use 'unsafeCopy' instead. 134 basicUnsafeCopy :: PrimMonad m => v (PrimState m) a -- ^ target 135 -> v (PrimState m) a -- ^ source 136 -> m () 137 138 -- | Move the contents of a vector. The two vectors may overlap. This method 139 -- should not be called directly, use 'unsafeMove' instead. 140 basicUnsafeMove :: PrimMonad m => v (PrimState m) a -- ^ target 141 -> v (PrimState m) a -- ^ source 142 -> m () 143 144 -- | Grow a vector by the given number of elements. This method should not be 145 -- called directly, use 'unsafeGrow' instead. 146 basicUnsafeGrow :: PrimMonad m => v (PrimState m) a -> Int 147 -> m (v (PrimState m) a) 148 149 {-# INLINE basicUnsafeReplicate #-} 150 basicUnsafeReplicate n x 151 = do 152 v <- basicUnsafeNew n 153 basicSet v x 154 return v 155 156 {-# INLINE basicClear #-} 157 basicClear _ = return () 158 159 {-# INLINE basicSet #-} 160 basicSet !v x 161 | n == 0 = return () 162 | otherwise = do 163 basicUnsafeWrite v 0 x 164 do_set 1 165 where 166 !n = basicLength v 167 168 do_set i | 2*i < n = do basicUnsafeCopy (basicUnsafeSlice i i v) 169 (basicUnsafeSlice 0 i v) 170 do_set (2*i) 171 | otherwise = basicUnsafeCopy (basicUnsafeSlice i (n-i) v) 172 (basicUnsafeSlice 0 (n-i) v) 173 174 {-# INLINE basicUnsafeCopyPointer #-} 175 basicUnsafeCopyPointer !dst !src = do_copy 0 src 176 where 177 do_copy !i p | Just (x,q) <- MBundle.pget p = do 178 basicUnsafeWrite dst i x 179 do_copy (i+1) q 180 | otherwise = return () 181 182 {-# INLINE basicUnsafeCopy #-} 183 basicUnsafeCopy !dst !src = do_copy 0 184 where 185 !n = basicLength src 186 187 do_copy i | i < n = do 188 x <- basicUnsafeRead src i 189 basicUnsafeWrite dst i x 190 do_copy (i+1) 191 | otherwise = return () 192 193 {-# INLINE basicUnsafeMove #-} 194 basicUnsafeMove !dst !src 195 | basicOverlaps dst src = do 196 srcCopy <- clone src 197 basicUnsafeCopy dst srcCopy 198 | otherwise = basicUnsafeCopy dst src 199 200 {-# INLINE basicUnsafeGrow #-} 201 basicUnsafeGrow v by 202 = do 203 v' <- basicUnsafeNew (n+by) 204 basicUnsafeCopy (basicUnsafeSlice 0 n v') v 205 return v' 206 where 207 n = basicLength v 208-} 209 210-- ------------------ 211-- Internal functions 212-- ------------------ 213 214unsafeAppend1 :: (PrimMonad m, MVector v a) 215 => v (PrimState m) a -> Int -> a -> m (v (PrimState m) a) 216{-# INLINE_INNER unsafeAppend1 #-} 217 -- NOTE: The case distinction has to be on the outside because 218 -- GHC creates a join point for the unsafeWrite even when everything 219 -- is inlined. This is bad because with the join point, v isn't getting 220 -- unboxed. 221unsafeAppend1 v i x 222 | i < length v = do 223 unsafeWrite v i x 224 return v 225 | otherwise = do 226 v' <- enlarge v 227 INTERNAL_CHECK(checkIndex) "unsafeAppend1" i (length v') 228 $ unsafeWrite v' i x 229 return v' 230 231unsafePrepend1 :: (PrimMonad m, MVector v a) 232 => v (PrimState m) a -> Int -> a -> m (v (PrimState m) a, Int) 233{-# INLINE_INNER unsafePrepend1 #-} 234unsafePrepend1 v i x 235 | i /= 0 = do 236 let i' = i-1 237 unsafeWrite v i' x 238 return (v, i') 239 | otherwise = do 240 (v', j) <- enlargeFront v 241 let i' = j-1 242 INTERNAL_CHECK(checkIndex) "unsafePrepend1" i' (length v') 243 $ unsafeWrite v' i' x 244 return (v', i') 245 246mstream :: (PrimMonad m, MVector v a) => v (PrimState m) a -> Stream m a 247{-# INLINE mstream #-} 248mstream v = v `seq` n `seq` (Stream.unfoldrM get 0) 249 where 250 n = length v 251 252 {-# INLINE_INNER get #-} 253 get i | i < n = do x <- unsafeRead v i 254 return $ Just (x, i+1) 255 | otherwise = return $ Nothing 256 257fill :: (PrimMonad m, MVector v a) 258 => v (PrimState m) a -> Stream m a -> m (v (PrimState m) a) 259{-# INLINE fill #-} 260fill v s = v `seq` do 261 n' <- Stream.foldM put 0 s 262 return $ unsafeSlice 0 n' v 263 where 264 {-# INLINE_INNER put #-} 265 put i x = do 266 INTERNAL_CHECK(checkIndex) "fill" i (length v) 267 $ unsafeWrite v i x 268 return (i+1) 269 270transform 271 :: (PrimMonad m, MVector v a) 272 => (Stream m a -> Stream m a) -> v (PrimState m) a -> m (v (PrimState m) a) 273{-# INLINE_FUSED transform #-} 274transform f v = fill v (f (mstream v)) 275 276mstreamR :: (PrimMonad m, MVector v a) => v (PrimState m) a -> Stream m a 277{-# INLINE mstreamR #-} 278mstreamR v = v `seq` n `seq` (Stream.unfoldrM get 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 289fillR :: (PrimMonad m, MVector v a) 290 => v (PrimState m) a -> Stream m a -> m (v (PrimState m) a) 291{-# INLINE fillR #-} 292fillR v s = v `seq` do 293 i <- Stream.foldM put n s 294 return $ unsafeSlice i (n-i) v 295 where 296 n = length v 297 298 {-# INLINE_INNER put #-} 299 put i x = do 300 unsafeWrite v j x 301 return j 302 where 303 j = i-1 304 305transformR 306 :: (PrimMonad m, MVector v a) 307 => (Stream m a -> Stream m a) -> v (PrimState m) a -> m (v (PrimState m) a) 308{-# INLINE_FUSED transformR #-} 309transformR f v = fillR v (f (mstreamR v)) 310 311-- | Create a new mutable vector and fill it with elements from the 'Bundle'. 312-- The vector will grow exponentially if the maximum size of the 'Bundle' is 313-- unknown. 314unstream :: (PrimMonad m, MVector v a) 315 => Bundle u a -> m (v (PrimState m) a) 316-- NOTE: replace INLINE_FUSED by INLINE? (also in unstreamR) 317{-# INLINE_FUSED unstream #-} 318unstream s = munstream (Bundle.lift s) 319 320-- | Create a new mutable vector and fill it with elements from the monadic 321-- stream. The vector will grow exponentially if the maximum size of the stream 322-- is unknown. 323munstream :: (PrimMonad m, MVector v a) 324 => MBundle m u a -> m (v (PrimState m) a) 325{-# INLINE_FUSED munstream #-} 326munstream s = case upperBound (MBundle.size s) of 327 Just n -> munstreamMax s n 328 Nothing -> munstreamUnknown s 329 330-- FIXME: I can't think of how to prevent GHC from floating out 331-- unstreamUnknown. That is bad because SpecConstr then generates two 332-- specialisations: one for when it is called from unstream (it doesn't know 333-- the shape of the vector) and one for when the vector has grown. To see the 334-- problem simply compile this: 335-- 336-- fromList = Data.Vector.Unboxed.unstream . Bundle.fromList 337-- 338-- I'm not sure this still applies (19/04/2010) 339 340munstreamMax :: (PrimMonad m, MVector v a) 341 => MBundle m u a -> Int -> m (v (PrimState m) a) 342{-# INLINE munstreamMax #-} 343munstreamMax s n 344 = do 345 v <- INTERNAL_CHECK(checkLength) "munstreamMax" n 346 $ unsafeNew n 347 let put i x = do 348 INTERNAL_CHECK(checkIndex) "munstreamMax" i n 349 $ unsafeWrite v i x 350 return (i+1) 351 n' <- MBundle.foldM' put 0 s 352 return $ INTERNAL_CHECK(checkSlice) "munstreamMax" 0 n' n 353 $ unsafeSlice 0 n' v 354 355munstreamUnknown :: (PrimMonad m, MVector v a) 356 => MBundle m u a -> m (v (PrimState m) a) 357{-# INLINE munstreamUnknown #-} 358munstreamUnknown s 359 = do 360 v <- unsafeNew 0 361 (v', n) <- MBundle.foldM put (v, 0) s 362 return $ INTERNAL_CHECK(checkSlice) "munstreamUnknown" 0 n (length v') 363 $ unsafeSlice 0 n v' 364 where 365 {-# INLINE_INNER put #-} 366 put (v,i) x = do 367 v' <- unsafeAppend1 v i x 368 return (v',i+1) 369 370 371 372 373 374 375 376-- | Create a new mutable vector and fill it with elements from the 'Bundle'. 377-- The vector will grow exponentially if the maximum size of the 'Bundle' is 378-- unknown. 379vunstream :: (PrimMonad m, V.Vector v a) 380 => Bundle v a -> m (V.Mutable v (PrimState m) a) 381-- NOTE: replace INLINE_FUSED by INLINE? (also in unstreamR) 382{-# INLINE_FUSED vunstream #-} 383vunstream s = vmunstream (Bundle.lift s) 384 385-- | Create a new mutable vector and fill it with elements from the monadic 386-- stream. The vector will grow exponentially if the maximum size of the stream 387-- is unknown. 388vmunstream :: (PrimMonad m, V.Vector v a) 389 => MBundle m v a -> m (V.Mutable v (PrimState m) a) 390{-# INLINE_FUSED vmunstream #-} 391vmunstream s = case upperBound (MBundle.size s) of 392 Just n -> vmunstreamMax s n 393 Nothing -> vmunstreamUnknown s 394 395-- FIXME: I can't think of how to prevent GHC from floating out 396-- unstreamUnknown. That is bad because SpecConstr then generates two 397-- specialisations: one for when it is called from unstream (it doesn't know 398-- the shape of the vector) and one for when the vector has grown. To see the 399-- problem simply compile this: 400-- 401-- fromList = Data.Vector.Unboxed.unstream . Bundle.fromList 402-- 403-- I'm not sure this still applies (19/04/2010) 404 405vmunstreamMax :: (PrimMonad m, V.Vector v a) 406 => MBundle m v a -> Int -> m (V.Mutable v (PrimState m) a) 407{-# INLINE vmunstreamMax #-} 408vmunstreamMax s n 409 = do 410 v <- INTERNAL_CHECK(checkLength) "munstreamMax" n 411 $ unsafeNew n 412 let {-# INLINE_INNER copyChunk #-} 413 copyChunk i (Chunk m f) = 414 INTERNAL_CHECK(checkSlice) "munstreamMax.copyChunk" i m (length v) $ do 415 f (basicUnsafeSlice i m v) 416 return (i+m) 417 418 n' <- Stream.foldlM' copyChunk 0 (MBundle.chunks s) 419 return $ INTERNAL_CHECK(checkSlice) "munstreamMax" 0 n' n 420 $ unsafeSlice 0 n' v 421 422vmunstreamUnknown :: (PrimMonad m, V.Vector v a) 423 => MBundle m v a -> m (V.Mutable v (PrimState m) a) 424{-# INLINE vmunstreamUnknown #-} 425vmunstreamUnknown s 426 = do 427 v <- unsafeNew 0 428 (v', n) <- Stream.foldlM copyChunk (v,0) (MBundle.chunks s) 429 return $ INTERNAL_CHECK(checkSlice) "munstreamUnknown" 0 n (length v') 430 $ unsafeSlice 0 n v' 431 where 432 {-# INLINE_INNER copyChunk #-} 433 copyChunk (v,i) (Chunk n f) 434 = do 435 let j = i+n 436 v' <- if basicLength v < j 437 then unsafeGrow v (delay_inline max (enlarge_delta v) (j - basicLength v)) 438 else return v 439 INTERNAL_CHECK(checkSlice) "munstreamUnknown.copyChunk" i n (length v') 440 $ f (basicUnsafeSlice i n v') 441 return (v',j) 442 443 444 445 446-- | Create a new mutable vector and fill it with elements from the 'Bundle' 447-- from right to left. The vector will grow exponentially if the maximum size 448-- of the 'Bundle' is unknown. 449unstreamR :: (PrimMonad m, MVector v a) 450 => Bundle u a -> m (v (PrimState m) a) 451-- NOTE: replace INLINE_FUSED by INLINE? (also in unstream) 452{-# INLINE_FUSED unstreamR #-} 453unstreamR s = munstreamR (Bundle.lift s) 454 455-- | Create a new mutable vector and fill it with elements from the monadic 456-- stream from right to left. The vector will grow exponentially if the maximum 457-- size of the stream is unknown. 458munstreamR :: (PrimMonad m, MVector v a) 459 => MBundle m u a -> m (v (PrimState m) a) 460{-# INLINE_FUSED munstreamR #-} 461munstreamR s = case upperBound (MBundle.size s) of 462 Just n -> munstreamRMax s n 463 Nothing -> munstreamRUnknown s 464 465munstreamRMax :: (PrimMonad m, MVector v a) 466 => MBundle m u a -> Int -> m (v (PrimState m) a) 467{-# INLINE munstreamRMax #-} 468munstreamRMax s n 469 = do 470 v <- INTERNAL_CHECK(checkLength) "munstreamRMax" n 471 $ unsafeNew n 472 let put i x = do 473 let i' = i-1 474 INTERNAL_CHECK(checkIndex) "munstreamRMax" i' n 475 $ unsafeWrite v i' x 476 return i' 477 i <- MBundle.foldM' put n s 478 return $ INTERNAL_CHECK(checkSlice) "munstreamRMax" i (n-i) n 479 $ unsafeSlice i (n-i) v 480 481munstreamRUnknown :: (PrimMonad m, MVector v a) 482 => MBundle m u a -> m (v (PrimState m) a) 483{-# INLINE munstreamRUnknown #-} 484munstreamRUnknown s 485 = do 486 v <- unsafeNew 0 487 (v', i) <- MBundle.foldM put (v, 0) s 488 let n = length v' 489 return $ INTERNAL_CHECK(checkSlice) "unstreamRUnknown" i (n-i) n 490 $ unsafeSlice i (n-i) v' 491 where 492 {-# INLINE_INNER put #-} 493 put (v,i) x = unsafePrepend1 v i x 494 495-- Length 496-- ------ 497 498-- | Length of the mutable vector. 499length :: MVector v a => v s a -> Int 500{-# INLINE length #-} 501length = basicLength 502 503-- | Check whether the vector is empty 504null :: MVector v a => v s a -> Bool 505{-# INLINE null #-} 506null v = length v == 0 507 508-- Extracting subvectors 509-- --------------------- 510 511-- | Yield a part of the mutable vector without copying it. The vector must 512-- contain at least @i+n@ elements. 513slice :: MVector v a 514 => Int -- ^ @i@ starting index 515 -> Int -- ^ @n@ length 516 -> v s a 517 -> v s a 518{-# INLINE slice #-} 519slice i n v = BOUNDS_CHECK(checkSlice) "slice" i n (length v) 520 $ unsafeSlice i n v 521 522take :: MVector v a => Int -> v s a -> v s a 523{-# INLINE take #-} 524take n v = unsafeSlice 0 (min (max n 0) (length v)) v 525 526drop :: MVector v a => Int -> v s a -> v s a 527{-# INLINE drop #-} 528drop n v = unsafeSlice (min m n') (max 0 (m - n')) v 529 where 530 n' = max n 0 531 m = length v 532 533{-# INLINE splitAt #-} 534splitAt :: MVector v a => Int -> v s a -> (v s a, v s a) 535splitAt n v = ( unsafeSlice 0 m v 536 , unsafeSlice m (max 0 (len - n')) v 537 ) 538 where 539 m = min n' len 540 n' = max n 0 541 len = length v 542 543init :: MVector v a => v s a -> v s a 544{-# INLINE init #-} 545init v = slice 0 (length v - 1) v 546 547tail :: MVector v a => v s a -> v s a 548{-# INLINE tail #-} 549tail v = slice 1 (length v - 1) v 550 551-- | Yield a part of the mutable vector without copying it. No bounds checks 552-- are performed. 553unsafeSlice :: MVector v a => Int -- ^ starting index 554 -> Int -- ^ length of the slice 555 -> v s a 556 -> v s a 557{-# INLINE unsafeSlice #-} 558unsafeSlice i n v = UNSAFE_CHECK(checkSlice) "unsafeSlice" i n (length v) 559 $ basicUnsafeSlice i n v 560 561unsafeInit :: MVector v a => v s a -> v s a 562{-# INLINE unsafeInit #-} 563unsafeInit v = unsafeSlice 0 (length v - 1) v 564 565unsafeTail :: MVector v a => v s a -> v s a 566{-# INLINE unsafeTail #-} 567unsafeTail v = unsafeSlice 1 (length v - 1) v 568 569unsafeTake :: MVector v a => Int -> v s a -> v s a 570{-# INLINE unsafeTake #-} 571unsafeTake n v = unsafeSlice 0 n v 572 573unsafeDrop :: MVector v a => Int -> v s a -> v s a 574{-# INLINE unsafeDrop #-} 575unsafeDrop n v = unsafeSlice n (length v - n) v 576 577-- Overlapping 578-- ----------- 579 580-- | Check whether two vectors overlap. 581overlaps :: MVector v a => v s a -> v s a -> Bool 582{-# INLINE overlaps #-} 583overlaps = basicOverlaps 584 585-- Initialisation 586-- -------------- 587 588-- | Create a mutable vector of the given length. 589new :: (PrimMonad m, MVector v a) => Int -> m (v (PrimState m) a) 590{-# INLINE new #-} 591new n = BOUNDS_CHECK(checkLength) "new" n 592 $ unsafeNew n >>= \v -> basicInitialize v >> return v 593 594-- | Create a mutable vector of the given length. The vector content 595-- should be presumed uninitialized. However exact semantics depends 596-- on vector implementation. For example unboxed and storable 597-- vectors will create vector filled with whatever underlying memory 598-- buffer happens to contain, while boxed vector's elements are 599-- initialized to bottoms which will throw exception when evaluated. 600-- 601-- @since 0.4 602unsafeNew :: (PrimMonad m, MVector v a) => Int -> m (v (PrimState m) a) 603{-# INLINE unsafeNew #-} 604unsafeNew 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. 609replicate :: (PrimMonad m, MVector v a) => Int -> a -> m (v (PrimState m) a) 610{-# INLINE replicate #-} 611replicate 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. 615replicateM :: (PrimMonad m, MVector v a) => Int -> m a -> m (v (PrimState m) a) 616{-# INLINE replicateM #-} 617replicateM n m = munstream (MBundle.replicateM n m) 618 619-- | Create a copy of a mutable vector. 620clone :: (PrimMonad m, MVector v a) => v (PrimState m) a -> m (v (PrimState m) a) 621{-# INLINE clone #-} 622clone 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 not be 631-- negative otherwise error is thrown. Semantics of this function is exactly the 632-- same as `unsafeGrow`, except that it will initialize the newly 633-- allocated memory first. 634-- 635-- It is important to note that mutating the returned vector will not affect the 636-- vector that was used as a source. In other words it does not, nor will it 637-- ever have the semantics of @realloc@ from C. 638-- 639-- > grow mv 0 === clone mv 640-- 641-- @since 0.4.0 642grow :: (PrimMonad m, MVector v a) 643 => v (PrimState m) a -> Int -> m (v (PrimState m) a) 644{-# INLINE grow #-} 645grow v by = BOUNDS_CHECK(checkLength) "grow" by 646 $ do vnew <- unsafeGrow v by 647 basicInitialize $ basicUnsafeSlice (length v) by vnew 648 return vnew 649 650-- | Same as `grow`, except that it copies data towards the end of the newly 651-- allocated vector making extra space available at the beginning. 652-- 653-- @since 0.11.0.0 654growFront :: (PrimMonad m, MVector v a) 655 => v (PrimState m) a -> Int -> m (v (PrimState m) a) 656{-# INLINE growFront #-} 657growFront v by = BOUNDS_CHECK(checkLength) "growFront" by 658 $ do vnew <- unsafeGrowFront v by 659 basicInitialize $ basicUnsafeSlice 0 by vnew 660 return vnew 661 662enlarge_delta :: MVector v a => v s a -> Int 663enlarge_delta v = max (length v) 1 664 665-- | Grow a vector logarithmically 666enlarge :: (PrimMonad m, MVector v a) 667 => v (PrimState m) a -> m (v (PrimState m) a) 668{-# INLINE enlarge #-} 669enlarge v = do vnew <- unsafeGrow v by 670 basicInitialize $ basicUnsafeSlice (length v) by vnew 671 return vnew 672 where 673 by = enlarge_delta v 674 675enlargeFront :: (PrimMonad m, MVector v a) 676 => v (PrimState m) a -> m (v (PrimState m) a, Int) 677{-# INLINE enlargeFront #-} 678enlargeFront v = do 679 v' <- unsafeGrowFront v by 680 basicInitialize $ basicUnsafeSlice 0 by v' 681 return (v', by) 682 where 683 by = enlarge_delta v 684 685-- | Grow a vector by allocating a new mutable vector of the same size plus the 686-- the given number of elements and copying all the data over to the new vector 687-- starting at its beginning. The newly allocated memory is not initialized and 688-- the extra space at the end will likely contain garbage data or uninitialzed 689-- error. Use `unsafeGrowFront` to make the extra space available in the front 690-- of the new vector. 691-- 692-- It is important to note that mutating the returned vector will not affect 693-- elements of the vector that was used as a source. In other words it does not, 694-- nor will it ever have the semantics of @realloc@ from C. Keep in mind, 695-- however, that values themselves can be of a mutable type 696-- (eg. `Foreign.Ptr.Ptr`), in which case it would be possible to affect values 697-- stored in both vectors. 698-- 699-- > unsafeGrow mv 0 === clone mv 700-- 701-- @since 0.4.0 702unsafeGrow :: 703 (PrimMonad m, MVector v a) 704 => v (PrimState m) a 705 -- ^ A mutable vector to copy the data from. 706 -> Int 707 -- ^ Number of elements to grow the vector by. It must be non-negative but 708 -- this is not checked. 709 -> m (v (PrimState m) a) 710{-# INLINE unsafeGrow #-} 711unsafeGrow v n = UNSAFE_CHECK(checkLength) "unsafeGrow" n 712 $ basicUnsafeGrow v n 713 714-- | Same as `unsafeGrow`, except that it copies data towards the end of the 715-- newly allocated vector making extra space available at the beginning. 716-- 717-- @since 0.11.0.0 718unsafeGrowFront :: (PrimMonad m, MVector v a) 719 => v (PrimState m) a -> Int -> m (v (PrimState m) a) 720{-# INLINE unsafeGrowFront #-} 721unsafeGrowFront v by = UNSAFE_CHECK(checkLength) "unsafeGrowFront" by 722 $ do 723 let n = length v 724 v' <- basicUnsafeNew (by+n) 725 basicUnsafeCopy (basicUnsafeSlice by n v') v 726 return v' 727 728-- Restricting memory usage 729-- ------------------------ 730 731-- | Reset all elements of the vector to some undefined value, clearing all 732-- references to external objects. This is usually a noop for unboxed vectors. 733clear :: (PrimMonad m, MVector v a) => v (PrimState m) a -> m () 734{-# INLINE clear #-} 735clear = basicClear 736 737-- Accessing individual elements 738-- ----------------------------- 739 740-- | Yield the element at the given position. 741read :: (PrimMonad m, MVector v a) => v (PrimState m) a -> Int -> m a 742{-# INLINE read #-} 743read v i = BOUNDS_CHECK(checkIndex) "read" i (length v) 744 $ unsafeRead v i 745 746-- | Replace the element at the given position. 747write :: (PrimMonad m, MVector v a) => v (PrimState m) a -> Int -> a -> m () 748{-# INLINE write #-} 749write v i x = BOUNDS_CHECK(checkIndex) "write" i (length v) 750 $ unsafeWrite v i x 751 752-- | Modify the element at the given position. 753modify :: (PrimMonad m, MVector v a) => v (PrimState m) a -> (a -> a) -> Int -> m () 754{-# INLINE modify #-} 755modify v f i = BOUNDS_CHECK(checkIndex) "modify" i (length v) 756 $ unsafeModify v f i 757 758-- | Swap the elements at the given positions. 759swap :: (PrimMonad m, MVector v a) => v (PrimState m) a -> Int -> Int -> m () 760{-# INLINE swap #-} 761swap v i j = BOUNDS_CHECK(checkIndex) "swap" i (length v) 762 $ BOUNDS_CHECK(checkIndex) "swap" j (length v) 763 $ unsafeSwap v i j 764 765-- | Replace the element at the given position and return the old element. 766exchange :: (PrimMonad m, MVector v a) => v (PrimState m) a -> Int -> a -> m a 767{-# INLINE exchange #-} 768exchange v i x = BOUNDS_CHECK(checkIndex) "exchange" i (length v) 769 $ unsafeExchange v i x 770 771-- | Yield the element at the given position. No bounds checks are performed. 772unsafeRead :: (PrimMonad m, MVector v a) => v (PrimState m) a -> Int -> m a 773{-# INLINE unsafeRead #-} 774unsafeRead v i = UNSAFE_CHECK(checkIndex) "unsafeRead" i (length v) 775 $ basicUnsafeRead v i 776 777-- | Replace the element at the given position. No bounds checks are performed. 778unsafeWrite :: (PrimMonad m, MVector v a) 779 => v (PrimState m) a -> Int -> a -> m () 780{-# INLINE unsafeWrite #-} 781unsafeWrite v i x = UNSAFE_CHECK(checkIndex) "unsafeWrite" i (length v) 782 $ basicUnsafeWrite v i x 783 784-- | Modify the element at the given position. No bounds checks are performed. 785unsafeModify :: (PrimMonad m, MVector v a) => v (PrimState m) a -> (a -> a) -> Int -> m () 786{-# INLINE unsafeModify #-} 787unsafeModify v f i = UNSAFE_CHECK(checkIndex) "unsafeModify" i (length v) 788 $ basicUnsafeRead v i >>= \x -> 789 basicUnsafeWrite v i (f x) 790 791-- | Swap the elements at the given positions. No bounds checks are performed. 792unsafeSwap :: (PrimMonad m, MVector v a) 793 => v (PrimState m) a -> Int -> Int -> m () 794{-# INLINE unsafeSwap #-} 795unsafeSwap v i j = UNSAFE_CHECK(checkIndex) "unsafeSwap" i (length v) 796 $ UNSAFE_CHECK(checkIndex) "unsafeSwap" j (length v) 797 $ do 798 x <- unsafeRead v i 799 y <- unsafeRead v j 800 unsafeWrite v i y 801 unsafeWrite v j x 802 803-- | Replace the element at the given position and return the old element. No 804-- bounds checks are performed. 805unsafeExchange :: (PrimMonad m, MVector v a) 806 => v (PrimState m) a -> Int -> a -> m a 807{-# INLINE unsafeExchange #-} 808unsafeExchange v i x = UNSAFE_CHECK(checkIndex) "unsafeExchange" i (length v) 809 $ do 810 y <- unsafeRead v i 811 unsafeWrite v i x 812 return y 813 814-- Filling and copying 815-- ------------------- 816 817-- | Set all elements of the vector to the given value. 818set :: (PrimMonad m, MVector v a) => v (PrimState m) a -> a -> m () 819{-# INLINE set #-} 820set = basicSet 821 822-- | Copy a vector. The two vectors must have the same length and may not 823-- overlap. 824copy :: (PrimMonad m, MVector v a) => v (PrimState m) a -- ^ target 825 -> v (PrimState m) a -- ^ source 826 -> m () 827{-# INLINE copy #-} 828copy dst src = BOUNDS_CHECK(check) "copy" "overlapping vectors" 829 (not (dst `overlaps` src)) 830 $ BOUNDS_CHECK(check) "copy" "length mismatch" 831 (length dst == length src) 832 $ unsafeCopy dst src 833 834-- | Move the contents of a vector. The two vectors must have the same 835-- length. 836-- 837-- If the vectors do not overlap, then this is equivalent to 'copy'. 838-- Otherwise, the copying is performed as if the source vector were 839-- copied to a temporary vector and then the temporary vector was copied 840-- to the target vector. 841move :: (PrimMonad m, MVector v a) 842 => v (PrimState m) a -- ^ target 843 -> v (PrimState m) a -- ^ source 844 -> m () 845{-# INLINE move #-} 846move dst src = BOUNDS_CHECK(check) "move" "length mismatch" 847 (length dst == length src) 848 $ unsafeMove dst src 849 850-- | Copy a vector. The two vectors must have the same length and may not 851-- overlap. This is not checked. 852unsafeCopy :: (PrimMonad m, MVector v a) => v (PrimState m) a -- ^ target 853 -> v (PrimState m) a -- ^ source 854 -> m () 855{-# INLINE unsafeCopy #-} 856unsafeCopy dst src = UNSAFE_CHECK(check) "unsafeCopy" "length mismatch" 857 (length dst == length src) 858 $ UNSAFE_CHECK(check) "unsafeCopy" "overlapping vectors" 859 (not (dst `overlaps` src)) 860 $ (dst `seq` src `seq` basicUnsafeCopy dst src) 861 862-- | Move the contents of a vector. The two vectors must have the same 863-- length, but this is not checked. 864-- 865-- If the vectors do not overlap, then this is equivalent to 'unsafeCopy'. 866-- Otherwise, the copying is performed as if the source vector were 867-- copied to a temporary vector and then the temporary vector was copied 868-- to the target vector. 869unsafeMove :: (PrimMonad m, MVector v a) => v (PrimState m) a -- ^ target 870 -> v (PrimState m) a -- ^ source 871 -> m () 872{-# INLINE unsafeMove #-} 873unsafeMove dst src = UNSAFE_CHECK(check) "unsafeMove" "length mismatch" 874 (length dst == length src) 875 $ (dst `seq` src `seq` basicUnsafeMove dst src) 876 877-- Permutations 878-- ------------ 879 880accum :: (PrimMonad m, MVector v a) 881 => (a -> b -> a) -> v (PrimState m) a -> Bundle u (Int, b) -> m () 882{-# INLINE accum #-} 883accum f !v s = Bundle.mapM_ upd s 884 where 885 {-# INLINE_INNER upd #-} 886 upd (i,b) = do 887 a <- BOUNDS_CHECK(checkIndex) "accum" i n 888 $ unsafeRead v i 889 unsafeWrite v i (f a b) 890 891 !n = length v 892 893update :: (PrimMonad m, MVector v a) 894 => v (PrimState m) a -> Bundle u (Int, a) -> m () 895{-# INLINE update #-} 896update !v s = Bundle.mapM_ upd s 897 where 898 {-# INLINE_INNER upd #-} 899 upd (i,b) = BOUNDS_CHECK(checkIndex) "update" i n 900 $ unsafeWrite v i b 901 902 !n = length v 903 904unsafeAccum :: (PrimMonad m, MVector v a) 905 => (a -> b -> a) -> v (PrimState m) a -> Bundle u (Int, b) -> m () 906{-# INLINE unsafeAccum #-} 907unsafeAccum f !v s = Bundle.mapM_ upd s 908 where 909 {-# INLINE_INNER upd #-} 910 upd (i,b) = do 911 a <- UNSAFE_CHECK(checkIndex) "accum" i n 912 $ unsafeRead v i 913 unsafeWrite v i (f a b) 914 915 !n = length v 916 917unsafeUpdate :: (PrimMonad m, MVector v a) 918 => v (PrimState m) a -> Bundle u (Int, a) -> m () 919{-# INLINE unsafeUpdate #-} 920unsafeUpdate !v s = Bundle.mapM_ upd s 921 where 922 {-# INLINE_INNER upd #-} 923 upd (i,b) = UNSAFE_CHECK(checkIndex) "accum" i n 924 $ unsafeWrite v i b 925 926 !n = length v 927 928reverse :: (PrimMonad m, MVector v a) => v (PrimState m) a -> m () 929{-# INLINE reverse #-} 930reverse !v = reverse_loop 0 (length v - 1) 931 where 932 reverse_loop i j | i < j = do 933 unsafeSwap v i j 934 reverse_loop (i + 1) (j - 1) 935 reverse_loop _ _ = return () 936 937unstablePartition :: forall m v a. (PrimMonad m, MVector v a) 938 => (a -> Bool) -> v (PrimState m) a -> m Int 939{-# INLINE unstablePartition #-} 940unstablePartition f !v = from_left 0 (length v) 941 where 942 -- NOTE: GHC 6.10.4 panics without the signatures on from_left and 943 -- from_right 944 from_left :: Int -> Int -> m Int 945 from_left i j 946 | i == j = return i 947 | otherwise = do 948 x <- unsafeRead v i 949 if f x 950 then from_left (i+1) j 951 else from_right i (j-1) 952 953 from_right :: Int -> Int -> m Int 954 from_right i j 955 | i == j = return i 956 | otherwise = do 957 x <- unsafeRead v j 958 if f x 959 then do 960 y <- unsafeRead v i 961 unsafeWrite v i x 962 unsafeWrite v j y 963 from_left (i+1) j 964 else from_right i (j-1) 965 966unstablePartitionBundle :: (PrimMonad m, MVector v a) 967 => (a -> Bool) -> Bundle u a -> m (v (PrimState m) a, v (PrimState m) a) 968{-# INLINE unstablePartitionBundle #-} 969unstablePartitionBundle f s 970 = case upperBound (Bundle.size s) of 971 Just n -> unstablePartitionMax f s n 972 Nothing -> partitionUnknown f s 973 974unstablePartitionMax :: (PrimMonad m, MVector v a) 975 => (a -> Bool) -> Bundle u a -> Int 976 -> m (v (PrimState m) a, v (PrimState m) a) 977{-# INLINE unstablePartitionMax #-} 978unstablePartitionMax f s n 979 = do 980 v <- INTERNAL_CHECK(checkLength) "unstablePartitionMax" n 981 $ unsafeNew n 982 let {-# INLINE_INNER put #-} 983 put (i, j) x 984 | f x = do 985 unsafeWrite v i x 986 return (i+1, j) 987 | otherwise = do 988 unsafeWrite v (j-1) x 989 return (i, j-1) 990 991 (i,j) <- Bundle.foldM' put (0, n) s 992 return (unsafeSlice 0 i v, unsafeSlice j (n-j) v) 993 994partitionBundle :: (PrimMonad m, MVector v a) 995 => (a -> Bool) -> Bundle u a -> m (v (PrimState m) a, v (PrimState m) a) 996{-# INLINE partitionBundle #-} 997partitionBundle f s 998 = case upperBound (Bundle.size s) of 999 Just n -> partitionMax f s n 1000 Nothing -> partitionUnknown f s 1001 1002partitionMax :: (PrimMonad m, MVector v a) 1003 => (a -> Bool) -> Bundle u a -> Int -> m (v (PrimState m) a, v (PrimState m) a) 1004{-# INLINE partitionMax #-} 1005partitionMax f s n 1006 = do 1007 v <- INTERNAL_CHECK(checkLength) "unstablePartitionMax" n 1008 $ unsafeNew n 1009 1010 let {-# INLINE_INNER put #-} 1011 put (i,j) x 1012 | f x = do 1013 unsafeWrite v i x 1014 return (i+1,j) 1015 1016 | otherwise = let j' = j-1 in 1017 do 1018 unsafeWrite v j' x 1019 return (i,j') 1020 1021 (i,j) <- Bundle.foldM' put (0,n) s 1022 INTERNAL_CHECK(check) "partitionMax" "invalid indices" (i <= j) 1023 $ return () 1024 let l = unsafeSlice 0 i v 1025 r = unsafeSlice j (n-j) v 1026 reverse r 1027 return (l,r) 1028 1029partitionUnknown :: (PrimMonad m, MVector v a) 1030 => (a -> Bool) -> Bundle u a -> m (v (PrimState m) a, v (PrimState m) a) 1031{-# INLINE partitionUnknown #-} 1032partitionUnknown f s 1033 = do 1034 v1 <- unsafeNew 0 1035 v2 <- unsafeNew 0 1036 (v1', n1, v2', n2) <- Bundle.foldM' put (v1, 0, v2, 0) s 1037 INTERNAL_CHECK(checkSlice) "partitionUnknown" 0 n1 (length v1') 1038 $ INTERNAL_CHECK(checkSlice) "partitionUnknown" 0 n2 (length v2') 1039 $ return (unsafeSlice 0 n1 v1', unsafeSlice 0 n2 v2') 1040 where 1041 -- NOTE: The case distinction has to be on the outside because 1042 -- GHC creates a join point for the unsafeWrite even when everything 1043 -- is inlined. This is bad because with the join point, v isn't getting 1044 -- unboxed. 1045 {-# INLINE_INNER put #-} 1046 put (v1, i1, v2, i2) x 1047 | f x = do 1048 v1' <- unsafeAppend1 v1 i1 x 1049 return (v1', i1+1, v2, i2) 1050 | otherwise = do 1051 v2' <- unsafeAppend1 v2 i2 x 1052 return (v1, i1, v2', i2+1) 1053 1054 1055partitionWithBundle :: (PrimMonad m, MVector v a, MVector v b, MVector v c) 1056 => (a -> Either b c) -> Bundle u a -> m (v (PrimState m) b, v (PrimState m) c) 1057{-# INLINE partitionWithBundle #-} 1058partitionWithBundle f s 1059 = case upperBound (Bundle.size s) of 1060 Just n -> partitionWithMax f s n 1061 Nothing -> partitionWithUnknown f s 1062 1063partitionWithMax :: (PrimMonad m, MVector v a, MVector v b, MVector v c) 1064 => (a -> Either b c) -> Bundle u a -> Int -> m (v (PrimState m) b, v (PrimState m) c) 1065{-# INLINE partitionWithMax #-} 1066partitionWithMax f s n 1067 = do 1068 v1 <- unsafeNew n 1069 v2 <- unsafeNew n 1070 let {-# INLINE_INNER put #-} 1071 put (i1, i2) x = case f x of 1072 Left b -> do 1073 unsafeWrite v1 i1 b 1074 return (i1+1, i2) 1075 Right c -> do 1076 unsafeWrite v2 i2 c 1077 return (i1, i2+1) 1078 (n1, n2) <- Bundle.foldM' put (0, 0) s 1079 INTERNAL_CHECK(checkSlice) "partitionEithersMax" 0 n1 (length v1) 1080 $ INTERNAL_CHECK(checkSlice) "partitionEithersMax" 0 n2 (length v2) 1081 $ return (unsafeSlice 0 n1 v1, unsafeSlice 0 n2 v2) 1082 1083partitionWithUnknown :: forall m v u a b c. 1084 (PrimMonad m, MVector v a, MVector v b, MVector v c) 1085 => (a -> Either b c) -> Bundle u a -> m (v (PrimState m) b, v (PrimState m) c) 1086{-# INLINE partitionWithUnknown #-} 1087partitionWithUnknown f s 1088 = do 1089 v1 <- unsafeNew 0 1090 v2 <- unsafeNew 0 1091 (v1', n1, v2', n2) <- Bundle.foldM' put (v1, 0, v2, 0) s 1092 INTERNAL_CHECK(checkSlice) "partitionEithersUnknown" 0 n1 (length v1') 1093 $ INTERNAL_CHECK(checkSlice) "partitionEithersUnknown" 0 n2 (length v2') 1094 $ return (unsafeSlice 0 n1 v1', unsafeSlice 0 n2 v2') 1095 where 1096 put :: (v (PrimState m) b, Int, v (PrimState m) c, Int) 1097 -> a 1098 -> m (v (PrimState m) b, Int, v (PrimState m) c, Int) 1099 {-# INLINE_INNER put #-} 1100 put (v1, i1, v2, i2) x = case f x of 1101 Left b -> do 1102 v1' <- unsafeAppend1 v1 i1 b 1103 return (v1', i1+1, v2, i2) 1104 Right c -> do 1105 v2' <- unsafeAppend1 v2 i2 c 1106 return (v1, i1, v2', i2+1) 1107 1108{- 1109http://en.wikipedia.org/wiki/Permutation#Algorithms_to_generate_permutations 1110 1111The following algorithm generates the next permutation lexicographically after 1112a given permutation. It changes the given permutation in-place. 1113 11141. Find the largest index k such that a[k] < a[k + 1]. If no such index exists, 1115 the permutation is the last permutation. 11162. Find the largest index l greater than k such that a[k] < a[l]. 11173. Swap the value of a[k] with that of a[l]. 11184. Reverse the sequence from a[k + 1] up to and including the final element a[n] 1119-} 1120 1121-- | Compute the next (lexicographically) permutation of given vector in-place. 1122-- Returns False when input is the last permutation 1123nextPermutation :: (PrimMonad m,Ord e,MVector v e) => v (PrimState m) e -> m Bool 1124nextPermutation v 1125 | dim < 2 = return False 1126 | otherwise = do 1127 val <- unsafeRead v 0 1128 (k,l) <- loop val (-1) 0 val 1 1129 if k < 0 1130 then return False 1131 else unsafeSwap v k l >> 1132 reverse (unsafeSlice (k+1) (dim-k-1) v) >> 1133 return True 1134 where loop !kval !k !l !prev !i 1135 | i == dim = return (k,l) 1136 | otherwise = do 1137 cur <- unsafeRead v i 1138 -- TODO: make tuple unboxed 1139 let (kval',k') = if prev < cur then (prev,i-1) else (kval,k) 1140 l' = if kval' < cur then i else l 1141 loop kval' k' l' cur (i+1) 1142 dim = length v 1143