Move wasm_import_module_map handling into shims/wasi

Author: bjorn3

src/machine.rs

@@ -1243,24 +1243,16 @@ impl<'tcx> Machine<'tcx> for MiriMachine<'tcx> {
             // foreign function
             // Any needed call to `goto_block` will be performed by `emulate_foreign_item`.
             let args = ecx.copy_fn_args(args); // FIXME: Should `InPlace` arguments be reset to uninit?
-
-            let link_name = if ecx.tcx.sess.target.is_like_wasm
-                && let Some(module) =
-                    ecx.tcx.wasm_import_module_map(instance.def_id().krate).get(&instance.def_id())
-            {
-                // Adapted from https://github.com/rust-lang/rust/blob/90b65889799733f21ebdf59d96411aa531c5900a/compiler/rustc_codegen_llvm/src/attributes.rs#L549-L562
-                let codegen_fn_attrs = ecx.tcx.codegen_instance_attrs(instance.def);
-                let name = codegen_fn_attrs
-                    .symbol_name
-                    .unwrap_or_else(|| ecx.tcx.item_name(instance.def_id()));
-                // $$$ is unlikely to occur in either the import module name or item name, so use it
-                // as a separator here. It will be split again in emulate_foreign_item_inner for wasi.
-                Symbol::intern(&format!("{}$$${}", module, name))
-            } else {
-                Symbol::intern(ecx.tcx.symbol_name(instance).name)
-            };
-
-            return ecx.emulate_foreign_item(link_name, abi, &args, dest, ret, unwind);
+            let link_name = Symbol::intern(ecx.tcx.symbol_name(instance).name);
+            return ecx.emulate_foreign_item(
+                Some(instance),
+                link_name,
+                abi,
+                &args,
+                dest,
+                ret,
+                unwind,
+            );
         }
 
         if ecx.machine.data_race.as_genmc_ref().is_some()

src/shims/foreign_items.rs

@@ -43,6 +43,7 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
     /// is delegated to another function.
     fn emulate_foreign_item(
         &mut self,
+        instance: Option<Instance<'tcx>>,
         link_name: Symbol,
         abi: &FnAbi<'tcx, Ty<'tcx>>,
         args: &[OpTy<'tcx>],
@@ -73,7 +74,7 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
         let dest = this.force_allocation(dest)?;
 
         // The rest either implements the logic, or falls back to `lookup_exported_symbol`.
-        match this.emulate_foreign_item_inner(link_name, abi, args, &dest)? {
+        match this.emulate_foreign_item_inner(instance, link_name, abi, args, &dest)? {
             EmulateItemResult::NeedsReturn => {
                 trace!("{:?}", this.dump_place(&dest.clone().into()));
                 this.return_to_block(ret)?;
@@ -118,7 +119,7 @@ pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
         ret: Option<mir::BasicBlock>,
         unwind: mir::UnwindAction,
     ) -> InterpResult<'tcx> {
-        let res = self.emulate_foreign_item(sym.0, abi, args, dest, ret, unwind)?;
+        let res = self.emulate_foreign_item(None, sym.0, abi, args, dest, ret, unwind)?;
         assert!(res.is_none(), "DynSyms that delegate are not supported");
         interp_ok(())
     }
@@ -248,6 +249,7 @@ impl<'tcx> EvalContextExtPriv<'tcx> for crate::MiriInterpCx<'tcx> {}
 trait EvalContextExtPriv<'tcx>: crate::MiriInterpCxExt<'tcx> {
     fn emulate_foreign_item_inner(
         &mut self,
+        instance: Option<Instance<'tcx>>,
         link_name: Symbol,
         abi: &FnAbi<'tcx, Ty<'tcx>>,
         args: &[OpTy<'tcx>],
@@ -848,7 +850,7 @@ trait EvalContextExtPriv<'tcx>: crate::MiriInterpCxExt<'tcx> {
                         ),
                     "wasi" =>
                         shims::wasi::foreign_items::EvalContextExt::emulate_foreign_item_inner(
-                            this, link_name, abi, args, dest,
+                            this, instance, link_name, abi, args, dest,
                         ),
                     "windows" =>
                         shims::windows::foreign_items::EvalContextExt::emulate_foreign_item_inner(

src/shims/wasi/foreign_items.rs

@@ -1,5 +1,5 @@
 use rustc_abi::CanonAbi;
-use rustc_middle::ty::Ty;
+use rustc_middle::ty::{Instance, Ty};
 use rustc_span::Symbol;
 use rustc_target::callconv::FnAbi;
 
@@ -14,27 +14,38 @@ impl<'tcx> EvalContextExt<'tcx> for crate::MiriInterpCx<'tcx> {}
 pub trait EvalContextExt<'tcx>: crate::MiriInterpCxExt<'tcx> {
     fn emulate_foreign_item_inner(
         &mut self,
+        instance: Option<Instance<'tcx>>,
         link_name: Symbol,
         abi: &FnAbi<'tcx, Ty<'tcx>>,
         args: &[OpTy<'tcx>],
         dest: &MPlaceTy<'tcx>,
     ) -> InterpResult<'tcx, EmulateItemResult> {
         let this = self.eval_context_mut();
 
-        let (interface, name) = if let Some((module, name)) = link_name.as_str().split_once("$$$") {
+        let (interface, name) = if let Some(instance) = instance
+            && let Some(module) =
+                this.tcx.wasm_import_module_map(instance.def_id().krate).get(&instance.def_id())
+        {
+            // Adapted from https://github.com/rust-lang/rust/blob/90b65889799733f21ebdf59d96411aa531c5900a/compiler/rustc_codegen_llvm/src/attributes.rs#L549-L562
+            let codegen_fn_attrs = this.tcx.codegen_instance_attrs(instance.def);
+            let name = codegen_fn_attrs
+                .symbol_name
+                .unwrap_or_else(|| this.tcx.item_name(instance.def_id()));
+
             // According to the component model, the version should be matched as semver, but for
             // simplicity we strip the version entirely for now. Once we support wasm-wasip3 it may
             // become actually important to match on the version, but for now it shouldn't matter.
-            let (module, _version) = module
+            let (interface, _version) = module
                 .split_once('@')
                 .ok_or_else(|| err_unsup_format!("module name {module} must contain a version"))?;
-            (Some(module), name)
+
+            (Some(interface), name)
         } else {
             // This item is provided by wasi-libc, not imported from the wasi runtime
-            (None, link_name.as_str())
+            (None, link_name)
         };
 
-        match (interface, name) {
+        match (interface, name.as_str()) {
             // Allocation
             (None, "posix_memalign") => {
                 let [memptr, align, size] =