vec::with_capacity: do one alloc for non-managed

1 files changed, 20 insertions, 1 deletions

diff --git a/src/libstd/vec.rs b/src/libstd/vec.rs
index 2d730b7a6a2..0b29caacb7a 100644
--- a/src/libstd/vec.rs
+++ b/src/libstd/vec.rs
@@ -27,7 +27,7 @@ use option::{None, Option, Some};
 use ptr::to_unsafe_ptr;
 use ptr;
 use ptr::RawPtr;
-use rt::global_heap::realloc_raw;
+use rt::global_heap::{malloc_raw, realloc_raw};
 use sys;
 use sys::size_of;
 use uint;
@@ -101,12 +101,31 @@ pub fn to_owned<T:Copy>(t: &[T]) -> ~[T] {
 }
 
 /// Creates a new vector with a capacity of `capacity`
+#[cfg(stage0)]
 pub fn with_capacity<T>(capacity: uint) -> ~[T] {
     let mut vec = ~[];
     vec.reserve(capacity);
     vec
 }
 
+/// Creates a new vector with a capacity of `capacity`
+#[cfg(not(stage0))]
+pub fn with_capacity<T>(capacity: uint) -> ~[T] {
+    unsafe {
+        if contains_managed::<T>() {
+            let mut vec = ~[];
+            vec.reserve(capacity);
+            vec
+        } else {
+            let alloc = capacity * sys::nonzero_size_of::<T>();
+            let ptr = malloc_raw(alloc + size_of::<raw::VecRepr>()) as *mut raw::VecRepr;
+            (*ptr).unboxed.alloc = alloc;
+            (*ptr).unboxed.fill = 0;
+            cast::transmute(ptr)
+        }
+    }
+}
+
 /**
  * Builds a vector by calling a provided function with an argument
  * function that pushes an element to the back of a vector.