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

Author: snf

Diff for: src/liballoc/allocator.rs

@@ -373,6 +373,24 @@ impl fmt::Display for CannotReallocInPlace {
     }
 }
 
+/// Augments `AllocErr` with a CapacityOverflow variant.
+#[derive(Clone, PartialEq, Eq, Debug)]
+#[unstable(feature = "try_reserve", reason = "new API", issue="48043")]
+pub enum CollectionAllocErr {
+    /// Error due to the computed capacity exceeding the collection's maximum
+    /// (usually `isize::MAX` bytes).
+    CapacityOverflow,
+    /// Error due to the allocator (see the `AllocErr` type's docs).
+    AllocErr(AllocErr),
+}
+
+#[unstable(feature = "try_reserve", reason = "new API", issue="48043")]
+impl From<AllocErr> for CollectionAllocErr {
+    fn from(err: AllocErr) -> Self {
+        CollectionAllocErr::AllocErr(err)
+    }
+}
+
 /// An implementation of `Alloc` can allocate, reallocate, and
 /// deallocate arbitrary blocks of data described via `Layout`.
 ///

Diff for: src/liballoc/lib.rs

@@ -117,6 +117,7 @@
 #![feature(staged_api)]
 #![feature(str_internals)]
 #![feature(trusted_len)]
+#![feature(try_reserve)]
 #![feature(unboxed_closures)]
 #![feature(unicode)]
 #![feature(unsize)]

Diff for: src/liballoc/raw_vec.rs

@@ -15,6 +15,8 @@ use core::ptr::{self, Unique};
 use core::slice;
 use heap::{Alloc, Layout, Heap};
 use super::boxed::Box;
+use super::allocator::CollectionAllocErr;
+use super::allocator::CollectionAllocErr::*;
 
 /// A low-level utility for more ergonomically allocating, reallocating, and deallocating
 /// a buffer of memory on the heap without having to worry about all the corner cases
