Make inlay hints work in attributed items

crates/hir/src/semantics.rs

@@ -16,8 +16,9 @@ use itertools::Itertools;
 use rustc_hash::{FxHashMap, FxHashSet};
 use smallvec::{smallvec, SmallVec};
 use syntax::{
+    algo::skip_trivia_token,
     ast::{self, GenericParamsOwner, LoopBodyOwner},
-    match_ast, AstNode, SyntaxNode, SyntaxNodePtr, SyntaxToken, TextRange, TextSize,
+    match_ast, AstNode, Direction, SyntaxNode, SyntaxNodePtr, SyntaxToken, TextRange, TextSize,
 };
 
 use crate::{
@@ -184,6 +185,11 @@ impl<'db, DB: HirDatabase> Semantics<'db, DB> {
         self.imp.descend_into_macros(token)
     }
 
+    /// Maps a node down by mapping its first and last token down.
+    pub fn descend_node_into_attributes<N: AstNode>(&self, node: N) -> SmallVec<[N; 1]> {
+        self.imp.descend_node_into_attributes(node)
+    }
+
     pub fn hir_file_for(&self, syntax_node: &SyntaxNode) -> HirFileId {
         self.imp.find_file(syntax_node.clone()).file_id
     }
@@ -192,6 +198,10 @@ impl<'db, DB: HirDatabase> Semantics<'db, DB> {
         self.imp.original_range(node)
     }
 
+    pub fn original_range_opt(&self, node: &SyntaxNode) -> Option<FileRange> {
+        self.imp.original_range_opt(node)
+    }
+
     pub fn diagnostics_display_range(&self, diagnostics: InFile<SyntaxNodePtr>) -> FileRange {
         self.imp.diagnostics_display_range(diagnostics)
     }
@@ -471,16 +481,69 @@ impl<'db> SemanticsImpl<'db> {
         )
     }
 
+    // This might not be the correct way to due this, but it works for now
+    fn descend_node_into_attributes<N: AstNode>(&self, node: N) -> SmallVec<[N; 1]> {
+        let mut res = smallvec![];
+        let tokens = (|| {
+            let first = skip_trivia_token(node.syntax().first_token()?, Direction::Next)?;
+            let last = skip_trivia_token(node.syntax().last_token()?, Direction::Prev)?;
+            Some((first, last))
+        })();
+        let (first, last) = match tokens {
+            Some(it) => it,
+            None => return res,
+        };
+
+        if first == last {
+            self.descend_into_macros_impl(first, |InFile { value, .. }| {
+                if let Some(node) = value.ancestors().find_map(N::cast) {
+                    res.push(node)
+                }
+            });
+        } else {
+            // Descend first and last token, then zip them to look for the node they belong to
+            let mut scratch: SmallVec<[_; 1]> = smallvec![];
+            self.descend_into_macros_impl(first, |token| {
+                scratch.push(token);
+            });
+
+            let mut scratch = scratch.into_iter();
+            self.descend_into_macros_impl(last, |InFile { value: last, file_id: last_fid }| {
+                if let Some(InFile { value: first, file_id: first_fid }) = scratch.next() {
+                    if first_fid == last_fid {
+                        if let Some(p) = first.parent() {
+                            let range = first.text_range().cover(last.text_range());
+                            let node = find_root(&p)
+                                .covering_element(range)
+                                .ancestors()
+                                .take_while(|it| it.text_range() == range)
+                                .find_map(N::cast);
+                            if let Some(node) = node {
+                                res.push(node);
+                            }
+                        }
+                    }
+                }
+            });
+        }
+        res
+    }
+
     fn descend_into_macros(&self, token: SyntaxToken) -> SmallVec<[SyntaxToken; 1]> {
+        let mut res = smallvec![];
+        self.descend_into_macros_impl(token, |InFile { value, .. }| res.push(value));
+        res
+    }
+
+    fn descend_into_macros_impl(&self, token: SyntaxToken, mut f: impl FnMut(InFile<SyntaxToken>)) {
         let _p = profile::span("descend_into_macros");
         let parent = match token.parent() {
             Some(it) => it,
-            None => return smallvec![token],
+            None => return,
         };
         let sa = self.analyze(&parent);
         let mut queue = vec![InFile::new(sa.file_id, token)];
         let mut cache = self.expansion_info_cache.borrow_mut();
-        let mut res = smallvec![];
         // Remap the next token in the queue into a macro call its in, if it is not being remapped
         // either due to not being in a macro-call or because its unused push it into the result vec,
         // otherwise push the remapped tokens back into the queue as they can potentially be remapped again.
@@ -546,10 +609,9 @@ impl<'db> SemanticsImpl<'db> {
             .is_none();
 
             if was_not_remapped {
-                res.push(token.value)
+                f(token)
             }
         }
-        res
     }
 
     // Note this return type is deliberate as [`find_nodes_at_offset_with_descend`] wants to stop
