Add fallible allocation methods to cranelift_bitset::CompoundBitSet

Cargo.toml

@@ -158,6 +158,7 @@ members = [
   "cranelift/serde",
   "crates/bench-api",
   "crates/c-api/artifact",
+  "crates/core",
   "crates/environ/fuzz",
   "crates/error",
   "crates/misc/component-async-tests",
@@ -263,6 +264,7 @@ wasmtime-wast = { path = "crates/wast", version = "=42.0.0" }
 # that these are internal unsupported crates for external use. These exist as
 # part of the project organization of other public crates in Wasmtime and are
 # otherwise not supported in terms of CVEs for example.
+wasmtime-core = { path = "crates/core", version = "=42.0.0", package = 'wasmtime-internal-core' }
 wasmtime-error = { path = "crates/error", version = "=42.0.0", package = 'wasmtime-internal-error' }
 wasmtime-wmemcheck = { path = "crates/wmemcheck", version = "=42.0.0", package = 'wasmtime-internal-wmemcheck' }
 wasmtime-c-api-macros = { path = "crates/c-api-macros", version = "=42.0.0", package = 'wasmtime-internal-c-api-macros' }

cranelift/bitset/Cargo.toml

@@ -16,6 +16,8 @@ workspace = true
 arbitrary = { workspace = true, optional = true }
 serde = { workspace = true, optional = true }
 serde_derive = { workspace = true, optional = true }
+wasmtime-error = { workspace = true }
+wasmtime-core = { workspace = true }
 
 [features]
 enable-serde = ["dep:serde", "dep:serde_derive"]

cranelift/bitset/src/compound.rs

@@ -3,6 +3,7 @@
 use crate::scalar::{self, ScalarBitSet, ScalarBitSetStorage};
 use alloc::boxed::Box;
 use core::{cmp, iter, mem};
+use wasmtime_error::OutOfMemory;
 
 /// A large bit set backed by dynamically-sized storage.
 ///
@@ -83,6 +84,11 @@ impl<T: ScalarBitSetStorage> CompoundBitSet<T> {
     ///
     /// The actual capacity reserved may be greater than that requested.
     ///
+    /// # Panics
+    ///
+    /// Panics on allocation failure. Use [`CompoundBitSet::try_with_capacity`]
+    /// to handle allocation failure.
+    ///
     /// # Example
     ///
     /// ```
@@ -95,9 +101,16 @@ impl<T: ScalarBitSetStorage> CompoundBitSet<T> {
     /// ```
     #[inline]
     pub fn with_capacity(capacity: usize) -> Self {
+        Self::try_with_capacity(capacity).unwrap()
+    }
+
+    /// Like [`CompoundBitSet::with_capacity`] but returns an error on
+    /// allocation failure.
+    #[inline]
+    pub fn try_with_capacity(capacity: usize) -> Result<Self, OutOfMemory> {
         let mut bitset = Self::default();
-        bitset.ensure_capacity(capacity);
-        bitset
+        bitset.try_ensure_capacity(capacity)?;
+        Ok(bitset)
     }
 
     /// Get the number of elements in this bitset.
@@ -221,6 +234,11 @@ impl<T: ScalarBitSetStorage> CompoundBitSet<T> {
     /// where `i < n` is guaranteed to succeed without growing the bitset's
     /// backing storage.
     ///
+    /// # Panics
+    ///
+    /// Panics on allocation failure. Use
+    /// [`CompoundBitSet::try_ensure_capacity`] to handle allocation failure.
+    ///
     /// # Example
     ///
     /// ```
@@ -243,31 +261,53 @@ impl<T: ScalarBitSetStorage> CompoundBitSet<T> {
     /// ```
     #[inline]
     pub fn ensure_capacity(&mut self, n: usize) {
+        self.try_ensure_capacity(n).unwrap()
+    }
+
+    /// Like [`CompoundBitSet::ensure_capacity`] but returns an error on
+    /// allocation failure.
+    #[inline]
+    pub fn try_ensure_capacity(&mut self, n: usize) -> Result<(), OutOfMemory> {
         // Subtract one from the capacity to get the maximum bit that we might
         // set. If `n` is 0 then nothing need be done as no capacity needs to be
         // allocated.
         let (word, _bit) = Self::word_and_bit(match n.checked_sub(1) {
-            None => return,
+            None => return Ok(()),
             Some(n) => n,
         });
-        if word >= self.elems.len() {
-            assert!(word < usize::try_from(isize::MAX).unwrap());
 
-            let delta = word - self.elems.len();
-            let to_grow = delta + 1;
+        if word < self.elems.len() {
+            // Already have capacity.
+            return Ok(());
+        }
+
+        // Need to allocate additional capacity.
+
+        assert!(word < usize::try_from(isize::MAX).unwrap());
 
-            // Amortize the cost of growing.
-            let to_grow = cmp::max(to_grow, self.elems.len() * 2);
-            // Don't make ridiculously small allocations.
-            let to_grow = cmp::max(to_grow, 4);
+        let delta = word - self.elems.len();
+        let to_grow = delta + 1;
 
-            let new_elems = self
-                .elems
+        // Amortize the cost of growing by at least growing another
+        // `self.elems.len()`, so the new length is double the old length.
+        let to_grow = cmp::max(to_grow, self.elems.len());
+        // Don't make ridiculously small allocations.
+        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(
+            new_len,
+            self.elems
                 .iter()
                 .copied()
-                .chain(iter::repeat(ScalarBitSet::new()).take(to_grow))
-                .collect::<Box<[_]>>();
-            self.elems = new_elems;
+                .chain(iter::repeat(ScalarBitSet::new())),
+        ) {
+            Ok(new_elems) => {
+                self.elems = new_elems;
+                Ok(())
+            }
+            Err(wasmtime_core::alloc::BoxedSliceFromIterError::Oom(oom)) => Err(oom),
+            Err(wasmtime_core::alloc::BoxedSliceFromIterError::TooFewItems) => unreachable!(),
         }
     }
 

crates/core/Cargo.toml

@@ -0,0 +1,18 @@
+[package]
+authors.workspace = true
+description = "INTERNAL: Wasmtime's core utilities and helpers with minimal dependencies"
+edition.workspace = true
+license = "Apache-2.0 WITH LLVM-exception"
+name = "wasmtime-internal-core"
+rust-version.workspace = true
+version.workspace = true
+
+[dependencies]
+# This crate should have *very* minimal dependencies! Don't add new ones
+# lightly.
+wasmtime-error = { workspace = true }
+
+[lints]
+workspace = true
+
+[features]

crates/core/src/alloc.rs

@@ -0,0 +1,30 @@
+//! Low-level allocation and OOM-handling utilities.
+
+mod arc;
+mod boxed;
+mod try_new;
+
+pub use boxed::{BoxedSliceFromIterError, new_boxed_slice_from_iter};
+pub use try_new::{TryNew, try_new};
+
+use core::{alloc::Layout, ptr::NonNull};
+use wasmtime_error::OutOfMemory;
+
+/// Try to allocate a block of memory that fits the given layout, or return an
+/// `OutOfMemory` error.
+///
+/// # Safety
+///
+/// Same as `alloc::alloc::alloc`: layout must have non-zero size.
+#[inline]
+unsafe fn try_alloc(layout: Layout) -> Result<NonNull<u8>, OutOfMemory> {
+    // Safety: same as our safety conditions.
+    debug_assert!(layout.size() > 0);
+    let ptr = unsafe { std_alloc::alloc::alloc(layout) };
+
+    if let Some(ptr) = NonNull::new(ptr) {
+        Ok(ptr)
+    } else {
+        Err(OutOfMemory::new(layout.size()))
+    }
+}

crates/core/src/alloc/arc.rs

@@ -0,0 +1,28 @@
+use super::TryNew;
+use std_alloc::sync::Arc;
+use wasmtime_error::OutOfMemory;
+
+/// XXX: Stable Rust doesn't actually give us any method to build fallible
+/// allocation for `Arc<T>`, so this is only actually fallible when using
+/// nightly Rust and setting `RUSTFLAGS="--cfg arc_try_new"`.
+impl<T> TryNew for Arc<T> {
+    type Value = T;
+
+    #[inline]
+    fn try_new(value: T) -> Result<Self, OutOfMemory>
+    where
+        Self: Sized,
+    {
+        #[cfg(arc_try_new)]
+        return Arc::try_new(value).map_err(|_| {
+            // We don't have access to the exact size of the inner `Arc`
+            // allocation, but (at least at one point) it was made up of a
+            // strong ref count, a weak ref count, and the inner value.
+            let bytes = core::mem::size_of::<(usize, usize, T)>();
+            OutOfMemory::new(bytes)
+        });
+
+        #[cfg(not(arc_try_new))]
+        return Ok(Arc::new(value));
+    }
+}