diff --git a/Cargo.lock b/Cargo.lock index 5d731343966a..9990d17cea79 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -734,6 +734,7 @@ name = "cranelift-bforest" version = "0.130.0" dependencies = [ "cranelift-entity", + "wasmtime-internal-core", ] [[package]] @@ -4789,6 +4790,7 @@ version = "0.0.0" dependencies = [ "arbitrary", "backtrace", + "cranelift-bforest", "cranelift-bitset", "env_logger 0.11.5", "futures", diff --git a/cranelift/bforest/Cargo.toml b/cranelift/bforest/Cargo.toml index 1f8884fbdd49..584556c80003 100644 --- a/cranelift/bforest/Cargo.toml +++ b/cranelift/bforest/Cargo.toml @@ -17,3 +17,4 @@ workspace = true [dependencies] cranelift-entity = { workspace = true } +wasmtime-core = { workspace = true } diff --git a/cranelift/bforest/src/map.rs b/cranelift/bforest/src/map.rs index 67446eab78cf..82ca968cdcdd 100644 --- a/cranelift/bforest/src/map.rs +++ b/cranelift/bforest/src/map.rs @@ -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(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>( + &mut self, + key: K, + value: V, + forest: &mut MapForest, + comp: &C, + ) -> Result, 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>( &mut self, @@ -345,12 +357,17 @@ 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 { + 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, 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. @@ -358,9 +375,9 @@ where 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) } } } diff --git a/cranelift/bforest/src/path.rs b/cranelift/bforest/src/path.rs index 660999306482..bcf50239d25d 100644 --- a/cranelift/bforest/src/path.rs +++ b/cranelift/bforest/src/path.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 Path { /// 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) -> Node { + pub fn insert( + &mut self, + key: F::Key, + value: F::Value, + pool: &mut NodePool, + ) -> Result { 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 Path { /// 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) { + fn split_and_insert( + &mut self, + mut key: F::Key, + value: F::Value, + pool: &mut NodePool, + ) -> 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 Path { 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 Path { 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 fmt::Display for Path { #[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::::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::::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(); diff --git a/cranelift/bforest/src/pool.rs b/cranelift/bforest/src/pool.rs index 9e3f57fccedd..d4c78fb8c53f 100644 --- a/cranelift/bforest/src/pool.rs +++ b/cranelift/bforest/src/pool.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 { @@ -30,7 +31,7 @@ impl NodePool { } /// Allocate a new node containing `data`. - pub fn alloc_node(&mut self, data: NodeData) -> Node { + pub fn alloc_node(&mut self, data: NodeData) -> Result { debug_assert!(!data.is_free(), "can't allocate free node"); match self.freelist { Some(node) => { @@ -40,11 +41,12 @@ impl NodePool { _ => 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)) } } } diff --git a/cranelift/bforest/src/set.rs b/cranelift/bforest/src/set.rs index 73e791628ee2..fecc149f6157 100644 --- a/cranelift/bforest/src/set.rs +++ b/cranelift/bforest/src/set.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(PhantomData); @@ -112,6 +113,16 @@ where self.cursor(forest, comp).insert(key) } + /// Like `insert` but returns an error on allocation failure. + pub fn try_insert>( + &mut self, + key: K, + forest: &mut SetForest, + comp: &C, + ) -> Result { + 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 { 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) } } } diff --git a/crates/fuzzing/Cargo.toml b/crates/fuzzing/Cargo.toml index 2c042f408ac8..044a629e8c5b 100644 --- a/crates/fuzzing/Cargo.toml +++ b/crates/fuzzing/Cargo.toml @@ -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 diff --git a/crates/fuzzing/tests/oom.rs b/crates/fuzzing/tests/oom.rs index 814424d28e4c..dc3f5b1ce3c2 100644 --- a/crates/fuzzing/tests/oom.rs +++ b/crates/fuzzing/tests/oom.rs @@ -991,3 +991,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(()) + }) +}