Make Wasmtime's FuncKey one-to-one with Cranelift's ir::UserExternalName (#11415)

* Make Wasmtime's `FuncKey` one-to-one with Cranelift's `ir::UserExternalName`

`FuncKey`, which used to be called `CompileKey`, is now one-to-one with
`cranelift_codegen::ir::UserExternalName`, and is used for not just identifying
compilation objects but also relocations and call-graph edges. This allows us to
determine the `StaticModuleIndex` and `DefinedFuncIndex` pair for any
`cranelift_codegen::ir::FuncRef`, regardless of inlining depth, which fixes some
fuzz bugs on OSS-Fuzz.

This continues pushing on the idea that Wasmtime's compilation orchestration and
linking should be relatively agnostic to the kinds of things it is actually
compiling and linking, allowing us to tweak, add, and remove new kinds of
`FuncKey`s more easily. Adding a new `FuncKey` should not require modifying
relocation resolution, for example, just a little bit of code to run the
associated compilation and optionally some code to extract metadata into our
final artifacts for querying at runtime. Everything in between should Just
Continue Working. We still aren't all the way there yet, but this does bring us
a little bit closer.

Finally, in Cranelift's inlining pass, this adds a check that a block is
inserted in the layout before attempting to remove it from the layout, which
would otherwise cause panics. This was triggered by multi-level inlining and
now-unreachable blocks in the inner callees.

I'll note that this does update basically all of the disas tests, or at least
nearly all of them that make function calls. This is because the namespace/index
numbering pair changed slightly to align with `FuncKey`, but that should pretty
much be the only changes.

* remove debug info from panic message, it is only available in some `cfg`s

* fill out module doc comment

* Fix compilation without `component-model` feature

* Fix some more cfg compilations

* cargo fmt

* fix a wrong `&dyn Any` auto coercion; add helpful debug logging and assertions for this kind of thing

Author: fitzgen

Cargo.lock

@@ -445,7 +445,7 @@ fn inline_one(
     // terminator, so do a quick pass over the inlined blocks and remove any
     // empty blocks from the caller's layout.
     for block in entity_map.iter_inlined_blocks(func) {
-        if func.layout.first_inst(block).is_none() {
+        if func.layout.is_block_inserted(block) && func.layout.first_inst(block).is_none() {
             func.layout.remove_block(block);
         }
     }

cranelift/codegen/src/inline.rs

@@ -9,12 +9,12 @@ use cranelift_codegen::ir::condcodes::IntCC;
 use cranelift_codegen::ir::{self, InstBuilder, MemFlags, Value};
 use cranelift_codegen::isa::{CallConv, TargetIsa};
 use cranelift_frontend::FunctionBuilder;
-use wasmtime_environ::fact::PREPARE_CALL_FIXED_PARAMS;
 use wasmtime_environ::{CompiledFunctionBody, component::*};
 use wasmtime_environ::{
     EntityRef, HostCall, ModuleInternedTypeIndex, PtrSize, TrapSentinel, Tunables, WasmFuncType,
     WasmValType,
 };
+use wasmtime_environ::{FuncKey, fact::PREPARE_CALL_FIXED_PARAMS};
 
 struct TrampolineCompiler<'a> {
     compiler: &'a Compiler,
@@ -1281,7 +1281,7 @@ impl ComponentCompiler for Compiler {
         &self,
         component: &ComponentTranslation,
         types: &ComponentTypesBuilder,
-        index: TrampolineIndex,
+        key: FuncKey,
         tunables: &Tunables,
         _symbol: &str,
     ) -> Result<AllCallFunc<CompiledFunctionBody>> {
@@ -1292,7 +1292,7 @@ impl ComponentCompiler for Compiler {
                 &mut compiler,
                 &component.component,
                 types,
-                index,
+                key.unwrap_component_trampoline(),
                 abi,
                 tunables,
             );
@@ -1325,12 +1325,12 @@ impl ComponentCompiler for Compiler {
                 );
             }
 
-            c.translate(&component.trampolines[index]);
+            c.translate(&component.trampolines[key.unwrap_component_trampoline()]);
             c.builder.finalize();
             compiler.cx.abi = Some(abi);
 
             Ok(CompiledFunctionBody {
-                code: Box::new(Some(compiler.cx)),
+                code: super::box_dyn_any_compiler_context(Some(compiler.cx)),
                 needs_gc_heap: false,
             })
         };

crates/cranelift/src/compiler.rs

@@ -498,8 +498,10 @@ impl AddressTransform {
         >,
         wasm_file: wasmtime_environ::WasmFileInfo,
     ) -> Self {
+        use cranelift_entity::EntityRef;
+
         let mut translations = wasmtime_environ::PrimaryMap::new();
-        let mut translation = wasmtime_environ::ModuleTranslation::default();
+        let mut translation = wasmtime_environ::ModuleTranslation::new(StaticModuleIndex::new(0));
         translation.debuginfo.wasm_file = wasm_file;
         translation
             .module

crates/cranelift/src/compiler/component.rs

@@ -121,16 +121,20 @@ fn read_dwarf_package_from_bytes<'data>(
     let mut validator = wasmparser::Validator::new();
     let parser = wasmparser::Parser::new(0);
     let mut types = wasmtime_environ::ModuleTypesBuilder::new(&validator);
-    let translation =
-        match wasmtime_environ::ModuleEnvironment::new(tunables, &mut validator, &mut types)
-            .translate(parser, dwp_bytes)
-        {
-            Ok(translation) => translation,
-            Err(e) => {
-                log::warn!("failed to parse wasm dwarf package: {e:?}");
-                return None;
-            }
-        };
+    let translation = match wasmtime_environ::ModuleEnvironment::new(
+        tunables,
+        &mut validator,
+        &mut types,
+        StaticModuleIndex::from_u32(0),
+    )
+    .translate(parser, dwp_bytes)
+    {
+        Ok(translation) => translation,
+        Err(e) => {
+            log::warn!("failed to parse wasm dwarf package: {e:?}");
+            return None;
+        }
+    };
 
     match load_dwp(translation, buffer) {
         Ok(package) => Some(package),

crates/cranelift/src/debug/transform/address_transform.rs

@@ -23,10 +23,10 @@ use std::mem;
 use wasmparser::{Operator, WasmFeatures};
 use wasmtime_environ::{
     BuiltinFunctionIndex, DataIndex, DefinedFuncIndex, ElemIndex, EngineOrModuleTypeIndex,
-    FuncIndex, GlobalIndex, IndexType, Memory, MemoryIndex, Module, ModuleInternedTypeIndex,
-    ModuleTranslation, ModuleTypesBuilder, PtrSize, Table, TableIndex, TripleExt, Tunables,
-    TypeConvert, TypeIndex, VMOffsets, WasmCompositeInnerType, WasmFuncType, WasmHeapTopType,
-    WasmHeapType, WasmRefType, WasmResult, WasmValType,
+    FuncIndex, FuncKey, GlobalIndex, IndexType, Memory, MemoryIndex, Module,
+    ModuleInternedTypeIndex, ModuleTranslation, ModuleTypesBuilder, PtrSize, Table, TableIndex,
+    TripleExt, Tunables, TypeConvert, TypeIndex, VMOffsets, WasmCompositeInnerType, WasmFuncType,
+    WasmHeapTopType, WasmHeapType, WasmRefType, WasmResult, WasmValType,
 };
 use wasmtime_environ::{FUNCREF_INIT_BIT, FUNCREF_MASK};
 use wasmtime_math::f64_cvt_to_int_bounds;
@@ -53,17 +53,17 @@ impl BuiltinFunctions {
         }
     }
 
-    fn load_builtin(&mut self, func: &mut Function, index: BuiltinFunctionIndex) -> ir::FuncRef {
-        let cache = &mut self.builtins[index.index() as usize];
+    fn load_builtin(&mut self, func: &mut Function, builtin: BuiltinFunctionIndex) -> ir::FuncRef {
+        let cache = &mut self.builtins[builtin.index() as usize];
         if let Some(f) = cache {
             return *f;
         }
-        let signature = func.import_signature(self.types.wasm_signature(index));
-        let name =
-            ir::ExternalName::User(func.declare_imported_user_function(ir::UserExternalName {
-                namespace: crate::NS_WASMTIME_BUILTIN,
-                index: index.index(),
-            }));
+        let signature = func.import_signature(self.types.wasm_signature(builtin));
+        let key = FuncKey::WasmToBuiltinTrampoline(builtin);
+        let (namespace, index) = key.into_raw_parts();
+        let name = ir::ExternalName::User(
+            func.declare_imported_user_function(ir::UserExternalName { namespace, index }),
+        );
         let f = func.import_function(ir::ExtFuncData {
             name,
             signature,
@@ -1265,15 +1265,18 @@ impl FuncEnvironment<'_> {
         def_func_index: DefinedFuncIndex,
     ) -> ir::FuncRef {
         let func_index = self.module.func_index(def_func_index);
+
         let ty = self.module.functions[func_index]
             .signature
             .unwrap_module_type_index();
         let signature = self.get_or_create_interned_sig_ref(func, ty);
-        let name =
-            ir::ExternalName::User(func.declare_imported_user_function(ir::UserExternalName {
-                namespace: crate::NS_WASM_FUNC,
-                index: func_index.as_u32(),
-            }));
+
+        let key = FuncKey::DefinedWasmFunction(self.translation.module_index, def_func_index);
+        let (namespace, index) = key.into_raw_parts();
+        let name = ir::ExternalName::User(
+            func.declare_imported_user_function(ir::UserExternalName { namespace, index }),
+        );
+
         func.import_function(ir::ExtFuncData {
             name,
             signature,
@@ -1288,15 +1291,20 @@ impl FuncEnvironment<'_> {
     ) -> ir::FuncRef {
         assert!(self.module.is_imported_function(func_index));
         assert!(self.translation.known_imported_functions[func_index].is_some());
+
         let ty = self.module.functions[func_index]
             .signature
             .unwrap_module_type_index();
         let signature = self.get_or_create_interned_sig_ref(func, ty);
-        let name =
-            ir::ExternalName::User(func.declare_imported_user_function(ir::UserExternalName {
-                namespace: crate::NS_WASM_FUNC,
-                index: func_index.as_u32(),
-            }));
+
+        let (module, def_func_index) =
+            self.translation.known_imported_functions[func_index].unwrap();
+        let key = FuncKey::DefinedWasmFunction(module, def_func_index);
+        let (namespace, index) = key.into_raw_parts();
+        let name = ir::ExternalName::User(
+            func.declare_imported_user_function(ir::UserExternalName { namespace, index }),
+        );
+
         func.import_function(ir::ExtFuncData {
             name,
             signature,

crates/cranelift/src/debug/transform/mod.rs

@@ -25,8 +25,8 @@ use cranelift_entity::PrimaryMap;
 
 use target_lexicon::Architecture;
 use wasmtime_environ::{
-    BuiltinFunctionIndex, FlagValue, FuncIndex, RelocationTarget, Trap, TrapInformation, Tunables,
-    WasmFuncType, WasmHeapTopType, WasmHeapType, WasmValType,
+    BuiltinFunctionIndex, FlagValue, FuncKey, Trap, TrapInformation, Tunables, WasmFuncType,
+    WasmHeapTopType, WasmHeapType, WasmValType,
 };
 
 pub use builder::builder;
@@ -223,29 +223,13 @@ fn reference_type(wasm_ht: WasmHeapType, pointer_type: ir::Type) -> ir::Type {
 
 // List of namespaces which are processed in `mach_reloc_to_reloc` below.
 
-/// Namespace corresponding to wasm functions, the index is the index of the
-/// defined function that's being referenced.
-pub const NS_WASM_FUNC: u32 = 0;
-
-/// Namespace for builtin function trampolines. The index is the index of the
-/// builtin that's being referenced. These trampolines invoke the real host
-/// function through an indirect function call loaded by the `VMContext`.
-pub const NS_WASMTIME_BUILTIN: u32 = 1;
-
-/// Namespace used to when a call from Pulley to the host is being made. This is
-/// used with a `colocated: false` name to trigger codegen for a special opcode
-/// for pulley-to-host communication. The index of the functions used in this
-/// namespace correspond to the function signature of `for_each_host_signature!`
-/// in the pulley_interpreter crate.
-pub const NS_PULLEY_HOSTCALL: u32 = 2;
-
 /// A record of a relocation to perform.
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct Relocation {
     /// The relocation code.
     pub reloc: binemit::Reloc,
     /// Relocation target.
-    pub reloc_target: RelocationTarget,
+    pub reloc_target: FuncKey,
     /// The offset where to apply the relocation.
     pub offset: binemit::CodeOffset,
     /// The addend to add to the relocation value.
@@ -312,14 +296,7 @@ fn mach_reloc_to_reloc(
     let reloc_target = match *target {
         FinalizedRelocTarget::ExternalName(ExternalName::User(user_func_ref)) => {
             let name = &name_map[user_func_ref];
-            match name.namespace {
-                NS_WASM_FUNC => RelocationTarget::Wasm(FuncIndex::from_u32(name.index)),
-                NS_WASMTIME_BUILTIN => {
-                    RelocationTarget::Builtin(BuiltinFunctionIndex::from_u32(name.index))
-                }
-                NS_PULLEY_HOSTCALL => RelocationTarget::PulleyHostcall(name.index),
-                _ => panic!("unknown namespace {}", name.namespace),
-            }
+            FuncKey::from_raw_parts(name.namespace, name.index)
         }
         FinalizedRelocTarget::ExternalName(ExternalName::LibCall(libcall)) => {
             // We should have avoided any code that needs this style of libcalls

crates/cranelift/src/func_environ.rs

@@ -13,7 +13,7 @@
 //! function body, the imported wasm function do not. The trampolines symbol
 //! names have format "_trampoline_N", where N is `SignatureIndex`.
 
-use crate::{CompiledFunction, RelocationTarget};
+use crate::CompiledFunction;
 use anyhow::Result;
 use cranelift_codegen::TextSectionBuilder;
 use cranelift_codegen::isa::unwind::{UnwindInfo, systemv};
@@ -23,8 +23,8 @@ use gimli::write::{Address, EhFrame, EndianVec, FrameTable, Writer};
 use object::write::{Object, SectionId, StandardSegment, Symbol, SymbolId, SymbolSection};
 use object::{Architecture, SectionFlags, SectionKind, SymbolFlags, SymbolKind, SymbolScope};
 use std::ops::Range;
-use wasmtime_environ::obj;
 use wasmtime_environ::{Compiler, TripleExt};
+use wasmtime_environ::{FuncKey, obj};
 
 const TEXT_SECTION_NAME: &[u8] = b".text";
 
@@ -120,7 +120,7 @@ impl<'a> ModuleTextBuilder<'a> {
         &mut self,
         name: &str,
         compiled_func: &'a CompiledFunction,
-        resolve_reloc_target: impl Fn(wasmtime_environ::RelocationTarget) -> usize,
+        resolve_reloc_target: impl Fn(wasmtime_environ::FuncKey) -> usize,
     ) -> (SymbolId, Range<u64>) {
         let body = compiled_func.buffer.data();
         let alignment = compiled_func.alignment;
@@ -146,65 +146,49 @@ impl<'a> ModuleTextBuilder<'a> {
 
         for r in compiled_func.relocations() {
             let reloc_offset = off + u64::from(r.offset);
-            match r.reloc_target {
-                // Relocations against user-defined functions means that this is
-                // a relocation against a module-local function, typically a
-                // call between functions. The `text` field is given priority to
-                // resolve this relocation before we actually emit an object
-                // file, but if it can't handle it then we pass through the
-                // relocation.
-                RelocationTarget::Wasm(_) | RelocationTarget::Builtin(_) => {
-                    let target = resolve_reloc_target(r.reloc_target);
-                    if self
-                        .text
-                        .resolve_reloc(reloc_offset, r.reloc, r.addend, target)
-                    {
-                        continue;
-                    }
-
-                    // At this time it's expected that all relocations are
-                    // handled by `text.resolve_reloc`, and anything that isn't
-                    // handled is a bug in `text.resolve_reloc` or something
-                    // transitively there. If truly necessary, though, then this
-                    // loop could also be updated to forward the relocation to
-                    // the final object file as well.
-                    panic!(
-                        "unresolved relocation could not be processed against \
-                         {:?}: {r:?}",
-                        r.reloc_target,
-                    );
-                }
 
-                // This relocation is used to fill in which hostcall id is
-                // desired within the `call_indirect_host` opcode of Pulley
-                // itself. The relocation target is the start of the instruction
-                // and the goal is to insert the static signature number, `n`,
-                // into the instruction.
-                //
-                // At this time the instruction looks like:
-                //
-                //      +------+------+------+------+
-                //      | OP   | OP_EXTENDED |  N   |
-                //      +------+------+------+------+
-                //
-                // This 4-byte encoding has `OP` indicating this is an "extended
-                // opcode" where `OP_EXTENDED` is a 16-bit extended opcode.
-                // The `N` byte is the index of the signature being called and
-                // is what's b eing filled in.
-                //
-                // See the `test_call_indirect_host_width` in
-                // `pulley/tests/all.rs` for this guarantee as well.
-                RelocationTarget::PulleyHostcall(n) => {
-                    #[cfg(feature = "pulley")]
-                    {
-                        use pulley_interpreter::encode::Encode;
-                        assert_eq!(pulley_interpreter::CallIndirectHost::WIDTH, 4);
-                    }
-                    let byte = u8::try_from(n).unwrap();
-                    self.text.write(reloc_offset + 3, &[byte]);
+            // This relocation is used to fill in which hostcall id is
+            // desired within the `call_indirect_host` opcode of Pulley
+            // itself. The relocation target is the start of the instruction
+            // and the goal is to insert the static signature number, `n`,
+            // into the instruction.
+            //
+            // At this time the instruction looks like:
+            //
+            //      +------+------+------+------+
+            //      | OP   | OP_EXTENDED |  N   |
+            //      +------+------+------+------+
+            //
+            // This 4-byte encoding has `OP` indicating this is an "extended
+            // opcode" where `OP_EXTENDED` is a 16-bit extended opcode.
+            // The `N` byte is the index of the signature being called and
+            // is what's b eing filled in.
+            //
+            // See the `test_call_indirect_host_width` in
+            // `pulley/tests/all.rs` for this guarantee as well.
+            if let FuncKey::PulleyHostCall(host_call) = r.reloc_target {
+                #[cfg(feature = "pulley")]
+                {
+                    use pulley_interpreter::encode::Encode;
+                    assert_eq!(pulley_interpreter::CallIndirectHost::WIDTH, 4);
                 }
-            };
+                let n = host_call.index();
+                let byte = u8::try_from(n).unwrap();
+                self.text.write(reloc_offset + 3, &[byte]);
+                continue;
+            }
+
+            let target = resolve_reloc_target(r.reloc_target);
+            if self
+                .text
+                .resolve_reloc(reloc_offset, r.reloc, r.addend, target)
+            {
+                continue;
+            }
+
+            panic!("failed to resolve relocation: {r:?} -> {target}");
         }
+
         (symbol_id, off..off + body_len)
     }