From 9789e88cfe7f50ac46d04175e32df4329d75024e Mon Sep 17 00:00:00 2001
From: Michael Goulet <michael@errs.io>
Date: Sun, 11 Feb 2024 19:17:42 +0000
Subject: [PATCH] Check that the ABI of the instance we are inlining is correct

---
 compiler/rustc_mir_transform/src/inline.rs | 12 ++++++++
 tests/ui/mir/inline-wrong-abi.rs           | 32 ++++++++++++++++++++++
 tests/ui/mir/inline-wrong-abi.stderr       | 12 ++++++++
 3 files changed, 56 insertions(+)
 create mode 100644 tests/ui/mir/inline-wrong-abi.rs
 create mode 100644 tests/ui/mir/inline-wrong-abi.stderr

diff --git a/compiler/rustc_mir_transform/src/inline.rs b/compiler/rustc_mir_transform/src/inline.rs
index be19bd8349ed4..956d855ab81a2 100644
--- a/compiler/rustc_mir_transform/src/inline.rs
+++ b/compiler/rustc_mir_transform/src/inline.rs
@@ -2,6 +2,7 @@
 use crate::deref_separator::deref_finder;
 use rustc_attr::InlineAttr;
 use rustc_const_eval::transform::validate::validate_types;
+use rustc_hir::def::DefKind;
 use rustc_hir::def_id::DefId;
 use rustc_index::bit_set::BitSet;
 use rustc_index::Idx;
@@ -384,6 +385,17 @@ impl<'tcx> Inliner<'tcx> {
                 }
 
                 let fn_sig = self.tcx.fn_sig(def_id).instantiate(self.tcx, args);
+
+                // Additionally, check that the body that we're inlining actually agrees
+                // with the ABI of the trait that the item comes from.
+                if let InstanceDef::Item(instance_def_id) = callee.def
+                    && self.tcx.def_kind(instance_def_id) == DefKind::AssocFn
+                    && let instance_fn_sig = self.tcx.fn_sig(instance_def_id).skip_binder()
+                    && instance_fn_sig.abi() != fn_sig.abi()
+                {
+                    return None;
+                }
+
                 let source_info = SourceInfo { span: fn_span, ..terminator.source_info };
 
                 return Some(CallSite { callee, fn_sig, block: bb, source_info });
diff --git a/tests/ui/mir/inline-wrong-abi.rs b/tests/ui/mir/inline-wrong-abi.rs
new file mode 100644
index 0000000000000..1f61a5dcfbe6d
--- /dev/null
+++ b/tests/ui/mir/inline-wrong-abi.rs
@@ -0,0 +1,32 @@
+// compile-flags: -Zpolymorphize=on -Zinline-mir=yes -Zmir-opt-level=0
+
+#![feature(fn_traits, unboxed_closures)]
+struct Foo<T>(T);
+
+impl<T: Copy> Fn<()> for Foo<T> {
+    extern "C" fn call(&self, _: ()) -> T {
+        //~^ ERROR method `call` has an incompatible type for trait
+        match *self {
+            Foo(t) => t,
+        }
+    }
+}
+
+impl<T: Copy> FnMut<()> for Foo<T> {
+    extern "rust-call" fn call_mut(&mut self, _: ()) -> T {
+        self.call(())
+    }
+}
+
+impl<T: Copy> FnOnce<()> for Foo<T> {
+    type Output = T;
+
+    extern "rust-call" fn call_once(self, _: ()) -> T {
+        self.call(())
+    }
+}
+
+fn main() {
+    let t: u8 = 1;
+    println!("{}", Foo(t)());
+}
diff --git a/tests/ui/mir/inline-wrong-abi.stderr b/tests/ui/mir/inline-wrong-abi.stderr
new file mode 100644
index 0000000000000..e44899eb92a5c
--- /dev/null
+++ b/tests/ui/mir/inline-wrong-abi.stderr
@@ -0,0 +1,12 @@
+error[E0053]: method `call` has an incompatible type for trait
+  --> $DIR/inline-wrong-abi.rs:7:5
+   |
+LL |     extern "C" fn call(&self, _: ()) -> T {
+   |     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ expected "rust-call" fn, found "C" fn
+   |
+   = note: expected signature `extern "rust-call" fn(&Foo<_>, ()) -> _`
+              found signature `extern "C" fn(&Foo<_>, ()) -> _`
+
+error: aborting due to 1 previous error
+
+For more information about this error, try `rustc --explain E0053`.