Refactor ModuleRegistry to not rely on image layout (#13483)

Prior to this commit the methods of `ModuleRegistry` implicitly relied
on the in-memory layout of compiled modules/components, specifically
that defined functions for modules were the only interesting functions
and they're all clumped in groups. This reliance is surfaced in the
`module_and_code_by_pc` method which was called for various operations
such as when throwing an exception or looking at the stack for GC roots.
In these scenarios, however, there's no need for a `Module` to be used
and instead a bland `CodeMemory` can be used instead.

This commit refactors the internals to no longer look at the
`finished_functions` of a module and also avoid a `BTreeMap` within the
`ModuleRegistry`. By doing so all methods are now fully agnostic to the
internal layout of the code image of a module or component, which in
turn empowers refactoring/moving various items around as needed. This is
needed for an upcoming change I'm making, for example. Internally most
users have moved over to `store_code_by_pc` since the call sites didn't
need a module but instead "just some metadata". The main holdout is the
logic to generate a trap frame from a stack walk. This is enabled by
refactoring the previous `BTreeMap` to instead hold a
`CompiledFunctionsTable`, used to translate a pc to a `FuncKey`, and
then a mapping of `StaticModuleIndex` to `RegisteredModuleId`. This adds
up to the ability to be able to go from an arbitrary pc to a module for
a defined wasm function.

Author: alexcrichton

crates/wasmtime/src/runtime/code.rs

@@ -215,12 +215,6 @@ impl EngineCode {
         self.original_code.exception_tables()
     }
 
-    /// Returns the encoded frame-tables section to pass to
-    /// `wasmtime_environ::FrameTable::parse`.
-    pub fn frame_tables(&self) -> &[u8] {
-        self.original_code.frame_tables()
-    }
-
     /// Returns the data in the `ELF_NAME_DATA` section.
     #[inline]
     pub fn func_name_data(&self) -> &[u8] {
@@ -262,6 +256,12 @@ impl EngineCode {
     pub(crate) fn module_memory_image_source(&self) -> &Arc<impl ModuleMemoryImageSource> {
         &self.original_code
     }
+
+    /// See [`CodeMemory::frame_table`].
+    #[cfg(feature = "debug")]
+    pub(crate) fn frame_table<'a>(&'a self) -> Option<wasmtime_environ::FrameTable<'a>> {
+        self.original_code.frame_table()
+    }
 }
 
 impl Drop for EngineCode {
@@ -430,23 +430,6 @@ impl<'a> ModuleWithCode<'a> {
         self.module
     }
 
-    /// Provide the StoreCode wrapped in this tuple.
-    pub fn store_code(&self) -> &'a StoreCode {
-        self.store_code
-    }
-
-    /// Returns an iterator over all functions defined within this module with
-    /// their index and their raw pointer.
-    #[inline]
-    pub fn finished_functions(
-        &self,
-    ) -> impl ExactSizeIterator<Item = (DefinedFuncIndex, &[u8])> + '_ {
-        self.module
-            .env_module()
-            .defined_func_indices()
-            .map(|i| (i, self.finished_function(i)))
-    }
-
     /// Returns the slice in the text section of the function that
     /// `index` points to.
     #[inline]
@@ -474,20 +457,4 @@ impl<'a> ModuleWithCode<'a> {
             .array_to_wasm_trampoline_range(def_func_index)?;
         Some(&self.store_code.text()[range])
     }
-
-    /// Get the text offset (relative PC) for a given absolute PC in
-    /// this module.
-    #[cfg(feature = "gc")]
-    pub(crate) fn text_offset(&self, pc: usize) -> Option<u32> {
-        StoreCodePC::offset_of(self.store_code.text_range(), pc)
-            .map(|offset| u32::try_from(offset).expect("Module larger than 4GiB"))
-    }
-
-    /// Lookup the stack map at a program counter value.
-    #[cfg(feature = "gc")]
-    pub(crate) fn lookup_stack_map(&self, pc: usize) -> Option<wasmtime_environ::StackMap<'_>> {
-        let text_offset = self.text_offset(pc)?;
-        let info = self.module.engine_code().stack_map_data();
-        wasmtime_environ::StackMap::lookup(text_offset, info)
-    }
 }

crates/wasmtime/src/runtime/code_memory.rs

@@ -633,6 +633,22 @@ impl CodeMemory {
         let mmap = self.mmap.deep_clone()?;
         Self::new(engine, mmap)
     }
+
+    /// Obtain a frame-table parser on this module's frame state slot
+    /// (debug instrumentation) metadata.
+    #[cfg(feature = "debug")]
+    pub(crate) fn frame_table(&self) -> Option<wasmtime_environ::FrameTable<'_>> {
+        let data = self.frame_tables();
+        if data.is_empty() {
+            None
+        } else {
+            let orig_text = self.text();
+            Some(
+                wasmtime_environ::FrameTable::parse(data, orig_text)
+                    .expect("Frame tables were validated on module load"),
+            )
+        }
+    }
 }
 
 fn section_name<'a>(

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

@@ -933,6 +933,10 @@ impl Component {
             _ => unreachable!(),
         }
     }
+
+    pub(crate) fn index(&self) -> &Arc<CompiledFunctionsTable> {
+        &self.inner.index
+    }
 }
 
 /// A value which represents a known export of a component.

