From 68fc92242f1e6165f917cce8304c8cf35518f63a Mon Sep 17 00:00:00 2001
From: Adrian Palacios <accorell@amazon.com>
Date: Fri, 8 Mar 2024 19:11:59 +0000
Subject: [PATCH] Add `intrinsic_name` to get plain intrinsic name

---
 compiler/rustc_smir/src/rustc_smir/context.rs | 10 ++++++++++
 compiler/stable_mir/src/compiler_interface.rs |  1 +
 compiler/stable_mir/src/mir/mono.rs           | 11 +++++++++++
 tests/ui-fulldeps/stable-mir/check_defs.rs    | 14 +++++++++++++-
 4 files changed, 35 insertions(+), 1 deletion(-)

diff --git a/compiler/rustc_smir/src/rustc_smir/context.rs b/compiler/rustc_smir/src/rustc_smir/context.rs
index 540bc48354866..fed1d98d5e6ad 100644
--- a/compiler/rustc_smir/src/rustc_smir/context.rs
+++ b/compiler/rustc_smir/src/rustc_smir/context.rs
@@ -587,6 +587,16 @@ impl<'tcx> Context for TablesWrapper<'tcx> {
         }
     }
 
+    /// Retrieve the plain intrinsic name of an instance.
+    ///
+    /// This assumes that the instance is an intrinsic.
+    fn intrinsic_name(&self, def: InstanceDef) -> Symbol {
+        let tables = self.0.borrow_mut();
+        let instance = tables.instances[def];
+        let intrinsic = tables.tcx.intrinsic(instance.def_id()).unwrap();
+        intrinsic.name.to_string()
+    }
+
     fn ty_layout(&self, ty: Ty) -> Result<Layout, Error> {
         let mut tables = self.0.borrow_mut();
         let tcx = tables.tcx;
diff --git a/compiler/stable_mir/src/compiler_interface.rs b/compiler/stable_mir/src/compiler_interface.rs
index 0f7d8d7e083bf..1c51c175d81b5 100644
--- a/compiler/stable_mir/src/compiler_interface.rs
+++ b/compiler/stable_mir/src/compiler_interface.rs
@@ -183,6 +183,7 @@ pub trait Context {
     fn vtable_allocation(&self, global_alloc: &GlobalAlloc) -> Option<AllocId>;
     fn krate(&self, def_id: DefId) -> Crate;
     fn instance_name(&self, def: InstanceDef, trimmed: bool) -> Symbol;
+    fn intrinsic_name(&self, def: InstanceDef) -> Symbol;
 
     /// Return information about the target machine.
     fn target_info(&self) -> MachineInfo;
diff --git a/compiler/stable_mir/src/mir/mono.rs b/compiler/stable_mir/src/mir/mono.rs
index 97f57d2c7b359..38e5776c48cea 100644
--- a/compiler/stable_mir/src/mir/mono.rs
+++ b/compiler/stable_mir/src/mir/mono.rs
@@ -90,6 +90,17 @@ impl Instance {
         with(|context| context.instance_name(self.def, true))
     }
 
+    /// Retrieve the plain intrinsic name of an instance if it's an intrinsic.
+    ///
+    /// The plain name does not include type arguments (as `trimmed_name` does),
+    /// which is more convenient to match with intrinsic symbols.
+    pub fn intrinsic_name(&self) -> Option<Symbol> {
+        match self.kind {
+            InstanceKind::Intrinsic => Some(with(|context| context.intrinsic_name(self.def))),
+            InstanceKind::Item | InstanceKind::Virtual { .. } | InstanceKind::Shim => None,
+        }
+    }
+
     /// Resolve an instance starting from a function definition and generic arguments.
     pub fn resolve(def: FnDef, args: &GenericArgs) -> Result<Instance, crate::Error> {
         with(|context| {
diff --git a/tests/ui-fulldeps/stable-mir/check_defs.rs b/tests/ui-fulldeps/stable-mir/check_defs.rs
index 27b9b059c209b..5bb1053f18798 100644
--- a/tests/ui-fulldeps/stable-mir/check_defs.rs
+++ b/tests/ui-fulldeps/stable-mir/check_defs.rs
@@ -19,6 +19,7 @@ extern crate stable_mir;
 
 use std::assert_matches::assert_matches;
 use mir::{mono::Instance, TerminatorKind::*};
+use stable_mir::mir::mono::InstanceKind;
 use rustc_smir::rustc_internal;
 use stable_mir::ty::{RigidTy, TyKind, Ty, UintTy};
 use stable_mir::*;
@@ -35,9 +36,10 @@ fn test_stable_mir() -> ControlFlow<()> {
     assert_eq!(main_fn.trimmed_name(), "main");
 
     let instances = get_instances(main_fn.body().unwrap());
-    assert_eq!(instances.len(), 2);
+    assert_eq!(instances.len(), 3);
     test_fn(instances[0], "from_u32", "std::char::from_u32", "core");
     test_fn(instances[1], "Vec::<u8>::new", "std::vec::Vec::<u8>::new", "alloc");
+    test_fn(instances[2], "ctpop::<u8>", "std::intrinsics::ctpop::<u8>", "core");
     test_vec_new(instances[1]);
     ControlFlow::Continue(())
 }
@@ -48,6 +50,14 @@ fn test_fn(instance: Instance, expected_trimmed: &str, expected_qualified: &str,
     assert_eq!(&trimmed, expected_trimmed);
     assert_eq!(&qualified, expected_qualified);
 
+    if instance.kind == InstanceKind::Intrinsic {
+        let intrinsic = instance.intrinsic_name().unwrap();
+        let (trimmed_base, _trimmed_args) = trimmed.split_once("::").unwrap();
+        assert_eq!(intrinsic, trimmed_base);
+        return;
+    }
+    assert!(instance.intrinsic_name().is_none());
+
     let item = CrateItem::try_from(instance).unwrap();
     let trimmed = item.trimmed_name();
     let qualified = item.name();
@@ -119,10 +129,12 @@ fn generate_input(path: &str) -> std::io::Result<()> {
     write!(
         file,
         r#"
+        #![feature(core_intrinsics)]
 
         fn main() {{
             let _c = core::char::from_u32(99);
             let _v = Vec::<u8>::new();
+            let _i = std::intrinsics::ctpop::<u8>(0);
         }}
     "#
     )?;