From 798681cff2d59d42cfbf49d8274284155ce36c03 Mon Sep 17 00:00:00 2001 From: Nick Fitzgerald Date: Tue, 17 Feb 2026 11:45:31 -0800 Subject: [PATCH] Use OOM-handling `PrimaryMap` in `wasmtime_environ::Module` --- crates/environ/src/compile/module_environ.rs | 44 +++++++++++-------- crates/environ/src/module.rs | 38 +++++++++------- .../wasmtime/src/runtime/trampoline/memory.rs | 2 +- .../wasmtime/src/runtime/trampoline/table.rs | 2 +- crates/wasmtime/src/runtime/trampoline/tag.rs | 2 +- 5 files changed, 50 insertions(+), 38 deletions(-) diff --git a/crates/environ/src/compile/module_environ.rs b/crates/environ/src/compile/module_environ.rs index f6b15f22f356..61938af88442 100644 --- a/crates/environ/src/compile/module_environ.rs +++ b/crates/environ/src/compile/module_environ.rs @@ -281,7 +281,7 @@ impl<'a, 'data> ModuleEnvironment<'a, 'data> { log::trace!("interning {count} Wasm types"); let capacity = usize::try_from(count).unwrap(); - self.result.module.types.reserve(capacity); + self.result.module.types.reserve(capacity)?; self.types.reserve_wasm_signatures(capacity); // Iterate over each *rec group* -- not type -- defined in the @@ -318,9 +318,9 @@ impl<'a, 'data> ModuleEnvironment<'a, 'data> { let interned = self.types.intern_rec_group(validator_types, rec_group_id)?; let elems = self.types.rec_group_elements(interned); let len = elems.len(); - self.result.module.types.reserve(len); + self.result.module.types.reserve(len)?; for ty in elems { - self.result.module.types.push(ty.into()); + self.result.module.types.push(ty.into())?; } // Advance `type_index` to the start of the next rec group. @@ -381,7 +381,7 @@ impl<'a, 'data> ModuleEnvironment<'a, 'data> { self.validator.function_section(&functions)?; let cnt = usize::try_from(functions.count()).unwrap(); - self.result.module.functions.reserve_exact(cnt); + self.result.module.functions.reserve_exact(cnt)?; for entry in functions { let sigindex = entry?; @@ -394,13 +394,13 @@ impl<'a, 'data> ModuleEnvironment<'a, 'data> { Payload::TableSection(tables) => { self.validator.table_section(&tables)?; let cnt = usize::try_from(tables.count()).unwrap(); - self.result.module.tables.reserve_exact(cnt); + self.result.module.tables.reserve_exact(cnt)?; for entry in tables { let wasmparser::Table { ty, init } = entry?; let table = self.convert_table_type(&ty)?; self.result.module.needs_gc_heap |= table.ref_type.is_vmgcref_type(); - self.result.module.tables.push(table); + self.result.module.tables.push(table)?; let init = match init { wasmparser::TableInit::RefNull => TableInitialValue::Null { precomputed: collections::Vec::new(), @@ -417,7 +417,7 @@ impl<'a, 'data> ModuleEnvironment<'a, 'data> { .module .table_initialization .initial_values - .push(init); + .push(init)?; } } @@ -425,11 +425,11 @@ impl<'a, 'data> ModuleEnvironment<'a, 'data> { self.validator.memory_section(&memories)?; let cnt = usize::try_from(memories.count()).unwrap(); - self.result.module.memories.reserve_exact(cnt); + self.result.module.memories.reserve_exact(cnt)?; for entry in memories { let memory = entry?; - self.result.module.memories.push(memory.into()); + self.result.module.memories.push(memory.into())?; } } @@ -451,7 +451,7 @@ impl<'a, 'data> ModuleEnvironment<'a, 'data> { self.validator.global_section(&globals)?; let cnt = usize::try_from(globals.count()).unwrap(); - self.result.module.globals.reserve_exact(cnt); + self.result.module.globals.reserve_exact(cnt)?; for entry in globals { let wasmparser::Global { ty, init_expr } = entry?; @@ -460,8 +460,8 @@ impl<'a, 'data> ModuleEnvironment<'a, 'data> { self.flag_func_escaped(f); } let ty = self.convert_global_type(&ty)?; - self.result.module.globals.push(ty); - self.result.module.global_initializers.push(initializer); + self.result.module.globals.push(ty)?; + self.result.module.global_initializers.push(initializer)?; } } @@ -835,10 +835,18 @@ and for re-adding support for interface types you can see this issue: self.flag_func_escaped(func_index); func_index }), - EntityType::Table(ty) => EntityIndex::Table(self.result.module.tables.push(ty)), - EntityType::Memory(ty) => EntityIndex::Memory(self.result.module.memories.push(ty)), - EntityType::Global(ty) => EntityIndex::Global(self.result.module.globals.push(ty)), - EntityType::Tag(ty) => EntityIndex::Tag(self.result.module.tags.push(ty)), + EntityType::Table(ty) => { + EntityIndex::Table(self.result.module.tables.push(ty).panic_on_oom()) + } + EntityType::Memory(ty) => { + EntityIndex::Memory(self.result.module.memories.push(ty).panic_on_oom()) + } + EntityType::Global(ty) => { + EntityIndex::Global(self.result.module.globals.push(ty).panic_on_oom()) + } + EntityType::Tag(ty) => { + EntityIndex::Tag(self.result.module.tags.push(ty).panic_on_oom()) + } } } @@ -1099,7 +1107,7 @@ impl ModuleTranslation<'_> { // memory initialization image is built here from the page data and then // it's converted to a single initializer. let data = mem::replace(&mut self.data, Vec::new()); - let mut map = PrimaryMap::with_capacity(info.len()); + let mut map = collections::PrimaryMap::with_capacity(info.len()).panic_on_oom(); let mut module_data_size = 0u32; for (memory, info) in info.iter() { // Create the in-memory `image` which is the initialized contents of @@ -1183,7 +1191,7 @@ impl ModuleTranslation<'_> { } else { None }; - let idx = map.push(init); + let idx = map.push(init).panic_on_oom(); assert_eq!(idx, memory); module_data_size += len; } diff --git a/crates/environ/src/module.rs b/crates/environ/src/module.rs index f0f64710333d..05468109a5c0 100644 --- a/crates/environ/src/module.rs +++ b/crates/environ/src/module.rs @@ -82,7 +82,7 @@ pub enum MemoryInitialization { /// /// The offset, range base, and range end are all guaranteed to be page /// aligned to the page size passed in to `try_static_init`. - map: PrimaryMap>, + map: collections::PrimaryMap>, }, } @@ -226,7 +226,7 @@ pub struct TableInitialization { /// initialization. For example table initializers to a table that are all /// in-bounds will get removed from `segment` and moved into /// `initial_values` here. - pub initial_values: PrimaryMap, + pub initial_values: collections::PrimaryMap, /// Element segments present in the initial wasm module which are executed /// at instantiation time. @@ -325,7 +325,7 @@ pub struct Module { pub passive_data_map: BTreeMap>, /// Types declared in the wasm module. - pub types: PrimaryMap, + pub types: collections::PrimaryMap, /// Number of imported or aliased functions in the module. pub num_imported_funcs: usize, @@ -353,22 +353,22 @@ pub struct Module { pub num_escaped_funcs: usize, /// Types of functions, imported and local. - pub functions: PrimaryMap, + pub functions: collections::PrimaryMap, /// WebAssembly tables. - pub tables: PrimaryMap, + pub tables: collections::PrimaryMap, /// WebAssembly linear memory plans. - pub memories: PrimaryMap, + pub memories: collections::PrimaryMap, /// WebAssembly global variables. - pub globals: PrimaryMap, + pub globals: collections::PrimaryMap, /// WebAssembly global initializers for locally-defined globals. - pub global_initializers: PrimaryMap, + pub global_initializers: collections::PrimaryMap, /// WebAssembly exception and control tags. - pub tags: PrimaryMap, + pub tags: collections::PrimaryMap, } /// Initialization routines for creating an instance, encompassing imports, @@ -605,10 +605,12 @@ impl Module { ) -> TagIndex { let signature = signature.into(); let exception = exception.into(); - self.tags.push(Tag { - signature, - exception, - }) + self.tags + .push(Tag { + signature, + exception, + }) + .panic_on_oom() } /// Appends a new function to this module with the given type information, @@ -616,10 +618,12 @@ impl Module { /// they escape yet. pub fn push_function(&mut self, signature: impl Into) -> FuncIndex { let signature = signature.into(); - self.functions.push(FunctionType { - signature, - func_ref: FuncRefIndex::reserved_value(), - }) + self.functions + .push(FunctionType { + signature, + func_ref: FuncRefIndex::reserved_value(), + }) + .panic_on_oom() } /// Returns an iterator over all of the defined function indices in this diff --git a/crates/wasmtime/src/runtime/trampoline/memory.rs b/crates/wasmtime/src/runtime/trampoline/memory.rs index bc68d43569cc..175e77b137e1 100644 --- a/crates/wasmtime/src/runtime/trampoline/memory.rs +++ b/crates/wasmtime/src/runtime/trampoline/memory.rs @@ -33,7 +33,7 @@ pub async fn create_memory( // Create a memory, though it will never be used for constructing a memory // with an allocator: instead the memories are either preallocated (i.e., // shared memory) or allocated manually below. - let memory_id = module.memories.push(*memory_ty.wasmtime_memory()); + let memory_id = module.memories.push(*memory_ty.wasmtime_memory())?; // Since we have only associated a single memory with the "frankenstein" // instance, it will be exported at index 0. diff --git a/crates/wasmtime/src/runtime/trampoline/table.rs b/crates/wasmtime/src/runtime/trampoline/table.rs index 7842ae5a3cd4..b32cac5bfd99 100644 --- a/crates/wasmtime/src/runtime/trampoline/table.rs +++ b/crates/wasmtime/src/runtime/trampoline/table.rs @@ -21,7 +21,7 @@ pub async fn create_table( wasmtime_table.ref_type ); - let table_id = module.tables.push(wasmtime_table); + let table_id = module.tables.push(wasmtime_table)?; // TODO: can this `exports.insert` get removed? let name = module.strings.insert("")?; diff --git a/crates/wasmtime/src/runtime/trampoline/tag.rs b/crates/wasmtime/src/runtime/trampoline/tag.rs index a42ac757686b..91de74e999db 100644 --- a/crates/wasmtime/src/runtime/trampoline/tag.rs +++ b/crates/wasmtime/src/runtime/trampoline/tag.rs @@ -26,7 +26,7 @@ pub fn create_tag(store: &mut StoreOpaque, ty: &TagType) -> Result { let tag_id = module.tags.push(Tag { signature: EngineOrModuleTypeIndex::Engine(func_ty.index()), exception: EngineOrModuleTypeIndex::Engine(exn_ty.index()), - }); + })?; let name = module.strings.insert("")?; module.exports.insert(name, EntityIndex::Tag(tag_id))?;