b26d0c32ff9c3e420e582bf828d2f9d872588286
[darcs-mirrors/vector.git] / Data / Vector / MVector.hs
1 {-# LANGUAGE MultiParamTypeClasses #-}
2 -- |
3 -- Module : Data.Vector.MVector
4 -- Copyright : (c) Roman Leshchinskiy 2008
5 -- License : BSD-style
6 --
7 -- Maintainer : rl@cse.unsw.edu.au
8 -- Stability : experimental
9 -- Portability : non-portable
10 --
11 -- Generic interface to mutable vectors
12 --
13
14 module Data.Vector.MVector (
15 MVectorPure(..), MVector(..),
16
17 slice, new, newWith, read, write, copy, grow, unstream
18 ) where
19
20 import qualified Data.Vector.Stream as Stream
21 import Data.Vector.Stream ( Stream )
22 import Data.Vector.Stream.Size
23
24 import Control.Monad.ST ( ST )
25 import Control.Exception ( assert )
26
27 import GHC.Float (
28 double2Int, int2Double
29 )
30
31 import Prelude hiding ( length, read )
32
33 gROWTH_FACTOR :: Double
34 gROWTH_FACTOR = 1.5
35
36 -- | Basic pure functions on mutable vectors
37 class MVectorPure v a where
38 -- | Length of the mutable vector
39 length :: v a -> Int
40
41 -- | Yield a part of the mutable vector without copying it. No range checks!
42 unsafeSlice :: v a -> Int -- ^ starting index
43 -> Int -- ^ length of the slice
44 -> v a
45
46 -- Check whether two vectors overlap.
47 overlaps :: v a -> v a -> Bool
48
49 -- | Class of mutable vectors. The type @m@ is the monad in which the mutable
50 -- vector can be transformed and @a@ is the type of elements.
51 --
52 class (Monad m, MVectorPure v a) => MVector v m a where
53 -- | Create a mutable vector of the given length. Length is not checked!
54 unsafeNew :: Int -> m (v a)
55
56 -- | Create a mutable vector of the given length and fill it with an
57 -- initial value. Length is not checked!
58 unsafeNewWith :: Int -> a -> m (v a)
59
60 -- | Yield the element at the given position. Index is not checked!
61 unsafeRead :: v a -> Int -> m a
62
63 -- | Replace the element at the given position. Index is not checked!
64 unsafeWrite :: v a -> Int -> a -> m ()
65
66 -- | Write the value at each position.
67 set :: v a -> a -> m ()
68
69 -- | Copy a vector. The two vectors may not overlap. This is not checked!
70 unsafeCopy :: v a -- ^ target
71 -> v a -- ^ source
72 -> m ()
73
74 -- | Grow a vector by the given number of elements. The length is not
75 -- checked!
76 unsafeGrow :: v a -> Int -> m (v a)
77
78 {-# INLINE unsafeNewWith #-}
79 unsafeNewWith n x = do
80 v <- unsafeNew n
81 set v x
82 return v
83
84 {-# INLINE set #-}
85 set v x = do_set 0
86 where
87 n = length v
88
89 do_set i | i < n = do
90 unsafeWrite v i x
91 do_set (i+1)
92 | otherwise = return ()
93
94 {-# INLINE unsafeCopy #-}
95 unsafeCopy dst src = do_copy 0
96 where
97 n = length src
98
99 do_copy i | i < n = do
100 x <- unsafeRead src i
101 unsafeWrite dst i x
102 do_copy (i+1)
103 | otherwise = return ()
104
105 {-# INLINE unsafeGrow #-}
106 unsafeGrow v by = do
107 v' <- unsafeNew (n+by)
108 unsafeCopy (unsafeSlice v' 0 n) v
109 return v'
110 where
111 n = length v
112
113 -- | Test whether the index is valid for the vector
114 inBounds :: MVectorPure v a => v a -> Int -> Bool
115 {-# INLINE inBounds #-}
116 inBounds v i = i >= 0 && i < length v
117
118 -- | Yield a part of the mutable vector without copying it. Safer version of
119 -- 'unsafeSlice'.
120 slice :: MVectorPure v a => v a -> Int -> Int -> v a
121 {-# INLINE slice #-}
122 slice v i n = assert (i >=0 && n >= 0 && i+n <= length v)
123 $ unsafeSlice v i n
124
125 -- | Create a mutable vector of the given length. Safer version of
126 -- 'unsafeNew'.
127 new :: MVector v m a => Int -> m (v a)
128 {-# INLINE new #-}
129 new n = assert (n >= 0) $ unsafeNew n
130
131 -- | Create a mutable vector of the given length and fill it with an
132 -- initial value. Safer version of 'unsafeNewWith'.
133 newWith :: MVector v m a => Int -> a -> m (v a)
134 {-# INLINE newWith #-}
135 newWith n x = assert (n >= 0) $ unsafeNewWith n x
136
137 -- | Yield the element at the given position. Safer version of 'unsafeRead'.
138 read :: MVector v m a => v a -> Int -> m a
139 {-# INLINE read #-}
140 read v i = assert (inBounds v i) $ unsafeRead v i
141
142 -- | Replace the element at the given position. Safer version of
143 -- 'unsafeWrite'.
144 write :: MVector v m a => v a -> Int -> a -> m ()
145 {-# INLINE write #-}
146 write v i x = assert (inBounds v i) $ unsafeWrite v i x
147
148 -- | Copy a vector. The two vectors may not overlap. Safer version of
149 -- 'unsafeCopy'.
150 copy :: MVector v m a => v a -> v a -> m ()
151 {-# INLINE copy #-}
152 copy dst src = assert (not (dst `overlaps` src) && length dst == length src)
153 $ unsafeCopy dst src
154
155 -- | Grow a vector by the given number of elements. Safer version of
156 -- 'unsafeGrow'.
157 grow :: MVector v m a => v a -> Int -> m (v a)
158 {-# INLINE grow #-}
159 grow v by = assert (by >= 0)
160 $ unsafeGrow v by
161
162
163 -- | Create a new mutable vector and fill it with elements from the 'Stream'.
164 -- The vector will grow logarithmically if the 'Size' hint of the 'Stream' is
165 -- inexact.
166 unstream :: MVector v m a => Stream a -> m (v a)
167 {-# INLINE unstream #-}
168 unstream s = case upperBound (Stream.size s) of
169 Just n -> unstreamMax s n
170 Nothing -> unstreamUnknown s
171
172 unstreamMax :: MVector v m a => Stream a -> Int -> m (v a)
173 {-# INLINE unstreamMax #-}
174 unstreamMax s n
175 = do
176 v <- new n
177 let put i x = do { write v i x; return (i+1) }
178 n' <- Stream.foldM put 0 s
179 return $ slice v 0 n'
180
181 unstreamUnknown :: MVector v m a => Stream a -> m (v a)
182 {-# INLINE unstreamUnknown #-}
183 unstreamUnknown s
184 = do
185 v <- new 0
186 (v', n) <- Stream.foldM put (v, 0) s
187 return $ slice v' 0 n
188 where
189 {-# INLINE put #-}
190 put (v, i) x = do
191 v' <- enlarge v i
192 unsafeWrite v' i x
193 return (v', i+1)
194
195 {-# INLINE enlarge #-}
196 enlarge v i | i < length v = return v
197 | otherwise = unsafeGrow v
198 . max 1
199 . double2Int
200 $ int2Double (length v) * gROWTH_FACTOR
201