diff --git a/cranelift/bitset/src/compound.rs b/cranelift/bitset/src/compound.rs index ed5991821fa9..5ebe10bf1f7a 100644 --- a/cranelift/bitset/src/compound.rs +++ b/cranelift/bitset/src/compound.rs @@ -295,7 +295,7 @@ impl CompoundBitSet { let to_grow = cmp::max(to_grow, 4); let new_len = self.elems.len() + to_grow; - match wasmtime_core::alloc::new_boxed_slice_from_iter( + match wasmtime_core::alloc::new_boxed_slice_from_iter_with_len( new_len, self.elems .iter() @@ -306,8 +306,10 @@ impl CompoundBitSet { self.elems = new_elems; Ok(()) } - Err(wasmtime_core::alloc::BoxedSliceFromIterError::Oom(oom)) => Err(oom), - Err(wasmtime_core::alloc::BoxedSliceFromIterError::TooFewItems) => unreachable!(), + Err(wasmtime_core::alloc::BoxedSliceFromIterWithLenError::Oom(oom)) => Err(oom), + Err(wasmtime_core::alloc::BoxedSliceFromIterWithLenError::TooFewItems) => { + unreachable!() + } } } diff --git a/crates/core/src/alloc.rs b/crates/core/src/alloc.rs index e5260512e1de..695c8bb39836 100644 --- a/crates/core/src/alloc.rs +++ b/crates/core/src/alloc.rs @@ -5,7 +5,11 @@ mod boxed; mod try_new; mod vec; -pub use boxed::{BoxedSliceFromIterError, new_boxed_slice_from_iter}; +pub use boxed::{ + BoxedSliceFromFallibleIterError, BoxedSliceFromIterWithLenError, + new_boxed_slice_from_fallible_iter, new_boxed_slice_from_iter, + new_boxed_slice_from_iter_with_len, +}; pub use try_new::{TryNew, try_new}; pub use vec::Vec; diff --git a/crates/core/src/alloc/boxed.rs b/crates/core/src/alloc/boxed.rs index 840e063fc735..c0d179214ac6 100644 --- a/crates/core/src/alloc/boxed.rs +++ b/crates/core/src/alloc/boxed.rs @@ -1,6 +1,9 @@ -use super::{TryNew, try_alloc}; +use super::{TryNew, Vec, try_alloc}; use crate::error::OutOfMemory; -use core::{alloc::Layout, mem::MaybeUninit}; +use core::{ + alloc::Layout, + mem::{self, MaybeUninit}, +}; use std_alloc::boxed::Box; /// Allocate an `Box>` with uninitialized contents, returning @@ -39,63 +42,153 @@ impl TryNew for Box { } } -fn new_uninit_boxed_slice(len: usize) -> Result]>, OutOfMemory> { - let layout = Layout::array::>(len) - .map_err(|_| OutOfMemory::new(len.saturating_mul(core::mem::size_of::())))?; +use boxed_slice_builder::BoxedSliceBuilder; +mod boxed_slice_builder { + use super::*; - if layout.size() == 0 { - // NB: no actual allocation takes place when boxing zero-sized - // types. - return Ok(Box::new_uninit_slice(len)); + /// Builder for constructing and initalizing a boxed slice. + /// + /// Also acts as an RAII guard to handle dropping the already-initialized + /// elements when we get too few items or an iterator panics during + /// construction. + pub struct BoxedSliceBuilder { + vec: Vec, } - // Safety: we just ensured that the new length is non-zero. - debug_assert_ne!(layout.size(), 0); - let ptr = unsafe { try_alloc(layout)? }; + impl BoxedSliceBuilder { + pub fn new(len: usize) -> Result { + let mut vec = Vec::new(); + vec.reserve_exact(len)?; + Ok(BoxedSliceBuilder { vec }) + } + + pub fn init_len(&self) -> usize { + self.vec.len() + } + + pub fn capacity(&self) -> usize { + self.vec.capacity() + } + + pub fn push(&mut self, value: T) -> Result<(), OutOfMemory> { + self.vec.push(value) + } + + /// Finish this builder and take its boxed slice out. + /// + /// Panics if `self.init_len() != self.capacity()`. Call + /// `self.shrink_to_fit()` if necessary. + pub fn finish(mut self) -> Box<[T]> { + assert_eq!(self.init_len(), self.capacity()); + let vec = mem::take(&mut self.vec); + mem::forget(self); + let (ptr, len, cap) = vec.into_raw_parts(); + debug_assert_eq!(len, cap); + let ptr = core::ptr::slice_from_raw_parts_mut(ptr, len); + unsafe { Box::from_raw(ptr) } + } - let ptr = ptr.cast::>().as_ptr(); - let ptr = core::ptr::slice_from_raw_parts_mut(ptr, len); + /// Shrink this builder's allocation such that `self.init_len() == + /// self.capacity()`. + pub fn shrink_to_fit(&mut self) -> Result<(), OutOfMemory> { + if self.init_len() == self.capacity() { + return Ok(()); + } + + let len = self.init_len(); + let cap = self.capacity(); + let vec = mem::take(&mut self.vec); + + let old_layout = Layout::array::(cap).expect( + "already have an allocation with this layout so should be able to recreate it", + ); + let new_layout = Layout::array::(len) + .expect("if `cap` is fine for an array layout, then `len` must be as well"); + debug_assert_eq!(old_layout.align(), new_layout.align()); + + // Handle zero-sized reallocations, since the global `realloc` function + // does not. + if new_layout.size() == 0 { + debug_assert!(mem::size_of::() == 0 || len == 0); + if len == 0 { + debug_assert_eq!(self.capacity(), 0); + debug_assert_eq!(self.init_len(), 0); + } else { + debug_assert_eq!(mem::size_of::(), 0); + let ptr = core::ptr::dangling_mut::(); + debug_assert!(!ptr.is_null()); + debug_assert!(ptr.is_aligned()); + // Safety: T's dangling pointer is always non-null and aligned. + self.vec = unsafe { Vec::from_raw_parts(ptr, len, len) }; + } + debug_assert_eq!(self.capacity(), self.init_len()); + return Ok(()); + } - // Safety: `ptr` points to a memory block that is valid for - // `[MaybeUninit; len]` and which was allocated by the global memory - // allocator. - let boxed = unsafe { Box::from_raw(ptr) }; - Ok(boxed) + let (ptr, _len, _cap) = vec.into_raw_parts(); + debug_assert_eq!(len, _len); + debug_assert_eq!(cap, _cap); + + // Safety: `ptr` was allocated by the global allocator, its memory block + // is described by `old_layout`, the new size is non-zero, and the new + // size will not overflow `isize::MAX` when rounded up to the layout's + // alignment (this is checked in the construction of `new_layout`). + let new_ptr = unsafe { + std_alloc::alloc::realloc(ptr.cast::(), old_layout, new_layout.size()) + }; + + // Update `self` based on whether the reallocation succeeded or not, + // either inserting the new vec or reconstructing and replacing the + // old one. + if new_ptr.is_null() { + // Safety: The allocation failed so we retain ownership of `ptr`, + // which was a valid vec and we can safely make it a vec again. + self.vec = unsafe { Vec::from_raw_parts(ptr, len, cap) }; + Err(OutOfMemory::new(new_layout.size())) + } else { + let new_ptr = new_ptr.cast::(); + // Safety: The allocation succeeded, `new_ptr` was reallocated by + // the global allocator and points to a valid boxed slice of length + // `len`. + self.vec = unsafe { Vec::from_raw_parts(new_ptr, len, len) }; + debug_assert_eq!(self.capacity(), self.init_len()); + Ok(()) + } + } + } } /// An error returned by [`new_boxed_slice_from_iter`]. #[derive(Debug)] -pub enum BoxedSliceFromIterError { +pub enum BoxedSliceFromIterWithLenError { /// The iterator did not yield enough items to fill the boxed slice. TooFewItems, /// Failed to allocate space for the boxed slice. Oom(OutOfMemory), } -impl From for BoxedSliceFromIterError { +impl From for BoxedSliceFromIterWithLenError { fn from(oom: OutOfMemory) -> Self { Self::Oom(oom) } } -impl core::fmt::Display for BoxedSliceFromIterError { +impl core::fmt::Display for BoxedSliceFromIterWithLenError { fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { match self { - BoxedSliceFromIterError::TooFewItems => { + Self::TooFewItems => { f.write_str("The iterator did not yield enough items to fill the boxed slice") } - BoxedSliceFromIterError::Oom(_) => { - f.write_str("Failed to allocate space for the boxed slice") - } + Self::Oom(_) => f.write_str("Failed to allocate space for the boxed slice"), } } } -impl core::error::Error for BoxedSliceFromIterError { +impl core::error::Error for BoxedSliceFromIterWithLenError { fn cause(&self) -> Option<&dyn core::error::Error> { match self { - BoxedSliceFromIterError::TooFewItems => None, - BoxedSliceFromIterError::Oom(oom) => Some(oom), + Self::TooFewItems => None, + Self::Oom(oom) => Some(oom), } } } @@ -107,68 +200,109 @@ impl core::error::Error for BoxedSliceFromIterError { /// /// The iterator is dropped after `len` elements have been yielded, this /// function does not check that the iterator yields exactly `len` elements. -pub fn new_boxed_slice_from_iter( +pub fn new_boxed_slice_from_iter_with_len( len: usize, iter: impl IntoIterator, -) -> Result, BoxedSliceFromIterError> { - /// RAII guard to handle dropping the initialized elements of the boxed - /// slice in the cases where we get too few items or the iterator panics. - struct DropGuard { - boxed: Box<[MaybeUninit]>, - init_len: usize, +) -> Result, BoxedSliceFromIterWithLenError> { + let mut builder = BoxedSliceBuilder::new(len)?; + assert_eq!(len, builder.capacity()); + + for elem in iter.into_iter().take(len) { + builder.push(elem).expect("reserved capacity"); } - impl Drop for DropGuard { - fn drop(&mut self) { - debug_assert!(self.init_len <= self.boxed.len()); + if builder.init_len() < builder.capacity() { + return Err(BoxedSliceFromIterWithLenError::TooFewItems); + } - if !core::mem::needs_drop::() { - return; - } + debug_assert_eq!(builder.init_len(), builder.capacity()); + Ok(builder.finish()) +} - for elem in self.boxed.iter_mut().take(self.init_len) { - // Safety: the elements in `self.boxed[..self.init_len]` are - // valid and initialized and will not be used again. - unsafe { - core::ptr::drop_in_place(elem.as_mut_ptr()); - } - } +/// An error returned by [`new_boxed_slice_from_fallible_iter`]. +#[derive(Debug)] +pub enum BoxedSliceFromFallibleIterError { + /// The fallible iterator produced an error. + IterError(E), + /// Failed to allocate space for the boxed slice. + Oom(OutOfMemory), +} + +impl From for BoxedSliceFromFallibleIterError { + fn from(oom: OutOfMemory) -> Self { + Self::Oom(oom) + } +} + +impl core::fmt::Display for BoxedSliceFromFallibleIterError { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + match self { + Self::IterError(_) => f.write_str("The fallible iterator produced an error"), + Self::Oom(_) => f.write_str("Failed to allocate space for the boxed slice"), } } +} - let mut guard = DropGuard { - boxed: new_uninit_boxed_slice(len)?, - init_len: 0, - }; - assert_eq!(len, guard.boxed.len()); +impl core::error::Error for BoxedSliceFromFallibleIterError +where + E: core::error::Error, +{ + fn cause(&self) -> Option<&dyn core::error::Error> { + match self { + Self::IterError(e) => Some(e), + Self::Oom(oom) => Some(oom), + } + } +} - for (i, elem) in iter.into_iter().enumerate().take(len) { - debug_assert!(i < len); - debug_assert_eq!(guard.init_len, i); - guard.boxed[i].write(elem); - guard.init_len += 1; +impl BoxedSliceFromFallibleIterError { + /// Flatten this error into its inner OOM. + pub fn flatten(self) -> OutOfMemory { + match self { + Self::IterError(oom) | Self::Oom(oom) => oom, + } } +} - debug_assert!(guard.init_len <= guard.boxed.len()); +/// Create a `Box<[T]>` from the given iterator's `Result` items. +/// +/// Returns an error on allocation failure or if an iterator item is an `Err`. +pub fn new_boxed_slice_from_fallible_iter( + iter: impl IntoIterator>, +) -> Result, BoxedSliceFromFallibleIterError> { + let iter = iter.into_iter(); - if guard.init_len < guard.boxed.len() { - return Err(BoxedSliceFromIterError::TooFewItems); - } + let (min, max) = iter.size_hint(); + let len = max.unwrap_or_else(|| min); - debug_assert_eq!(guard.init_len, guard.boxed.len()); + let mut builder = BoxedSliceBuilder::new(len)?; + assert_eq!(len, builder.capacity()); - // Take the boxed slice out of the guard. - let boxed = { - guard.init_len = 0; - let boxed = core::mem::take(&mut guard.boxed); - core::mem::forget(guard); - boxed - }; + for result in iter { + let elem = result.map_err(BoxedSliceFromFallibleIterError::IterError)?; + builder.push(elem)?; + } - // Safety: we initialized all elements. - let boxed = unsafe { boxed.assume_init() }; + debug_assert!(builder.init_len() <= builder.capacity()); + builder.shrink_to_fit()?; + debug_assert_eq!(builder.init_len(), builder.capacity()); - Ok(boxed) + Ok(builder.finish()) +} + +/// Create a `Box<[T]>` from the given iterator's elements. +/// +/// Returns an error on allocation failure. +pub fn new_boxed_slice_from_iter( + iter: impl IntoIterator, +) -> Result, OutOfMemory> { + let iter = iter + .into_iter() + .map(Result::::Ok); + new_boxed_slice_from_fallible_iter(iter).map_err(|e| match e { + BoxedSliceFromFallibleIterError::Oom(oom) => oom, + BoxedSliceFromFallibleIterError::IterError(_) => unreachable!(), + }) } #[cfg(test)] @@ -199,20 +333,20 @@ mod tests { } #[test] - fn new_boxed_slice_from_iter_smoke_test() { - let slice = new_boxed_slice_from_iter(3, [42, 36, 1337]).unwrap(); + fn new_boxed_slice_from_iter_with_len_smoke_test() { + let slice = new_boxed_slice_from_iter_with_len(3, [42, 36, 1337]).unwrap(); assert_eq!(&*slice, &[42, 36, 1337]); } #[test] - fn new_boxed_slice_from_iter_with_too_few_elems() { + fn new_boxed_slice_from_iter_with_len_with_too_few_elems() { let (a_dropped, a) = SetFlagOnDrop::new(); let (b_dropped, b) = SetFlagOnDrop::new(); let (c_dropped, c) = SetFlagOnDrop::new(); - match new_boxed_slice_from_iter(4, [a, b, c]) { - Err(BoxedSliceFromIterError::TooFewItems) => {} - Ok(_) | Err(BoxedSliceFromIterError::Oom(_)) => unreachable!(), + match new_boxed_slice_from_iter_with_len(4, [a, b, c]) { + Err(BoxedSliceFromIterWithLenError::TooFewItems) => {} + Ok(_) | Err(BoxedSliceFromIterWithLenError::Oom(_)) => unreachable!(), } assert!(a_dropped.get()); @@ -221,12 +355,12 @@ mod tests { } #[test] - fn new_boxed_slice_from_iter_with_too_many_elems() { + fn new_boxed_slice_from_iter_with_len_with_too_many_elems() { let (a_dropped, a) = SetFlagOnDrop::new(); let (b_dropped, b) = SetFlagOnDrop::new(); let (c_dropped, c) = SetFlagOnDrop::new(); - let slice = new_boxed_slice_from_iter(2, [a, b, c]).unwrap(); + let slice = new_boxed_slice_from_iter_with_len(2, [a, b, c]).unwrap(); assert!(!a_dropped.get()); assert!(!b_dropped.get()); @@ -238,4 +372,26 @@ mod tests { assert!(b_dropped.get()); assert!(c_dropped.get()); } + + #[test] + fn new_boxed_slice_from_iter_smoke_test() { + let slice = new_boxed_slice_from_iter([10, 20, 30]).unwrap(); + assert_eq!(&*slice, &[10, 20, 30]); + } + + #[test] + fn new_boxed_slice_from_fallible_iter_smoke_test() { + let slice = + new_boxed_slice_from_fallible_iter::<_, &str>([Ok(10), Ok(20), Ok(30)]).unwrap(); + assert_eq!(&*slice, &[10, 20, 30]); + } + + #[test] + fn new_boxed_slice_from_fallible_iter_error() { + let result = new_boxed_slice_from_fallible_iter::<_, u32>([Ok(10), Ok(20), Err(30)]); + let Err(BoxedSliceFromFallibleIterError::IterError(err)) = result else { + panic!("unexpected result: {result:?}"); + }; + assert_eq!(err, 30); + } } diff --git a/crates/core/src/alloc/vec.rs b/crates/core/src/alloc/vec.rs index fc126892bb1b..82bdc8af6908 100644 --- a/crates/core/src/alloc/vec.rs +++ b/crates/core/src/alloc/vec.rs @@ -1,6 +1,6 @@ use crate::error::OutOfMemory; use core::{ - fmt, + fmt, mem, ops::{Deref, DerefMut, Index, IndexMut}, }; use std_alloc::vec::Vec as StdVec; @@ -47,7 +47,7 @@ impl Vec { OutOfMemory::new( self.len() .saturating_add(additional) - .saturating_mul(core::mem::size_of::()), + .saturating_mul(mem::size_of::()), ) }) } @@ -89,6 +89,7 @@ impl Vec { let ptr = self.as_mut_ptr(); let len = self.len(); let cap = self.capacity(); + mem::forget(self); (ptr, len, cap) }