Add insertion methods with fallible allocation to cranelift_bforest::{Map,Set} (#12685)

Author: fitzgen

Cargo.lock

@@ -17,3 +17,4 @@ workspace = true
 
 [dependencies]
 cranelift-entity = { workspace = true }
+wasmtime-core = { workspace = true }

cranelift/bforest/Cargo.toml

@@ -7,6 +7,7 @@ use alloc::string::String;
 #[cfg(test)]
 use core::fmt;
 use core::marker::PhantomData;
+use wasmtime_core::{alloc::PanicOnOom as _, error::OutOfMemory};
 
 /// Tag type defining forest types for a map.
 struct MapTypes<K, V>(PhantomData<(K, V)>);
@@ -135,6 +136,17 @@ where
         self.cursor(forest, comp).insert(key, value)
     }
 
+    /// Like `insert` but returns an error on allocation failure.
+    pub fn try_insert<C: Comparator<K>>(
+        &mut self,
+        key: K,
+        value: V,
+        forest: &mut MapForest<K, V>,
+        comp: &C,
+    ) -> Result<Option<V>, OutOfMemory> {
+        self.cursor(forest, comp).try_insert(key, value)
+    }
+
     /// Remove `key` from the map and return the removed value for `key`, if any.
     pub fn remove<C: Comparator<K>>(
         &mut self,
@@ -345,22 +357,27 @@ where
     ///
     /// If `key` is already present, replace the existing with `value` and return the old value.
     pub fn insert(&mut self, key: K, value: V) -> Option<V> {
+        self.try_insert(key, value).panic_on_oom()
+    }
+
+    /// Like `insert` but returns an error on allocation failure.
+    pub fn try_insert(&mut self, key: K, value: V) -> Result<Option<V>, OutOfMemory> {
         match self.root.expand() {
             None => {
-                let root = self.pool.alloc_node(NodeData::leaf(key, value));
+                let root = self.pool.alloc_node(NodeData::leaf(key, value))?;
                 *self.root = root.into();
                 self.path.set_root_node(root);
-                None
+                Ok(None)
             }
             Some(root) => {
                 // TODO: Optimize the case where `self.path` is already at the correct insert pos.
                 let old = self.path.find(key, root, self.pool, self.comp);
                 if old.is_some() {
                     *self.path.value_mut(self.pool) = value;
                 } else {
-                    *self.root = self.path.insert(key, value, self.pool).into();
+                    *self.root = self.path.insert(key, value, self.pool)?.into();
                 }
-                old
+                Ok(old)
             }
         }
     }

cranelift/bforest/src/map.rs

@@ -4,6 +4,7 @@ use super::node::Removed;
 use super::{Comparator, Forest, MAX_PATH, Node, NodeData, NodePool, slice_insert, slice_shift};
 use core::borrow::Borrow;
 use core::marker::PhantomData;
+use wasmtime_core::error::OutOfMemory;
 
 #[cfg(test)]
 use core::fmt;
@@ -260,11 +261,16 @@ impl<F: Forest> Path<F> {
     /// The current position must be the correct insertion location for the key.
     /// This function does not check for duplicate keys. Use `find` or similar for that.
     /// Returns the new root node.
-    pub fn insert(&mut self, key: F::Key, value: F::Value, pool: &mut NodePool<F>) -> Node {
+    pub fn insert(
+        &mut self,
+        key: F::Key,
+        value: F::Value,
+        pool: &mut NodePool<F>,
+    ) -> Result<Node, OutOfMemory> {
         if !self.try_leaf_insert(key, value, pool) {
-            self.split_and_insert(key, value, pool);
+            self.split_and_insert(key, value, pool)?;
         }
-        self.node[0]
+        Ok(self.node[0])
     }
 
     /// Try to insert `key, value` at the current position, but fail and return false if the leaf
@@ -282,7 +288,12 @@ impl<F: Forest> Path<F> {
 
     /// Split the current leaf node and then insert `key, value`.
     /// This should only be used if `try_leaf_insert()` fails.
-    fn split_and_insert(&mut self, mut key: F::Key, value: F::Value, pool: &mut NodePool<F>) {
+    fn split_and_insert(
+        &mut self,
+        mut key: F::Key,
+        value: F::Value,
+        pool: &mut NodePool<F>,
+    ) -> Result<(), OutOfMemory> {
         let orig_root = self.node[0];
 
         // Loop invariant: We need to split the node at `level` and then retry a failed insertion.
@@ -294,7 +305,7 @@ impl<F: Forest> Path<F> {
             let mut node = self.node[level];
             let mut entry = self.entry[level].into();
             split = pool[node].split(entry);
-            let rhs_node = pool.alloc_node(split.rhs_data);
+            let rhs_node = pool.alloc_node(split.rhs_data)?;
 
             // Should the path be moved to the new RHS node?
             // Prefer the smaller node if we're right in the middle.
@@ -347,18 +358,19 @@ impl<F: Forest> Path<F> {
                     if node == rhs_node {
                         self.entry[level - 1] += 1;
                     }
-                    return;
+                    return Ok(());
                 }
             }
         }
 
         // If we get here we have split the original root node and need to add an extra level.
         let rhs_node = ins_node.expect("empty path");
-        let root = pool.alloc_node(NodeData::inner(orig_root, key, rhs_node));
+        let root = pool.alloc_node(NodeData::inner(orig_root, key, rhs_node))?;
         let entry = if self.node[0] == rhs_node { 1 } else { 0 };
         self.size += 1;
         slice_insert(&mut self.node[0..self.size], 0, root);
         slice_insert(&mut self.entry[0..self.size], 0, entry);
+        Ok(())
     }
 
     /// Remove the key-value pair at the current position and advance the path to the next
@@ -699,6 +711,8 @@ impl<F: Forest> fmt::Display for Path<F> {
 
 #[cfg(test)]
 mod tests {
+    use wasmtime_core::alloc::PanicOnOom;
+
     use super::*;
     use core::cmp::Ordering;
 
@@ -731,7 +745,7 @@ mod tests {
     fn search_single_leaf() {
         // Testing Path::new() for trees with a single leaf node.
         let mut pool = NodePool::<TF>::new();
-        let root = pool.alloc_node(NodeData::leaf(10, 'a'));
+        let root = pool.alloc_node(NodeData::leaf(10, 'a')).panic_on_oom();
         let mut p = Path::default();
         let comp = TC();
 
@@ -784,9 +798,11 @@ mod tests {
     fn search_single_inner() {
         // Testing Path::new() for trees with a single inner node and two leaves.
         let mut pool = NodePool::<TF>::new();
-        let leaf1 = pool.alloc_node(NodeData::leaf(10, 'a'));
-        let leaf2 = pool.alloc_node(NodeData::leaf(20, 'b'));
-        let root = pool.alloc_node(NodeData::inner(leaf1, 20, leaf2));
+        let leaf1 = pool.alloc_node(NodeData::leaf(10, 'a')).panic_on_oom();
+        let leaf2 = pool.alloc_node(NodeData::leaf(20, 'b')).panic_on_oom();
+        let root = pool
+            .alloc_node(NodeData::inner(leaf1, 20, leaf2))
+            .panic_on_oom();
         let mut p = Path::default();
         let comp = TC();
 

cranelift/bforest/src/path.rs

@@ -7,6 +7,7 @@ use crate::entity::PrimaryMap;
 #[cfg(test)]
 use core::fmt;
 use core::ops::{Index, IndexMut};
+use wasmtime_core::error::OutOfMemory;
 
 /// A pool of nodes, including a free list.
 pub(super) struct NodePool<F: Forest> {
@@ -30,7 +31,7 @@ impl<F: Forest> NodePool<F> {
     }
 
     /// Allocate a new node containing `data`.
-    pub fn alloc_node(&mut self, data: NodeData<F>) -> Node {
+    pub fn alloc_node(&mut self, data: NodeData<F>) -> Result<Node, OutOfMemory> {
         debug_assert!(!data.is_free(), "can't allocate free node");
         match self.freelist {
             Some(node) => {
@@ -40,11 +41,12 @@ impl<F: Forest> NodePool<F> {
                     _ => panic!("Invalid {} on free list", node),
                 }
                 self.nodes[node] = data;
-                node
+                Ok(node)
             }
             None => {
                 // The free list is empty. Allocate a new node.
-                self.nodes.push(data)
+                self.nodes.try_reserve(1)?;
+                Ok(self.nodes.push(data))
             }
         }
     }

cranelift/bforest/src/pool.rs

@@ -7,6 +7,7 @@ use alloc::string::String;
 #[cfg(test)]
 use core::fmt;
 use core::marker::PhantomData;
+use wasmtime_core::{alloc::PanicOnOom, error::OutOfMemory};
 
 /// Tag type defining forest types for a set.
 struct SetTypes<K>(PhantomData<K>);
@@ -112,6 +113,16 @@ where
         self.cursor(forest, comp).insert(key)
     }
 
+    /// Like `insert` but returns an error on allocation failure.
+    pub fn try_insert<C: Comparator<K>>(
+        &mut self,
+        key: K,
+        forest: &mut SetForest<K>,
+        comp: &C,
+    ) -> Result<bool, OutOfMemory> {
+        self.cursor(forest, comp).try_insert(key)
+    }
+
     /// Remove `key` from the set and return true.
     ///
     /// If `key` was not present in the set, return false.
@@ -277,20 +288,25 @@ where
     /// If `elem` is already present, don't change the set, place the cursor at `goto(elem)`, and
     /// return false.
     pub fn insert(&mut self, elem: K) -> bool {
+        self.try_insert(elem).panic_on_oom()
+    }
+
+    /// Like `insert` but returns an error on allocation failure.
+    pub fn try_insert(&mut self, elem: K) -> Result<bool, OutOfMemory> {
         match self.root.expand() {
             None => {
-                let root = self.pool.alloc_node(NodeData::leaf(elem, SetValue()));
+                let root = self.pool.alloc_node(NodeData::leaf(elem, SetValue()))?;
                 *self.root = root.into();
                 self.path.set_root_node(root);
-                true
+                Ok(true)
             }
             Some(root) => {
                 // TODO: Optimize the case where `self.path` is already at the correct insert pos.
                 if self.path.find(elem, root, self.pool, self.comp).is_none() {
-                    *self.root = self.path.insert(elem, SetValue(), self.pool).into();
-                    true
+                    *self.root = self.path.insert(elem, SetValue(), self.pool)?.into();
+                    Ok(true)
                 } else {
-                    false
+                    Ok(false)
                 }
             }
         }

cranelift/bforest/src/set.rs

@@ -71,6 +71,7 @@ v8 = "139.0.0"
 wat = { workspace = true }
 rand = { version = "0.8.0", features = ["small_rng"] }
 wasmtime-environ = { workspace = true }
+cranelift-bforest = { workspace = true }
 cranelift-bitset = { workspace = true }
 
 # Only enable the `build-libinterpret` feature when fuzzing is enabled, enabling

crates/fuzzing/Cargo.toml

@@ -1036,3 +1036,35 @@ fn index_map_shift_insert() -> Result<()> {
             Ok(())
         })
 }
+
+#[test]
+fn bforest_map() -> Result<()> {
+    use cranelift_bforest::*;
+    OomTest::new().test(|| {
+        let mut forest = MapForest::new();
+        let mut map = Map::new();
+        for i in 0..100 {
+            map.try_insert(Key(i), i, &mut forest, &())?;
+        }
+        for i in 0..100 {
+            assert_eq!(map.get(Key(i), &forest, &()), Some(i));
+        }
+        Ok(())
+    })
+}
+
+#[test]
+fn bforest_set() -> Result<()> {
+    use cranelift_bforest::*;
+    OomTest::new().test(|| {
+        let mut forest = SetForest::new();
+        let mut set = Set::new();
+        for i in 0..100 {
+            set.try_insert(Key(i), &mut forest, &())?;
+        }
+        for i in 0..100 {
+            assert!(set.contains(Key(i), &forest, &()));
+        }
+        Ok(())
+    })
+}