32fb5fd1c3864992f713c6e1df53ff1a23bbcdb5
[darcs-mirrors/vector.git] / Data / Vector / Unboxed / Mutable.hs
1 {-# LANGUAGE BangPatterns #-}
2
3 module Data.Vector.Unboxed.Mutable (
4 Vector,
5
6 new, length, slice, read, write, unstream, fill, fillIndexed,
7 dataOf
8 ) where
9
10 import qualified Data.Vector.Unboxed.Prim as Prim
11 import Data.Vector.Unboxed.Unbox ( Unbox )
12
13 import Data.Vector.Stream.Size ( upperBound )
14 import qualified Data.Vector.Stream as Stream
15 import Data.Vector.Stream ( Stream )
16
17 import Control.Exception ( assert )
18 import Control.Monad.ST ( ST )
19
20 import Prelude hiding ( length, read )
21
22 data Vector s a = Vector {-# UNPACK #-} !Int -- ^ start
23 {-# UNPACK #-} !Int -- ^ length
24 {-# UNPACK #-} !(Prim.MutableVector s a) -- ^ data
25
26 dataOf :: Vector s a -> (Prim.MutableVector s a, Int, Int)
27 {-# INLINE dataOf #-}
28 dataOf (Vector i n v) = (v, i, n)
29
30 new :: Unbox a => Int -> ST s (Vector s a)
31 {-# INLINE new #-}
32 new n = assert (n >= 0)
33 $ Vector 0 n `fmap` Prim.new n
34
35 length :: Unbox a => Vector s a -> Int
36 {-# INLINE length #-}
37 length (Vector _ n _) = n
38
39 slice :: Unbox a => Vector s a -> Int -> Int -> Vector s a
40 {-# INLINE slice #-}
41 slice (Vector i n v) j m
42 = assert (j + m <= n && j >= 0 && m >= 0)
43 $ Vector (i+j) m v
44
45 slicel :: Unbox a => Vector s a -> Int -> Vector s a
46 {-# INLINE slicel #-}
47 slicel (Vector i n v) m
48 = assert (m <= n && m >= 0)
49 $ Vector i m v
50
51 read :: Unbox a => Vector s a -> Int -> ST s a
52 {-# INLINE read #-}
53 read (Vector i n v) j
54 = assert (j < n)
55 $ Prim.read v (i+j)
56
57 write :: Unbox a => Vector s a -> Int -> a -> ST s ()
58 {-# INLINE write #-}
59 write (Vector i n v) j x
60 = assert (j < n)
61 $ Prim.write v (i+j) x
62
63 unstream :: Unbox a => Stream a -> ST s (Vector s a)
64 {-# INLINE unstream #-}
65 unstream s = case upperBound (Stream.size s) of
66 Just n -> unstream_known s n
67 Nothing -> unstream_unknown s
68
69 gROWTH_FACTOR :: Double
70 gROWTH_FACTOR = 1.6
71
72 unstream_known :: Unbox a => Stream a -> Int -> ST s (Vector s a)
73 {-# INLINE unstream_known #-}
74 unstream_known s n
75 = do
76 v <- new n
77 n' <- fill v s
78 return $ slice v 0 n'
79
80 unstream_unknown :: Unbox a => Stream a -> ST s (Vector s a)
81 {-# INLINE unstream_unknown #-}
82 unstream_unknown s
83 = do
84 v <- Prim.new 0
85 (w, n, _) <- Stream.foldM put (v, 0, 0) s
86 return $ Vector 0 n w
87 where
88 {-# INLINE put #-}
89 put (v, i, n) x = do
90 (v', n') <- enlarge v i n
91 Prim.write v' i x
92 return (v', i+1, n')
93
94 {-# INLINE enlarge #-}
95 enlarge v i n | i < n = return (v, n)
96 | otherwise = Prim.grow v n gROWTH_FACTOR
97
98 fill :: Unbox a => Vector s a -> Stream a -> ST s Int
99 {-# INLINE fill #-}
100 fill !v s = Stream.foldM put 0 s
101 where
102 {-# INLINE put #-}
103 put i x = do { write v i x; return (i+1) }
104
105 fillIndexed :: Unbox a => Vector s a -> Stream (Int, a) -> ST s ()
106 {-# INLINE fillIndexed #-}
107 fillIndexed !v s = Stream.mapM_ put s
108 where
109 {-# INLINE put #-}
110 put (i,x) = write v i x
111
112 copyTo :: Unbox a => Vector s a -> Vector s a -> ST s ()
113 {-# INLINE copyTo #-}
114 copyTo !v !w = assert (length v == length w)
115 $ copy_loop 0
116 where
117 n = length v
118
119 copy_loop i | i < n = do
120 x <- read v i
121 write w i x
122 copy_loop (i+1)
123 | otherwise = return ()
124
125 clone :: Unbox a => Vector s a -> ST s (Vector s a)
126 {-# INLINE clone #-}
127 clone v = do
128 w <- new (length v)
129 v `copyTo` w
130 return w
131