Refactor String usage in wasmtime_environ::Module (#12565)

* Refactor `String` usage in `wasmtime_environ::Module`

Do not hold regular `String`s; instead use our own OOM-handling
`wasmtime_core::alloc::String` or, even better, an interned `Atom` from the
`StringPool`.

Also avoid `IndexMap`, as it doesn't handle OOMs.

* Use OOM-handling `IndexMap` in `wasmtime_environ::Module`

* review feedback

Author: fitzgen

crates/environ/src/compile/module_environ.rs

@@ -1,4 +1,4 @@
-use crate::error::{Result, bail};
+use crate::error::{OutOfMemory, Result, bail};
 use crate::module::{
     FuncRefIndex, Initializer, MemoryInitialization, MemoryInitializer, Module, TableSegment,
     TableSegmentElements,
@@ -7,10 +7,10 @@ use crate::prelude::*;
 use crate::{
     ConstExpr, ConstOp, DataIndex, DefinedFuncIndex, ElemIndex, EngineOrModuleTypeIndex,
     EntityIndex, EntityType, FuncIndex, FuncKey, GlobalIndex, IndexType, InitMemory, MemoryIndex,
-    ModuleInternedTypeIndex, ModuleTypesBuilder, PrimaryMap, SizeOverflow, StaticMemoryInitializer,
-    StaticModuleIndex, TableIndex, TableInitialValue, Tag, TagIndex, Tunables, TypeConvert,
-    TypeIndex, WasmError, WasmHeapTopType, WasmHeapType, WasmResult, WasmValType,
-    WasmparserTypeConverter,
+    ModuleInternedTypeIndex, ModuleTypesBuilder, PanicOnOom as _, PrimaryMap, SizeOverflow,
+    StaticMemoryInitializer, StaticModuleIndex, TableIndex, TableInitialValue, Tag, TagIndex,
+    Tunables, TypeConvert, TypeIndex, WasmError, WasmHeapTopType, WasmHeapType, WasmResult,
+    WasmValType, WasmparserTypeConverter,
 };
 use cranelift_entity::SecondaryMap;
 use cranelift_entity::packed_option::ReservedValue;
@@ -373,7 +373,7 @@ impl<'a, 'data> ModuleEnvironment<'a, 'data> {
                             bail!("custom-descriptors proposal not implemented yet");
                         }
                     };
-                    self.declare_import(import.module, import.name, ty);
+                    self.declare_import(import.module, import.name, ty)?;
                 }
             }
 
@@ -469,7 +469,7 @@ impl<'a, 'data> ModuleEnvironment<'a, 'data> {
                 self.validator.export_section(&exports)?;
 
                 let cnt = usize::try_from(exports.count()).unwrap();
-                self.result.module.exports.reserve(cnt);
+                self.result.module.exports.reserve(cnt)?;
 
                 for entry in exports {
                     let wasmparser::Export { name, kind, index } = entry?;
@@ -484,10 +484,8 @@ impl<'a, 'data> ModuleEnvironment<'a, 'data> {
                         ExternalKind::Global => EntityIndex::Global(GlobalIndex::from_u32(index)),
                         ExternalKind::Tag => EntityIndex::Tag(TagIndex::from_u32(index)),
                     };
-                    self.result
-                        .module
-                        .exports
-                        .insert(String::from(name), entity);
+                    let name = self.result.module.strings.insert(name)?;
+                    self.result.module.exports.insert(name, entity)?;
                 }
             }
 
@@ -812,13 +810,19 @@ and for re-adding support for interface types you can see this issue:
     /// When the module linking proposal is disabled, however, disregard this
     /// logic and instead work directly with two-level imports since no
     /// instances are defined.
