Use OOM-handling PrimaryMap in wasmtime_environ::Module (#12606)

Author: fitzgen

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,19 +417,19 @@ impl<'a, 'data> ModuleEnvironment<'a, 'data> {
                         .module
                         .table_initialization
                         .initial_values
-                        .push(init);
+                        .push(init)?;
                 }
             }
 
             Payload::MemorySection(memories) => {
                 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;
         }

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<MemoryIndex, Option<StaticMemoryInitializer>>,
+        map: collections::PrimaryMap<MemoryIndex, Option<StaticMemoryInitializer>>,
     },
 }
 
@@ -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<DefinedTableIndex, TableInitialValue>,
+    pub initial_values: collections::PrimaryMap<DefinedTableIndex, TableInitialValue>,
 
     /// 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<DataIndex, Range<u32>>,
 
     /// Types declared in the wasm module.
-    pub types: PrimaryMap<TypeIndex, EngineOrModuleTypeIndex>,
+    pub types: collections::PrimaryMap<TypeIndex, EngineOrModuleTypeIndex>,
 
     /// 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<FuncIndex, FunctionType>,
+    pub functions: collections::PrimaryMap<FuncIndex, FunctionType>,
 
     /// WebAssembly tables.
-    pub tables: PrimaryMap<TableIndex, Table>,
+    pub tables: collections::PrimaryMap<TableIndex, Table>,
 
     /// WebAssembly linear memory plans.
-    pub memories: PrimaryMap<MemoryIndex, Memory>,
+    pub memories: collections::PrimaryMap<MemoryIndex, Memory>,
 
     /// WebAssembly global variables.
-    pub globals: PrimaryMap<GlobalIndex, Global>,
+    pub globals: collections::PrimaryMap<GlobalIndex, Global>,
 
     /// WebAssembly global initializers for locally-defined globals.
-    pub global_initializers: PrimaryMap<DefinedGlobalIndex, ConstExpr>,
+    pub global_initializers: collections::PrimaryMap<DefinedGlobalIndex, ConstExpr>,
 
     /// WebAssembly exception and control tags.
-    pub tags: PrimaryMap<TagIndex, Tag>,
+    pub tags: collections::PrimaryMap<TagIndex, Tag>,
 }
 
 /// Initialization routines for creating an instance, encompassing imports,
@@ -605,21 +605,25 @@ 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,
     /// used for functions that either don't escape or aren't certain whether
     /// they escape yet.
     pub fn push_function(&mut self, signature: impl Into<EngineOrModuleTypeIndex>) -> 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

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.

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("")?;

crates/wasmtime/src/runtime/trampoline/tag.rs

@@ -26,7 +26,7 @@ pub fn create_tag(store: &mut StoreOpaque, ty: &TagType) -> Result<InstanceId> {
     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))?;