{-# LANGUAGE BangPatterns #-} {-# LANGUAGE CPP #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeFamilies #-} -- | -- Module : Data.Vector.Generic.Mutable -- Copyright : (c) Roman Leshchinskiy 2008-2010 -- Alexey Kuleshevich 2020-2022 -- Aleksey Khudyakov 2020-2022 -- Andrew Lelechenko 2020-2022 -- License : BSD-style -- -- Maintainer : Haskell Libraries Team -- Stability : experimental -- Portability : non-portable -- -- Generic interface to mutable vectors. module Data.Vector.Generic.Mutable ( -- * Class of mutable vector types MVector(..), -- * Accessors -- ** Length information length, null, -- ** Extracting subvectors slice, init, tail, take, drop, splitAt, unsafeSlice, unsafeInit, unsafeTail, unsafeTake, unsafeDrop, -- ** Overlapping overlaps, -- * Construction -- ** Initialisation new, unsafeNew, replicate, replicateM, generate, generateM, clone, -- ** Growing grow, unsafeGrow, growFront, unsafeGrowFront, -- ** Restricting memory usage clear, -- * Accessing individual elements read, readMaybe, write, modify, modifyM, swap, exchange, unsafeRead, unsafeWrite, unsafeModify, unsafeModifyM, unsafeSwap, unsafeExchange, -- * Folds mapM_, imapM_, forM_, iforM_, foldl, foldl', foldM, foldM', foldr, foldr', foldrM, foldrM', ifoldl, ifoldl', ifoldM, ifoldM', ifoldr, ifoldr', ifoldrM, ifoldrM', -- * Modifying vectors nextPermutation, -- ** Filling and copying set, copy, move, unsafeCopy, unsafeMove, -- * Internal operations mstream, mstreamR, unstream, unstreamR, vunstream, munstream, munstreamR, transform, transformR, fill, fillR, unsafeAccum, accum, unsafeUpdate, update, reverse, unstablePartition, unstablePartitionBundle, partitionBundle, partitionWithBundle, -- * Re-exports PrimMonad, PrimState, RealWorld ) where import Data.Vector.Generic.Mutable.Base import qualified Data.Vector.Generic.Base as V import qualified Data.Vector.Fusion.Bundle as Bundle import Data.Vector.Fusion.Bundle ( Bundle, MBundle, Chunk(..) ) import qualified Data.Vector.Fusion.Bundle.Monadic as MBundle import Data.Vector.Fusion.Stream.Monadic ( Stream ) import qualified Data.Vector.Fusion.Stream.Monadic as Stream import Data.Vector.Fusion.Bundle.Size import Data.Vector.Fusion.Util ( delay_inline ) import Data.Vector.Internal.Check import Control.Monad.Primitive ( PrimMonad(..), RealWorld, stToPrim ) import Prelude hiding ( length, null, replicate, reverse, map, read, take, drop, splitAt, init, tail, mapM_, foldr, foldl ) #include "vector.h" -- ------------------ -- Internal functions -- ------------------ unsafeAppend1 :: (PrimMonad m, MVector v a) => v (PrimState m) a -> Int -> a -> m (v (PrimState m) a) {-# INLINE_INNER unsafeAppend1 #-} -- NOTE: The case distinction has to be on the outside because -- GHC creates a join point for the unsafeWrite even when everything -- is inlined. This is bad because with the join point, v isn't getting -- unboxed. unsafeAppend1 v i x | i < length v = do unsafeWrite v i x return v | otherwise = do v' <- enlarge v checkIndex Internal i (length v') $ unsafeWrite v' i x return v' unsafePrepend1 :: (PrimMonad m, MVector v a) => v (PrimState m) a -> Int -> a -> m (v (PrimState m) a, Int) {-# INLINE_INNER unsafePrepend1 #-} unsafePrepend1 v i x | i /= 0 = do let i' = i-1 unsafeWrite v i' x return (v, i') | otherwise = do (v', j) <- enlargeFront v let i' = j-1 checkIndex Internal i' (length v') $ unsafeWrite v' i' x return (v', i') mstream :: (PrimMonad m, MVector v a) => v (PrimState m) a -> Stream m a {-# INLINE mstream #-} mstream v = v `seq` n `seq` Stream.unfoldrM get 0 where n = length v {-# INLINE_INNER get #-} get i | i < n = do x <- unsafeRead v i return $ Just (x, i+1) | otherwise = return Nothing fill :: (PrimMonad m, MVector v a) => v (PrimState m) a -> Stream m a -> m (v (PrimState m) a) {-# INLINE fill #-} fill v s = v `seq` do n' <- Stream.foldM put 0 s return $ unsafeSlice 0 n' v where {-# INLINE_INNER put #-} put i x = do checkIndex Internal i (length v) $ unsafeWrite v i x return (i+1) transform :: (PrimMonad m, MVector v a) => (Stream m a -> Stream m a) -> v (PrimState m) a -> m (v (PrimState m) a) {-# INLINE_FUSED transform #-} transform f v = fill v (f (mstream v)) mstreamR :: (PrimMonad m, MVector v a) => v (PrimState m) a -> Stream m a {-# INLINE mstreamR #-} mstreamR v = v `seq` n `seq` Stream.unfoldrM get n where n = length v {-# INLINE_INNER get #-} get i | j >= 0 = do x <- unsafeRead v j return $ Just (x,j) | otherwise = return Nothing where j = i-1 fillR :: (PrimMonad m, MVector v a) => v (PrimState m) a -> Stream m a -> m (v (PrimState m) a) {-# INLINE fillR #-} fillR v s = v `seq` do i <- Stream.foldM put n s return $ unsafeSlice i (n-i) v where n = length v {-# INLINE_INNER put #-} put i x = do unsafeWrite v j x return j where j = i-1 transformR :: (PrimMonad m, MVector v a) => (Stream m a -> Stream m a) -> v (PrimState m) a -> m (v (PrimState m) a) {-# INLINE_FUSED transformR #-} transformR f v = fillR v (f (mstreamR v)) -- | Create a new mutable vector and fill it with elements from the 'Bundle'. -- The vector will grow exponentially if the maximum size of the 'Bundle' is -- unknown. unstream :: (PrimMonad m, MVector v a) => Bundle u a -> m (v (PrimState m) a) -- NOTE: replace INLINE_FUSED by INLINE? (also in unstreamR) {-# INLINE_FUSED unstream #-} unstream s = munstream (Bundle.lift s) -- | Create a new mutable vector and fill it with elements from the monadic -- stream. The vector will grow exponentially if the maximum size of the stream -- is unknown. munstream :: (PrimMonad m, MVector v a) => MBundle m u a -> m (v (PrimState m) a) {-# INLINE_FUSED munstream #-} munstream s = case upperBound (MBundle.size s) of Just n -> munstreamMax s n Nothing -> munstreamUnknown s munstreamMax :: (PrimMonad m, MVector v a) => MBundle m u a -> Int -> m (v (PrimState m) a) {-# INLINE munstreamMax #-} munstreamMax s n = do v <- checkLength Internal n $ unsafeNew n let put i x = do checkIndex Internal i n $ unsafeWrite v i x return (i+1) n' <- MBundle.foldM' put 0 s return $ checkSlice Internal 0 n' n $ unsafeSlice 0 n' v munstreamUnknown :: (PrimMonad m, MVector v a) => MBundle m u a -> m (v (PrimState m) a) {-# INLINE munstreamUnknown #-} munstreamUnknown s = do v <- unsafeNew 0 (v', n) <- MBundle.foldM put (v, 0) s return $ checkSlice Internal 0 n (length v') $ unsafeSlice 0 n v' where {-# INLINE_INNER put #-} put (v,i) x = do v' <- unsafeAppend1 v i x return (v',i+1) -- | Create a new mutable vector and fill it with elements from the 'Bundle'. -- The vector will grow exponentially if the maximum size of the 'Bundle' is -- unknown. vunstream :: (PrimMonad m, V.Vector v a) => Bundle v a -> m (V.Mutable v (PrimState m) a) -- NOTE: replace INLINE_FUSED by INLINE? (also in unstreamR) {-# INLINE_FUSED vunstream #-} vunstream s = vmunstream (Bundle.lift s) -- | Create a new mutable vector and fill it with elements from the monadic -- stream. The vector will grow exponentially if the maximum size of the stream -- is unknown. vmunstream :: (PrimMonad m, V.Vector v a) => MBundle m v a -> m (V.Mutable v (PrimState m) a) {-# INLINE_FUSED vmunstream #-} vmunstream s = case upperBound (MBundle.size s) of Just n -> vmunstreamMax s n Nothing -> vmunstreamUnknown s vmunstreamMax :: (PrimMonad m, V.Vector v a) => MBundle m v a -> Int -> m (V.Mutable v (PrimState m) a) {-# INLINE vmunstreamMax #-} vmunstreamMax s n = do v <- checkLength Internal n $ unsafeNew n let {-# INLINE_INNER copyChunk #-} copyChunk i (Chunk m f) = checkSlice Internal i m (length v) $ do f (basicUnsafeSlice i m v) return (i+m) n' <- Stream.foldlM' copyChunk 0 (MBundle.chunks s) return $ checkSlice Internal 0 n' n $ unsafeSlice 0 n' v vmunstreamUnknown :: (PrimMonad m, V.Vector v a) => MBundle m v a -> m (V.Mutable v (PrimState m) a) {-# INLINE vmunstreamUnknown #-} vmunstreamUnknown s = do v <- unsafeNew 0 (v', n) <- Stream.foldlM copyChunk (v,0) (MBundle.chunks s) return $ checkSlice Internal 0 n (length v') $ unsafeSlice 0 n v' where {-# INLINE_INNER copyChunk #-} copyChunk (v,i) (Chunk n f) = do let j = i+n v' <- if basicLength v < j then unsafeGrow v (delay_inline max (enlarge_delta v) (j - basicLength v)) else return v checkSlice Internal i n (length v') $ f (basicUnsafeSlice i n v') return (v',j) -- | Create a new mutable vector and fill it with elements from the 'Bundle' -- from right to left. The vector will grow exponentially if the maximum size -- of the 'Bundle' is unknown. unstreamR :: (PrimMonad m, MVector v a) => Bundle u a -> m (v (PrimState m) a) -- NOTE: replace INLINE_FUSED by INLINE? (also in unstream) {-# INLINE_FUSED unstreamR #-} unstreamR s = munstreamR (Bundle.lift s) -- | Create a new mutable vector and fill it with elements from the monadic -- stream from right to left. The vector will grow exponentially if the maximum -- size of the stream is unknown. munstreamR :: (PrimMonad m, MVector v a) => MBundle m u a -> m (v (PrimState m) a) {-# INLINE_FUSED munstreamR #-} munstreamR s = case upperBound (MBundle.size s) of Just n -> munstreamRMax s n Nothing -> munstreamRUnknown s munstreamRMax :: (PrimMonad m, MVector v a) => MBundle m u a -> Int -> m (v (PrimState m) a) {-# INLINE munstreamRMax #-} munstreamRMax s n = do v <- checkLength Internal n $ unsafeNew n let put i x = do let i' = i-1 checkIndex Internal i' n $ unsafeWrite v i' x return i' i <- MBundle.foldM' put n s return $ checkSlice Internal i (n-i) n $ unsafeSlice i (n-i) v munstreamRUnknown :: (HasCallStack, PrimMonad m, MVector v a) => MBundle m u a -> m (v (PrimState m) a) {-# INLINE munstreamRUnknown #-} munstreamRUnknown s = do v <- unsafeNew 0 (v', i) <- MBundle.foldM put (v, 0) s let n = length v' return $ checkSlice Internal i (n-i) n $ unsafeSlice i (n-i) v' where {-# INLINE_INNER put #-} put (v,i) x = unsafePrepend1 v i x -- Length -- ------ -- | Length of the mutable vector. length :: MVector v a => v s a -> Int {-# INLINE length #-} length = basicLength -- | Check whether the vector is empty. null :: MVector v a => v s a -> Bool {-# INLINE null #-} null v = length v == 0 -- Extracting subvectors -- --------------------- -- | Yield a part of the mutable vector without copying it. The vector must -- contain at least @i+n@ elements. slice :: (HasCallStack, MVector v a) => Int -- ^ @i@ starting index -> Int -- ^ @n@ length -> v s a -> v s a {-# INLINE slice #-} slice i n v = checkSlice Bounds i n (length v) $ unsafeSlice i n v -- | Take the @n@ first elements of the mutable vector without making a -- copy. For negative @n@, the empty vector is returned. If @n@ is larger -- than the vector's length, the vector is returned unchanged. take :: MVector v a => Int -> v s a -> v s a {-# INLINE take #-} take n v = unsafeSlice 0 (min (max n 0) (length v)) v -- | Drop the @n@ first element of the mutable vector without making a -- copy. For negative @n@, the vector is returned unchanged. If @n@ is -- larger than the vector's length, the empty vector is returned. drop :: MVector v a => Int -> v s a -> v s a {-# INLINE drop #-} drop n v = unsafeSlice (min m n') (max 0 (m - n')) v where n' = max n 0 m = length v -- | /O(1)/ Split the mutable vector into the first @n@ elements -- and the remainder, without copying. -- -- Note that @'splitAt' n v@ is equivalent to @('take' n v, 'drop' n v)@, -- but slightly more efficient. splitAt :: MVector v a => Int -> v s a -> (v s a, v s a) {-# INLINE splitAt #-} splitAt n v = ( unsafeSlice 0 m v , unsafeSlice m (max 0 (len - n')) v ) where m = min n' len n' = max n 0 len = length v -- | Drop the last element of the mutable vector without making a copy. -- If the vector is empty, an exception is thrown. init :: MVector v a => v s a -> v s a {-# INLINE init #-} init v = slice 0 (length v - 1) v -- | Drop the first element of the mutable vector without making a copy. -- If the vector is empty, an exception is thrown. tail :: MVector v a => v s a -> v s a {-# INLINE tail #-} tail v = slice 1 (length v - 1) v -- | Yield a part of the mutable vector without copying it. No bounds checks -- are performed. unsafeSlice :: MVector v a => Int -- ^ starting index -> Int -- ^ length of the slice -> v s a -> v s a {-# INLINE unsafeSlice #-} unsafeSlice i n v = checkSlice Unsafe i n (length v) $ basicUnsafeSlice i n v -- | Same as 'init', but doesn't do range checks. unsafeInit :: MVector v a => v s a -> v s a {-# INLINE unsafeInit #-} unsafeInit v = unsafeSlice 0 (length v - 1) v -- | Same as 'tail', but doesn't do range checks. unsafeTail :: MVector v a => v s a -> v s a {-# INLINE unsafeTail #-} unsafeTail v = unsafeSlice 1 (length v - 1) v -- | Unsafe variant of 'take'. If @n@ is out of range, it will -- simply create an invalid slice that likely violate memory safety. unsafeTake :: MVector v a => Int -> v s a -> v s a {-# INLINE unsafeTake #-} unsafeTake n v = unsafeSlice 0 n v -- | Unsafe variant of 'drop'. If @n@ is out of range, it will -- simply create an invalid slice that likely violate memory safety. unsafeDrop :: MVector v a => Int -> v s a -> v s a {-# INLINE unsafeDrop #-} unsafeDrop n v = unsafeSlice n (length v - n) v -- Overlapping -- ----------- -- | Check whether two vectors overlap. overlaps :: MVector v a => v s a -> v s a -> Bool {-# INLINE overlaps #-} overlaps = basicOverlaps -- Initialisation -- -------------- -- | Create a mutable vector of the given length. new :: (HasCallStack, PrimMonad m, MVector v a) => Int -> m (v (PrimState m) a) {-# INLINE new #-} new n = checkLength Bounds n $ stToPrim $ unsafeNew n >>= \v -> basicInitialize v >> return v -- | Create a mutable vector of the given length. The vector content -- should be assumed to be uninitialized. However, the exact semantics depend -- on the vector implementation. For example, unboxed and storable -- vectors will create a vector filled with whatever the underlying memory -- buffer happens to contain, while boxed vector's elements are -- initialized to bottoms which will throw exception when evaluated. -- -- @since 0.4 unsafeNew :: (PrimMonad m, MVector v a) => Int -> m (v (PrimState m) a) {-# INLINE unsafeNew #-} unsafeNew n = checkLength Unsafe n $ stToPrim $ basicUnsafeNew n -- | Create a mutable vector of the given length (0 if the length is negative) -- and fill it with an initial value. replicate :: (PrimMonad m, MVector v a) => Int -> a -> m (v (PrimState m) a) {-# INLINE replicate #-} replicate n x = stToPrim $ basicUnsafeReplicate (delay_inline max 0 n) x -- | Create a mutable vector of the given length (0 if the length is negative) -- and fill it with values produced by repeatedly executing the monadic action. replicateM :: (PrimMonad m, MVector v a) => Int -> m a -> m (v (PrimState m) a) {-# INLINE replicateM #-} replicateM n m = munstream (MBundle.replicateM n m) -- | /O(n)/ Create a mutable vector of the given length (0 if the length is negative) -- and fill it with the results of applying the function to each index. -- Iteration starts at index 0. -- -- @since 0.12.3.0 generate :: (PrimMonad m, MVector v a) => Int -> (Int -> a) -> m (v (PrimState m) a) {-# INLINE generate #-} generate n f = stToPrim $ generateM n (return . f) -- | /O(n)/ Create a mutable vector of the given length (0 if the length is -- negative) and fill it with the results of applying the monadic function to each -- index. Iteration starts at index 0. -- -- @since 0.12.3.0 generateM :: (PrimMonad m, MVector v a) => Int -> (Int -> m a) -> m (v (PrimState m) a) {-# INLINE generateM #-} generateM n f | n <= 0 = new 0 | otherwise = do vec <- new n let loop i | i >= n = return vec | otherwise = do unsafeWrite vec i =<< f i loop (i + 1) loop 0 -- | Create a copy of a mutable vector. clone :: (PrimMonad m, MVector v a) => v (PrimState m) a -> m (v (PrimState m) a) {-# INLINE clone #-} clone v = do v' <- unsafeNew (length v) unsafeCopy v' v return v' -- Growing -- ------- -- | Grow a vector by the given number of elements. The number must not be -- negative, otherwise an exception is thrown. The semantics of this function -- are exactly the same as of 'unsafeGrow', except that it will initialize the newly -- allocated memory first. -- -- It is important to note that mutating the returned vector will not affect the -- vector that was used as a source. In other words, it does not, nor will it -- ever have the semantics of @realloc@ from C. -- -- > grow mv 0 === clone mv -- -- @since 0.4.0 grow :: (HasCallStack, PrimMonad m, MVector v a) => v (PrimState m) a -> Int -> m (v (PrimState m) a) {-# INLINE grow #-} grow v by = checkLength Bounds by $ stToPrim $ do vnew <- unsafeGrow v by basicInitialize $ basicUnsafeSlice (length v) by vnew return vnew -- | Same as 'grow', except that it copies data towards the end of the newly -- allocated vector, making extra space available at the beginning. -- -- @since 0.11.0.0 growFront :: (HasCallStack, PrimMonad m, MVector v a) => v (PrimState m) a -> Int -> m (v (PrimState m) a) {-# INLINE growFront #-} growFront v by = checkLength Bounds by $ stToPrim $ do vnew <- unsafeGrowFront v by basicInitialize $ basicUnsafeSlice 0 by vnew return vnew enlarge_delta :: MVector v a => v s a -> Int enlarge_delta v = max (length v) 1 -- | Grow a vector logarithmically. enlarge :: (PrimMonad m, MVector v a) => v (PrimState m) a -> m (v (PrimState m) a) {-# INLINE enlarge #-} enlarge v = stToPrim $ do vnew <- unsafeGrow v by basicInitialize $ basicUnsafeSlice (length v) by vnew return vnew where by = enlarge_delta v enlargeFront :: (PrimMonad m, MVector v a) => v (PrimState m) a -> m (v (PrimState m) a, Int) {-# INLINE enlargeFront #-} enlargeFront v = stToPrim $ do v' <- unsafeGrowFront v by basicInitialize $ basicUnsafeSlice 0 by v' return (v', by) where by = enlarge_delta v -- | Grow a vector by allocating a new mutable vector of the same size plus the -- the given number of elements and copying all the data over to the new vector, -- starting at its beginning. The newly allocated memory is not initialized and -- the extra space at the end will likely contain garbage data or bottoms. -- Use 'unsafeGrowFront' to make the extra space available in the front -- of the new vector. -- -- It is important to note that mutating the returned vector will not affect -- elements of the vector that was used as a source. In other words, it does not, -- nor will it ever have the semantics of @realloc@ from C. Keep in mind, -- however, that values themselves can be of a mutable type -- (eg. 'Foreign.Ptr.Ptr'), in which case it would be possible to affect values -- stored in both vectors. -- -- > unsafeGrow mv 0 === clone mv -- -- @since 0.4.0 unsafeGrow :: (PrimMonad m, MVector v a) => v (PrimState m) a -- ^ mutable vector to copy from -> Int -- ^ number of elements to grow the vector by (must be non-negative, but -- this is not checked) -> m (v (PrimState m) a) {-# INLINE unsafeGrow #-} unsafeGrow v n = checkLength Unsafe n $ stToPrim $ basicUnsafeGrow v n -- | Same as 'unsafeGrow', except that it copies data towards the end of the -- newly allocated vector, making extra space available at the beginning. -- -- @since 0.11.0.0 unsafeGrowFront :: (PrimMonad m, MVector v a) => v (PrimState m) a -> Int -> m (v (PrimState m) a) {-# INLINE unsafeGrowFront #-} unsafeGrowFront v by = checkLength Unsafe by $ stToPrim $ do let n = length v v' <- basicUnsafeNew (by+n) basicUnsafeCopy (basicUnsafeSlice by n v') v return v' -- Restricting memory usage -- ------------------------ -- | Reset all elements of the vector to some undefined value, clearing all -- references to external objects. This is usually a noop for unboxed vectors. clear :: (PrimMonad m, MVector v a) => v (PrimState m) a -> m () {-# INLINE clear #-} clear = stToPrim . basicClear -- Accessing individual elements -- ----------------------------- -- | Yield the element at the given position. Will throw an exception if -- the index is out of range. -- -- ==== __Examples__ -- -- >>> import qualified Data.Vector.Mutable as MV -- >>> v <- MV.generate 10 (\x -> x*x) -- >>> MV.read v 3 -- 9 read :: (HasCallStack, PrimMonad m, MVector v a) => v (PrimState m) a -> Int -> m a {-# INLINE read #-} read v i = checkIndex Bounds i (length v) $ unsafeRead v i -- | Yield the element at the given position. Returns 'Nothing' if -- the index is out of range. -- -- @since 0.13 -- -- ==== __Examples__ -- -- >>> import qualified Data.Vector.Mutable as MV -- >>> v <- MV.generate 10 (\x -> x*x) -- >>> MV.readMaybe v 3 -- Just 9 -- >>> MV.readMaybe v 13 -- Nothing readMaybe :: (PrimMonad m, MVector v a) => v (PrimState m) a -> Int -> m (Maybe a) {-# INLINE readMaybe #-} readMaybe v i | i `inRange` (length v) = Just <$> unsafeRead v i | otherwise = pure Nothing -- | Replace the element at the given position. write :: (HasCallStack, PrimMonad m, MVector v a) => v (PrimState m) a -> Int -> a -> m () {-# INLINE write #-} write v i x = checkIndex Bounds i (length v) $ unsafeWrite v i x -- | Modify the element at the given position. modify :: (HasCallStack, PrimMonad m, MVector v a) => v (PrimState m) a -> (a -> a) -> Int -> m () {-# INLINE modify #-} modify v f i = checkIndex Bounds i (length v) $ unsafeModify v f i -- | Modify the element at the given position using a monadic function. -- -- @since 0.12.3.0 modifyM :: (HasCallStack, PrimMonad m, MVector v a) => v (PrimState m) a -> (a -> m a) -> Int -> m () {-# INLINE modifyM #-} modifyM v f i = checkIndex Bounds i (length v) $ unsafeModifyM v f i -- | Swap the elements at the given positions. swap :: (HasCallStack, PrimMonad m, MVector v a) => v (PrimState m) a -> Int -> Int -> m () {-# INLINE swap #-} swap v i j = checkIndex Bounds i (length v) $ checkIndex Bounds j (length v) $ unsafeSwap v i j -- | Replace the element at the given position and return the old element. exchange :: (HasCallStack, PrimMonad m, MVector v a) => v (PrimState m) a -> Int -> a -> m a {-# INLINE exchange #-} exchange v i x = checkIndex Bounds i (length v) $ unsafeExchange v i x -- | Yield the element at the given position. No bounds checks are performed. unsafeRead :: (PrimMonad m, MVector v a) => v (PrimState m) a -> Int -> m a {-# INLINE unsafeRead #-} unsafeRead v i = checkIndex Unsafe i (length v) $ stToPrim $ basicUnsafeRead v i -- | Replace the element at the given position. No bounds checks are performed. unsafeWrite :: (PrimMonad m, MVector v a) => v (PrimState m) a -> Int -> a -> m () {-# INLINE unsafeWrite #-} unsafeWrite v i x = checkIndex Unsafe i (length v) $ stToPrim $ basicUnsafeWrite v i x -- | Modify the element at the given position. No bounds checks are performed. unsafeModify :: (PrimMonad m, MVector v a) => v (PrimState m) a -> (a -> a) -> Int -> m () {-# INLINE unsafeModify #-} unsafeModify v f i = checkIndex Unsafe i (length v) $ stToPrim $ basicUnsafeRead v i >>= \x -> basicUnsafeWrite v i (f x) -- | Modify the element at the given position using a monadic -- function. No bounds checks are performed. -- -- @since 0.12.3.0 unsafeModifyM :: (PrimMonad m, MVector v a) => v (PrimState m) a -> (a -> m a) -> Int -> m () {-# INLINE unsafeModifyM #-} unsafeModifyM v f i = checkIndex Unsafe i (length v) $ stToPrim . basicUnsafeWrite v i =<< f =<< stToPrim (basicUnsafeRead v i) -- | Swap the elements at the given positions. No bounds checks are performed. unsafeSwap :: (PrimMonad m, MVector v a) => v (PrimState m) a -> Int -> Int -> m () {-# INLINE unsafeSwap #-} unsafeSwap v i j = checkIndex Unsafe i (length v) $ checkIndex Unsafe j (length v) $ stToPrim $ do x <- unsafeRead v i y <- unsafeRead v j unsafeWrite v i y unsafeWrite v j x -- | Replace the element at the given position and return the old element. No -- bounds checks are performed. unsafeExchange :: (PrimMonad m, MVector v a) => v (PrimState m) a -> Int -> a -> m a {-# INLINE unsafeExchange #-} unsafeExchange v i x = checkIndex Unsafe i (length v) $ stToPrim $ do y <- unsafeRead v i unsafeWrite v i x return y -- Folds -- ----- forI_ :: (Monad m, MVector v a) => v (PrimState m) a -> (Int -> m b) -> m () {-# INLINE forI_ #-} forI_ v f = loop 0 where loop i | i >= n = return () | otherwise = f i >> loop (i + 1) n = length v -- | /O(n)/ Apply the monadic action to every element of the vector, discarding the results. -- -- @since 0.12.3.0 mapM_ :: (PrimMonad m, MVector v a) => (a -> m b) -> v (PrimState m) a -> m () {-# INLINE mapM_ #-} mapM_ f v = forI_ v $ \i -> f =<< unsafeRead v i -- | /O(n)/ Apply the monadic action to every element of the vector and its index, discarding the results. -- -- @since 0.12.3.0 imapM_ :: (PrimMonad m, MVector v a) => (Int -> a -> m b) -> v (PrimState m) a -> m () {-# INLINE imapM_ #-} imapM_ f v = forI_ v $ \i -> f i =<< unsafeRead v i -- | /O(n)/ Apply the monadic action to every element of the vector, -- discarding the results. It's the same as @flip mapM_@. -- -- @since 0.12.3.0 forM_ :: (PrimMonad m, MVector v a) => v (PrimState m) a -> (a -> m b) -> m () {-# INLINE forM_ #-} forM_ = flip mapM_ -- | /O(n)/ Apply the monadic action to every element of the vector -- and its index, discarding the results. It's the same as @flip imapM_@. -- -- @since 0.12.3.0 iforM_ :: (PrimMonad m, MVector v a) => v (PrimState m) a -> (Int -> a -> m b) -> m () {-# INLINE iforM_ #-} iforM_ = flip imapM_ -- | /O(n)/ Pure left fold. -- -- @since 0.12.3.0 foldl :: (PrimMonad m, MVector v a) => (b -> a -> b) -> b -> v (PrimState m) a -> m b {-# INLINE foldl #-} foldl f = ifoldl (\b _ -> f b) -- | /O(n)/ Pure left fold with strict accumulator. -- -- @since 0.12.3.0 foldl' :: (PrimMonad m, MVector v a) => (b -> a -> b) -> b -> v (PrimState m) a -> m b {-# INLINE foldl' #-} foldl' f = ifoldl' (\b _ -> f b) -- | /O(n)/ Pure left fold using a function applied to each element and its index. -- -- @since 0.12.3.0 ifoldl :: (PrimMonad m, MVector v a) => (b -> Int -> a -> b) -> b -> v (PrimState m) a -> m b {-# INLINE ifoldl #-} ifoldl f b0 v = stToPrim $ ifoldM (\b i a -> return $ f b i a) b0 v -- | /O(n)/ Pure left fold with strict accumulator using a function applied to each element and its index. -- -- @since 0.12.3.0 ifoldl' :: (PrimMonad m, MVector v a) => (b -> Int -> a -> b) -> b -> v (PrimState m) a -> m b {-# INLINE ifoldl' #-} ifoldl' f b0 v = stToPrim $ ifoldM' (\b i a -> return $ f b i a) b0 v -- | /O(n)/ Pure right fold. -- -- @since 0.12.3.0 foldr :: (PrimMonad m, MVector v a) => (a -> b -> b) -> b -> v (PrimState m) a -> m b {-# INLINE foldr #-} foldr f = ifoldr (const f) -- | /O(n)/ Pure right fold with strict accumulator. -- -- @since 0.12.3.0 foldr' :: (PrimMonad m, MVector v a) => (a -> b -> b) -> b -> v (PrimState m) a -> m b {-# INLINE foldr' #-} foldr' f = ifoldr' (const f) -- | /O(n)/ Pure right fold using a function applied to each element and its index. -- -- @since 0.12.3.0 ifoldr :: (PrimMonad m, MVector v a) => (Int -> a -> b -> b) -> b -> v (PrimState m) a -> m b {-# INLINE ifoldr #-} ifoldr f b0 v = stToPrim $ ifoldrM (\i a b -> return $ f i a b) b0 v -- | /O(n)/ Pure right fold with strict accumulator using a function applied -- to each element and its index. -- -- @since 0.12.3.0 ifoldr' :: (PrimMonad m, MVector v a) => (Int -> a -> b -> b) -> b -> v (PrimState m) a -> m b {-# INLINE ifoldr' #-} ifoldr' f b0 v = stToPrim $ ifoldrM' (\i a b -> return $ f i a b) b0 v -- | /O(n)/ Monadic fold. -- -- @since 0.12.3.0 foldM :: (PrimMonad m, MVector v a) => (b -> a -> m b) -> b -> v (PrimState m) a -> m b {-# INLINE foldM #-} foldM f = ifoldM (\x _ -> f x) -- | /O(n)/ Monadic fold with strict accumulator. -- -- @since 0.12.3.0 foldM' :: (PrimMonad m, MVector v a) => (b -> a -> m b) -> b -> v (PrimState m) a -> m b {-# INLINE foldM' #-} foldM' f = ifoldM' (\x _ -> f x) -- | /O(n)/ Monadic fold using a function applied to each element and its index. -- -- @since 0.12.3.0 ifoldM :: (PrimMonad m, MVector v a) => (b -> Int -> a -> m b) -> b -> v (PrimState m) a -> m b {-# INLINE ifoldM #-} ifoldM f b0 v = loop 0 b0 where loop i b | i >= n = return b | otherwise = do a <- unsafeRead v i loop (i + 1) =<< f b i a n = length v -- | /O(n)/ Monadic fold with strict accumulator using a function applied to each element and its index. -- -- @since 0.12.3.0 ifoldM' :: (PrimMonad m, MVector v a) => (b -> Int -> a -> m b) -> b -> v (PrimState m) a -> m b {-# INLINE ifoldM' #-} ifoldM' f b0 v = loop 0 b0 where loop i !b | i >= n = return b | otherwise = do a <- unsafeRead v i loop (i + 1) =<< f b i a n = length v -- | /O(n)/ Monadic right fold. -- -- @since 0.12.3.0 foldrM :: (PrimMonad m, MVector v a) => (a -> b -> m b) -> b -> v (PrimState m) a -> m b {-# INLINE foldrM #-} foldrM f = ifoldrM (const f) -- | /O(n)/ Monadic right fold with strict accumulator. -- -- @since 0.12.3.0 foldrM' :: (PrimMonad m, MVector v a) => (a -> b -> m b) -> b -> v (PrimState m) a -> m b {-# INLINE foldrM' #-} foldrM' f = ifoldrM' (const f) -- | /O(n)/ Monadic right fold using a function applied to each element and its index. -- -- @since 0.12.3.0 ifoldrM :: (PrimMonad m, MVector v a) => (Int -> a -> b -> m b) -> b -> v (PrimState m) a -> m b {-# INLINE ifoldrM #-} ifoldrM f b0 v = loop (n-1) b0 where loop i b | i < 0 = return b | otherwise = do a <- unsafeRead v i loop (i - 1) =<< f i a b n = length v -- | /O(n)/ Monadic right fold with strict accumulator using a function applied -- to each element and its index. -- -- @since 0.12.3.0 ifoldrM' :: (PrimMonad m, MVector v a) => (Int -> a -> b -> m b) -> b -> v (PrimState m) a -> m b {-# INLINE ifoldrM' #-} ifoldrM' f b0 v = loop (n-1) b0 where loop i !b | i < 0 = return b | otherwise = do a <- unsafeRead v i loop (i - 1) =<< f i a b n = length v -- Filling and copying -- ------------------- -- | Set all elements of the vector to the given value. set :: (PrimMonad m, MVector v a) => v (PrimState m) a -> a -> m () {-# INLINE set #-} set v = stToPrim . basicSet v -- | Copy a vector. The two vectors must have the same length and may not -- overlap. copy :: (HasCallStack, PrimMonad m, MVector v a) => v (PrimState m) a -- ^ target -> v (PrimState m) a -- ^ source -> m () {-# INLINE copy #-} copy dst src = check Bounds "overlapping vectors" (not (dst `overlaps` src)) $ check Bounds "length mismatch" (length dst == length src) $ unsafeCopy dst src -- | Move the contents of a vector. The two vectors must have the same -- length. -- -- If the vectors do not overlap, then this is equivalent to 'copy'. -- Otherwise, the copying is performed as if the source vector were -- copied to a temporary vector and then the temporary vector was copied -- to the target vector. move :: (HasCallStack, PrimMonad m, MVector v a) => v (PrimState m) a -- ^ target -> v (PrimState m) a -- ^ source -> m () {-# INLINE move #-} move dst src = check Bounds "length mismatch" (length dst == length src) $ unsafeMove dst src -- | Copy a vector. The two vectors must have the same length and may not -- overlap, but this is not checked. unsafeCopy :: (PrimMonad m, MVector v a) => v (PrimState m) a -- ^ target -> v (PrimState m) a -- ^ source -> m () {-# INLINE unsafeCopy #-} unsafeCopy dst src = check Unsafe "length mismatch" (length dst == length src) $ check Unsafe "overlapping vectors" (not (dst `overlaps` src)) $ dst `seq` src `seq` stToPrim (basicUnsafeCopy dst src) -- | Move the contents of a vector. The two vectors must have the same -- length, but this is not checked. -- -- If the vectors do not overlap, then this is equivalent to 'unsafeCopy'. -- Otherwise, the copying is performed as if the source vector were -- copied to a temporary vector and then the temporary vector was copied -- to the target vector. unsafeMove :: (PrimMonad m, MVector v a) => v (PrimState m) a -- ^ target -> v (PrimState m) a -- ^ source -> m () {-# INLINE unsafeMove #-} unsafeMove dst src = check Unsafe "length mismatch" (length dst == length src) $ dst `seq` src `seq` stToPrim (basicUnsafeMove dst src) accum :: forall m v a b u. (HasCallStack, PrimMonad m, MVector v a) => (a -> b -> a) -> v (PrimState m) a -> Bundle u (Int, b) -> m () {-# INLINE accum #-} accum f !v s = Bundle.mapM_ upd s where {-# INLINE_INNER upd #-} upd :: HasCallStack => (Int, b) -> m () upd (i,b) = do a <- checkIndex Bounds i n $ unsafeRead v i unsafeWrite v i (f a b) !n = length v update :: forall m v a u. (HasCallStack, PrimMonad m, MVector v a) => v (PrimState m) a -> Bundle u (Int, a) -> m () {-# INLINE update #-} update !v s = Bundle.mapM_ upd s where {-# INLINE_INNER upd #-} upd :: HasCallStack => (Int, a) -> m () upd (i,b) = checkIndex Bounds i n $ unsafeWrite v i b !n = length v unsafeAccum :: (PrimMonad m, MVector v a) => (a -> b -> a) -> v (PrimState m) a -> Bundle u (Int, b) -> m () {-# INLINE unsafeAccum #-} unsafeAccum f !v s = Bundle.mapM_ upd s where {-# INLINE_INNER upd #-} upd (i,b) = do a <- checkIndex Unsafe i n $ unsafeRead v i unsafeWrite v i (f a b) !n = length v unsafeUpdate :: (PrimMonad m, MVector v a) => v (PrimState m) a -> Bundle u (Int, a) -> m () {-# INLINE unsafeUpdate #-} unsafeUpdate !v s = Bundle.mapM_ upd s where {-# INLINE_INNER upd #-} upd (i,b) = checkIndex Unsafe i n $ unsafeWrite v i b !n = length v reverse :: (PrimMonad m, MVector v a) => v (PrimState m) a -> m () {-# INLINE reverse #-} reverse !v = reverse_loop 0 (length v - 1) where reverse_loop i j | i < j = do unsafeSwap v i j reverse_loop (i + 1) (j - 1) reverse_loop _ _ = return () unstablePartition :: forall m v a. (PrimMonad m, MVector v a) => (a -> Bool) -> v (PrimState m) a -> m Int {-# INLINE unstablePartition #-} unstablePartition f !v = from_left 0 (length v) where -- NOTE: GHC 6.10.4 panics without the signatures on from_left and -- from_right from_left :: Int -> Int -> m Int from_left i j | i == j = return i | otherwise = do x <- unsafeRead v i if f x then from_left (i+1) j else from_right i (j-1) from_right :: Int -> Int -> m Int from_right i j | i == j = return i | otherwise = do x <- unsafeRead v j if f x then do y <- unsafeRead v i unsafeWrite v i x unsafeWrite v j y from_left (i+1) j else from_right i (j-1) unstablePartitionBundle :: (PrimMonad m, MVector v a) => (a -> Bool) -> Bundle u a -> m (v (PrimState m) a, v (PrimState m) a) {-# INLINE unstablePartitionBundle #-} unstablePartitionBundle f s = case upperBound (Bundle.size s) of Just n -> unstablePartitionMax f s n Nothing -> partitionUnknown f s unstablePartitionMax :: (PrimMonad m, MVector v a) => (a -> Bool) -> Bundle u a -> Int -> m (v (PrimState m) a, v (PrimState m) a) {-# INLINE unstablePartitionMax #-} unstablePartitionMax f s n = do v <- checkLength Internal n $ unsafeNew n let {-# INLINE_INNER put #-} put (i, j) x | f x = do unsafeWrite v i x return (i+1, j) | otherwise = do unsafeWrite v (j-1) x return (i, j-1) (i,j) <- Bundle.foldM' put (0, n) s return (unsafeSlice 0 i v, unsafeSlice j (n-j) v) partitionBundle :: (PrimMonad m, MVector v a) => (a -> Bool) -> Bundle u a -> m (v (PrimState m) a, v (PrimState m) a) {-# INLINE partitionBundle #-} partitionBundle f s = case upperBound (Bundle.size s) of Just n -> partitionMax f s n Nothing -> partitionUnknown f s partitionMax :: (PrimMonad m, MVector v a) => (a -> Bool) -> Bundle u a -> Int -> m (v (PrimState m) a, v (PrimState m) a) {-# INLINE partitionMax #-} partitionMax f s n = do v <- checkLength Internal n $ unsafeNew n let {-# INLINE_INNER put #-} put (i,j) x | f x = do unsafeWrite v i x return (i+1,j) | otherwise = let j' = j-1 in do unsafeWrite v j' x return (i,j') (i,j) <- Bundle.foldM' put (0,n) s check Internal "invalid indices" (i <= j) $ return () let l = unsafeSlice 0 i v r = unsafeSlice j (n-j) v reverse r return (l,r) partitionUnknown :: (PrimMonad m, MVector v a) => (a -> Bool) -> Bundle u a -> m (v (PrimState m) a, v (PrimState m) a) {-# INLINE partitionUnknown #-} partitionUnknown f s = do v1 <- unsafeNew 0 v2 <- unsafeNew 0 (v1', n1, v2', n2) <- Bundle.foldM' put (v1, 0, v2, 0) s checkSlice Internal 0 n1 (length v1') $ checkSlice Internal 0 n2 (length v2') $ return (unsafeSlice 0 n1 v1', unsafeSlice 0 n2 v2') where -- NOTE: The case distinction has to be on the outside because -- GHC creates a join point for the unsafeWrite even when everything -- is inlined. This is bad because with the join point, v isn't getting -- unboxed. {-# INLINE_INNER put #-} put (v1, i1, v2, i2) x | f x = do v1' <- unsafeAppend1 v1 i1 x return (v1', i1+1, v2, i2) | otherwise = do v2' <- unsafeAppend1 v2 i2 x return (v1, i1, v2', i2+1) partitionWithBundle :: (PrimMonad m, MVector v a, MVector v b, MVector v c) => (a -> Either b c) -> Bundle u a -> m (v (PrimState m) b, v (PrimState m) c) {-# INLINE partitionWithBundle #-} partitionWithBundle f s = case upperBound (Bundle.size s) of Just n -> partitionWithMax f s n Nothing -> partitionWithUnknown f s partitionWithMax :: (PrimMonad m, MVector v a, MVector v b, MVector v c) => (a -> Either b c) -> Bundle u a -> Int -> m (v (PrimState m) b, v (PrimState m) c) {-# INLINE partitionWithMax #-} partitionWithMax f s n = do v1 <- unsafeNew n v2 <- unsafeNew n let {-# INLINE_INNER put #-} put (i1, i2) x = case f x of Left b -> do unsafeWrite v1 i1 b return (i1+1, i2) Right c -> do unsafeWrite v2 i2 c return (i1, i2+1) (n1, n2) <- Bundle.foldM' put (0, 0) s checkSlice Internal 0 n1 (length v1) $ checkSlice Internal 0 n2 (length v2) $ return (unsafeSlice 0 n1 v1, unsafeSlice 0 n2 v2) partitionWithUnknown :: forall m v u a b c. (PrimMonad m, MVector v a, MVector v b, MVector v c) => (a -> Either b c) -> Bundle u a -> m (v (PrimState m) b, v (PrimState m) c) {-# INLINE partitionWithUnknown #-} partitionWithUnknown f s = do v1 <- unsafeNew 0 v2 <- unsafeNew 0 (v1', n1, v2', n2) <- Bundle.foldM' put (v1, 0, v2, 0) s checkSlice Internal 0 n1 (length v1') $ checkSlice Internal 0 n2 (length v2') $ return (unsafeSlice 0 n1 v1', unsafeSlice 0 n2 v2') where put :: (v (PrimState m) b, Int, v (PrimState m) c, Int) -> a -> m (v (PrimState m) b, Int, v (PrimState m) c, Int) {-# INLINE_INNER put #-} put (v1, i1, v2, i2) x = case f x of Left b -> do v1' <- unsafeAppend1 v1 i1 b return (v1', i1+1, v2, i2) Right c -> do v2' <- unsafeAppend1 v2 i2 c return (v1, i1, v2', i2+1) -- Modifying vectors -- ----------------- {- http://en.wikipedia.org/wiki/Permutation#Algorithms_to_generate_permutations The following algorithm generates the next permutation lexicographically after a given permutation. It changes the given permutation in-place. 1. Find the largest index k such that a[k] < a[k + 1]. If no such index exists, the permutation is the last permutation. 2. Find the largest index l greater than k such that a[k] < a[l]. 3. Swap the value of a[k] with that of a[l]. 4. Reverse the sequence from a[k + 1] up to and including the final element a[n] -} -- | Compute the (lexicographically) next permutation of the given vector in-place. -- Returns False when the input is the last permutation. nextPermutation :: (PrimMonad m,Ord e,MVector v e) => v (PrimState m) e -> m Bool nextPermutation v | dim < 2 = return False | otherwise = do val <- unsafeRead v 0 (k,l) <- loop val (-1) 0 val 1 if k < 0 then return False else unsafeSwap v k l >> reverse (unsafeSlice (k+1) (dim-k-1) v) >> return True where loop !kval !k !l !prev !i | i == dim = return (k,l) | otherwise = do cur <- unsafeRead v i -- TODO: make tuple unboxed let (kval',k') = if prev < cur then (prev,i-1) else (kval,k) l' = if kval' < cur then i else l loop kval' k' l' cur (i+1) dim = length v