crates/wasmtime/src/runtime/debug.rs

@@ -645,10 +645,10 @@ impl FrameDataCache {
                 // module that actually contains the physical PC
                 // (i.e., the outermost function that inlined the
                 // others).
-                let (module, frames) = VirtualFrame::decode(registry, frame.pc());
+                let (store_code, frames) = VirtualFrame::decode(registry, frame.pc());
                 let frames = frames
                     .into_iter()
-                    .map(|frame| FrameData::compute(frame, &module))
+                    .map(|frame| FrameData::compute(frame, store_code))
                     .collect::<Vec<_>>();
                 v.insert(frames)
             }
@@ -672,18 +672,17 @@ struct VirtualFrame {
 impl VirtualFrame {
     /// Return virtual frames corresponding to a physical frame, from
     /// outermost to innermost.
-    fn decode(registry: &ModuleRegistry, pc: usize) -> (Module, Vec<VirtualFrame>) {
-        let (module_with_code, pc) = registry
-            .module_and_code_by_pc(pc)
+    fn decode(registry: &ModuleRegistry, pc: usize) -> (&StoreCode, Vec<VirtualFrame>) {
+        let (store_code, pc) = registry
+            .store_code_by_pc(pc)
             .expect("Wasm frame PC does not correspond to a module");
-        let module = module_with_code.module();
-        let table = module.frame_table().unwrap();
+        let table = store_code.code_memory().frame_table().unwrap();
         let pc = u32::try_from(pc).expect("PC offset too large");
         let program_points = table.find_program_point(pc, FrameInstPos::Post)
             .expect("There must be a program point record in every frame when debug instrumentation is enabled");
 
         (
-            module.clone(),
+            store_code,
             program_points
                 .map(|(wasm_pc, frame_descriptor, stack_shape)| VirtualFrame {
                     wasm_pc,
@@ -719,8 +718,8 @@ struct FrameData {
 }
 
 impl FrameData {
-    fn compute(frame: VirtualFrame, module: &Module) -> Self {
-        let frame_table = module.frame_table().unwrap();
+    fn compute(frame: VirtualFrame, store_code: &StoreCode) -> Self {
+        let frame_table = store_code.code_memory().frame_table().unwrap();
         // Parse the frame descriptor.
         let (data, slot_to_fp_offset) = frame_table
             .frame_descriptor(frame.frame_descriptor)

crates/wasmtime/src/runtime/instantiate.rs

@@ -267,6 +267,10 @@ impl CompiledModule {
         self.engine_code
             .wasm_bytecode_for_module(self.module.module_index)
     }
+
+    pub(crate) fn index(&self) -> &Arc<CompiledFunctionsTable> {
+        &self.index
+    }
 }
 
 #[cfg(feature = "addr2line")]

crates/wasmtime/src/runtime/module.rs

@@ -18,14 +18,10 @@ use core::ptr::NonNull;
 #[cfg(feature = "std")]
 use std::{fs::File, path::Path};
 use wasmparser::{Parser, ValidPayload, Validator};
-#[cfg(feature = "debug")]
-use wasmtime_environ::FrameTable;
 use wasmtime_environ::{
     CompiledFunctionsTable, CompiledModuleInfo, EntityIndex, HostPtr, ModuleTypes, ObjectKind,
     StaticModuleIndex, TypeTrace, VMOffsets, VMSharedTypeIndex, WasmChecksum,
 };
-#[cfg(feature = "gc")]
-use wasmtime_unwinder::ExceptionTable;
 mod registry;
 
 pub use registry::*;
@@ -1178,27 +1174,10 @@ impl Module {
         Ok(images)
     }
 
-    /// Obtain an exception-table parser on this module's exception metadata.
-    #[cfg(feature = "gc")]
-    pub(crate) fn exception_table<'a>(&'a self) -> ExceptionTable<'a> {
-        ExceptionTable::parse(self.inner.code.exception_tables())
-            .expect("Exception tables were validated on module load")
-    }
-
-    /// Obtain a frame-table parser on this module's frame state slot
-    /// (debug instrumentation) metadata.
+    /// See [`CodeMemory::frame_table`].
     #[cfg(feature = "debug")]
-    pub(crate) fn frame_table<'a>(&'a self) -> Option<FrameTable<'a>> {
-        let data = self.inner.code.frame_tables();
-        if data.is_empty() {
-            None
-        } else {
-            let orig_text = self.inner.code.text();
-            Some(
-                FrameTable::parse(data, orig_text)
-                    .expect("Frame tables were validated on module load"),
-            )
-        }
+    pub(crate) fn frame_table<'a>(&'a self) -> Option<wasmtime_environ::FrameTable<'a>> {
+        self.inner.code.frame_table()
     }
 
     /// Is this `Module` the same as another?
@@ -1216,6 +1195,10 @@ impl Module {
     pub fn same(a: &Module, b: &Module) -> bool {
         Arc::ptr_eq(&a.inner, &b.inner)
     }
+
+    pub(crate) fn index(&self) -> &Arc<CompiledFunctionsTable> {
+        &self.inner.module.index()
+    }
 }
 
 /// Describes a function for a given module.