@@ -580,6 +642,11 @@ impl<'db> SemanticsImpl<'db> {
         node.as_ref().original_file_range(self.db.upcast())
     }
 
+    fn original_range_opt(&self, node: &SyntaxNode) -> Option<FileRange> {
+        let node = self.find_file(node.clone());
+        node.as_ref().original_file_range_opt(self.db.upcast())
+    }
+
     fn diagnostics_display_range(&self, src: InFile<SyntaxNodePtr>) -> FileRange {
         let root = self.db.parse_or_expand(src.file_id).unwrap();
         let node = src.value.to_node(&root);

crates/ide/src/inlay_hints.rs

@@ -1,7 +1,7 @@
 use either::Either;
 use hir::{known, Callable, HasVisibility, HirDisplay, Semantics, TypeInfo};
-use ide_db::helpers::FamousDefs;
 use ide_db::RootDatabase;
+use ide_db::{base_db::FileRange, helpers::FamousDefs};
 use stdx::to_lower_snake_case;
 use syntax::{
     ast::{self, ArgListOwner, AstNode, NameOwner},
@@ -79,7 +79,7 @@ pub(crate) fn inlay_hints(
                 _ => (),
             }
         } else if let Some(it) = ast::IdentPat::cast(node.clone()) {
-            get_bind_pat_hints(&mut res, &sema, config, it);
+            get_bind_pat_hints(&mut res, &sema, config, &it);
         }
     }
     res
@@ -99,7 +99,9 @@ fn get_chaining_hints(
         return None;
     }
 
-    let krate = sema.scope(expr.syntax()).module().map(|it| it.krate());
+    let descended = sema.descend_node_into_attributes(expr.clone()).pop();
+    let desc_expr = descended.as_ref().unwrap_or(expr);
+    let krate = sema.scope(desc_expr.syntax()).module().map(|it| it.krate());
     let famous_defs = FamousDefs(sema, krate);
 
     let mut tokens = expr
