diff --git a/src/liballoc/allocator.rs b/src/liballoc/allocator.rs index 55e8c0b430f..fdc4efc66b9 100644 --- a/src/liballoc/allocator.rs +++ b/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 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 --git a/src/liballoc/lib.rs b/src/liballoc/lib.rs index 3f306784558..b93e128d508 100644 --- a/src/liballoc/lib.rs +++ b/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 --git a/src/liballoc/raw_vec.rs b/src/liballoc/raw_vec.rs index 621e1906961..229ae54d747 100644 --- a/src/liballoc/raw_vec.rs +++ b/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 RawVec { let elem_size = mem::size_of::(); 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 RawVec { 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 RawVec { // 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 RawVec { /// # 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,16 +417,15 @@ impl RawVec { // 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::(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::(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; @@ -430,26 +433,34 @@ impl RawVec { } 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 { + // 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 RawVec { /// # 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,17 +526,15 @@ impl RawVec { // 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::(new_cap).ok_or(CapacityOverflow)?; + + // FIXME: may crash and burn on over-reserve + alloc_guard(new_layout.size())?; - let new_layout = match Layout::array::(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; @@ -532,15 +542,22 @@ impl RawVec { } 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 RawVec { 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 RawVec { let ptr = self.ptr() as *mut _; let new_layout = Layout::new::().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 { // all 4GB in user-space. e.g. PAE or x32 #[inline] -fn alloc_guard(alloc_size: usize) { - if mem::size_of::() < 8 { - assert!(alloc_size <= ::core::isize::MAX as usize, - "capacity overflow"); +fn alloc_guard(alloc_size: usize) -> Result<(), CollectionAllocErr> { + if mem::size_of::() < 8 && alloc_size > ::core::isize::MAX as usize { + Err(CapacityOverflow) + } else { + Ok(()) } } - #[cfg(test)] mod tests { use super::*; diff --git a/src/liballoc/string.rs b/src/liballoc/string.rs index 370fb6b4e89..dcc81417346 100644 --- a/src/liballoc/string.rs +++ b/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 { + /// 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 { + /// 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 --git a/src/liballoc/tests/lib.rs b/src/liballoc/tests/lib.rs index 168dbb2ce9b..285cba0270c 100644 --- a/src/liballoc/tests/lib.rs +++ b/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)] diff --git a/src/liballoc/tests/string.rs b/src/liballoc/tests/string.rs index ef6f5e10a72..d1e746ea43b 100644 --- a/src/liballoc/tests/string.rs +++ b/src/liballoc/tests/string.rs @@ -9,6 +9,9 @@ // except according to those terms. use std::borrow::Cow; +use std::collections::CollectionAllocErr::*; +use std::mem::size_of; +use std::{usize, isize}; pub trait IntoCow<'a, B: ?Sized> where B: ToOwned { fn into_cow(self) -> Cow<'a, B>; @@ -504,3 +507,163 @@ fn test_into_boxed_str() { let ys = xs.into_boxed_str(); assert_eq!(&*ys, "hello my name is bob"); } + +#[test] +fn test_reserve_exact() { + // This is all the same as test_reserve + + let mut s = String::new(); + assert_eq!(s.capacity(), 0); + + s.reserve_exact(2); + assert!(s.capacity() >= 2); + + for _i in 0..16 { + s.push('0'); + } + + assert!(s.capacity() >= 16); + s.reserve_exact(16); + assert!(s.capacity() >= 32); + + s.push('0'); + + s.reserve_exact(16); + assert!(s.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:: (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::() < 8; + + { + // Note: basic stuff is checked by test_reserve + let mut empty_string: String = String::new(); + + // Check isize::MAX doesn't count as an overflow + if let Err(CapacityOverflow) = empty_string.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_string.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_string.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_string.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_string.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_string.try_reserve(MAX_USIZE) { + } else { panic!("usize::MAX should trigger an OOM!") } + } + } + + + { + // Same basic idea, but with non-zero len + let mut ten_bytes: String = String::from("0123456789"); + + 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!") } + } + +} + +#[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::() < 8; + + { + let mut empty_string: String = String::new(); + + if let Err(CapacityOverflow) = empty_string.try_reserve_exact(MAX_CAP) { + panic!("isize::MAX shouldn't trigger an overflow!"); + } + if let Err(CapacityOverflow) = empty_string.try_reserve_exact(MAX_CAP) { + panic!("isize::MAX shouldn't trigger an overflow!"); + } + + if guards_against_isize { + if let Err(CapacityOverflow) = empty_string.try_reserve_exact(MAX_CAP + 1) { + } else { panic!("isize::MAX + 1 should trigger an overflow!") } + + if let Err(CapacityOverflow) = empty_string.try_reserve_exact(MAX_USIZE) { + } else { panic!("usize::MAX should trigger an overflow!") } + } else { + if let Err(AllocErr(_)) = empty_string.try_reserve_exact(MAX_CAP + 1) { + } else { panic!("isize::MAX + 1 should trigger an OOM!") } + + if let Err(AllocErr(_)) = empty_string.try_reserve_exact(MAX_USIZE) { + } else { panic!("usize::MAX should trigger an OOM!") } + } + } + + + { + let mut ten_bytes: String = String::from("0123456789"); + + 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!") } + } + +} 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:: (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::() < 8; + + { + // Note: basic stuff is checked by test_reserve + let mut empty_bytes: Vec = 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 = 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 = 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::() < 8; + + { + let mut empty_bytes: Vec = 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 = 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 = 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!") } + } + +} diff --git a/src/liballoc/tests/vec_deque.rs b/src/liballoc/tests/vec_deque.rs index f2935c05d4f..fc1a0b624a5 100644 --- a/src/liballoc/tests/vec_deque.rs +++ b/src/liballoc/tests/vec_deque.rs @@ -11,6 +11,9 @@ use std::collections::VecDeque; use std::fmt::Debug; use std::collections::vec_deque::{Drain}; +use std::collections::CollectionAllocErr::*; +use std::mem::size_of; +use std::{usize, isize}; use self::Taggy::*; use self::Taggypar::*; @@ -1022,3 +1025,208 @@ fn test_placement_in() { } assert_eq!(buf, [5,4,3,1,2,6]); } + +#[test] +fn test_reserve_exact_2() { + // This is all the same as test_reserve + + let mut v = VecDeque::new(); + + v.reserve_exact(2); + assert!(v.capacity() >= 2); + + for i in 0..16 { + v.push_back(i); + } + + assert!(v.capacity() >= 16); + v.reserve_exact(16); + assert!(v.capacity() >= 32); + + v.push_back(16); + + v.reserve_exact(16); + assert!(v.capacity() >= 48) +} + +#[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:: (to get bytes) + + const MAX_CAP: usize = (isize::MAX as usize + 1) / 2 - 1; + 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::() < 8; + + { + // Note: basic stuff is checked by test_reserve + let mut empty_bytes: VecDeque = VecDeque::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 is an OOM + // VecDeque starts with capacity 7, always adds 1 to the capacity + // and also rounds the number to next power of 2 so this is the + // furthest we can go without triggering CapacityOverflow + if let Err(AllocErr(_)) = empty_bytes.try_reserve(MAX_CAP) { + } else { panic!("isize::MAX + 1 should trigger an OOM!") } + } + } + + + { + // Same basic idea, but with non-zero len + let mut ten_bytes: VecDeque = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10].into_iter().collect(); + + 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: VecDeque = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10].into_iter().collect(); + + 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 + 1) / 2 - 1; + const MAX_USIZE: usize = usize::MAX; + + let guards_against_isize = size_of::() < 8; + + { + let mut empty_bytes: VecDeque = VecDeque::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 { + // Check isize::MAX is an OOM + // VecDeque starts with capacity 7, always adds 1 to the capacity + // and also rounds the number to next power of 2 so this is the + // furthest we can go without triggering CapacityOverflow + if let Err(AllocErr(_)) = empty_bytes.try_reserve_exact(MAX_CAP) { + } else { panic!("isize::MAX + 1 should trigger an OOM!") } + } + } + + + { + let mut ten_bytes: VecDeque = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10].into_iter().collect(); + + 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: VecDeque = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10].into_iter().collect(); + + 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!") } + } + +} diff --git a/src/liballoc/vec.rs b/src/liballoc/vec.rs index 1bb2bed463b..953f95876be 100644 --- a/src/liballoc/vec.rs +++ b/src/liballoc/vec.rs @@ -86,6 +86,7 @@ use borrow::Cow; use boxed::Box; use raw_vec::RawVec; use super::range::RangeArgument; +use super::allocator::CollectionAllocErr; use Bound::{Excluded, Included, Unbounded}; /// A contiguous growable array type, written `Vec` but pronounced 'vector'. @@ -489,6 +490,83 @@ impl Vec { self.buf.reserve_exact(self.len, additional); } + /// Tries to reserve capacity for at least `additional` more elements to be inserted + /// in the given `Vec`. 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: &[u32]) -> Result, CollectionAllocErr> { + /// let mut output = Vec::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.extend(data.iter().map(|&val| { + /// val * 2 + 5 // very complicated + /// })); + /// + /// Ok(output) + /// } + /// # process_data(&[1, 2, 3]).expect("why is the test harness OOMing on 12 bytes?"); + /// ``` + #[unstable(feature = "try_reserve", reason = "new API", issue="48043")] + pub fn try_reserve(&mut self, additional: usize) -> Result<(), CollectionAllocErr> { + self.buf.try_reserve(self.len, additional) + } + + /// Tries to reserves the minimum capacity for exactly `additional` more elements to + /// be inserted in the given `Vec`. 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: &[u32]) -> Result, CollectionAllocErr> { + /// let mut output = Vec::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.extend(data.iter().map(|&val| { + /// val * 2 + 5 // very complicated + /// })); + /// + /// Ok(output) + /// } + /// # process_data(&[1, 2, 3]).expect("why is the test harness OOMing on 12 bytes?"); + /// ``` + #[unstable(feature = "try_reserve", reason = "new API", issue="48043")] + pub fn try_reserve_exact(&mut self, additional: usize) -> Result<(), CollectionAllocErr> { + self.buf.try_reserve_exact(self.len, additional) + } + /// Shrinks the capacity of the vector as much as possible. /// /// It will drop down as close as possible to the length but the allocator diff --git a/src/liballoc/vec_deque.rs b/src/liballoc/vec_deque.rs index 68add3cbd51..0658777f0a0 100644 --- a/src/liballoc/vec_deque.rs +++ b/src/liballoc/vec_deque.rs @@ -31,6 +31,7 @@ use core::cmp; use raw_vec::RawVec; +use super::allocator::CollectionAllocErr; use super::range::RangeArgument; use Bound::{Excluded, Included, Unbounded}; use super::vec::Vec; @@ -566,6 +567,97 @@ impl VecDeque { } } + /// Tries to reserves the minimum capacity for exactly `additional` more elements to + /// be inserted in the given `VecDeque`. 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; + /// use std::collections::VecDeque; + /// + /// fn process_data(data: &[u32]) -> Result, CollectionAllocErr> { + /// let mut output = VecDeque::new(); + /// + /// // Pre-reserve the memory, exiting if we can't + /// output.try_reserve_exact(data.len())?; + /// + /// // Now we know this can't OOM in the middle of our complex work + /// output.extend(data.iter().map(|&val| { + /// val * 2 + 5 // very complicated + /// })); + /// + /// Ok(output) + /// } + /// # process_data(&[1, 2, 3]).expect("why is the test harness OOMing on 12 bytes?"); + /// ``` + #[unstable(feature = "try_reserve", reason = "new API", issue="48043")] + pub fn try_reserve_exact(&mut self, additional: usize) -> Result<(), CollectionAllocErr> { + self.try_reserve(additional) + } + + /// Tries to reserve capacity for at least `additional` more elements to be inserted + /// in the given `VecDeque`. 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; + /// use std::collections::VecDeque; + /// + /// fn process_data(data: &[u32]) -> Result, CollectionAllocErr> { + /// let mut output = VecDeque::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.extend(data.iter().map(|&val| { + /// val * 2 + 5 // very complicated + /// })); + /// + /// Ok(output) + /// } + /// # process_data(&[1, 2, 3]).expect("why is the test harness OOMing on 12 bytes?"); + /// ``` + #[unstable(feature = "try_reserve", reason = "new API", issue="48043")] + pub fn try_reserve(&mut self, additional: usize) -> Result<(), CollectionAllocErr> { + let old_cap = self.cap(); + let used_cap = self.len() + 1; + let new_cap = used_cap.checked_add(additional) + .and_then(|needed_cap| needed_cap.checked_next_power_of_two()) + .ok_or(CollectionAllocErr::CapacityOverflow)?; + + if new_cap > old_cap { + self.buf.try_reserve_exact(used_cap, new_cap - used_cap)?; + unsafe { + self.handle_cap_increase(old_cap); + } + } + Ok(()) + } + /// Shrinks the capacity of the `VecDeque` as much as possible. /// /// It will drop down as close as possible to the length but the allocator may still inform the diff --git a/src/libstd/collections/hash/map.rs b/src/libstd/collections/hash/map.rs index 6f4528a0e24..b18b38ec302 100644 --- a/src/libstd/collections/hash/map.rs +++ b/src/libstd/collections/hash/map.rs @@ -11,6 +11,8 @@ use self::Entry::*; use self::VacantEntryState::*; +use alloc::heap::{Heap, Alloc}; +use alloc::allocator::CollectionAllocErr; use cell::Cell; use borrow::Borrow; use cmp::max; @@ -42,21 +44,28 @@ impl DefaultResizePolicy { /// provide that capacity, accounting for maximum loading. The raw capacity /// is always zero or a power of two. #[inline] - fn raw_capacity(&self, len: usize) -> usize { + fn try_raw_capacity(&self, len: usize) -> Result { if len == 0 { - 0 + Ok(0) } else { // 1. Account for loading: `raw_capacity >= len * 1.1`. // 2. Ensure it is a power of two. // 3. Ensure it is at least the minimum size. - let mut raw_cap = len * 11 / 10; - assert!(raw_cap >= len, "raw_cap overflow"); - raw_cap = raw_cap.checked_next_power_of_two().expect("raw_capacity overflow"); + let mut raw_cap = len.checked_mul(11) + .map(|l| l / 10) + .and_then(|l| l.checked_next_power_of_two()) + .ok_or(CollectionAllocErr::CapacityOverflow)?; + raw_cap = max(MIN_NONZERO_RAW_CAPACITY, raw_cap); - raw_cap + Ok(raw_cap) } } + #[inline] + fn raw_capacity(&self, len: usize) -> usize { + self.try_raw_capacity(len).expect("raw_capacity overflow") + } + /// The capacity of the given raw capacity. #[inline] fn capacity(&self, raw_cap: usize) -> usize { @@ -775,17 +784,45 @@ impl HashMap /// ``` #[stable(feature = "rust1", since = "1.0.0")] pub fn reserve(&mut self, additional: usize) { + match self.try_reserve(additional) { + Err(CollectionAllocErr::CapacityOverflow) => panic!("capacity overflow"), + Err(CollectionAllocErr::AllocErr(e)) => Heap.oom(e), + Ok(()) => { /* yay */ } + } + } + + /// Tries to reserve capacity for at least `additional` more elements to be inserted + /// in the given `HashMap`. The collection may reserve more space to avoid + /// frequent reallocations. + /// + /// # Errors + /// + /// If the capacity overflows, or the allocator reports a failure, then an error + /// is returned. + /// + /// # Examples + /// + /// ``` + /// #![feature(try_reserve)] + /// use std::collections::HashMap; + /// let mut map: HashMap<&str, isize> = HashMap::new(); + /// map.try_reserve(10).expect("why is the test harness OOMing on 10 bytes?"); + /// ``` + #[unstable(feature = "try_reserve", reason = "new API", issue="48043")] + pub fn try_reserve(&mut self, additional: usize) -> Result<(), CollectionAllocErr> { let remaining = self.capacity() - self.len(); // this can't overflow if remaining < additional { - let min_cap = self.len().checked_add(additional).expect("reserve overflow"); - let raw_cap = self.resize_policy.raw_capacity(min_cap); - self.resize(raw_cap); + let min_cap = self.len().checked_add(additional) + .ok_or(CollectionAllocErr::CapacityOverflow)?; + let raw_cap = self.resize_policy.try_raw_capacity(min_cap)?; + self.try_resize(raw_cap)?; } else if self.table.tag() && remaining <= self.len() { // Probe sequence is too long and table is half full, // resize early to reduce probing length. let new_capacity = self.table.capacity() * 2; - self.resize(new_capacity); + self.try_resize(new_capacity)?; } + Ok(()) } /// Resizes the internal vectors to a new capacity. It's your @@ -795,15 +832,15 @@ impl HashMap /// 2) Ensure `new_raw_cap` is a power of two or zero. #[inline(never)] #[cold] - fn resize(&mut self, new_raw_cap: usize) { + fn try_resize(&mut self, new_raw_cap: usize) -> Result<(), CollectionAllocErr> { assert!(self.table.size() <= new_raw_cap); assert!(new_raw_cap.is_power_of_two() || new_raw_cap == 0); - let mut old_table = replace(&mut self.table, RawTable::new(new_raw_cap)); + let mut old_table = replace(&mut self.table, RawTable::try_new(new_raw_cap)?); let old_size = old_table.size(); if old_table.size() == 0 { - return; + return Ok(()); } let mut bucket = Bucket::head_bucket(&mut old_table); @@ -838,6 +875,7 @@ impl HashMap } assert_eq!(self.table.size(), old_size); + Ok(()) } /// Shrinks the capacity of the map as much as possible. It will drop @@ -2717,6 +2755,9 @@ mod test_map { use cell::RefCell; use rand::{thread_rng, Rng}; use panic; + use realstd::collections::CollectionAllocErr::*; + use realstd::mem::size_of; + use realstd::usize; #[test] fn test_zero_capacities() { @@ -3651,4 +3692,33 @@ mod test_map { let _ = panic::catch_unwind(panic::AssertUnwindSafe(|| { hm.entry(0) <- makepanic(); })); assert_eq!(hm.len(), 0); } + + #[test] + fn test_try_reserve() { + + let mut empty_bytes: HashMap = HashMap::new(); + + const MAX_USIZE: usize = usize::MAX; + + // HashMap and RawTables use complicated size calculations + // hashes_size is sizeof(HashUint) * capacity; + // pairs_size is sizeof((K. V)) * capacity; + // alignment_hashes_size is 8 + // alignment_pairs size is 4 + let size_of_multiplier = (size_of::() + size_of::<(u8, u8)>()).next_power_of_two(); + // The following formula is used to calculate the new capacity + let max_no_ovf = ((MAX_USIZE / 11) * 10) / size_of_multiplier - 1; + + if let Err(CapacityOverflow) = empty_bytes.try_reserve(MAX_USIZE) { + } else { panic!("usize::MAX should trigger an overflow!"); } + + if size_of::() < 8 { + if let Err(CapacityOverflow) = empty_bytes.try_reserve(max_no_ovf) { + } else { panic!("isize::MAX + 1 should trigger a CapacityOverflow!") } + } else { + if let Err(AllocErr(_)) = empty_bytes.try_reserve(max_no_ovf) { + } else { panic!("isize::MAX + 1 should trigger an OOM!") } + } + } + } diff --git a/src/libstd/collections/hash/table.rs b/src/libstd/collections/hash/table.rs index 73bd5747c10..8e78dc546c6 100644 --- a/src/libstd/collections/hash/table.rs +++ b/src/libstd/collections/hash/table.rs @@ -17,6 +17,7 @@ use mem::{align_of, size_of, needs_drop}; use mem; use ops::{Deref, DerefMut}; use ptr::{self, Unique, NonNull}; +use alloc::allocator::CollectionAllocErr; use self::BucketState::*; @@ -741,14 +742,15 @@ fn test_offset_calculation() { impl RawTable { /// Does not initialize the buckets. The caller should ensure they, /// at the very least, set every hash to EMPTY_BUCKET. - unsafe fn new_uninitialized(capacity: usize) -> RawTable { + /// Returns an error if it cannot allocate or capacity overflows. + unsafe fn try_new_uninitialized(capacity: usize) -> Result, CollectionAllocErr> { if capacity == 0 { - return RawTable { + return Ok(RawTable { size: 0, capacity_mask: capacity.wrapping_sub(1), hashes: TaggedHashUintPtr::new(EMPTY as *mut HashUint), marker: marker::PhantomData, - }; + }); } // No need for `checked_mul` before a more restrictive check performed @@ -768,25 +770,38 @@ impl RawTable { align_of::(), pairs_size, align_of::<(K, V)>()); - assert!(!oflo, "capacity overflow"); + if oflo { + return Err(CollectionAllocErr::CapacityOverflow); + } // One check for overflow that covers calculation and rounding of size. - let size_of_bucket = size_of::().checked_add(size_of::<(K, V)>()).unwrap(); - assert!(size >= - capacity.checked_mul(size_of_bucket) - .expect("capacity overflow"), - "capacity overflow"); + let size_of_bucket = size_of::().checked_add(size_of::<(K, V)>()) + .ok_or(CollectionAllocErr::CapacityOverflow)?; + let capacity_mul_size_of_bucket = capacity.checked_mul(size_of_bucket); + if capacity_mul_size_of_bucket.is_none() || size < capacity_mul_size_of_bucket.unwrap() { + return Err(CollectionAllocErr::CapacityOverflow); + } - let buffer = Heap.alloc(Layout::from_size_align(size, alignment).unwrap()) - .unwrap_or_else(|e| Heap.oom(e)); + let buffer = Heap.alloc(Layout::from_size_align(size, alignment) + .ok_or(CollectionAllocErr::CapacityOverflow)?)?; let hashes = buffer as *mut HashUint; - RawTable { + Ok(RawTable { capacity_mask: capacity.wrapping_sub(1), size: 0, hashes: TaggedHashUintPtr::new(hashes), marker: marker::PhantomData, + }) + } + + /// Does not initialize the buckets. The caller should ensure they, + /// at the very least, set every hash to EMPTY_BUCKET. + unsafe fn new_uninitialized(capacity: usize) -> RawTable { + match Self::try_new_uninitialized(capacity) { + Err(CollectionAllocErr::CapacityOverflow) => panic!("capacity overflow"), + Err(CollectionAllocErr::AllocErr(e)) => Heap.oom(e), + Ok(table) => { table } } } @@ -809,13 +824,23 @@ impl RawTable { } } + /// Tries to create a new raw table from a given capacity. If it cannot allocate, + /// it returns with AllocErr. + pub fn try_new(capacity: usize) -> Result, CollectionAllocErr> { + unsafe { + let ret = RawTable::try_new_uninitialized(capacity)?; + ptr::write_bytes(ret.hashes.ptr(), 0, capacity); + Ok(ret) + } + } + /// Creates a new raw table from a given capacity. All buckets are /// initially empty. pub fn new(capacity: usize) -> RawTable { - unsafe { - let ret = RawTable::new_uninitialized(capacity); - ptr::write_bytes(ret.hashes.ptr(), 0, capacity); - ret + match Self::try_new(capacity) { + Err(CollectionAllocErr::CapacityOverflow) => panic!("capacity overflow"), + Err(CollectionAllocErr::AllocErr(e)) => Heap.oom(e), + Ok(table) => { table } } } diff --git a/src/libstd/collections/mod.rs b/src/libstd/collections/mod.rs index e9a150f34a5..be88f4e268a 100644 --- a/src/libstd/collections/mod.rs +++ b/src/libstd/collections/mod.rs @@ -438,6 +438,9 @@ pub use self::hash_set::HashSet; #[stable(feature = "rust1", since = "1.0.0")] pub use alloc::range; +#[unstable(feature = "try_reserve", reason = "new API", issue="48043")] +pub use alloc::allocator::CollectionAllocErr; + mod hash; #[stable(feature = "rust1", since = "1.0.0")] diff --git a/src/libstd/lib.rs b/src/libstd/lib.rs index da15941374d..ccc5373acc7 100644 --- a/src/libstd/lib.rs +++ b/src/libstd/lib.rs @@ -314,6 +314,7 @@ #![feature(thread_local)] #![feature(toowned_clone_into)] #![feature(try_from)] +#![feature(try_reserve)] #![feature(unboxed_closures)] #![feature(unicode)] #![feature(untagged_unions)] diff --git a/src/test/ui/feature-gate-try_reserve.rs b/src/test/ui/feature-gate-try_reserve.rs new file mode 100644 index 00000000000..9322dbd272f --- /dev/null +++ b/src/test/ui/feature-gate-try_reserve.rs @@ -0,0 +1,14 @@ +// Copyright 2016 The Rust Project Developers. See the COPYRIGHT +// file at the top-level directory of this distribution and at +// http://rust-lang.org/COPYRIGHT. +// +// Licensed under the Apache License, Version 2.0 or the MIT license +// , at your +// option. This file may not be copied, modified, or distributed +// except according to those terms. + +fn main() { + let v = Vec::new(); + v.try_reserve(10); //~ ERROR: use of unstable library feature 'try_reserve' +} diff --git a/src/test/ui/feature-gate-try_reserve.stderr b/src/test/ui/feature-gate-try_reserve.stderr new file mode 100644 index 00000000000..b1fef61dd24 --- /dev/null +++ b/src/test/ui/feature-gate-try_reserve.stderr @@ -0,0 +1,11 @@ +error[E0658]: use of unstable library feature 'try_reserve': new API (see issue #48043) + --> $DIR/feature-gate-try_reserve.rs:13:7 + | +LL | v.try_reserve(10); //~ ERROR: use of unstable library feature 'try_reserve' + | ^^^^^^^^^^^ + | + = help: add #![feature(try_reserve)] to the crate attributes to enable + +error: aborting due to previous error + +If you want more information on this error, try using "rustc --explain E0658"