1{-# LANGUAGE CPP, DeriveDataTypeable, MultiParamTypeClasses, FlexibleInstances, ScopedTypeVariables #-}
2
3-- |
4-- Module      : Data.Vector.Primitive.Mutable
5-- Copyright   : (c) Roman Leshchinskiy 2008-2010
6-- License     : BSD-style
7--
8-- Maintainer  : Roman Leshchinskiy <rl@cse.unsw.edu.au>
9-- Stability   : experimental
10-- Portability : non-portable
11--
12-- Mutable primitive vectors.
13--
14
15module Data.Vector.Primitive.Mutable (
16  -- * Mutable vectors of primitive types
17  MVector(..), IOVector, STVector, Prim,
18
19  -- * Accessors
20
21  -- ** Length information
22  length, null,
23
24  -- ** Extracting subvectors
25  slice, init, tail, take, drop, splitAt,
26  unsafeSlice, unsafeInit, unsafeTail, unsafeTake, unsafeDrop,
27
28  -- ** Overlapping
29  overlaps,
30
31  -- * Construction
32
33  -- ** Initialisation
34  new, unsafeNew, replicate, replicateM, generate, generateM, clone,
35
36  -- ** Growing
37  grow, unsafeGrow,
38
39  -- ** Restricting memory usage
40  clear,
41
42  -- * Accessing individual elements
43  read, write, modify, modifyM, swap, exchange,
44  unsafeRead, unsafeWrite, unsafeModify, unsafeModifyM, unsafeSwap, unsafeExchange,
45
46  -- * Folds
47  mapM_, imapM_, forM_, iforM_,
48  foldl, foldl', foldM, foldM',
49  foldr, foldr', foldrM, foldrM',
50  ifoldl, ifoldl', ifoldM, ifoldM',
51  ifoldr, ifoldr', ifoldrM, ifoldrM',
52
53  -- * Modifying vectors
54  nextPermutation,
55
56  -- ** Filling and copying
57  set, copy, move, unsafeCopy, unsafeMove
58) where
59
60import qualified Data.Vector.Generic.Mutable as G
61import           Data.Primitive.ByteArray
62import           Data.Primitive ( Prim, sizeOf )
63import           Data.Word ( Word8 )
64import           Control.Monad.Primitive
65import           Control.Monad ( liftM )
66
67import Control.DeepSeq ( NFData(rnf)
68#if MIN_VERSION_deepseq(1,4,3)
69                       , NFData1(liftRnf)
70#endif
71                       )
72
73import Prelude hiding ( length, null, replicate, reverse, map, read,
74                        take, drop, splitAt, init, tail, foldr, foldl, mapM_ )
75
76import Data.Typeable ( Typeable )
77
78-- Data.Vector.Internal.Check is unnecessary
79#define NOT_VECTOR_MODULE
80#include "vector.h"
81
82-- | Mutable vectors of primitive types.
83data MVector s a = MVector {-# UNPACK #-} !Int
84                           {-# UNPACK #-} !Int
85                           {-# UNPACK #-} !(MutableByteArray s) -- ^ offset, length, underlying mutable byte array
86        deriving ( Typeable )
87
88type IOVector = MVector RealWorld
89type STVector s = MVector s
90
91instance NFData (MVector s a) where
92  rnf (MVector _ _ _) = ()
93
94#if MIN_VERSION_deepseq(1,4,3)
95instance NFData1 (MVector s) where
96  liftRnf _ (MVector _ _ _) = ()
97#endif
98
99instance Prim a => G.MVector MVector a where
100  basicLength (MVector _ n _) = n
101  basicUnsafeSlice j m (MVector i _ arr)
102    = MVector (i+j) m arr
103
104  {-# INLINE basicOverlaps #-}
105  basicOverlaps (MVector i m arr1) (MVector j n arr2)
106    = sameMutableByteArray arr1 arr2
107      && (between i j (j+n) || between j i (i+m))
108    where
109      between x y z = x >= y && x < z
110
111  {-# INLINE basicUnsafeNew #-}
112  basicUnsafeNew n
113    | n < 0 = error $ "Primitive.basicUnsafeNew: negative length: " ++ show n
114    | n > mx = error $ "Primitive.basicUnsafeNew: length to large: " ++ show n
115    | otherwise = MVector 0 n `liftM` newByteArray (n * size)
116    where
117      size = sizeOf (undefined :: a)
118      mx = maxBound `div` size :: Int
119
120  {-# INLINE basicInitialize #-}
121  basicInitialize (MVector off n v) =
122      setByteArray v (off * size) (n * size) (0 :: Word8)
123    where
124      size = sizeOf (undefined :: a)
125
126
127  {-# INLINE basicUnsafeRead #-}
128  basicUnsafeRead (MVector i _ arr) j = readByteArray arr (i+j)
129
130  {-# INLINE basicUnsafeWrite #-}
131  basicUnsafeWrite (MVector i _ arr) j x = writeByteArray arr (i+j) x
132
133  {-# INLINE basicUnsafeCopy #-}
134  basicUnsafeCopy (MVector i n dst) (MVector j _ src)
135    = copyMutableByteArray dst (i*sz) src (j*sz) (n*sz)
136    where
137      sz = sizeOf (undefined :: a)
138
139  {-# INLINE basicUnsafeMove #-}
140  basicUnsafeMove (MVector i n dst) (MVector j _ src)
141    = moveByteArray dst (i*sz) src (j*sz) (n * sz)
142    where
143      sz = sizeOf (undefined :: a)
144
145  {-# INLINE basicSet #-}
146  basicSet (MVector i n arr) x = setByteArray arr i n x
147
148-- Length information
149-- ------------------
150
151-- | Length of the mutable vector.
152length :: Prim a => MVector s a -> Int
153{-# INLINE length #-}
154length = G.length
155
156-- | Check whether the vector is empty
157null :: Prim a => MVector s a -> Bool
158{-# INLINE null #-}
159null = G.null
160
161-- Extracting subvectors
162-- ---------------------
163
164-- | Yield a part of the mutable vector without copying it. The vector must
165-- contain at least @i+n@ elements.
166slice :: Prim a
167      => Int  -- ^ @i@ starting index
168      -> Int  -- ^ @n@ length
169      -> MVector s a
170      -> MVector s a
171{-# INLINE slice #-}
172slice = G.slice
173
174take :: Prim a => Int -> MVector s a -> MVector s a
175{-# INLINE take #-}
176take = G.take
177
178drop :: Prim a => Int -> MVector s a -> MVector s a
179{-# INLINE drop #-}
180drop = G.drop
181
182splitAt :: Prim a => Int -> MVector s a -> (MVector s a, MVector s a)
183{-# INLINE splitAt #-}
184splitAt = G.splitAt
185
186init :: Prim a => MVector s a -> MVector s a
187{-# INLINE init #-}
188init = G.init
189
190tail :: Prim a => MVector s a -> MVector s a
191{-# INLINE tail #-}
192tail = G.tail
193
194-- | Yield a part of the mutable vector without copying it. No bounds checks
195-- are performed.
196unsafeSlice :: Prim a
197            => Int  -- ^ starting index
198            -> Int  -- ^ length of the slice
199            -> MVector s a
200            -> MVector s a
201{-# INLINE unsafeSlice #-}
202unsafeSlice = G.unsafeSlice
203
204unsafeTake :: Prim a => Int -> MVector s a -> MVector s a
205{-# INLINE unsafeTake #-}
206unsafeTake = G.unsafeTake
207
208unsafeDrop :: Prim a => Int -> MVector s a -> MVector s a
209{-# INLINE unsafeDrop #-}
210unsafeDrop = G.unsafeDrop
211
212unsafeInit :: Prim a => MVector s a -> MVector s a
213{-# INLINE unsafeInit #-}
214unsafeInit = G.unsafeInit
215
216unsafeTail :: Prim a => MVector s a -> MVector s a
217{-# INLINE unsafeTail #-}
218unsafeTail = G.unsafeTail
219
220-- Overlapping
221-- -----------
222
223-- | Check whether two vectors overlap.
224overlaps :: Prim a => MVector s a -> MVector s a -> Bool
225{-# INLINE overlaps #-}
226overlaps = G.overlaps
227
228-- Initialisation
229-- --------------
230
231-- | Create a mutable vector of the given length.
232new :: (PrimMonad m, Prim a) => Int -> m (MVector (PrimState m) a)
233{-# INLINE new #-}
234new = G.new
235
236-- | Create a mutable vector of the given length. The vector content
237--   is uninitialized, which means it is filled with whatever underlying memory
238--   buffer happens to contain.
239--
240-- @since 0.5
241unsafeNew :: (PrimMonad m, Prim a) => Int -> m (MVector (PrimState m) a)
242{-# INLINE unsafeNew #-}
243unsafeNew = G.unsafeNew
244
245-- | Create a mutable vector of the given length (0 if the length is negative)
246-- and fill it with an initial value.
247replicate :: (PrimMonad m, Prim a) => Int -> a -> m (MVector (PrimState m) a)
248{-# INLINE replicate #-}
249replicate = G.replicate
250
251-- | Create a mutable vector of the given length (0 if the length is negative)
252-- and fill it with values produced by repeatedly executing the monadic action.
253replicateM :: (PrimMonad m, Prim a) => Int -> m a -> m (MVector (PrimState m) a)
254{-# INLINE replicateM #-}
255replicateM = G.replicateM
256
257-- | /O(n)/ Create a mutable vector of the given length (0 if the length is negative)
258-- and fill it with the results of applying the function to each index.
259--
260-- @since 0.12.3.0
261generate :: (PrimMonad m, Prim a) => Int -> (Int -> a) -> m (MVector (PrimState m) a)
262{-# INLINE generate #-}
263generate = G.generate
264
265-- | /O(n)/ Create a mutable vector of the given length (0 if the length is
266-- negative) and fill it with the results of applying the monadic function to each
267-- index. Iteration starts at index 0.
268--
269-- @since 0.12.3.0
270generateM :: (PrimMonad m, Prim a) => Int -> (Int -> m a) -> m (MVector (PrimState m) a)
271{-# INLINE generateM #-}
272generateM = G.generateM
273
274-- | Create a copy of a mutable vector.
275clone :: (PrimMonad m, Prim a)
276      => MVector (PrimState m) a -> m (MVector (PrimState m) a)
277{-# INLINE clone #-}
278clone = G.clone
279
280-- Growing
281-- -------
282
283-- | Grow a primitive vector by the given number of elements. The number must be
284-- non-negative. Same semantics as in `G.grow` for generic vector.
285--
286-- ====__Examples__
287--
288-- >>> import qualified Data.Vector.Primitive as VP
289-- >>> import qualified Data.Vector.Primitive.Mutable as MVP
290-- >>> mv <- VP.thaw $ VP.fromList ([10, 20, 30] :: [Int])
291-- >>> mv' <- MVP.grow mv 2
292--
293-- Extra memory at the end of the newly allocated vector is initialized to 0
294-- bytes, which for `Prim` instance will usually correspond to some default
295-- value for a particular type, eg. @0@ for @Int@, @\NUL@ for @Char@,
296-- etc. However, if `unsafeGrow` was used instead this would not have been
297-- guaranteed and some garbage would be there instead:
298--
299-- >>> VP.unsafeFreeze mv'
300-- [10,20,30,0,0]
301--
302-- Having the extra space we can write new values in there:
303--
304-- >>> MVP.write mv' 3 999
305-- >>> VP.unsafeFreeze mv'
306-- [10,20,30,999,0]
307--
308-- It is important to note that the source mutable vector is not affected when
309-- the newly allocated one is mutated.
310--
311-- >>> MVP.write mv' 2 888
312-- >>> VP.unsafeFreeze mv'
313-- [10,20,888,999,0]
314-- >>> VP.unsafeFreeze mv
315-- [10,20,30]
316--
317-- @since 0.5
318grow :: (PrimMonad m, Prim a)
319              => MVector (PrimState m) a -> Int -> m (MVector (PrimState m) a)
320{-# INLINE grow #-}
321grow = G.grow
322
323-- | Grow a vector by the given number of elements. The number must be non-negative but
324-- this is not checked. Same semantics as in `G.unsafeGrow` for generic vector.
325--
326-- @since 0.5
327unsafeGrow :: (PrimMonad m, Prim a)
328               => MVector (PrimState m) a -> Int -> m (MVector (PrimState m) a)
329{-# INLINE unsafeGrow #-}
330unsafeGrow = G.unsafeGrow
331
332-- Restricting memory usage
333-- ------------------------
334
335-- | Reset all elements of the vector to some undefined value, clearing all
336-- references to external objects. This is usually a noop for unboxed vectors.
337clear :: (PrimMonad m, Prim a) => MVector (PrimState m) a -> m ()
338{-# INLINE clear #-}
339clear = G.clear
340
341-- Accessing individual elements
342-- -----------------------------
343
344-- | Yield the element at the given position.
345read :: (PrimMonad m, Prim a) => MVector (PrimState m) a -> Int -> m a
346{-# INLINE read #-}
347read = G.read
348
349-- | Replace the element at the given position.
350write :: (PrimMonad m, Prim a) => MVector (PrimState m) a -> Int -> a -> m ()
351{-# INLINE write #-}
352write = G.write
353
354-- | Modify the element at the given position.
355modify :: (PrimMonad m, Prim a) => MVector (PrimState m) a -> (a -> a) -> Int -> m ()
356{-# INLINE modify #-}
357modify = G.modify
358
359-- | Modify the element at the given position using a monadic function.
360--
361-- @since 0.12.3.0
362modifyM :: (PrimMonad m, Prim a) => MVector (PrimState m) a -> (a -> m a) -> Int -> m ()
363{-# INLINE modifyM #-}
364modifyM = G.modifyM
365
366-- | Swap the elements at the given positions.
367swap :: (PrimMonad m, Prim a) => MVector (PrimState m) a -> Int -> Int -> m ()
368{-# INLINE swap #-}
369swap = G.swap
370
371-- | Replace the element at the given position and return the old element.
372exchange :: (PrimMonad m, Prim a) => MVector (PrimState m) a -> Int -> a -> m a
373{-# INLINE exchange #-}
374exchange = G.exchange
375
376-- | Yield the element at the given position. No bounds checks are performed.
377unsafeRead :: (PrimMonad m, Prim a) => MVector (PrimState m) a -> Int -> m a
378{-# INLINE unsafeRead #-}
379unsafeRead = G.unsafeRead
380
381-- | Replace the element at the given position. No bounds checks are performed.
382unsafeWrite
383    :: (PrimMonad m, Prim a) =>  MVector (PrimState m) a -> Int -> a -> m ()
384{-# INLINE unsafeWrite #-}
385unsafeWrite = G.unsafeWrite
386
387-- | Modify the element at the given position. No bounds checks are performed.
388unsafeModify :: (PrimMonad m, Prim a) => MVector (PrimState m) a -> (a -> a) -> Int -> m ()
389{-# INLINE unsafeModify #-}
390unsafeModify = G.unsafeModify
391
392-- | Modify the element at the given position using a monadic
393-- function. No bounds checks are performed.
394--
395-- @since 0.12.3.0
396unsafeModifyM :: (PrimMonad m, Prim a) => MVector (PrimState m) a -> (a -> m a) -> Int -> m ()
397{-# INLINE unsafeModifyM #-}
398unsafeModifyM = G.unsafeModifyM
399
400-- | Swap the elements at the given positions. No bounds checks are performed.
401unsafeSwap
402    :: (PrimMonad m, Prim a) => MVector (PrimState m) a -> Int -> Int -> m ()
403{-# INLINE unsafeSwap #-}
404unsafeSwap = G.unsafeSwap
405
406-- | Replace the element at the given position and return the old element. No
407-- bounds checks are performed.
408unsafeExchange :: (PrimMonad m, Prim a) => MVector (PrimState m) a -> Int -> a -> m a
409{-# INLINE unsafeExchange #-}
410unsafeExchange = G.unsafeExchange
411
412-- Filling and copying
413-- -------------------
414
415-- | Set all elements of the vector to the given value.
416set :: (PrimMonad m, Prim a) => MVector (PrimState m) a -> a -> m ()
417{-# INLINE set #-}
418set = G.set
419
420-- | Copy a vector. The two vectors must have the same length and may not
421-- overlap.
422copy :: (PrimMonad m, Prim a)
423     => MVector (PrimState m) a   -- ^ target
424     -> MVector (PrimState m) a   -- ^ source
425     -> m ()
426{-# INLINE copy #-}
427copy = G.copy
428
429-- | Copy a vector. The two vectors must have the same length and may not
430-- overlap. This is not checked.
431unsafeCopy :: (PrimMonad m, Prim a)
432           => MVector (PrimState m) a   -- ^ target
433           -> MVector (PrimState m) a   -- ^ source
434           -> m ()
435{-# INLINE unsafeCopy #-}
436unsafeCopy = G.unsafeCopy
437
438-- | Move the contents of a vector. The two vectors must have the same
439-- length.
440--
441-- If the vectors do not overlap, then this is equivalent to 'copy'.
442-- Otherwise, the copying is performed as if the source vector were
443-- copied to a temporary vector and then the temporary vector was copied
444-- to the target vector.
445move :: (PrimMonad m, Prim a)
446     => MVector (PrimState m) a   -- ^ target
447     -> MVector (PrimState m) a   -- ^ source
448     -> m ()
449{-# INLINE move #-}
450move = G.move
451
452-- | Move the contents of a vector. The two vectors must have the same
453-- length, but this is not checked.
454--
455-- If the vectors do not overlap, then this is equivalent to 'unsafeCopy'.
456-- Otherwise, the copying is performed as if the source vector were
457-- copied to a temporary vector and then the temporary vector was copied
458-- to the target vector.
459unsafeMove :: (PrimMonad m, Prim a)
460                          => MVector (PrimState m) a   -- ^ target
461                          -> MVector (PrimState m) a   -- ^ source
462                          -> m ()
463{-# INLINE unsafeMove #-}
464unsafeMove = G.unsafeMove
465
466-- | Compute the next (lexicographically) permutation of given vector in-place.
467--   Returns False when input is the last permutation
468nextPermutation :: (PrimMonad m,Ord e,Prim e) => MVector (PrimState m) e -> m Bool
469{-# INLINE nextPermutation #-}
470nextPermutation = G.nextPermutation
471
472
473-- Folds
474-- -----
475
476-- | /O(n)/ Apply the monadic action to every element of the vector, discarding the results.
477--
478-- @since 0.12.3.0
479mapM_ :: (PrimMonad m, Prim a) => (a -> m b) -> MVector (PrimState m) a -> m ()
480{-# INLINE mapM_ #-}
481mapM_ = G.mapM_
482
483-- | /O(n)/ Apply the monadic action to every element of the vector and its index, discarding the results.
484--
485-- @since 0.12.3.0
486imapM_ :: (PrimMonad m, Prim a) => (Int -> a -> m b) -> MVector (PrimState m) a -> m ()
487{-# INLINE imapM_ #-}
488imapM_ = G.imapM_
489
490-- | /O(n)/ Apply the monadic action to every element of the vector,
491-- discarding the results. It's same as the @flip mapM_@.
492--
493-- @since 0.12.3.0
494forM_ :: (PrimMonad m, Prim a) => MVector (PrimState m) a -> (a -> m b) -> m ()
495{-# INLINE forM_ #-}
496forM_ = G.forM_
497
498-- | /O(n)/ Apply the monadic action to every element of the vector
499-- and its index, discarding the results. It's same as the @flip imapM_@.
500--
501-- @since 0.12.3.0
502iforM_ :: (PrimMonad m, Prim a) => MVector (PrimState m) a -> (Int -> a -> m b) -> m ()
503{-# INLINE iforM_ #-}
504iforM_ = G.iforM_
505
506-- | /O(n)/ Pure left fold.
507--
508-- @since 0.12.3.0
509foldl :: (PrimMonad m, Prim a) => (b -> a -> b) -> b -> MVector (PrimState m) a -> m b
510{-# INLINE foldl #-}
511foldl = G.foldl
512
513-- | /O(n)/ Pure left fold with strict accumulator.
514--
515-- @since 0.12.3.0
516foldl' :: (PrimMonad m, Prim a) => (b -> a -> b) -> b -> MVector (PrimState m) a -> m b
517{-# INLINE foldl' #-}
518foldl' = G.foldl'
519
520-- | /O(n)/ Pure left fold (function applied to each element and its index).
521--
522-- @since 0.12.3.0
523ifoldl :: (PrimMonad m, Prim a) => (b -> Int -> a -> b) -> b -> MVector (PrimState m) a -> m b
524{-# INLINE ifoldl #-}
525ifoldl = G.ifoldl
526
527-- | /O(n)/ Pure left fold with strict accumulator (function applied to each element and its index).
528--
529-- @since 0.12.3.0
530ifoldl' :: (PrimMonad m, Prim a) => (b -> Int -> a -> b) -> b -> MVector (PrimState m) a -> m b
531{-# INLINE ifoldl' #-}
532ifoldl' = G.ifoldl'
533
534-- | /O(n)/ Pure right fold.
535--
536-- @since 0.12.3.0
537foldr :: (PrimMonad m, Prim a) => (a -> b -> b) -> b -> MVector (PrimState m) a -> m b
538{-# INLINE foldr #-}
539foldr = G.foldr
540
541-- | /O(n)/ Pure right fold with strict accumulator.
542--
543-- @since 0.12.3.0
544foldr' :: (PrimMonad m, Prim a) => (a -> b -> b) -> b -> MVector (PrimState m) a -> m b
545{-# INLINE foldr' #-}
546foldr' = G.foldr'
547
548-- | /O(n)/ Pure right fold (function applied to each element and its index).
549--
550-- @since 0.12.3.0
551ifoldr :: (PrimMonad m, Prim a) => (Int -> a -> b -> b) -> b -> MVector (PrimState m) a -> m b
552{-# INLINE ifoldr #-}
553ifoldr = G.ifoldr
554
555-- | /O(n)/ Pure right fold with strict accumulator (function applied
556-- to each element and its index).
557--
558-- @since 0.12.3.0
559ifoldr' :: (PrimMonad m, Prim a) => (Int -> a -> b -> b) -> b -> MVector (PrimState m) a -> m b
560{-# INLINE ifoldr' #-}
561ifoldr' = G.ifoldr'
562
563-- | /O(n)/ Monadic fold.
564--
565-- @since 0.12.3.0
566foldM :: (PrimMonad m, Prim a) => (b -> a -> m b) -> b -> MVector (PrimState m) a -> m b
567{-# INLINE foldM #-}
568foldM = G.foldM
569
570-- | /O(n)/ Monadic fold with strict accumulator.
571--
572-- @since 0.12.3.0
573foldM' :: (PrimMonad m, Prim a) => (b -> a -> m b) -> b -> MVector (PrimState m) a -> m b
574{-# INLINE foldM' #-}
575foldM' = G.foldM'
576
577-- | /O(n)/ Monadic fold (action applied to each element and its index).
578--
579-- @since 0.12.3.0
580ifoldM :: (PrimMonad m, Prim a) => (b -> Int -> a -> m b) -> b -> MVector (PrimState m) a -> m b
581{-# INLINE ifoldM #-}
582ifoldM = G.ifoldM
583
584-- | /O(n)/ Monadic fold with strict accumulator (action applied to each element and its index).
585--
586-- @since 0.12.3.0
587ifoldM' :: (PrimMonad m, Prim a) => (b -> Int -> a -> m b) -> b -> MVector (PrimState m) a -> m b
588{-# INLINE ifoldM' #-}
589ifoldM' = G.ifoldM'
590
591-- | /O(n)/ Monadic right fold.
592--
593-- @since 0.12.3.0
594foldrM :: (PrimMonad m, Prim a) => (a -> b -> m b) -> b -> MVector (PrimState m) a -> m b
595{-# INLINE foldrM #-}
596foldrM = G.foldrM
597
598-- | /O(n)/ Monadic right fold with strict accumulator.
599--
600-- @since 0.12.3.0
601foldrM' :: (PrimMonad m, Prim a) => (a -> b -> m b) -> b -> MVector (PrimState m) a -> m b
602{-# INLINE foldrM' #-}
603foldrM' = G.foldrM'
604
605-- | /O(n)/ Monadic right fold (action applied to each element and its index).
606--
607-- @since 0.12.3.0
608ifoldrM :: (PrimMonad m, Prim a) => (Int -> a -> b -> m b) -> b -> MVector (PrimState m) a -> m b
609{-# INLINE ifoldrM #-}
610ifoldrM = G.ifoldrM
611
612-- | /O(n)/ Monadic right fold with strict accumulator (action applied
613-- to each element and its index).
614--
615-- @since 0.12.3.0
616ifoldrM' :: (PrimMonad m, Prim a) => (Int -> a -> b -> m b) -> b -> MVector (PrimState m) a -> m b
617{-# INLINE ifoldrM' #-}
618ifoldrM' = G.ifoldrM'
619