f9eb0045863150c146a04abd7fa6c1f2e4c88d14
[darcs-mirrors/vector.git] / Data / Vector / Prim.hs
1 {-# LANGUAGE MagicHash, UnboxedTuples #-}
2
3 module Data.Vector.Prim (
4 Vector, MutableVector,
5 new, new', unsafeFreeze, at, at', read, write, copy, grow
6 ) where
7
8 import GHC.Prim (
9 Array#, MutableArray#,
10 newArray#, readArray#, writeArray#, indexArray#, unsafeFreezeArray#
11 )
12 import GHC.Float (
13 double2Int, int2Double
14 )
15 import GHC.ST (
16 ST(..)
17 )
18 import GHC.Base (
19 Int(..)
20 )
21
22 import Prelude hiding ( read )
23
24 data Vector a = Vector (Array# a)
25 data MutableVector s a = MutableVector (MutableArray# s a)
26
27 new :: Int -> ST s (MutableVector s a)
28 {-# INLINE new #-}
29 new n = new' n (error "Data.Vector: uninitialised element")
30
31 new' :: Int -> a -> ST s (MutableVector s a)
32 {-# INLINE new' #-}
33 new' (I# n#) x = ST $ \s# ->
34 case newArray# n# x s# of
35 (# s2#, arr# #) -> (# s2#, MutableVector arr# #)
36
37 unsafeFreeze :: MutableVector s a -> ST s (Vector a)
38 {-# INLINE unsafeFreeze #-}
39 unsafeFreeze (MutableVector arr#) = ST $ \s# ->
40 case unsafeFreezeArray# arr# s# of
41 (# s2, frozen# #) -> (# s2, Vector frozen# #)
42
43 at :: Vector a -> Int -> a
44 {-# INLINE at #-}
45 at v i = at' v i id
46
47 at' :: Vector a -> Int -> (a -> b) -> b
48 {-# INLINE at' #-}
49 at' (Vector arr#) (I# n#) f = case indexArray# arr# n# of (# x #) -> f x
50
51 read :: MutableVector s a -> Int -> ST s a
52 {-# INLINE read #-}
53 read (MutableVector arr#) (I# n#) = ST $ readArray# arr# n#
54
55 write :: MutableVector s a -> Int -> a -> ST s ()
56 {-# INLINE write #-}
57 write (MutableVector arr#) (I# n#) x = ST $ \s# ->
58 case writeArray# arr# n# x s# of s2# -> (# s2#, () #)
59
60 copy :: MutableVector s a -> Int -> MutableVector s a -> Int -> Int -> ST s ()
61 {-# INLINE copy #-}
62 copy mv i mw j n = do_copy i j n
63 where
64 do_copy i j 0 = return ()
65 do_copy i j n = do
66 x <- read mw j
67 write mv i x
68 do_copy (i+1) (j+1) (n-1)
69
70 grow :: MutableVector s a -> Int -> Double -> ST s (MutableVector s a, Int)
71 {-# INLINE grow #-}
72 grow v n r
73 = do
74 w <- new m
75 copy w 0 v 0 n
76 return (w, m)
77 where
78 n' = double2Int (int2Double n * r)
79 m | n' <= n = n+1
80 | otherwise = n'
81