@@ -84,7 +86,7 @@ impl<T, A: Alloc> RawVec<T, A> {
             let elem_size = mem::size_of::<T>();
 
             let alloc_size = cap.checked_mul(elem_size).expect("capacity overflow");
-            alloc_guard(alloc_size);
+            alloc_guard(alloc_size).expect("capacity overflow");
 
             // handles ZSTs and `cap = 0` alike
             let ptr = if alloc_size == 0 {
@@ -308,7 +310,7 @@ impl<T, A: Alloc> RawVec<T, A> {
                     let new_cap = 2 * self.cap;
                     let new_size = new_cap * elem_size;
                     let new_layout = Layout::from_size_align_unchecked(new_size, cur.align());
-                    alloc_guard(new_size);
+                    alloc_guard(new_size).expect("capacity overflow");
                     let ptr_res = self.a.realloc(self.ptr.as_ptr() as *mut u8,
                                                  cur,
                                                  new_layout);
@@ -367,7 +369,7 @@ impl<T, A: Alloc> RawVec<T, A> {
             // overflow and the alignment is sufficiently small.
             let new_cap = 2 * self.cap;
             let new_size = new_cap * elem_size;
-            alloc_guard(new_size);
+            alloc_guard(new_size).expect("capacity overflow");
             let ptr = self.ptr() as *mut _;
             let new_layout = Layout::from_size_align_unchecked(new_size, old_layout.align());
             match self.a.grow_in_place(ptr, old_layout, new_layout) {
@@ -403,7 +405,9 @@ impl<T, A: Alloc> RawVec<T, A> {
     /// # Aborts
     ///
     /// Aborts on OOM
-    pub fn reserve_exact(&mut self, used_cap: usize, needed_extra_cap: usize) {
+    pub fn try_reserve_exact(&mut self, used_cap: usize, needed_extra_cap: usize)
+           -> Result<(), CollectionAllocErr> {
+
         unsafe {
             // NOTE: we don't early branch on ZSTs here because we want this
             // to actually catch "asking for more than usize::MAX" in that case.
@@ -413,43 +417,50 @@ impl<T, A: Alloc> RawVec<T, A> {
             // Don't actually need any more capacity.
             // Wrapping in case they gave a bad `used_cap`.
             if self.cap().wrapping_sub(used_cap) >= needed_extra_cap {
-                return;
+                return Ok(());
             }
 
             // Nothing we can really do about these checks :(
-            let new_cap = used_cap.checked_add(needed_extra_cap).expect("capacity overflow");
-            let new_layout = match Layout::array::<T>(new_cap) {
-                Some(layout) => layout,
-                None => panic!("capacity overflow"),
-            };
-            alloc_guard(new_layout.size());
+            let new_cap = used_cap.checked_add(needed_extra_cap).ok_or(CapacityOverflow)?;
+            let new_layout = Layout::array::<T>(new_cap).ok_or(CapacityOverflow)?;
+
+            alloc_guard(new_layout.size())?;
+
             let res = match self.current_layout() {
                 Some(layout) => {
                     let old_ptr = self.ptr.as_ptr() as *mut u8;
                     self.a.realloc(old_ptr, layout, new_layout)
                 }
                 None => self.a.alloc(new_layout),
             };
-            let uniq = match res {
-                Ok(ptr) => Unique::new_unchecked(ptr as *mut T),
-                Err(e) => self.a.oom(e),
-            };
-            self.ptr = uniq;
+
+            self.ptr = Unique::new_unchecked(res? as *mut T);
             self.cap = new_cap;
+
+            Ok(())
         }
     }
 
+    pub fn reserve_exact(&mut self, used_cap: usize, needed_extra_cap: usize) {
+        match self.try_reserve_exact(used_cap, needed_extra_cap) {
+            Err(CapacityOverflow) => panic!("capacity overflow"),
+            Err(AllocErr(e)) => self.a.oom(e),
+            Ok(()) => { /* yay */ }
+         }
+     }
+
     /// Calculates the buffer's new size given that it'll hold `used_cap +
     /// needed_extra_cap` elements. This logic is used in amortized reserve methods.
     /// Returns `(new_capacity, new_alloc_size)`.
-    fn amortized_new_size(&self, used_cap: usize, needed_extra_cap: usize) -> usize {
+    fn amortized_new_size(&self, used_cap: usize, needed_extra_cap: usize)
+        -> Result<usize, CollectionAllocErr> {
+
         // Nothing we can really do about these checks :(
-        let required_cap = used_cap.checked_add(needed_extra_cap)
-            .expect("capacity overflow");
+        let required_cap = used_cap.checked_add(needed_extra_cap).ok_or(CapacityOverflow)?;
         // Cannot overflow, because `cap <= isize::MAX`, and type of `cap` is `usize`.
         let double_cap = self.cap * 2;
         // `double_cap` guarantees exponential growth.
-        cmp::max(double_cap, required_cap)
+        Ok(cmp::max(double_cap, required_cap))
     }
 
     /// Ensures that the buffer contains at least enough space to hold
@@ -504,8 +515,9 @@ impl<T, A: Alloc> RawVec<T, A> {
     /// #   vector.push_all(&[1, 3, 5, 7, 9]);
     /// # }
     /// ```
-    pub fn reserve(&mut self, used_cap: usize, needed_extra_cap: usize) {
-        unsafe {
+    pub fn try_reserve(&mut self, used_cap: usize, needed_extra_cap: usize)
+        -> Result<(), CollectionAllocErr> {
+         unsafe {
             // NOTE: we don't early branch on ZSTs here because we want this
             // to actually catch "asking for more than usize::MAX" in that case.
             // If we make it past the first branch then we are guaranteed to
@@ -514,33 +526,38 @@ impl<T, A: Alloc> RawVec<T, A> {
             // Don't actually need any more capacity.
             // Wrapping in case they give a bad `used_cap`
             if self.cap().wrapping_sub(used_cap) >= needed_extra_cap {
-                return;
+               return Ok(());
             }
 
-            let new_cap = self.amortized_new_size(used_cap, needed_extra_cap);
+            let new_cap = self.amortized_new_size(used_cap, needed_extra_cap)?;
+            let new_layout = Layout::array::<T>(new_cap).ok_or(CapacityOverflow)?;
+
+             // FIXME: may crash and burn on over-reserve
+            alloc_guard(new_layout.size())?;
 
-            let new_layout = match Layout::array::<T>(new_cap) {
-                Some(layout) => layout,
-                None => panic!("capacity overflow"),
-            };
-            // FIXME: may crash and burn on over-reserve
-            alloc_guard(new_layout.size());
             let res = match self.current_layout() {
                 Some(layout) => {
                     let old_ptr = self.ptr.as_ptr() as *mut u8;
                     self.a.realloc(old_ptr, layout, new_layout)
                 }
                 None => self.a.alloc(new_layout),
             };
-            let uniq = match res {
-                Ok(ptr) => Unique::new_unchecked(ptr as *mut T),
-                Err(e) => self.a.oom(e),
-            };
-            self.ptr = uniq;
+
+            self.ptr = Unique::new_unchecked(res? as *mut T);
             self.cap = new_cap;
+
+            Ok(())
         }
     }
 
+    /// The same as try_reserve, but errors are lowered to a call to oom().
+    pub fn reserve(&mut self, used_cap: usize, needed_extra_cap: usize) {
+        match self.try_reserve(used_cap, needed_extra_cap) {
+            Err(CapacityOverflow) => panic!("capacity overflow"),
+            Err(AllocErr(e)) => self.a.oom(e),
+            Ok(()) => { /* yay */ }
+         }
+     }
     /// Attempts to ensure that the buffer contains at least enough space to hold
     /// `used_cap + needed_extra_cap` elements. If it doesn't already have
     /// enough capacity, will reallocate in place enough space plus comfortable slack
@@ -576,7 +593,8 @@ impl<T, A: Alloc> RawVec<T, A> {
                 return false;
             }
 
-            let new_cap = self.amortized_new_size(used_cap, needed_extra_cap);
+            let new_cap = self.amortized_new_size(used_cap, needed_extra_cap)
+                              .expect("capacity overflow");
 
             // Here, `cap < used_cap + needed_extra_cap <= new_cap`
             // (regardless of whether `self.cap - used_cap` wrapped).
@@ -585,7 +603,7 @@ impl<T, A: Alloc> RawVec<T, A> {
             let ptr = self.ptr() as *mut _;
             let new_layout = Layout::new::<T>().repeat(new_cap).unwrap().0;
             // FIXME: may crash and burn on over-reserve
-            alloc_guard(new_layout.size());
+            alloc_guard(new_layout.size()).expect("capacity overflow");
             match self.a.grow_in_place(ptr, old_layout, new_layout) {
                 Ok(_) => {
                     self.cap = new_cap;
@@ -709,14 +727,14 @@ unsafe impl<#[may_dangle] T, A: Alloc> Drop for RawVec<T, A> {
 // all 4GB in user-space. e.g. PAE or x32
 
 #[inline]
-fn alloc_guard(alloc_size: usize) {
-    if mem::size_of::<usize>() < 8 {
-        assert!(alloc_size <= ::core::isize::MAX as usize,
-                "capacity overflow");
+fn alloc_guard(alloc_size: usize) -> Result<(), CollectionAllocErr> {
+    if mem::size_of::<usize>() < 8 && alloc_size > ::core::isize::MAX as usize {
+        Err(CapacityOverflow)
+    } else {
+        Ok(())
     }
 }
 
-
 #[cfg(test)]
 mod tests {
     use super::*;

Diff for: src/liballoc/string.rs

@@ -71,6 +71,7 @@ use Bound::{Excluded, Included, Unbounded};
 use str::{self, from_boxed_utf8_unchecked, FromStr, Utf8Error, Chars};
 use vec::Vec;
 use boxed::Box;
+use super::allocator::CollectionAllocErr;
 
 /// A UTF-8 encoded, growable string.
 ///
@@ -920,6 +921,79 @@ impl String {
         self.vec.reserve_exact(additional)
     }
 
+    /// Tries to reserve capacity for at least `additional` more elements to be inserted
+    /// in the given `String`. The collection may reserve more space to avoid
+    /// frequent reallocations. After calling `reserve`, capacity will be
+    /// greater than or equal to `self.len() + additional`. Does nothing if
+    /// capacity is already sufficient.
+    ///
+    /// # Errors
+    ///
+    /// If the capacity overflows, or the allocator reports a failure, then an error
+    /// is returned.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(try_reserve)]
+    /// use std::collections::CollectionAllocErr;
+    ///
+    /// fn process_data(data: &str) -> Result<String, CollectionAllocErr> {
+    ///     let mut output = String::new();
+    ///
+    ///     // Pre-reserve the memory, exiting if we can't
+    ///     output.try_reserve(data.len())?;
+    ///
+    ///     // Now we know this can't OOM in the middle of our complex work
+    ///     output.push_str(data);
+    ///
+    ///     Ok(output)
+    /// }
+    /// # process_data("rust").expect("why is the test harness OOMing on 4 bytes?");
+    /// ```
+    #[unstable(feature = "try_reserve", reason = "new API", issue="48043")]
+    pub fn try_reserve(&mut self, additional: usize) -> Result<(), CollectionAllocErr> {
+        self.vec.try_reserve(additional)
+    }
+
+    /// Tries to reserves the minimum capacity for exactly `additional` more elements to
+    /// be inserted in the given `String`. After calling `reserve_exact`,
+    /// capacity will be greater than or equal to `self.len() + additional`.
+    /// Does nothing if the capacity is already sufficient.
+    ///
+    /// Note that the allocator may give the collection more space than it
+    /// requests. Therefore capacity can not be relied upon to be precisely
+    /// minimal. Prefer `reserve` if future insertions are expected.
+    ///
+    /// # Errors
+    ///
+    /// If the capacity overflows, or the allocator reports a failure, then an error
+    /// is returned.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(try_reserve)]
+    /// use std::collections::CollectionAllocErr;
+    ///
+    /// fn process_data(data: &str) -> Result<String, CollectionAllocErr> {
+    ///     let mut output = String::new();
+    ///
+    ///     // Pre-reserve the memory, exiting if we can't
+    ///     output.try_reserve(data.len())?;
+    ///
+    ///     // Now we know this can't OOM in the middle of our complex work
+    ///     output.push_str(data);
+    ///
+    ///     Ok(output)
+    /// }
+    /// # process_data("rust").expect("why is the test harness OOMing on 4 bytes?");
+    /// ```
+    #[unstable(feature = "try_reserve", reason = "new API", issue="48043")]
+    pub fn try_reserve_exact(&mut self, additional: usize) -> Result<(), CollectionAllocErr>  {
+        self.vec.try_reserve_exact(additional)
+    }
+
     /// Shrinks the capacity of this `String` to match its length.
     ///
     /// # Examples

Diff for: src/liballoc/tests/lib.rs

@@ -26,6 +26,7 @@
 #![feature(splice)]
 #![feature(str_escape)]
 #![feature(string_retain)]
+#![feature(try_reserve)]
 #![feature(unboxed_closures)]
 #![feature(unicode)]
 #![feature(exact_chunks)]