@@ -121,7 +123,7 @@ fn get_chaining_hints(
             next_next = tokens.next()?.kind();
         }
         if next_next == T![.] {
-            let ty = sema.type_of_expr(expr)?.original;
+            let ty = sema.type_of_expr(desc_expr)?.original;
             if ty.is_unknown() {
                 return None;
             }
@@ -133,7 +135,7 @@ fn get_chaining_hints(
                 }
             }
             acc.push(InlayHint {
-                range: sema.original_range(expr.syntax()).range,
+                range: expr.syntax().text_range(),
                 kind: InlayKind::ChainingHint,
                 label: hint_iterator(sema, &famous_defs, config, &ty).unwrap_or_else(|| {
                     ty.display_truncated(sema.db, config.max_length).to_string().into()
@@ -160,6 +162,8 @@ fn get_param_name_hints(
         .into_iter()
         .zip(arg_list.args())
         .filter_map(|((param, _ty), arg)| {
+            // Only annotate hints for expressions that exist in the original file
+            let range = sema.original_range_opt(arg.syntax())?;
             let param_name = match param? {
                 Either::Left(_) => "self".to_string(),
                 Either::Right(pat) => match pat {
@@ -167,11 +171,13 @@ fn get_param_name_hints(
                     _ => return None,
                 },
             };
-            Some((param_name, arg))
+            Some((param_name, arg, range))
         })
-        .filter(|(param_name, arg)| !should_hide_param_name_hint(sema, &callable, param_name, arg))
+        .filter(|(param_name, arg, _)| {
-        .map(|(param_name, arg)| InlayHint {
+            !should_hide_param_name_hint(sema, &callable, param_name, arg)
-            range: sema.original_range(arg.syntax()).range,
+        })
+        .map(|(param_name, _, FileRange { range, .. })| InlayHint {
+            range,
             kind: InlayKind::ParameterHint,
             label: param_name.into(),
         });
@@ -184,16 +190,18 @@ fn get_bind_pat_hints(
     acc: &mut Vec<InlayHint>,
     sema: &Semantics<RootDatabase>,
     config: &InlayHintsConfig,
-    pat: ast::IdentPat,
+    pat: &ast::IdentPat,
 ) -> Option<()> {
     if !config.type_hints {
         return None;
     }
 
-    let krate = sema.scope(pat.syntax()).module().map(|it| it.krate());
+    let descended = sema.descend_node_into_attributes(pat.clone()).pop();
+    let desc_pat = descended.as_ref().unwrap_or(pat);
+    let krate = sema.scope(desc_pat.syntax()).module().map(|it| it.krate());
     let famous_defs = FamousDefs(sema, krate);
 
-    let ty = sema.type_of_pat(&pat.clone().into())?.original;
+    let ty = sema.type_of_pat(&desc_pat.clone().into())?.original;
 
     if should_not_display_type_hint(sema, &pat, &ty) {
         return None;
@@ -201,8 +209,8 @@ fn get_bind_pat_hints(
 
     acc.push(InlayHint {
         range: match pat.name() {
-            Some(name) => sema.original_range(name.syntax()).range,
+            Some(name) => name.syntax().text_range(),
-            None => sema.original_range(pat.syntax()).range,
+            None => pat.syntax().text_range(),
         },
         kind: InlayKind::TypeHint,
         label: hint_iterator(sema, &famous_defs, config, &ty)
@@ -435,9 +443,13 @@ fn get_callable(
 ) -> Option<(hir::Callable, ast::ArgList)> {
     match expr {
         ast::Expr::CallExpr(expr) => {
+            let descended = sema.descend_node_into_attributes(expr.clone()).pop();
+            let expr = descended.as_ref().unwrap_or(expr);
             sema.type_of_expr(&expr.expr()?)?.original.as_callable(sema.db).zip(expr.arg_list())
         }
         ast::Expr::MethodCallExpr(expr) => {
+            let descended = sema.descend_node_into_attributes(expr.clone()).pop();
+            let expr = descended.as_ref().unwrap_or(expr);
             sema.resolve_method_call_as_callable(expr).zip(expr.arg_list())
         }
         _ => None,
@@ -1471,4 +1483,53 @@ fn main() {
             "#]],
         );
     }
+
+    #[test]
+    fn hints_in_attr_call() {
+        check_expect(
+            TEST_CONFIG,
+            r#"
+//- proc_macros: identity, input_replace
+struct Struct;
+impl Struct {
+    fn chain(self) -> Self {
+        self
+    }
+}
+#[proc_macros::identity]
+fn main() {
+    let strukt = Struct;
+    strukt
+        .chain()
+        .chain()
+        .chain();
+    Struct::chain(strukt);
+}
+"#,
+            expect![[r#"
+                [
+                    InlayHint {
+                        range: 124..130,
+                        kind: TypeHint,
+                        label: "Struct",
+                    },
+                    InlayHint {
+                        range: 145..185,
+                        kind: ChainingHint,
+                        label: "Struct",
+                    },
+                    InlayHint {
+                        range: 145..168,
+                        kind: ChainingHint,
+                        label: "Struct",
+                    },
+                    InlayHint {
+                        range: 222..228,
+                        kind: ParameterHint,
+                        label: "self",
+                    },
+                ]
+            "#]],
+        );
+    }
 }