vec: add `shrink_to_fit`

Closes #4960

1 files changed, 26 insertions, 0 deletions

diff --git a/src/libstd/vec.rs b/src/libstd/vec.rs
index 7c8046a64b2..0427eec2b05 100644
--- a/src/libstd/vec.rs
+++ b/src/libstd/vec.rs
@@ -1161,6 +1161,7 @@ pub trait OwnedVector<T> {
     fn reserve(&mut self, n: uint);
     fn reserve_at_least(&mut self, n: uint);
     fn capacity(&self) -> uint;
+    fn shrink_to_fit(&mut self);
 
     fn push(&mut self, t: T);
     unsafe fn push_fast(&mut self, t: T);
@@ -1254,6 +1255,7 @@ impl<T> OwnedVector<T> for ~[T] {
      *
      * * n - The number of elements to reserve space for
      */
+    #[inline]
     fn reserve_at_least(&mut self, n: uint) {
         self.reserve(uint::next_power_of_two(n));
     }
@@ -1272,6 +1274,17 @@ impl<T> OwnedVector<T> for ~[T] {
         }
     }
 
+    /// Shrink the capacity of the vector to match the length
+    fn shrink_to_fit(&mut self) {
+        unsafe {
+            let ptr: *mut *mut Vec<()> = cast::transmute(self);
+            let alloc = (**ptr).fill;
+            let size = alloc + sys::size_of::<Vec<()>>();
+            *ptr = realloc_raw(*ptr as *mut c_void, size) as *mut Vec<()>;
+            (**ptr).alloc = alloc;
+        }
+    }
+
     /// Append an element to a vector
     #[inline]
     fn push(&mut self, t: T) {
@@ -2327,6 +2340,7 @@ mod tests {
     use sys;
     use vec::*;
     use cmp::*;
+    use prelude::*;
 
     fn square(n: uint) -> uint { n * n }
 
@@ -3600,6 +3614,18 @@ mod tests {
         }
         assert!(cnt == 3);
     }
+
+    #[test]
+    fn test_shrink_to_fit() {
+        let mut xs = ~[0, 1, 2, 3];
+        for i in range(4, 100) {
+            xs.push(i)
+        }
+        assert_eq!(xs.capacity(), 128);
+        xs.shrink_to_fit();
+        assert_eq!(xs.capacity(), 100);
+        assert_eq!(xs, range(0, 100).to_owned_vec());
+    }
 }
 
 #[cfg(test)]