From 4d29038093fa5757f4a6109e2fee03952bd0ce4d Mon Sep 17 00:00:00 2001 From: Alexis Beingessner Date: Tue, 15 Aug 2017 15:43:12 -0400 Subject: [PATCH] add try_reserve and try_reserve_exact to Vec --- src/liballoc/allocator.rs | 19 ++++ src/liballoc/raw_vec.rs | 99 +++++++++------- src/liballoc/tests/lib.rs | 1 + src/liballoc/tests/vec.rs | 209 +++++++++++++++++++++++++++++++++- src/liballoc/vec.rs | 78 +++++++++++++ src/libstd/collections/mod.rs | 3 + src/libstd/lib.rs | 1 + 7 files changed, 370 insertions(+), 40 deletions(-) diff --git a/src/liballoc/allocator.rs b/src/liballoc/allocator.rs index 2b3df15f71618..7abe1a689e499 100644 --- a/src/liballoc/allocator.rs +++ b/src/liballoc/allocator.rs @@ -382,6 +382,25 @@ impl fmt::Display for AllocErr { } } +/// Augments `AllocErr` with a CapacityOverflow variant. +#[derive(Clone, PartialEq, Eq, Debug)] +#[unstable(feature = "try_reserve", reason = "new API", issue="0000")] +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="0000")] +impl From for CollectionAllocErr { + fn from(err: AllocErr) -> Self { + CollectionAllocErr::AllocErr(err) + } +} + + /// The `CannotReallocInPlace` error is used when `grow_in_place` or /// `shrink_in_place` were unable to reuse the given memory block for /// a requested layout. diff --git a/src/liballoc/raw_vec.rs b/src/liballoc/raw_vec.rs index 6090fc3942a56..eda47a656be5c 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,35 @@ 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_exact, but errors are lowered to a call to oom(). + 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,7 +516,8 @@ impl RawVec { /// # vector.push_all(&[1, 3, 5, 7, 9]); /// # } /// ``` - pub fn reserve(&mut self, used_cap: usize, needed_extra_cap: usize) { + 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. @@ -514,17 +527,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)?; - 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()); + alloc_guard(new_layout.size())?; + let res = match self.current_layout() { Some(layout) => { let old_ptr = self.ptr.as_ptr() as *mut u8; @@ -532,12 +543,20 @@ 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 */ } } } @@ -576,7 +595,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 +605,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,10 +729,11 @@ 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(()) } } diff --git a/src/liballoc/tests/lib.rs b/src/liballoc/tests/lib.rs index 27b23d14059f8..31d1b7e16de3a 100644 --- a/src/liballoc/tests/lib.rs +++ b/src/liballoc/tests/lib.rs @@ -26,6 +26,7 @@ #![feature(splice)] #![feature(str_escape)] #![feature(test)] +#![feature(try_reserve)] #![feature(unboxed_closures)] #![feature(unicode)] diff --git a/src/liballoc/tests/vec.rs b/src/liballoc/tests/vec.rs index 17f1229c2060e..c7e9aacf529ff 100644 --- a/src/liballoc/tests/vec.rs +++ b/src/liballoc/tests/vec.rs @@ -11,8 +11,9 @@ use std::ascii::AsciiExt; 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, @@ -79,6 +80,212 @@ fn test_reserve() { assert!(v.capacity() >= 33) } +#[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, 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, 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, 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, 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!") } + } + +} + #[test] fn test_extend() { let mut v = Vec::new(); diff --git a/src/liballoc/vec.rs b/src/liballoc/vec.rs index 5f68e59289d76..b1f1bfd53825e 100644 --- a/src/liballoc/vec.rs +++ b/src/liballoc/vec.rs @@ -85,6 +85,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'. @@ -485,6 +486,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="0000")] + 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="0000")] + 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/libstd/collections/mod.rs b/src/libstd/collections/mod.rs index b8a6a66eaa65d..b7525860d26a0 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="0000")] +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 880caa2ade5d8..91917c059abe9 100644 --- a/src/libstd/lib.rs +++ b/src/libstd/lib.rs @@ -309,6 +309,7 @@ #![feature(thread_local)] #![feature(toowned_clone_into)] #![feature(try_from)] +#![feature(try_reserve)] #![feature(unboxed_closures)] #![feature(unicode)] #![feature(unique)]