From 46068c9dafe8cfa763ef855ec21f577a1e058de1 Mon Sep 17 00:00:00 2001
From: Stepan Koltsov <stepan.koltsov@gmail.com>
Date: Sat, 31 Oct 2015 00:17:16 +0300
Subject: [PATCH] Fix excessive memory allocation in RawVec::reserve

Before this patch `reserve` function allocated twice as requested
amount elements (not twice as capacity).  It leaded to unnecessary
excessive memory usage in scenarios like this:

```
let mut v = Vec::new();
v.push(17);
v.extend(0..10);
println!("{}", v.capacity());
```

`Vec` allocated 22 elements, while it could allocate just 11.

`reserve` function must have a property of keeping `push` operation
cost (which calls `reserve`) `O(1)`. To achieve this `reserve` must
exponentialy grow its capacity when it does reallocation.

There's better strategy to implement `reserve`:

```
let new_capacity = max(current_capacity * 2, requested_capacity);
```

This strategy still guarantees that capacity grows at `O(1)` with
`reserve`, and fixes the issue with `extend`.

Patch imlpements this strategy.
---
 src/liballoc/raw_vec.rs | 52 ++++++++++++++++++++++++++++++++++++++---
 1 file changed, 49 insertions(+), 3 deletions(-)

diff --git a/src/liballoc/raw_vec.rs b/src/liballoc/raw_vec.rs
index eef8f3e94e07d..78deed0c84f7f 100644
--- a/src/liballoc/raw_vec.rs
+++ b/src/liballoc/raw_vec.rs
@@ -15,6 +15,7 @@ use heap;
 use super::oom;
 use super::boxed::Box;
 use core::ops::Drop;
+use core::cmp;
 use core;
 
 /// A low-level utility for more ergonomically allocating, reallocating, and deallocating a
@@ -360,9 +361,15 @@ impl<T> RawVec<T> {
             }
 
             // Nothing we can really do about these checks :(
-            let new_cap = used_cap.checked_add(needed_extra_cap)
-                                  .and_then(|cap| cap.checked_mul(2))
-                                  .expect("capacity overflow");
+            let required_cap = used_cap.checked_add(needed_extra_cap)
+                                       .expect("capacity overflow");
+
+            // Cannot overflow, because `cap <= isize::MAX`, and type of `cap` is `usize`.
+            let double_cap = self.cap * 2;
+
+            // `double_cap` guarantees exponential growth.
+            let new_cap = cmp::max(double_cap, required_cap);
+
             let new_alloc_size = new_cap.checked_mul(elem_size).expect("capacity overflow");
             // FIXME: may crash and burn on over-reserve
             alloc_guard(new_alloc_size);
@@ -486,3 +493,42 @@ fn alloc_guard(alloc_size: usize) {
                 "capacity overflow");
     }
 }
+
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn reserve_does_not_overallocate() {
+        {
+            let mut v: RawVec<u32> = RawVec::new();
+            // First `reserve` allocates like `reserve_exact`
+            v.reserve(0, 9);
+            assert_eq!(9, v.cap());
+        }
+
+        {
+            let mut v: RawVec<u32> = RawVec::new();
+            v.reserve(0, 7);
+            assert_eq!(7, v.cap());
+            // 97 if more than double of 7, so `reserve` should work
+            // like `reserve_exact`.
+            v.reserve(7, 90);
+            assert_eq!(97, v.cap());
+        }
+
+        {
+            let mut v: RawVec<u32> = RawVec::new();
+            v.reserve(0, 12);
+            assert_eq!(12, v.cap());
+            v.reserve(12, 3);
+            // 3 is less than half of 12, so `reserve` must grow
+            // exponentially. At the time of writing this test grow
+            // factor is 2, so new capacity is 24, however, grow factor
+            // of 1.5 is OK too. Hence `>= 18` in assert.
+            assert!(v.cap() >= 12 + 12 / 2);
+        }
+    }
+
+}