Use our OOM-handling Vec in wasmtime_environ::Module (#12584)

fitzgen · web-flow · commit c782855a9ab1 · 2026-02-13T21:32:53.000Z
diff --git a/crates/environ/src/compile/module_environ.rs b/crates/environ/src/compile/module_environ.rs
@@ -10,7 +10,7 @@ use crate::{
     ModuleInternedTypeIndex, ModuleTypesBuilder, PanicOnOom as _, PrimaryMap, SizeOverflow,
     StaticMemoryInitializer, StaticModuleIndex, TableIndex, TableInitialValue, Tag, TagIndex,
     Tunables, TypeConvert, TypeIndex, WasmError, WasmHeapTopType, WasmHeapType, WasmResult,
-    WasmValType, WasmparserTypeConverter,
+    WasmValType, WasmparserTypeConverter, collections,
 };
 use cranelift_entity::SecondaryMap;
 use cranelift_entity::packed_option::ReservedValue;
@@ -332,7 +332,7 @@ impl<'a, 'data> ModuleEnvironment<'a, 'data> {
                 self.validator.import_section(&imports)?;
 
                 let cnt = usize::try_from(imports.count()).unwrap();
-                self.result.module.initializers.reserve(cnt);
+                self.result.module.initializers.reserve(cnt)?;
 
                 for entry in imports.into_imports() {
                     let import = entry?;
@@ -403,7 +403,7 @@ impl<'a, 'data> ModuleEnvironment<'a, 'data> {
                     self.result.module.tables.push(table);
                     let init = match init {
                         wasmparser::TableInit::RefNull => TableInitialValue::Null {
-                            precomputed: Vec::new(),
+                            precomputed: collections::Vec::new(),
                         },
                         wasmparser::TableInit::Expr(expr) => {
                             let (init, escaped) = ConstExpr::from_wasmparser(self, expr)?;
@@ -547,21 +547,19 @@ impl<'a, 'data> ModuleEnvironment<'a, 'data> {
                             let (offset, escaped) = ConstExpr::from_wasmparser(self, offset_expr)?;
                             debug_assert!(escaped.is_empty());
 
-                            self.result
-                                .module
-                                .table_initialization
-                                .segments
-                                .push(TableSegment {
+                            self.result.module.table_initialization.segments.push(
+                                TableSegment {
                                     table_index,
                                     offset,
                                     elements,
-                                });
+                                },
+                            )?;
                         }
 
                         ElementKind::Passive => {
                             let elem_index = ElemIndex::from_u32(index as u32);
                             let index = self.result.module.passive_elements.len();
-                            self.result.module.passive_elements.push(elements);
+                            self.result.module.passive_elements.push(elements)?;
                             self.result
                                 .module
                                 .passive_elements_map
@@ -627,7 +625,7 @@ impl<'a, 'data> ModuleEnvironment<'a, 'data> {
                 };
 
                 let cnt = usize::try_from(data.count()).unwrap();
-                initializers.reserve_exact(cnt);
+                initializers.reserve_exact(cnt)?;
                 self.result.data.reserve_exact(cnt);
 
                 for (index, entry) in data.into_iter().enumerate() {
@@ -670,7 +668,7 @@ impl<'a, 'data> ModuleEnvironment<'a, 'data> {
                                 memory_index,
                                 offset,
                                 data: range,
-                            });
+                            })?;
                             self.result.data.push(data.into());
                         }
                         DataKind::Passive => {
@@ -821,7 +819,7 @@ and for re-adding support for interface types you can see this issue:
             name: self.result.module.strings.insert(module)?,
             field: self.result.module.strings.insert(field)?,
             index,
-        });
+        })?;
         Ok(())
     }
 
@@ -1224,7 +1222,7 @@ impl ModuleTranslation<'_> {
             if let TableInitialValue::Expr(expr) = init {
                 if let [ConstOp::RefFunc(f)] = expr.ops() {
                     *init = TableInitialValue::Null {
-                        precomputed: vec![*f; table_size as usize],
+                        precomputed: collections::vec![*f; table_size as usize].panic_on_oom(),
                     };
                 }
             }
@@ -1318,14 +1316,16 @@ impl ModuleTranslation<'_> {
             // it's large enough to hold the segment and then copying the
             // segment into the precomputed list.
             if precomputed.len() < top as usize {
-                precomputed.resize(top as usize, FuncIndex::reserved_value());
+                precomputed
+                    .resize(top as usize, FuncIndex::reserved_value())
+                    .panic_on_oom();
             }
             let dst = &mut precomputed[offset as usize..top as usize];
             dst.copy_from_slice(&function_elements);
 
             // advance the iterator to see the next segment
             let _ = segments.next();
         }
-        self.module.table_initialization.segments = segments.collect();
+        self.module.table_initialization.segments = segments.try_collect().panic_on_oom();
     }
 }
