Add method to get all attributes on a definition

5 files changed, 85 insertions, 30 deletions

diff --git a/compiler/rustc_smir/src/rustc_smir/context.rs b/compiler/rustc_smir/src/rustc_smir/context.rs
index 0e50241fce7..e23f4289e98 100644
--- a/compiler/rustc_smir/src/rustc_smir/context.rs
+++ b/compiler/rustc_smir/src/rustc_smir/context.rs
@@ -232,7 +232,7 @@ impl<'tcx> Context for TablesWrapper<'tcx> {
         &self,
         def_id: stable_mir::DefId,
         attr: &[stable_mir::Symbol],
-    ) -> Vec<stable_mir::ty::Attribute> {
+    ) -> Vec<stable_mir::crate_def::Attribute> {
         let mut tables = self.0.borrow_mut();
         let tcx = tables.tcx;
         let did = tables[def_id];
@@ -242,7 +242,28 @@ impl<'tcx> Context for TablesWrapper<'tcx> {
             .map(|attribute| {
                 let attr_str = rustc_ast_pretty::pprust::attribute_to_string(attribute);
                 let span = attribute.span;
-                stable_mir::ty::Attribute::new(attr_str, span.stable(&mut *tables))
+                stable_mir::crate_def::Attribute::new(attr_str, span.stable(&mut *tables))
+            })
+            .collect()
+    }
+
+    fn get_all_attrs(&self, def_id: stable_mir::DefId) -> Vec<stable_mir::crate_def::Attribute> {
+        let mut tables = self.0.borrow_mut();
+        let tcx = tables.tcx;
+        let did = tables[def_id];
+        let filter_fn = move |a: &&rustc_ast::ast::Attribute| {
+            matches!(a.kind, rustc_ast::ast::AttrKind::Normal(_))
+        };
+        let attrs_iter = if let Some(did) = did.as_local() {
+            tcx.hir().attrs(tcx.local_def_id_to_hir_id(did)).iter().filter(filter_fn)
+        } else {
+            tcx.item_attrs(did).iter().filter(filter_fn)
+        };
+        attrs_iter
+            .map(|attribute| {
+                let attr_str = rustc_ast_pretty::pprust::attribute_to_string(attribute);
+                let span = attribute.span;
+                stable_mir::crate_def::Attribute::new(attr_str, span.stable(&mut *tables))
             })
             .collect()
     }
diff --git a/compiler/stable_mir/src/compiler_interface.rs b/compiler/stable_mir/src/compiler_interface.rs
index a8f2cd7bcc3..5f2d9b96c73 100644
--- a/compiler/stable_mir/src/compiler_interface.rs
+++ b/compiler/stable_mir/src/compiler_interface.rs
@@ -6,12 +6,13 @@
 use std::cell::Cell;
 
 use crate::abi::{FnAbi, Layout, LayoutShape};
+use crate::crate_def::Attribute;
 use crate::mir::alloc::{AllocId, GlobalAlloc};
 use crate::mir::mono::{Instance, InstanceDef, StaticDef};
 use crate::mir::{BinOp, Body, Place, UnOp};
 use crate::target::MachineInfo;
 use crate::ty::{
-    AdtDef, AdtKind, Allocation, Attribute, ClosureDef, ClosureKind, FieldDef, FnDef, ForeignDef,
+    AdtDef, AdtKind, Allocation, ClosureDef, ClosureKind, FieldDef, FnDef, ForeignDef,
     ForeignItemKind, ForeignModule, ForeignModuleDef, GenericArgs, GenericPredicates, Generics,
     ImplDef, ImplTrait, IntrinsicDef, LineInfo, MirConst, PolyFnSig, RigidTy, Span, TraitDecl,
     TraitDef, Ty, TyConst, TyConstId, TyKind, UintTy, VariantDef,
@@ -55,9 +56,15 @@ pub trait Context {
     /// Returns the name of given `DefId`
     fn def_name(&self, def_id: DefId, trimmed: bool) -> Symbol;
 
-    /// Get all attributes with the given attribute name.
+    /// Return attributes with the given attribute name.
+    ///
+    /// Single segmented name like `#[inline]` is specified as `&["inline".to_string()]`.
+    /// Multi-segmented name like `#[rustfmt::skip]` is specified as `&["rustfmt".to_string(), "skip".to_string()]`.
     fn get_attrs_by_path(&self, def_id: DefId, attr: &[Symbol]) -> Vec<Attribute>;
 
+    /// Get all attributes of a definition.
+    fn get_all_attrs(&self, def_id: DefId) -> Vec<Attribute>;
+
     /// Returns printable, human readable form of `Span`
     fn span_to_string(&self, span: Span) -> String;
 
diff --git a/compiler/stable_mir/src/crate_def.rs b/compiler/stable_mir/src/crate_def.rs
index ee6214a47cf..d9b987c28a2 100644
--- a/compiler/stable_mir/src/crate_def.rs
+++ b/compiler/stable_mir/src/crate_def.rs
@@ -1,7 +1,7 @@
 //! Module that define a common trait for things that represent a crate definition,
 //! such as, a function, a trait, an enum, and any other definitions.
 
-use crate::ty::{Attribute, GenericArgs, Span, Ty};
+use crate::ty::{GenericArgs, Span, Ty};
 use crate::{with, Crate, Symbol};
 
 /// A unique identification number for each item accessible for the current compilation unit.
@@ -52,10 +52,19 @@ pub trait CrateDef {
     }
 
     /// Return attributes with the given attribute name.
+    ///
+    /// Single segmented name like `#[inline]` is specified as `&["inline".to_string()]`.
+    /// Multi-segmented name like `#[rustfmt::skip]` is specified as `&["rustfmt".to_string(), "skip".to_string()]`.
     fn attrs_by_path(&self, attr: &[Symbol]) -> Vec<Attribute> {
         let def_id = self.def_id();
         with(|cx| cx.get_attrs_by_path(def_id, attr))
     }
+
+    /// Return all attributes of this definition.
+    fn all_attrs(&self) -> Vec<Attribute> {
+        let def_id = self.def_id();
+        with(|cx| cx.get_all_attrs(def_id))
+    }
 }
 
 /// A trait that can be used to retrieve a definition's type.
@@ -75,6 +84,28 @@ pub trait CrateDefType: CrateDef {
     }
 }
 
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub struct Attribute {
+    value: String,
+    span: Span,
+}
+
+impl Attribute {
+    pub fn new(value: String, span: Span) -> Attribute {
+        Attribute { value, span }
+    }
+
+    /// Get the span of this attribute.
+    pub fn span(&self) -> Span {
+        self.span
+    }
+
+    /// Get the string representation of this attribute.
+    pub fn as_str(&self) -> &str {
+        &self.value
+    }
+}
+
 macro_rules! crate_def {
     ( $(#[$attr:meta])*
       $vis:vis $name:ident $(;)?
diff --git a/compiler/stable_mir/src/ty.rs b/compiler/stable_mir/src/ty.rs
index 46c163e5bf1..01e4f1d1f33 100644
--- a/compiler/stable_mir/src/ty.rs
+++ b/compiler/stable_mir/src/ty.rs
@@ -248,28 +248,6 @@ pub struct Placeholder<T> {
     pub bound: T,
 }
 
-#[derive(Clone, Debug, PartialEq, Eq)]
-pub struct Attribute {
-    value: String,
-    span: Span,
-}
-
-impl Attribute {
-    pub fn new(value: String, span: Span) -> Attribute {
-        Attribute { value, span }
-    }
-
-    /// Get the span of this attribute.
-    pub fn span(&self) -> Span {
-        self.span
-    }
-
-    /// Get the string representation of this attribute.
-    pub fn as_str(&self) -> &str {
-        &self.value
-    }
-}
-
 #[derive(Clone, Copy, PartialEq, Eq)]
 pub struct Span(usize);
 
diff --git a/tests/ui-fulldeps/stable-mir/check_attribute.rs b/tests/ui-fulldeps/stable-mir/check_attribute.rs
index 7fec0e1a90f..be52853a479 100644
--- a/tests/ui-fulldeps/stable-mir/check_attribute.rs
+++ b/tests/ui-fulldeps/stable-mir/check_attribute.rs
@@ -31,6 +31,7 @@ fn test_stable_mir() -> ControlFlow<()> {
     test_builtins(&items);
     test_derive(&items);
     test_tool(&items);
+    test_all_attrs(&items);
 
     ControlFlow::Continue(())
 }
@@ -73,14 +74,21 @@ fn test_tool(items: &CrateItems) {
     assert_eq!(clippy_attrs[0].as_str(), "#[clippy::cyclomatic_complexity = \"100\"]");
 }
 
+fn test_all_attrs(items: &CrateItems) {
+    let target_fn = *get_item(&items, "many_attrs").unwrap();
+    let all_attrs = target_fn.all_attrs();
+    assert_eq!(all_attrs[0].as_str(), "#[inline]");
+    assert_eq!(all_attrs[1].as_str(), "#[allow(unused_variables)]");
+    assert_eq!(all_attrs[2].as_str(), "#[allow(dead_code)]");
+    assert_eq!(all_attrs[3].as_str(), "#[allow(unused_imports)]");
+    assert_eq!(all_attrs[4].as_str(), "#![allow(clippy::filter_map)]");
+}
+
 
 fn get_item<'a>(
     items: &'a CrateItems,
     name: &str,
 ) -> Option<&'a stable_mir::CrateItem> {
-    for item in items {
-        println!("{:?}", item);
-    }
     items.iter().find(|crate_item| crate_item.name() == name)
 }
 
@@ -131,6 +139,16 @@ fn generate_input(path: &str) -> std::io::Result<()> {
         // A clippy tool attribute.
         #[clippy::cyclomatic_complexity = "100"]
         pub fn complex_fn() {{}}
+
+        // A function with many attributes.
+        #[inline]
+        #[allow(unused_variables)]
+        #[allow(dead_code)]
+        #[allow(unused_imports)]
+        fn many_attrs() {{
+            #![allow(clippy::filter_map)]
+            todo!()
+        }}
         "#
     )?;
     Ok(())