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