Define a fallible-allocation-only version of EntitySet in wasmtime-environ

cranelift/entity/Cargo.toml

@@ -19,6 +19,7 @@ workspace = true
 cranelift-bitset = { workspace=true }
 serde = { workspace = true, optional = true }
 serde_derive = { workspace = true, optional = true }
+wasmtime-error = { workspace = true }
 
 [features]
 enable-serde = ["serde", "serde_derive", "cranelift-bitset/enable-serde"]

cranelift/entity/src/set.rs

@@ -5,6 +5,7 @@ use crate::keys::Keys;
 use core::fmt;
 use core::marker::PhantomData;
 use cranelift_bitset::CompoundBitSet;
+use wasmtime_error::OutOfMemory;
 
 /// A set of `K` for densely indexed entity references.
 ///
@@ -67,12 +68,25 @@ where
         }
     }
 
+    /// Like `with_capacity` but returns an error on allocation failure.
+    pub fn try_with_capacity(capacity: usize) -> Result<Self, OutOfMemory> {
+        Ok(Self {
+            bitset: CompoundBitSet::try_with_capacity(capacity)?,
+            unused: PhantomData,
+        })
+    }
+
     /// Ensure that the set has enough capacity to hold `capacity` total
     /// elements.
     pub fn ensure_capacity(&mut self, capacity: usize) {
         self.bitset.ensure_capacity(capacity);
     }
 
+    /// Like `ensure_capacity` but returns an error on allocation failure.
+    pub fn try_ensure_capacity(&mut self, capacity: usize) -> Result<(), OutOfMemory> {
+        self.bitset.try_ensure_capacity(capacity)
+    }
+
     /// Get the element at `k` if it exists.
     pub fn contains(&self, k: K) -> bool {
         let index = k.index();
@@ -156,7 +170,7 @@ where
         self.bitset.remove(index)
     }
 
-    /// Removes and returns the entity from the set if it exists.
+    /// Removes and returns the highest-index entity from the set if it exists.
     pub fn pop(&mut self) -> Option<K> {
         let index = self.bitset.pop()?;
         Some(K::new(index))

crates/environ/src/collections.rs

@@ -1,3 +1,6 @@
 //! Fallible, OOM-handling collections.
 
+mod entity_set;
+
+pub use entity_set::EntitySet;
 pub use wasmtime_core::alloc::{TryNew, try_new};

crates/environ/src/collections/entity_set.rs

@@ -0,0 +1,85 @@
+use cranelift_entity::{EntityRef, Keys, SetIter};
+use wasmtime_error::OutOfMemory;
+
+/// Like `cranelift_entity::EntitySet` but enforces fallible allocation for all
+/// methods that allocate.
+#[derive(Debug, Default)]
+pub struct EntitySet<K>
+where
+    K: EntityRef,
+{
+    inner: cranelift_entity::EntitySet<K>,
+}
+
+impl<K> EntitySet<K>
+where
+    K: EntityRef,
+{
+    /// Create a new empty set.
+    pub fn new() -> Self {
+        EntitySet {
+            inner: Default::default(),
+        }
+    }
+
+    /// Creates a new empty set with the specified capacity.
+    pub fn with_capacity(capacity: usize) -> Result<Self, OutOfMemory> {
+        let mut set = Self::new();
+        set.inner.try_ensure_capacity(capacity)?;
+        Ok(set)
+    }
+
+    /// Ensure that there is enough capacity to hold `capacity` total elements.
+    pub fn ensure_capacity(&mut self, capacity: usize) -> Result<(), OutOfMemory> {
+        self.inner.try_ensure_capacity(capacity)
+    }
+
+    /// Is this set completely empty?
+    pub fn is_empty(&self) -> bool {
+        self.inner.is_empty()
+    }
+
+    /// Get the element at `k` if it exists.
+    pub fn contains(&self, k: K) -> bool {
+        self.inner.contains(k)
+    }
+
+    /// Remove all entries from this set.
+    pub fn clear(&mut self) {
+        self.inner.clear();
+    }
+
+    /// Iterate over all the keys up to the maximum in this set.
+    ///
+    /// This will yield intermediate keys on the way up to the max key, even if
+    /// they are not contained within the set.
+    pub fn keys(&self) -> Keys<K> {
+        self.inner.keys()
+    }
+
+    /// Iterate over the elements of this set.
+    pub fn iter(&self) -> SetIter<'_, K> {
+        self.inner.iter()
+    }
+
+    /// Insert the element at `k`.
+    ///
+    /// Returns `true` if `k` was not present in the set, i.e. this is a
+    /// newly-added element. Returns `false` otherwise.
+    pub fn insert(&mut self, k: K) -> Result<bool, OutOfMemory> {
+        self.inner.try_ensure_capacity(k.index())?;
+        Ok(self.inner.insert(k))
+    }
+
+    /// Remove `k` from this bitset.
+    ///
+    /// Returns whether `k` was previously in this set or not.
+    pub fn remove(&mut self, k: K) -> bool {
+        self.inner.remove(k)
+    }
+
+    /// Removes and returns the highest-index entity from the set if it exists.
+    pub fn pop(&mut self) -> Option<K> {
+        self.inner.pop()
+    }
+}

crates/environ/src/prelude.rs

@@ -19,7 +19,7 @@
 //!
 //! and then `use crate::*` works as usual.
 
-pub use crate::collections::{TryNew, try_new};
+pub use crate::collections::{EntitySet, TryNew, try_new};
 pub use crate::error::{Context, Error, Result, bail, ensure, format_err};
 pub use alloc::borrow::ToOwned;
 pub use alloc::boxed::Box;