diff --git a/crates/environ/src/module.rs b/crates/environ/src/module.rs
@@ -48,7 +48,7 @@ pub enum MemoryInitialization {
     /// data segments when the module is instantiated.
     ///
     /// This is the default memory initialization type.
-    Segmented(Vec<MemoryInitializer>),
+    Segmented(collections::Vec<MemoryInitializer>),
 
     /// Memory initialization is statically known and involves a single `memcpy`
     /// or otherwise simply making the defined data visible.
@@ -88,7 +88,7 @@ pub enum MemoryInitialization {
 
 impl Default for MemoryInitialization {
     fn default() -> Self {
-        Self::Segmented(Vec::new())
+        Self::Segmented(collections::Vec::new())
     }
 }
 
@@ -233,7 +233,7 @@ pub struct TableInitialization {
     ///
     /// These element segments are iterated over during instantiation to apply
     /// any segments that weren't already moved into `initial_values` above.
-    pub segments: Vec<TableSegment>,
+    pub segments: collections::Vec<TableSegment>,
 }
 
 /// Initial value for all elements in a table.
@@ -249,7 +249,7 @@ pub enum TableInitialValue {
         /// `FuncIndex::reserved_value()`. Note that this image is empty by
         /// default and may not encompass the entire span of the table in which
         /// case the elements are initialized to null.
-        precomputed: Vec<FuncIndex>,
+        precomputed: collections::Vec<FuncIndex>,
     },
     /// An arbitrary const expression.
     Expr(ConstExpr),
@@ -301,7 +301,7 @@ pub struct Module {
     pub name: Option<Atom>,
 
     /// All import records, in the order they are declared in the module.
-    pub initializers: Vec<Initializer>,
+    pub initializers: collections::Vec<Initializer>,
 
     /// Exported entities.
     pub exports: collections::IndexMap<Atom, EntityIndex>,
@@ -316,7 +316,7 @@ pub struct Module {
     pub memory_initialization: MemoryInitialization,
 
     /// WebAssembly passive elements.
-    pub passive_elements: Vec<TableSegmentElements>,
+    pub passive_elements: collections::Vec<TableSegmentElements>,
 
     /// The map from passive element index (element segment index space) to index in `passive_elements`.
     pub passive_elements_map: BTreeMap<ElemIndex, usize>,
diff --git a/crates/environ/src/types.rs b/crates/environ/src/types.rs
@@ -2,7 +2,6 @@ use crate::{Tunables, WasmResult, error::OutOfMemory, wasm_unsupported};
 use alloc::borrow::Cow;
 use alloc::boxed::Box;
 use core::{fmt, ops::Range};
-use cranelift_entity::entity_impl;
 use serde_derive::{Deserialize, Serialize};
 use smallvec::SmallVec;
 use wasmtime_core::alloc::{TryClone, TryCollect as _};
@@ -1536,67 +1535,80 @@ impl TypeTrace for WasmRecGroup {
     }
 }
 
+macro_rules! entity_impl_with_try_clone {
+    ( $ty:ident ) => {
+        cranelift_entity::entity_impl!($ty);
+
+        impl TryClone for $ty {
+            #[inline]
+            fn try_clone(&self) -> Result<Self, $crate::error::OutOfMemory> {
+                Ok(*self)
+            }
+        }
+    };
+}
+
 /// Index type of a function (imported or defined) inside the WebAssembly module.
 #[derive(Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord, Debug, Serialize, Deserialize)]
 pub struct FuncIndex(u32);
-entity_impl!(FuncIndex);
+entity_impl_with_try_clone!(FuncIndex);
 
 /// Index type of a defined function inside the WebAssembly module.
 #[derive(Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord, Debug, Serialize, Deserialize)]
 pub struct DefinedFuncIndex(u32);
-entity_impl!(DefinedFuncIndex);
+entity_impl_with_try_clone!(DefinedFuncIndex);
 
 /// Index type of a defined table inside the WebAssembly module.
 #[derive(Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord, Debug, Serialize, Deserialize)]
 pub struct DefinedTableIndex(u32);
-entity_impl!(DefinedTableIndex);
+entity_impl_with_try_clone!(DefinedTableIndex);
 
 /// Index type of a defined memory inside the WebAssembly module.
 #[derive(Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord, Debug, Serialize, Deserialize)]
 pub struct DefinedMemoryIndex(u32);
-entity_impl!(DefinedMemoryIndex);
+entity_impl_with_try_clone!(DefinedMemoryIndex);
 
 /// Index type of a defined memory inside the WebAssembly module.
 #[derive(Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord, Debug, Serialize, Deserialize)]
 pub struct OwnedMemoryIndex(u32);
-entity_impl!(OwnedMemoryIndex);
+entity_impl_with_try_clone!(OwnedMemoryIndex);
 
 /// Index type of a defined global inside the WebAssembly module.
 #[derive(Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord, Debug, Serialize, Deserialize)]
 pub struct DefinedGlobalIndex(u32);
-entity_impl!(DefinedGlobalIndex);
+entity_impl_with_try_clone!(DefinedGlobalIndex);
 
 /// Index type of a table (imported or defined) inside the WebAssembly module.
 #[derive(Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord, Debug, Serialize, Deserialize)]
 pub struct TableIndex(u32);
-entity_impl!(TableIndex);
+entity_impl_with_try_clone!(TableIndex);
 
 /// Index type of a global variable (imported or defined) inside the WebAssembly module.
 #[derive(Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord, Debug, Serialize, Deserialize)]
 pub struct GlobalIndex(u32);
-entity_impl!(GlobalIndex);
+entity_impl_with_try_clone!(GlobalIndex);
 
 /// Index type of a linear memory (imported or defined) inside the WebAssembly module.
 #[derive(Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord, Debug, Serialize, Deserialize)]
 pub struct MemoryIndex(u32);
-entity_impl!(MemoryIndex);
+entity_impl_with_try_clone!(MemoryIndex);
 
 /// Index type of a canonicalized recursive type group inside a WebAssembly
 /// module (as opposed to canonicalized within the whole engine).
 #[derive(Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord, Debug, Serialize, Deserialize)]
 pub struct ModuleInternedRecGroupIndex(u32);
-entity_impl!(ModuleInternedRecGroupIndex);
+entity_impl_with_try_clone!(ModuleInternedRecGroupIndex);
 
 /// Index type of a canonicalized recursive type group inside the whole engine
 /// (as opposed to canonicalized within just a single Wasm module).
 #[derive(Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord, Debug, Serialize, Deserialize)]
 pub struct EngineInternedRecGroupIndex(u32);
-entity_impl!(EngineInternedRecGroupIndex);
+entity_impl_with_try_clone!(EngineInternedRecGroupIndex);
 
 /// Index type of a type (imported or defined) inside the WebAssembly module.
 #[derive(Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord, Debug, Serialize, Deserialize)]
 pub struct TypeIndex(u32);
-entity_impl!(TypeIndex);
+entity_impl_with_try_clone!(TypeIndex);
 
 /// A canonicalized type index referencing a type within a single recursion
 /// group from another type within that same recursion group.
@@ -1605,7 +1617,7 @@ entity_impl!(TypeIndex);
 /// groups.
 #[derive(Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord, Debug, Serialize, Deserialize)]
 pub struct RecGroupRelativeTypeIndex(u32);
-entity_impl!(RecGroupRelativeTypeIndex);
+entity_impl_with_try_clone!(RecGroupRelativeTypeIndex);
 
 /// A canonicalized type index for a type within a single WebAssembly module.
 ///
@@ -1616,7 +1628,7 @@ entity_impl!(RecGroupRelativeTypeIndex);
 /// involve entities defined in different modules).
 #[derive(Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord, Debug, Serialize, Deserialize)]
 pub struct ModuleInternedTypeIndex(u32);
-entity_impl!(ModuleInternedTypeIndex);
+entity_impl_with_try_clone!(ModuleInternedTypeIndex);
 
 /// A canonicalized type index into an engine's shared type registry.
 ///
@@ -1628,7 +1640,7 @@ entity_impl!(ModuleInternedTypeIndex);
 #[repr(transparent)] // Used directly by JIT code.
 #[derive(Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord, Debug, Serialize, Deserialize)]
 pub struct VMSharedTypeIndex(u32);
-entity_impl!(VMSharedTypeIndex);
+entity_impl_with_try_clone!(VMSharedTypeIndex);
 
 impl VMSharedTypeIndex {
     /// Create a new `VMSharedTypeIndex`.
@@ -1659,30 +1671,30 @@ impl Default for VMSharedTypeIndex {
 /// Index type of a passive data segment inside the WebAssembly module.
 #[derive(Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord, Debug, Serialize, Deserialize)]
 pub struct DataIndex(u32);
-entity_impl!(DataIndex);
+entity_impl_with_try_clone!(DataIndex);
 
 /// Index type of a passive element segment inside the WebAssembly module.
 #[derive(Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord, Debug, Serialize, Deserialize)]
 pub struct ElemIndex(u32);
-entity_impl!(ElemIndex);
+entity_impl_with_try_clone!(ElemIndex);
 
 /// Index type of a defined tag inside the WebAssembly module.
 #[derive(Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord, Debug, Serialize, Deserialize)]
 pub struct DefinedTagIndex(u32);
-entity_impl!(DefinedTagIndex);
+entity_impl_with_try_clone!(DefinedTagIndex);
 
 /// Index type of an event inside the WebAssembly module.
 #[derive(Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord, Debug, Serialize, Deserialize)]
 pub struct TagIndex(u32);
-entity_impl!(TagIndex);
+entity_impl_with_try_clone!(TagIndex);
 
 /// Index into the global list of modules found within an entire component.
 ///
 /// Module translations are saved on the side to get fully compiled after
 /// the original component has finished being translated.
 #[derive(Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord, Debug, Serialize, Deserialize)]
 pub struct StaticModuleIndex(u32);
-entity_impl!(StaticModuleIndex);
+entity_impl_with_try_clone!(StaticModuleIndex);
 
 /// An index of an entity.
 #[derive(Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord, Debug, Serialize, Deserialize)]
