More flexible size hints
[darcs-mirrors/vector.git] / Data / Vector / Mutable.hs
index c958d4b..244d996 100644 (file)
@@ -3,12 +3,13 @@
 module Data.Vector.Mutable (
   Vector,
 
-  new, new', length, slice, read, write, fill,
+  new, new', length, slice, read, write, unstream, fill,
   dataOf
 ) where
 
 import qualified Data.Vector.Prim as Prim
 
+import           Data.Vector.Stream.Size ( upperBound )
 import qualified Data.Vector.Stream as Stream
 import           Data.Vector.Stream ( Stream )
 
@@ -56,6 +57,41 @@ write (Vector i n v) j x
   = assert (j < n)
   $ Prim.write v (i+j) x
 
+unstream :: Stream a -> ST s (Vector s a)
+{-# INLINE unstream #-}
+unstream s = case upperBound (Stream.size s) of
+               Just n  -> unstream_known   s n
+               Nothing -> unstream_unknown s
+
+gROWTH_FACTOR :: Double
+gROWTH_FACTOR = 1.6
+
+unstream_known :: Stream a -> Int -> ST s (Vector s a)
+{-# INLINE unstream_known #-}
+unstream_known s n
+  = do
+      v  <- new n
+      n' <- fill v s
+      return $ slice v 0 n'
+
+unstream_unknown :: Stream a -> ST s (Vector s a)
+{-# INLINE unstream_unknown #-}
+unstream_unknown s
+  = do
+      v <- Prim.new 0
+      (w, n, _) <- Stream.foldM put (v, 0, 0) s
+      return $ Vector 0 n w
+  where
+    {-# INLINE put #-}
+    put (v, i, n) x = do
+                        (v', n') <- enlarge v i n
+                        Prim.write v' i x
+                        return (v', i+1, n')
+
+    {-# INLINE enlarge #-}
+    enlarge v i n | i < n     = return (v, n)
+                  | otherwise = Prim.grow v n gROWTH_FACTOR
+
 fill :: Vector s a -> Stream a -> ST s Int
 {-# INLINE fill #-}
 fill !v s = Stream.foldM put 0 s