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