More flexible size hints
[darcs-mirrors/vector.git] / Data / Vector / Prim.hs
index df5c68f..f9eb004 100644 (file)
@@ -2,13 +2,16 @@
 
 module Data.Vector.Prim (
   Vector, MutableVector,
-  new, new', unsafeFreeze, at, at', read, write
+  new, new', unsafeFreeze, at, at', read, write, copy, grow
 ) where
 
 import GHC.Prim (
     Array#, MutableArray#,
     newArray#, readArray#, writeArray#, indexArray#, unsafeFreezeArray#
   )
+import GHC.Float (
+    double2Int, int2Double
+  )
 import GHC.ST (
     ST(..)
   )
@@ -54,3 +57,25 @@ write :: MutableVector s a -> Int -> a -> ST s ()
 write (MutableVector arr#) (I# n#) x = ST $ \s# ->
   case writeArray# arr# n# x s# of s2# -> (# s2#, () #)
 
+copy :: 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 :: 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'
+