-    fn declare_import(&mut self, module: &'data str, field: &'data str, ty: EntityType) {
+    fn declare_import(
+        &mut self,
+        module: &'data str,
+        field: &'data str,
+        ty: EntityType,
+    ) -> Result<(), OutOfMemory> {
         let index = self.push_type(ty);
         self.result.module.initializers.push(Initializer::Import {
-            name: module.to_owned(),
-            field: field.to_owned(),
+            name: self.result.module.strings.insert(module)?,
+            field: self.result.module.strings.insert(field)?,
             index,
         });
+        Ok(())
     }
 
     fn push_type(&mut self, ty: EntityType) -> EntityIndex {
@@ -877,7 +881,8 @@ and for re-adding support for interface types you can see this issue:
                     }
                 }
                 wasmparser::Name::Module { name, .. } => {
-                    self.result.module.name = Some(name.to_string());
+                    self.result.module.name =
+                        Some(self.result.module.strings.insert(name).panic_on_oom());
                     if self.tunables.debug_native {
                         self.result.debuginfo.name_section.module_name = Some(name);
                     }

crates/environ/src/component/translate/adapt.rs

@@ -252,8 +252,9 @@ impl<'data> Translator<'_, 'data> {
             // partitioned in-order so we're guaranteed to push the adapters
             // in-order here as well. (with an assert to double-check)
             for (adapter, name) in adapter_module.adapters.iter().zip(&names) {
-                let index = translation.module.exports[name];
-                let i = component.adapter_partitionings.push((module_id, index));
+                let name = translation.module.strings.get_atom(name).unwrap();
+                let export = translation.module.exports[&name];
+                let i = component.adapter_partitionings.push((module_id, export));
                 assert_eq!(i, *adapter);
             }
 

crates/environ/src/component/translate/inline.rs

@@ -1310,7 +1310,12 @@ impl<'a> Inliner<'a> {
                     ModuleInstanceDef::Instantiated(instance, module) => {
                         let (ty, item) = match &frame.modules[*module] {
                             ModuleDef::Static(idx, _ty) => {
-                                let entity = self.nested_modules[*idx].module.exports[*name];
+                                let name = self.nested_modules[*idx]
+                                    .module
+                                    .strings
+                                    .get_atom(name)
+                                    .unwrap();
+                                let entity = self.nested_modules[*idx].module.exports[&name];
                                 let ty = match entity {
                                     EntityIndex::Function(f) => {
                                         self.nested_modules[*idx].module.functions[f]
@@ -1481,7 +1486,12 @@ impl<'a> Inliner<'a> {
             ModuleInstanceDef::Instantiated(instance, module) => {
                 let item = match frame.modules[*module] {
                     ModuleDef::Static(idx, _ty) => {
-                        let entity = self.nested_modules[idx].module.exports[name];
+                        let name = self.nested_modules[idx]
+                            .module
+                            .strings
+                            .get_atom(name)
+                            .unwrap();
+                        let entity = self.nested_modules[idx].module.exports[&name];
                         ExportItem::Index(entity)
                     }
                     ModuleDef::Import(..) => ExportItem::Name(name.to_string()),

crates/environ/src/module.rs

@@ -294,14 +294,17 @@ pub struct Module {
     /// This module's index.
     pub module_index: StaticModuleIndex,
 
+    /// A pool of strings used in this module.
+    pub strings: StringPool,
+
     /// The name of this wasm module, often found in the wasm file.
-    pub name: Option<String>,
+    pub name: Option<Atom>,
 
     /// All import records, in the order they are declared in the module.
     pub initializers: Vec<Initializer>,
 
     /// Exported entities.
-    pub exports: IndexMap<String, EntityIndex>,
+    pub exports: collections::IndexMap<Atom, EntityIndex>,
 
     /// The module "start" function, if present.
     pub start_func: Option<FuncIndex>,
@@ -375,9 +378,9 @@ pub enum Initializer {
     /// An imported item is required to be provided.
     Import {
         /// Name of this import
-        name: String,
+        name: Atom,
         /// The field name projection of this import
-        field: String,
+        field: Atom,
         /// Where this import will be placed, which also has type information
         /// about the import.
         index: EntityIndex,
@@ -389,6 +392,7 @@ impl Module {
     pub fn new(module_index: StaticModuleIndex) -> Self {
         Self {
             module_index,
+            strings: Default::default(),
             name: Default::default(),
             initializers: Default::default(),
             exports: Default::default(),
@@ -563,17 +567,22 @@ impl Module {
     /// Returns an iterator of all the imports in this module, along with their
     /// module name, field name, and type that's being imported.
     pub fn imports(&self) -> impl ExactSizeIterator<Item = (&str, &str, EntityType)> {
+        let pool = &self.strings;
         self.initializers.iter().map(move |i| match i {
             Initializer::Import { name, field, index } => {
-                (name.as_str(), field.as_str(), self.type_of(*index))
+                (&pool[name], &pool[field], self.type_of(*index))
             }
         })
     }
 
     /// Get this module's `i`th import.
     pub fn import(&self, i: usize) -> Option<(&str, &str, EntityType)> {
         match self.initializers.get(i)? {
-            Initializer::Import { name, field, index } => Some((name, field, self.type_of(*index))),
+            Initializer::Import { name, field, index } => Some((
+                &self.strings[name],
+                &self.strings[field],
+                self.type_of(*index),
+            )),
         }
     }
 
@@ -670,6 +679,7 @@ impl TypeTrace for Module {
         // when adding new fields that might need re-canonicalization.
         let Self {
             module_index: _,
+            strings: _,
             name: _,
             initializers: _,
             exports: _,
@@ -721,6 +731,7 @@ impl TypeTrace for Module {
         // when adding new fields that might need re-canonicalization.
         let Self {
             module_index: _,
+            strings: _,
             name: _,
             initializers: _,
             exports: _,

crates/environ/src/string_pool.rs

@@ -101,6 +101,17 @@ impl core::ops::Index<Atom> for StringPool {
     }
 }
 
+// For convenience, to avoid `*atom` noise at call sites.
+impl core::ops::Index<&'_ Atom> for StringPool {
+    type Output = str;
+
+    #[inline]
+    #[track_caller]
+    fn index(&self, atom: &Atom) -> &Self::Output {
+        self.get(*atom).unwrap()
+    }
+}
+
 impl serde::ser::Serialize for StringPool {
     fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
     where
@@ -218,6 +229,20 @@ impl StringPool {
             None
         }
     }
+
+    /// Get the number of strings in this pool.
+    pub fn len(&self) -> usize {
+        self.strings.len()
+    }
+}
+
+impl Default for Atom {
+    #[inline]
+    fn default() -> Self {
+        Self {
+            index: NonZeroU32::MAX,
+        }
+    }
 }
 
 impl fmt::Debug for Atom {
@@ -228,6 +253,17 @@ impl fmt::Debug for Atom {
     }
 }
 
+// Allow using `Atom` in `SecondaryMap`s.
+impl crate::EntityRef for Atom {
+    fn new(index: usize) -> Self {
+        Atom::new(index)
+    }
+
+    fn index(self) -> usize {
+        Atom::index(&self)
+    }
+}
+
 impl serde::ser::Serialize for Atom {
     fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
     where

crates/wasmtime/src/runtime/component/instance.rs

@@ -632,7 +632,11 @@ where
         // during that phase so the actual instantiation of an `InstancePre`
         // skips all string lookups. This should probably only be
         // investigated if this becomes a performance issue though.
-        ExportItem::Name(name) => instance.env_module().exports[name],
+        ExportItem::Name(name) => {
+            let module = instance.env_module();
+            let name = module.strings.get_atom(name).unwrap();
+            module.exports[&name]
+        }
     };
     // SAFETY: the `store_id` owns this instance and all exports contained
     // within.

crates/wasmtime/src/runtime/component/matching.rs

@@ -119,8 +119,9 @@ impl TypeChecker<'_> {
         // Every export that is expected should be in the actual module we have
         for (name, expected) in expected.exports.iter() {
             let idx = actual
-                .exports
-                .get(name)
+                .strings
+                .get_atom(name)
+                .and_then(|atom| actual.exports.get(&atom))
                 .ok_or_else(|| format_err!("module export `{name}` not defined"))?;
             let actual = actual.type_of(*idx);
             matching::entity_ty(self.engine, expected, &actual)

crates/wasmtime/src/runtime/instance.rs

@@ -425,12 +425,12 @@ impl Instance {
         for (_name, entity) in module.exports.iter() {
             items.push(self._get_export(store, *entity));
         }
-        store[self.id]
-            .env_module()
+        let module = store[self.id].env_module();
+        module
             .exports
             .iter()
             .zip(items)
-            .map(|((name, _), item)| Export::new(name, item))
+            .map(|((name, _), item)| Export::new(&module.strings[name], item))
     }
 
     /// Looks up an exported [`Extern`] value by name.
@@ -452,7 +452,9 @@ impl Instance {
     /// mutable context.
     pub fn get_export(&self, mut store: impl AsContextMut, name: &str) -> Option<Extern> {
         let store = store.as_context_mut().0;
-        let entity = *store[self.id].env_module().exports.get(name)?;
+        let module = store[self.id].env_module();
+        let name = module.strings.get_atom(name)?;
+        let entity = *module.exports.get(&name)?;
         Some(self._get_export(store, entity))
     }
 

crates/wasmtime/src/runtime/module.rs

@@ -678,7 +678,9 @@ impl Module {
     /// # }
     /// ```
     pub fn name(&self) -> Option<&str> {
-        self.compiled_module().module().name.as_deref()
+        let module = self.compiled_module().module();
+        let name = module.name?;
+        Some(&module.strings[name])
     }
 
     /// Returns the list of imports that this [`Module`] has and must be
@@ -807,7 +809,12 @@ impl Module {
         let types = self.types();
         let engine = self.engine();
         module.exports.iter().map(move |(name, entity_index)| {
-            ExportType::new(name, module.type_of(*entity_index), types, engine)
+            ExportType::new(
+                &module.strings[name],
+                module.type_of(*entity_index),
+                types,
+                engine,
+            )
         })
     }
 
@@ -856,7 +863,8 @@ impl Module {
     /// ```
     pub fn get_export(&self, name: &str) -> Option<ExternType> {
         let module = self.compiled_module().module();
-        let entity_index = module.exports.get(name)?;
+        let name = module.strings.get_atom(name)?;
+        let entity_index = module.exports.get(&name)?;
         Some(ExternType::from_wasmtime(
             self.engine(),
             self.types(),
@@ -873,7 +881,8 @@ impl Module {
     pub fn get_export_index(&self, name: &str) -> Option<ModuleExport> {
         let compiled_module = self.compiled_module();
         let module = compiled_module.module();
-        let entity = *module.exports.get(name)?;
+        let name = module.strings.get_atom(name)?;
+        let entity = *module.exports.get(&name)?;
         Some(ModuleExport {
             module: self.id(),
             entity,

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

@@ -38,9 +38,10 @@ pub async fn create_memory(
     // Since we have only associated a single memory with the "frankenstein"
     // instance, it will be exported at index 0.
     debug_assert_eq!(memory_id.as_u32(), 0);
+    let name = module.strings.insert("")?;
     module
         .exports
-        .insert(String::new(), EntityIndex::Memory(memory_id));
+        .insert(name, EntityIndex::Memory(memory_id))?;
     let info = ModuleRuntimeInfo::bare(Arc::new(module))?;
 
     // We create an instance in the on-demand allocator when creating handles