implementing fallible allocation API (try_reserve) for Vec, String and HashMap

1 files changed, 208 insertions, 1 deletions

diff --git a/src/liballoc/tests/vec.rs b/src/liballoc/tests/vec.rs
index 9cfde5dcc73..3c17a401bba 100644
--- a/src/liballoc/tests/vec.rs
+++ b/src/liballoc/tests/vec.rs
@@ -10,8 +10,9 @@
 
 use std::borrow::Cow;
 use std::mem::size_of;
-use std::panic;
+use std::{usize, isize, panic};
 use std::vec::{Drain, IntoIter};
+use std::collections::CollectionAllocErr::*;
 
 struct DropCounter<'a> {
     count: &'a mut u32,
@@ -965,3 +966,209 @@ fn drain_filter_complex() {
         assert_eq!(vec, vec![1, 3, 5, 7, 9, 11, 13, 15, 17, 19]);
     }
 }
+
+#[test]
+fn test_reserve_exact() {
+    // This is all the same as test_reserve
+
+    let mut v = Vec::new();
+    assert_eq!(v.capacity(), 0);
+
+    v.reserve_exact(2);
+    assert!(v.capacity() >= 2);
+
+    for i in 0..16 {
+        v.push(i);
+    }
+
+    assert!(v.capacity() >= 16);
+    v.reserve_exact(16);
+    assert!(v.capacity() >= 32);
+
+    v.push(16);
+
+    v.reserve_exact(16);
+    assert!(v.capacity() >= 33)
+}
+
+#[test]
+fn test_try_reserve() {
+
+    // These are the interesting cases:
+    // * exactly isize::MAX should never trigger a CapacityOverflow (can be OOM)
+    // * > isize::MAX should always fail
+    //    * On 16/32-bit should CapacityOverflow
+    //    * On 64-bit should OOM
+    // * overflow may trigger when adding `len` to `cap` (in number of elements)
+    // * overflow may trigger when multiplying `new_cap` by size_of::<T> (to get bytes)
+
+    const MAX_CAP: usize = isize::MAX as usize;
+    const MAX_USIZE: usize = usize::MAX;
+
+    // On 16/32-bit, we check that allocations don't exceed isize::MAX,
+    // on 64-bit, we assume the OS will give an OOM for such a ridiculous size.
+    // Any platform that succeeds for these requests is technically broken with
+    // ptr::offset because LLVM is the worst.
+    let guards_against_isize = size_of::<usize>() < 8;
+
+    {
+        // Note: basic stuff is checked by test_reserve
+        let mut empty_bytes: Vec<u8> = Vec::new();
+
+        // Check isize::MAX doesn't count as an overflow
+        if let Err(CapacityOverflow) = empty_bytes.try_reserve(MAX_CAP) {
+            panic!("isize::MAX shouldn't trigger an overflow!");
+        }
+        // Play it again, frank! (just to be sure)
+        if let Err(CapacityOverflow) = empty_bytes.try_reserve(MAX_CAP) {
+            panic!("isize::MAX shouldn't trigger an overflow!");
+        }
+
+        if guards_against_isize {
+            // Check isize::MAX + 1 does count as overflow
+            if let Err(CapacityOverflow) = empty_bytes.try_reserve(MAX_CAP + 1) {
+            } else { panic!("isize::MAX + 1 should trigger an overflow!") }
+
+            // Check usize::MAX does count as overflow
+            if let Err(CapacityOverflow) = empty_bytes.try_reserve(MAX_USIZE) {
+            } else { panic!("usize::MAX should trigger an overflow!") }
+        } else {
+            // Check isize::MAX + 1 is an OOM
+            if let Err(AllocErr(_)) = empty_bytes.try_reserve(MAX_CAP + 1) {
+            } else { panic!("isize::MAX + 1 should trigger an OOM!") }
+
+            // Check usize::MAX is an OOM
+            if let Err(AllocErr(_)) = empty_bytes.try_reserve(MAX_USIZE) {
+            } else { panic!("usize::MAX should trigger an OOM!") }
+        }
+    }
+
+
+    {
+        // Same basic idea, but with non-zero len
+        let mut ten_bytes: Vec<u8> = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
+
+        if let Err(CapacityOverflow) = ten_bytes.try_reserve(MAX_CAP - 10) {
+            panic!("isize::MAX shouldn't trigger an overflow!");
+        }
+        if let Err(CapacityOverflow) = ten_bytes.try_reserve(MAX_CAP - 10) {
+            panic!("isize::MAX shouldn't trigger an overflow!");
+        }
+        if guards_against_isize {
+            if let Err(CapacityOverflow) = ten_bytes.try_reserve(MAX_CAP - 9) {
+            } else { panic!("isize::MAX + 1 should trigger an overflow!"); }
+        } else {
+            if let Err(AllocErr(_)) = ten_bytes.try_reserve(MAX_CAP - 9) {
+            } else { panic!("isize::MAX + 1 should trigger an OOM!") }
+        }
+        // Should always overflow in the add-to-len
+        if let Err(CapacityOverflow) = ten_bytes.try_reserve(MAX_USIZE) {
+        } else { panic!("usize::MAX should trigger an overflow!") }
+    }
+
+
+    {
+        // Same basic idea, but with interesting type size
+        let mut ten_u32s: Vec<u32> = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
+
+        if let Err(CapacityOverflow) = ten_u32s.try_reserve(MAX_CAP/4 - 10) {
+            panic!("isize::MAX shouldn't trigger an overflow!");
+        }
+        if let Err(CapacityOverflow) = ten_u32s.try_reserve(MAX_CAP/4 - 10) {
+            panic!("isize::MAX shouldn't trigger an overflow!");
+        }
+        if guards_against_isize {
+            if let Err(CapacityOverflow) = ten_u32s.try_reserve(MAX_CAP/4 - 9) {
+            } else { panic!("isize::MAX + 1 should trigger an overflow!"); }
+        } else {
+            if let Err(AllocErr(_)) = ten_u32s.try_reserve(MAX_CAP/4 - 9) {
+            } else { panic!("isize::MAX + 1 should trigger an OOM!") }
+        }
+        // Should fail in the mul-by-size
+        if let Err(CapacityOverflow) = ten_u32s.try_reserve(MAX_USIZE - 20) {
+        } else {
+            panic!("usize::MAX should trigger an overflow!");
+        }
+    }
+
+}
+
+#[test]
+fn test_try_reserve_exact() {
+
+    // This is exactly the same as test_try_reserve with the method changed.
+    // See that test for comments.
+
+    const MAX_CAP: usize = isize::MAX as usize;
+    const MAX_USIZE: usize = usize::MAX;
+
+    let guards_against_isize = size_of::<usize>() < 8;
+
+    {
+        let mut empty_bytes: Vec<u8> = Vec::new();
+
+        if let Err(CapacityOverflow) = empty_bytes.try_reserve_exact(MAX_CAP) {
+            panic!("isize::MAX shouldn't trigger an overflow!");
+        }
+        if let Err(CapacityOverflow) = empty_bytes.try_reserve_exact(MAX_CAP) {
+            panic!("isize::MAX shouldn't trigger an overflow!");
+        }
+
+        if guards_against_isize {
+            if let Err(CapacityOverflow) = empty_bytes.try_reserve_exact(MAX_CAP + 1) {
+            } else { panic!("isize::MAX + 1 should trigger an overflow!") }
+
+            if let Err(CapacityOverflow) = empty_bytes.try_reserve_exact(MAX_USIZE) {
+            } else { panic!("usize::MAX should trigger an overflow!") }
+        } else {
+            if let Err(AllocErr(_)) = empty_bytes.try_reserve_exact(MAX_CAP + 1) {
+            } else { panic!("isize::MAX + 1 should trigger an OOM!") }
+
+            if let Err(AllocErr(_)) = empty_bytes.try_reserve_exact(MAX_USIZE) {
+            } else { panic!("usize::MAX should trigger an OOM!") }
+        }
+    }
+
+
+    {
+        let mut ten_bytes: Vec<u8> = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
+
+        if let Err(CapacityOverflow) = ten_bytes.try_reserve_exact(MAX_CAP - 10) {
+            panic!("isize::MAX shouldn't trigger an overflow!");
+        }
+        if let Err(CapacityOverflow) = ten_bytes.try_reserve_exact(MAX_CAP - 10) {
+            panic!("isize::MAX shouldn't trigger an overflow!");
+        }
+        if guards_against_isize {
+            if let Err(CapacityOverflow) = ten_bytes.try_reserve_exact(MAX_CAP - 9) {
+            } else { panic!("isize::MAX + 1 should trigger an overflow!"); }
+        } else {
+            if let Err(AllocErr(_)) = ten_bytes.try_reserve_exact(MAX_CAP - 9) {
+            } else { panic!("isize::MAX + 1 should trigger an OOM!") }
+        }
+        if let Err(CapacityOverflow) = ten_bytes.try_reserve_exact(MAX_USIZE) {
+        } else { panic!("usize::MAX should trigger an overflow!") }
+    }
+
+
+    {
+        let mut ten_u32s: Vec<u32> = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10];
+
+        if let Err(CapacityOverflow) = ten_u32s.try_reserve_exact(MAX_CAP/4 - 10) {
+            panic!("isize::MAX shouldn't trigger an overflow!");
+        }
+        if let Err(CapacityOverflow) = ten_u32s.try_reserve_exact(MAX_CAP/4 - 10) {
+            panic!("isize::MAX shouldn't trigger an overflow!");
+        }
+        if guards_against_isize {
+            if let Err(CapacityOverflow) = ten_u32s.try_reserve_exact(MAX_CAP/4 - 9) {
+            } else { panic!("isize::MAX + 1 should trigger an overflow!"); }
+        } else {
+            if let Err(AllocErr(_)) = ten_u32s.try_reserve_exact(MAX_CAP/4 - 9) {
+            } else { panic!("isize::MAX + 1 should trigger an OOM!") }
+        }
+        if let Err(CapacityOverflow) = ten_u32s.try_reserve_exact(MAX_USIZE - 20) {
+        } else { panic!("usize::MAX should trigger an overflow!") }
+    }
+
+}