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