Provide proper function index and name in the FrameInfo (#824)

* fix function index

* Add function name to JITFunctionTag

* Add ModuleSyncString.

crates/api/src/trap.rs

@@ -42,7 +42,8 @@ impl Trap {
             if let Some(info) = wasmtime_runtime::jit_function_registry::find(pc) {
                 wasm_trace.push(FrameInfo {
                     func_index: info.func_index as u32,
-                    module_name: info.module_id.clone(),
+                    module_name: info.module_id.get().map(|s| s.to_string()),
+                    func_name: info.func_name.get().map(|s| s.to_string()),
                 })
             }
         }
@@ -87,6 +88,7 @@ impl std::error::Error for Trap {}
 pub struct FrameInfo {
     module_name: Option<String>,
     func_index: u32,
+    func_name: Option<String>,
 }
 
 impl FrameInfo {
@@ -109,4 +111,8 @@ impl FrameInfo {
     pub fn module_name(&self) -> Option<&str> {
         self.module_name.as_deref()
     }
+
+    pub fn func_name(&self) -> Option<&str> {
+        self.func_name.as_deref()
+    }
 }

crates/api/tests/traps.rs

@@ -69,8 +69,10 @@ fn test_trap_trace() -> Result<(), String> {
     assert_eq!(trace.len(), 2);
     assert_eq!(trace[0].module_name().unwrap(), "hello_mod");
     assert_eq!(trace[0].func_index(), 1);
+    assert_eq!(trace[0].func_name(), Some("hello"));
     assert_eq!(trace[1].module_name().unwrap(), "hello_mod");
     assert_eq!(trace[1].func_index(), 0);
+    assert_eq!(trace[1].func_name(), None);
     assert!(e.message().contains("unreachable"));
 
     Ok(())
@@ -114,9 +116,9 @@ fn test_trap_trace_cb() -> Result<(), String> {
     let trace = e.trace();
     assert_eq!(trace.len(), 2);
     assert_eq!(trace[0].module_name().unwrap(), "hello_mod");
-    assert_eq!(trace[0].func_index(), 1);
+    assert_eq!(trace[0].func_index(), 2);
     assert_eq!(trace[1].module_name().unwrap(), "hello_mod");
-    assert_eq!(trace[1].func_index(), 0);
+    assert_eq!(trace[1].func_index(), 1);
     assert_eq!(e.message(), "cb throw");
 
     Ok(())
@@ -149,6 +151,7 @@ fn test_trap_stack_overflow() -> Result<(), String> {
     for i in 0..trace.len() {
         assert_eq!(trace[i].module_name().unwrap(), "rec_mod");
         assert_eq!(trace[i].func_index(), 0);
+        assert_eq!(trace[1].func_name(), Some("run"));
     }
     assert!(e.message().contains("call stack exhausted"));
 

crates/environ/src/lib.rs

@@ -54,7 +54,7 @@ pub use crate::func_environ::BuiltinFunctionIndex;
 #[cfg(feature = "lightbeam")]
 pub use crate::lightbeam::Lightbeam;
 pub use crate::module::{
-    Export, MemoryPlan, MemoryStyle, Module, TableElements, TablePlan, TableStyle,
+    Export, MemoryPlan, MemoryStyle, Module, ModuleSyncString, TableElements, TablePlan, TableStyle,
 };
 pub use crate::module_environ::{
     translate_signature, DataInitializer, DataInitializerLocation, FunctionBodyData,

crates/environ/src/module.rs

@@ -3,7 +3,7 @@
 use crate::module_environ::FunctionBodyData;
 use crate::tunables::Tunables;
 use cranelift_codegen::ir;
-use cranelift_entity::{EntityRef, PrimaryMap};
+use cranelift_entity::{EntityRef, PrimaryMap, SecondaryMap};
 use cranelift_wasm::{
     DefinedFuncIndex, DefinedGlobalIndex, DefinedMemoryIndex, DefinedTableIndex, FuncIndex, Global,
     GlobalIndex, Memory, MemoryIndex, SignatureIndex, Table, TableIndex,
@@ -11,6 +11,7 @@ use cranelift_wasm::{
 use indexmap::IndexMap;
 use more_asserts::assert_ge;
 use std::hash::{Hash, Hasher};
+use std::sync::Arc;
 
 /// A WebAssembly table initializer.
 #[derive(Clone, Debug, Hash)]
@@ -128,6 +129,28 @@ impl TablePlan {
     }
 }
 
+/// Allows module strings to be cached as reused across
+/// multiple threads. Useful for debug/trace information.
+#[derive(Debug, Clone)]
+pub struct ModuleSyncString(Option<Arc<String>>);
+
+impl ModuleSyncString {
+    /// Gets optional string reference.
+    pub fn get(&self) -> Option<&str> {
+        self.0.as_deref().map(|s| s.as_str())
+    }
+    /// Constructs the string.
+    pub fn new(s: Option<&str>) -> Self {
+        ModuleSyncString(s.map(|s| Arc::new(s.to_string())))
+    }
+}
+
+impl Default for ModuleSyncString {
+    fn default() -> Self {
+        ModuleSyncString(None)
+    }
+}
+
 /// A translated WebAssembly module, excluding the function bodies and
 /// memory initializers.
 // WARNING: when modifying, make sure that `hash_for_cache` is still valid!
@@ -171,7 +194,10 @@ pub struct Module {
     pub table_elements: Vec<TableElements>,
 
     /// Module name.
-    pub name: Option<String>,
+    pub name: ModuleSyncString,
+
+    /// Function names.
+    pub func_names: SecondaryMap<FuncIndex, ModuleSyncString>,
 }
 
 impl Module {
@@ -190,7 +216,8 @@ impl Module {
             exports: IndexMap::new(),
             start_func: None,
             table_elements: Vec::new(),
-            name: None,
+            name: Default::default(),
+            func_names: SecondaryMap::new(),
         }
     }
 

crates/environ/src/module_environ.rs

@@ -1,5 +1,5 @@
 use crate::func_environ::FuncEnvironment;
-use crate::module::{Export, MemoryPlan, Module, TableElements, TablePlan};
+use crate::module::{Export, MemoryPlan, Module, ModuleSyncString, TableElements, TablePlan};
 use crate::tunables::Tunables;
 use cranelift_codegen::ir;
 use cranelift_codegen::ir::{AbiParam, ArgumentPurpose};
@@ -369,6 +369,11 @@ impl<'data> cranelift_wasm::ModuleEnvironment<'data> for ModuleEnvironment<'data
         });
         Ok(())
     }
+
+    fn declare_func_name(&mut self, func_index: FuncIndex, name: &'data str) -> WasmResult<()> {
+        self.result.module.func_names[func_index] = ModuleSyncString::new(Some(name));
+        Ok(())
+    }
 }
 
 /// Add environment-specific function parameters.

crates/jit/src/compiler.rs

@@ -166,9 +166,16 @@ impl Compiler {
             let body_len = compilation.get(i).body.len();
             self.jit_function_ranges
                 .push((ptr as usize, ptr as usize + body_len));
+            let func_index = module.func_index(i);
+            let func_name = module
+                .func_names
+                .get(func_index)
+                .cloned()
+                .unwrap_or_else(Default::default);
             let tag = jit_function_registry::JITFunctionTag {
                 module_id: module.name.clone(),
-                func_index: i.index(),
+                func_index: func_index.index(),
+                func_name,
             };
             jit_function_registry::register(ptr as usize, ptr as usize + body_len, tag);
         }

crates/jit/src/instantiate.rs

@@ -16,6 +16,7 @@ use wasmtime_environ::entity::{BoxedSlice, PrimaryMap};
 use wasmtime_environ::wasm::{DefinedFuncIndex, SignatureIndex};
 use wasmtime_environ::{
     CompileError, DataInitializer, DataInitializerLocation, Module, ModuleEnvironment,
+    ModuleSyncString,
 };
 use wasmtime_runtime::{
     Export, GdbJitImageRegistration, Imports, InstanceHandle, InstantiationError, VMFunctionBody,
@@ -76,7 +77,7 @@ impl<'data> RawCompiledModule<'data> {
             None
         };
 
-        translation.module.name = module_name.map(|s| s.to_string());
+        translation.module.name = ModuleSyncString::new(module_name);
 
         let (allocated_functions, jt_offsets, relocations, dbg_image) = compiler.compile(
             &translation.module,

crates/runtime/src/jit_function_registry.rs

@@ -3,6 +3,7 @@
 use lazy_static::lazy_static;
 use std::collections::BTreeMap;
 use std::sync::{Arc, RwLock};
+use wasmtime_environ::ModuleSyncString;
 
 lazy_static! {
     static ref REGISTRY: RwLock<JITFunctionRegistry> = RwLock::new(JITFunctionRegistry::default());
@@ -10,13 +11,14 @@ lazy_static! {
 
 #[derive(Clone)]
 pub struct JITFunctionTag {
-    pub module_id: Option<String>,
+    pub module_id: ModuleSyncString,
     pub func_index: usize,
+    pub func_name: ModuleSyncString,
 }
 
 impl std::fmt::Debug for JITFunctionTag {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        if let Some(ref module_id) = self.module_id {
+        if let Some(ref module_id) = self.module_id.get() {
             write!(f, "{}", module_id)?;
         } else {
             write!(f, "(module)")?;