More flexible size hints
[darcs-mirrors/vector.git] / Data / Vector / Unboxed / Prim.hs
index 4bc224c..12ca5cd 100644 (file)
@@ -3,7 +3,8 @@
 module Data.Vector.Unboxed.Prim (
   Unbox(..), Vector, MutableVector,
 
-  size, new, unsafeFreeze, at, read, write
+  size, new, unsafeFreeze, at, read, write,
+  copy, grow
 ) where
 
 import Data.Vector.Unboxed.Unbox
@@ -12,6 +13,9 @@ import GHC.Prim (
     ByteArray#, MutableByteArray#,
     newByteArray#, unsafeFreezeByteArray#,
   )
+import GHC.Float (
+    int2Double, double2Int
+  )
 import GHC.ST (
     ST(..)
   )
@@ -53,3 +57,28 @@ write :: Unbox a => MutableVector s a -> Int -> a -> ST s ()
 write (MutableVector arr#) (I# i#) x = ST $ \s# ->
   case write# arr# i# x s# of s2# -> (# s2#, () #)
 
+copy :: Unbox a => MutableVector s a -> Int
+                -> MutableVector s a -> Int -> Int
+                -> ST s ()
+{-# INLINE copy #-}
+copy mv i mw j n = do_copy i j n
+  where
+    do_copy i j 0 = return ()
+    do_copy i j n = do
+                      x <- read mw j
+                      write mv i x
+                      do_copy (i+1) (j+1) (n-1)
+
+grow :: Unbox a => MutableVector s a -> Int -> Double
+                -> ST s (MutableVector s a, Int)
+{-# INLINE grow #-}
+grow v n r
+  = do
+      w <- new m
+      copy w 0 v 0 n
+      return (w, m)
+  where
+    n' = double2Int (int2Double n * r)
+    m | n' <= n   = n+1
+      | otherwise = n'
+