From 225796a2df27b3a58bb163c7f59c51988b6dc9f3 Mon Sep 17 00:00:00 2001
From: Adwin White <adwinw01@gmail.com>
Date: Thu, 20 Jun 2024 11:08:12 +0800
Subject: [PATCH] Add method to get `FnAbi` of function pointer

---
 compiler/rustc_smir/src/rustc_smir/context.rs |  7 ++++++
 compiler/stable_mir/src/compiler_interface.rs |  3 +++
 compiler/stable_mir/src/ty.rs                 | 12 +++++++++-
 tests/ui-fulldeps/stable-mir/check_abi.rs     | 23 +++++++++++++++++++
 4 files changed, 44 insertions(+), 1 deletion(-)

diff --git a/compiler/rustc_smir/src/rustc_smir/context.rs b/compiler/rustc_smir/src/rustc_smir/context.rs
index 9afd507ce113a..dde5e30c3d071 100644
--- a/compiler/rustc_smir/src/rustc_smir/context.rs
+++ b/compiler/rustc_smir/src/rustc_smir/context.rs
@@ -533,6 +533,13 @@ impl<'tcx> Context for TablesWrapper<'tcx> {
         Ok(tables.fn_abi_of_instance(instance, List::empty())?.stable(&mut *tables))
     }
 
+    fn fn_ptr_abi(&self, fn_ptr: PolyFnSig) -> Result<FnAbi, Error> {
+        let mut tables = self.0.borrow_mut();
+        let tcx = tables.tcx;
+        let sig = fn_ptr.internal(&mut *tables, tcx);
+        Ok(tables.fn_abi_of_fn_ptr(sig, List::empty())?.stable(&mut *tables))
+    }
+
     fn instance_def_id(&self, def: InstanceDef) -> stable_mir::DefId {
         let mut tables = self.0.borrow_mut();
         let def_id = tables.instances[def].def_id();
diff --git a/compiler/stable_mir/src/compiler_interface.rs b/compiler/stable_mir/src/compiler_interface.rs
index 3e138e3c2e04e..44dbf549c1a73 100644
--- a/compiler/stable_mir/src/compiler_interface.rs
+++ b/compiler/stable_mir/src/compiler_interface.rs
@@ -215,6 +215,9 @@ pub trait Context {
     /// Get an instance ABI.
     fn instance_abi(&self, def: InstanceDef) -> Result<FnAbi, Error>;
 
+    /// Get the ABI of a function pointer.
+    fn fn_ptr_abi(&self, fn_ptr: PolyFnSig) -> Result<FnAbi, Error>;
+
     /// Get the layout of a type.
     fn ty_layout(&self, ty: Ty) -> Result<Layout, Error>;
 
diff --git a/compiler/stable_mir/src/ty.rs b/compiler/stable_mir/src/ty.rs
index 35927237281a0..8c120a96e75b1 100644
--- a/compiler/stable_mir/src/ty.rs
+++ b/compiler/stable_mir/src/ty.rs
@@ -2,7 +2,7 @@ use super::{
     mir::{Body, Mutability, Safety},
     with, DefId, Error, Symbol,
 };
-use crate::abi::Layout;
+use crate::abi::{FnAbi, Layout};
 use crate::crate_def::{CrateDef, CrateDefType};
 use crate::mir::alloc::{read_target_int, read_target_uint, AllocId};
 use crate::mir::mono::StaticDef;
@@ -996,6 +996,16 @@ pub struct AliasTerm {
 
 pub type PolyFnSig = Binder<FnSig>;
 
+impl PolyFnSig {
+    /// Compute a `FnAbi` suitable for indirect calls, i.e. to `fn` pointers.
+    ///
+    /// NB: this doesn't handle virtual calls - those should use `Instance::fn_abi`
+    /// instead, where the instance is an `InstanceKind::Virtual`.
+    pub fn fn_ptr_abi(self) -> Result<FnAbi, Error> {
+        with(|cx| cx.fn_ptr_abi(self))
+    }
+}
+
 #[derive(Clone, Debug, Eq, PartialEq)]
 pub struct FnSig {
     pub inputs_and_output: Vec<Ty>,
diff --git a/tests/ui-fulldeps/stable-mir/check_abi.rs b/tests/ui-fulldeps/stable-mir/check_abi.rs
index 359dd4146baf5..7518ea902ecf5 100644
--- a/tests/ui-fulldeps/stable-mir/check_abi.rs
+++ b/tests/ui-fulldeps/stable-mir/check_abi.rs
@@ -54,6 +54,21 @@ fn test_stable_mir() -> ControlFlow<()> {
     let variadic_fn = *get_item(&items, (ItemKind::Fn, "variadic_fn")).unwrap();
     check_variadic(variadic_fn);
 
+    // Extract function pointers.
+    let fn_ptr_holder = *get_item(&items, (ItemKind::Fn, "fn_ptr_holder")).unwrap();
+    let fn_ptr_holder_instance = Instance::try_from(fn_ptr_holder).unwrap();
+    let body = fn_ptr_holder_instance.body().unwrap();
+    let args = body.arg_locals();
+
+    // Test fn_abi of function pointer version.
+    let ptr_fn_abi = args[0].ty.kind().fn_sig().unwrap().fn_ptr_abi().unwrap();
+    assert_eq!(ptr_fn_abi, fn_abi);
+
+    // Test variadic_fn of function pointer version.
+    let ptr_variadic_fn_abi = args[1].ty.kind().fn_sig().unwrap().fn_ptr_abi().unwrap();
+    assert!(ptr_variadic_fn_abi.c_variadic);
+    assert_eq!(ptr_variadic_fn_abi.args.len(), 1);
+
     ControlFlow::Continue(())
 }
 
@@ -164,6 +179,14 @@ fn generate_input(path: &str) -> std::io::Result<()> {
         pub unsafe extern "C" fn variadic_fn(n: usize, mut args: ...) -> usize {{
             0
         }}
+
+        pub type ComplexFn = fn([u8; 0], char, NonZero<u8>) -> Result<usize, &'static str>;
+        pub type VariadicFn = unsafe extern "C" fn(usize, ...) -> usize;
+
+        pub fn fn_ptr_holder(complex_fn: ComplexFn, variadic_fn: VariadicFn) {{
+            // We only care about the signature.
+            todo!()
+        }}
         "#
     )?;
     Ok(())