Implement Mut.restream
[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 #include "phases.h"
15
16 module Data.Vector.MVector (
17 MVectorPure(..), MVector(..),
18
19 slice, new, newWith, read, write, copy, grow,
20 unstream, mstream, munstream,
21 update, reverse
22 ) where
23
24 import qualified Data.Vector.Stream as Stream
25 import Data.Vector.Stream ( Stream )
26 import Data.Vector.Stream.Size
27
28 import Control.Monad.ST ( ST )
29 import Control.Exception ( assert )
30
31 import GHC.Float (
32 double2Int, int2Double
33 )
34
35 import Prelude hiding ( length, reverse, map, read )
36
37 gROWTH_FACTOR :: Double
38 gROWTH_FACTOR = 1.5
39
40 -- | Basic pure functions on mutable vectors
41 class MVectorPure v a where
42 -- | Length of the mutable vector
43 length :: v a -> Int
44
45 -- | Yield a part of the mutable vector without copying it. No range checks!
46 unsafeSlice :: v a -> Int -- ^ starting index
47 -> Int -- ^ length of the slice
48 -> v a
49
50 -- Check whether two vectors overlap.
51 overlaps :: v a -> v a -> Bool
52
53 -- | Class of mutable vectors. The type @m@ is the monad in which the mutable
54 -- vector can be transformed and @a@ is the type of elements.
55 --
56 class (Monad m, MVectorPure v a) => MVector v m a where
57 -- | Create a mutable vector of the given length. Length is not checked!
58 unsafeNew :: Int -> m (v a)
59
60 -- | Create a mutable vector of the given length and fill it with an
61 -- initial value. Length is not checked!
62 unsafeNewWith :: Int -> a -> m (v a)
63
64 -- | Yield the element at the given position. Index is not checked!
65 unsafeRead :: v a -> Int -> m a
66
67 -- | Replace the element at the given position. Index is not checked!
68 unsafeWrite :: v a -> Int -> a -> m ()
69
70 -- | Clear all references to external objects
71 clear :: v a -> m ()
72
73 -- | Write the value at each position.
74 set :: v a -> a -> m ()
75
76 -- | Copy a vector. The two vectors may not overlap. This is not checked!
77 unsafeCopy :: v a -- ^ target
78 -> v a -- ^ source
79 -> m ()
80
81 -- | Grow a vector by the given number of elements. The length is not
82 -- checked!
83 unsafeGrow :: v a -> Int -> m (v a)
84
85 {-# INLINE unsafeNewWith #-}
86 unsafeNewWith n x = do
87 v <- unsafeNew n
88 set v x
89 return v
90
91 {-# INLINE set #-}
92 set v x = do_set 0
93 where
94 n = length v
95
96 do_set i | i < n = do
97 unsafeWrite v i x
98 do_set (i+1)
99 | otherwise = return ()
100
101 {-# INLINE unsafeCopy #-}
102 unsafeCopy dst src = do_copy 0
103 where
104 n = length src
105
106 do_copy i | i < n = do
107 x <- unsafeRead src i
108 unsafeWrite dst i x
109 do_copy (i+1)
110 | otherwise = return ()
111
112 {-# INLINE unsafeGrow #-}
113 unsafeGrow v by = do
114 v' <- unsafeNew (n+by)
115 unsafeCopy (unsafeSlice v' 0 n) v
116 return v'
117 where
118 n = length v
119
120 -- | Test whether the index is valid for the vector
121 inBounds :: MVectorPure v a => v a -> Int -> Bool
122 {-# INLINE inBounds #-}
123 inBounds v i = i >= 0 && i < length v
124
125 -- | Yield a part of the mutable vector without copying it. Safer version of
126 -- 'unsafeSlice'.
127 slice :: MVectorPure v a => v a -> Int -> Int -> v a
128 {-# INLINE slice #-}
129 slice v i n = assert (i >=0 && n >= 0 && i+n <= length v)
130 $ unsafeSlice v i n
131
132 -- | Create a mutable vector of the given length. Safer version of
133 -- 'unsafeNew'.
134 new :: MVector v m a => Int -> m (v a)
135 {-# INLINE new #-}
136 new n = assert (n >= 0) $ unsafeNew n
137
138 -- | Create a mutable vector of the given length and fill it with an
139 -- initial value. Safer version of 'unsafeNewWith'.
140 newWith :: MVector v m a => Int -> a -> m (v a)
141 {-# INLINE newWith #-}
142 newWith n x = assert (n >= 0) $ unsafeNewWith n x
143
144 -- | Yield the element at the given position. Safer version of 'unsafeRead'.
145 read :: MVector v m a => v a -> Int -> m a
146 {-# INLINE read #-}
147 read v i = assert (inBounds v i) $ unsafeRead v i
148
149 -- | Replace the element at the given position. Safer version of
150 -- 'unsafeWrite'.
151 write :: MVector v m a => v a -> Int -> a -> m ()
152 {-# INLINE write #-}
153 write v i x = assert (inBounds v i) $ unsafeWrite v i x
154
155 -- | Copy a vector. The two vectors may not overlap. Safer version of
156 -- 'unsafeCopy'.
157 copy :: MVector v m a => v a -> v a -> m ()
158 {-# INLINE copy #-}
159 copy dst src = assert (not (dst `overlaps` src) && length dst == length src)
160 $ unsafeCopy dst src
161
162 -- | Grow a vector by the given number of elements. Safer version of
163 -- 'unsafeGrow'.
164 grow :: MVector v m a => v a -> Int -> m (v a)
165 {-# INLINE grow #-}
166 grow v by = assert (by >= 0)
167 $ unsafeGrow v by
168
169 mstream :: MVector v m a => v a -> Stream (m a)
170 {-# INLINE mstream #-}
171 mstream v = v `seq` (Stream.unfold get 0 `Stream.sized` Exact n)
172 where
173 n = length v
174
175 {-# INLINE get #-}
176 get i | i < n = Just (unsafeRead v i, i+1)
177 | otherwise = Nothing
178
179 munstream :: MVector v m a => v a -> Stream (m a) -> m (v a)
180 {-# INLINE munstream #-}
181 munstream v s = v `seq` do
182 n' <- Stream.foldM put 0 s
183 return $ slice v 0 n'
184 where
185 put i m = do { write v i =<< m; return (i+1) }
186
187 -- | Create a new mutable vector and fill it with elements from the 'Stream'.
188 -- The vector will grow logarithmically if the 'Size' hint of the 'Stream' is
189 -- inexact.
190 unstream :: MVector v m a => Stream a -> m (v a)
191 {-# INLINE_STREAM unstream #-}
192 unstream s = case upperBound (Stream.size s) of
193 Just n -> unstreamMax s n
194 Nothing -> unstreamUnknown s
195
196 unstreamMax :: MVector v m a => Stream a -> Int -> m (v a)
197 {-# INLINE unstreamMax #-}
198 unstreamMax s n
199 = do
200 v <- new n
201 let put i x = do { write v i x; return (i+1) }
202 n' <- Stream.foldM put 0 s
203 return $ slice v 0 n'
204
205 unstreamUnknown :: MVector v m a => Stream a -> m (v a)
206 {-# INLINE unstreamUnknown #-}
207 unstreamUnknown s
208 = do
209 v <- new 0
210 (v', n) <- Stream.foldM put (v, 0) s
211 return $ slice v' 0 n
212 where
213 {-# INLINE put #-}
214 put (v, i) x = do
215 v' <- enlarge v i
216 unsafeWrite v' i x
217 return (v', i+1)
218
219 {-# INLINE enlarge #-}
220 enlarge v i | i < length v = return v
221 | otherwise = unsafeGrow v
222 . max 1
223 . double2Int
224 $ int2Double (length v) * gROWTH_FACTOR
225
226 update :: MVector v m a => v a -> Stream (Int, a) -> m ()
227 {-# INLINE update #-}
228 update v s = Stream.mapM_ put s
229 where
230 {-# INLINE put #-}
231 put (i, x) = write v i x
232
233 reverse :: MVector v m a => v a -> m ()
234 {-# INLINE reverse #-}
235 reverse v = reverse_loop 0 (length v - 1)
236 where
237 reverse_loop i j | i < j = do
238 x <- unsafeRead v i
239 y <- unsafeRead v j
240 unsafeWrite v i y
241 unsafeWrite v j x
242 reverse_loop _ _ = return ()
243