Merge pull request #21279 from J3m3/no-downmap

internal: add special `ErasedFileAstId` used for bypassing downmapping

Author: ChayimFriedman2

crates/hir-def/src/macro_expansion_tests/mod.rs

@@ -19,15 +19,18 @@ use std::{any::TypeId, iter, ops::Range, sync};
 use base_db::RootQueryDb;
 use expect_test::Expect;
 use hir_expand::{
-    AstId, InFile, MacroCallId, MacroCallKind, MacroKind,
+    AstId, ExpansionInfo, InFile, MacroCallId, MacroCallKind, MacroKind,
     builtin::quote::quote,
     db::ExpandDatabase,
     proc_macro::{ProcMacro, ProcMacroExpander, ProcMacroExpansionError, ProcMacroKind},
     span_map::SpanMapRef,
 };
 use intern::{Symbol, sym};
 use itertools::Itertools;
-use span::{Edition, ROOT_ERASED_FILE_AST_ID, Span, SpanAnchor, SyntaxContext};
+use span::{
+    Edition, NO_DOWNMAP_ERASED_FILE_AST_ID_MARKER, ROOT_ERASED_FILE_AST_ID, Span, SpanAnchor,
+    SyntaxContext,
+};
 use stdx::{format_to, format_to_acc};
 use syntax::{
     AstNode, AstPtr,
@@ -97,37 +100,6 @@ pub fn identity_when_valid(_attr: TokenStream, item: TokenStream) -> TokenStream
         },
     )];
 
-    fn resolve(
-        db: &dyn DefDatabase,
-        def_map: &DefMap,
-        ast_id: AstId<ast::MacroCall>,
-        ast_ptr: InFile<AstPtr<ast::MacroCall>>,
-    ) -> Option<MacroCallId> {
-        def_map.modules().find_map(|module| {
-            for decl in
-                module.1.scope.declarations().chain(module.1.scope.unnamed_consts().map(Into::into))
-            {
-                let body = match decl {
-                    ModuleDefId::FunctionId(it) => it.into(),
-                    ModuleDefId::ConstId(it) => it.into(),
-                    ModuleDefId::StaticId(it) => it.into(),
-                    _ => continue,
-                };
-
-                let (body, sm) = db.body_with_source_map(body);
-                if let Some(it) =
-                    body.blocks(db).find_map(|block| resolve(db, block.1, ast_id, ast_ptr))
-                {
-                    return Some(it);
-                }
-                if let Some((_, res)) = sm.macro_calls().find(|it| it.0 == ast_ptr) {
-                    return Some(res);
-                }
-            }
-            module.1.scope.macro_invoc(ast_id)
-        })
-    }
-
     let db = TestDB::with_files_extra_proc_macros(ra_fixture, extra_proc_macros);
     let krate = db.fetch_test_crate();
     let def_map = crate_def_map(&db, krate);
@@ -144,7 +116,7 @@ pub fn identity_when_valid(_attr: TokenStream, item: TokenStream) -> TokenStream
         let ast_id = db.ast_id_map(source.file_id).ast_id(&macro_call_node);
         let ast_id = InFile::new(source.file_id, ast_id);
         let ptr = InFile::new(source.file_id, AstPtr::new(&macro_call_node));
-        let macro_call_id = resolve(&db, def_map, ast_id, ptr)
+        let macro_call_id = resolve_macro_call_id(&db, def_map, ast_id, ptr)
             .unwrap_or_else(|| panic!("unable to find semantic macro call {macro_call_node}"));
         let expansion_result = db.parse_macro_expansion(macro_call_id);
         expansions.push((macro_call_node.clone(), expansion_result));
@@ -278,6 +250,38 @@ pub fn identity_when_valid(_attr: TokenStream, item: TokenStream) -> TokenStream
     expect.assert_eq(&expanded_text);
 }
 
+fn resolve_macro_call_id(
+    db: &dyn DefDatabase,
+    def_map: &DefMap,
+    ast_id: AstId<ast::MacroCall>,
+    ast_ptr: InFile<AstPtr<ast::MacroCall>>,
+) -> Option<MacroCallId> {
+    def_map.modules().find_map(|module| {
+        for decl in
+            module.1.scope.declarations().chain(module.1.scope.unnamed_consts().map(Into::into))
+        {
+            let body = match decl {
+                ModuleDefId::FunctionId(it) => it.into(),
+                ModuleDefId::ConstId(it) => it.into(),
+                ModuleDefId::StaticId(it) => it.into(),
+                _ => continue,
+            };
+
+            let (body, sm) = db.body_with_source_map(body);
+            if let Some(it) = body
+                .blocks(db)
+                .find_map(|block| resolve_macro_call_id(db, block.1, ast_id, ast_ptr))
+            {
+                return Some(it);
+            }
+            if let Some((_, res)) = sm.macro_calls().find(|it| it.0 == ast_ptr) {
+                return Some(res);
+            }
+        }
+        module.1.scope.macro_invoc(ast_id)
+    })
+}
+
 fn reindent(indent: IndentLevel, pp: String) -> String {
     if !pp.contains('\n') {
         return pp;
@@ -430,3 +434,47 @@ fn regression_20171() {
         Edition::CURRENT
     });
 }
+
+#[test]
+fn no_downmap() {
+    let fixture = r#"
+macro_rules! m {
+    ($func_name:ident) => {
+        fn $func_name() { todo!() }
+    };
+}
+m!(f);
+m!(g);
+    "#;
+
+    let (db, file_id) = TestDB::with_single_file(fixture);
+    let krate = file_id.krate(&db);
+    let def_map = crate_def_map(&db, krate);
+    let source = def_map[def_map.root].definition_source(&db);
+    let source_file = match source.value {
+        ModuleSource::SourceFile(it) => it,
+        ModuleSource::Module(_) | ModuleSource::BlockExpr(_) => panic!(),
+    };
+    let no_downmap_spans: Vec<_> = source_file
+        .syntax()
+        .descendants()
+        .map(|node| {
+            let mut span = db.real_span_map(file_id).span_for_range(node.text_range());
+            span.anchor.ast_id = NO_DOWNMAP_ERASED_FILE_AST_ID_MARKER;
+            span
+        })
+        .collect();
+
+    for macro_call_node in source_file.syntax().descendants().filter_map(ast::MacroCall::cast) {
+        let ast_id = db.ast_id_map(source.file_id).ast_id(&macro_call_node);
+        let ast_id = InFile::new(source.file_id, ast_id);
+        let ptr = InFile::new(source.file_id, AstPtr::new(&macro_call_node));
+        let macro_call_id = resolve_macro_call_id(&db, def_map, ast_id, ptr)
+            .unwrap_or_else(|| panic!("unable to find semantic macro call {macro_call_node}"));
+        let expansion_info = ExpansionInfo::new(&db, macro_call_id);
+        for &span in no_downmap_spans.iter() {
+            assert!(expansion_info.map_range_down(span).is_none());
+            assert!(expansion_info.map_range_down_exact(span).is_none());
+        }
+    }
+}

crates/hir-expand/src/lib.rs

@@ -37,7 +37,9 @@ use std::{hash::Hash, ops};
 
 use base_db::Crate;
 use either::Either;
-use span::{Edition, ErasedFileAstId, FileAstId, Span, SyntaxContext};
+use span::{
+    Edition, ErasedFileAstId, FileAstId, NO_DOWNMAP_ERASED_FILE_AST_ID_MARKER, Span, SyntaxContext,
+};
 use syntax::{
     SyntaxNode, SyntaxToken, TextRange, TextSize,
     ast::{self, AstNode},
@@ -854,6 +856,10 @@ impl ExpansionInfo {
         &self,
         span: Span,
     ) -> Option<InMacroFile<impl Iterator<Item = (SyntaxToken, SyntaxContext)> + '_>> {
+        if span.anchor.ast_id == NO_DOWNMAP_ERASED_FILE_AST_ID_MARKER {
+            return None;
+        }
+
         let tokens = self.exp_map.ranges_with_span_exact(span).flat_map(move |(range, ctx)| {
             self.expanded.value.covering_element(range).into_token().zip(Some(ctx))
         });
@@ -869,6 +875,10 @@ impl ExpansionInfo {
         &self,
         span: Span,
     ) -> Option<InMacroFile<impl Iterator<Item = (SyntaxToken, SyntaxContext)> + '_>> {
+        if span.anchor.ast_id == NO_DOWNMAP_ERASED_FILE_AST_ID_MARKER {
+            return None;
+        }
+
         let tokens = self.exp_map.ranges_with_span(span).flat_map(move |(range, ctx)| {
             self.expanded.value.covering_element(range).into_token().zip(Some(ctx))
         });

crates/span/src/ast_id.rs

@@ -48,6 +48,11 @@ pub const ROOT_ERASED_FILE_AST_ID: ErasedFileAstId =
 pub const FIXUP_ERASED_FILE_AST_ID_MARKER: ErasedFileAstId =
     ErasedFileAstId(pack_hash_index_and_kind(0, 0, ErasedFileAstIdKind::Fixup as u32));
 
+/// [`ErasedFileAstId`] used as the span for syntax nodes that should not be mapped down to
+/// macro expansion. Any `Span` containing this file id is to be considered fake.
+pub const NO_DOWNMAP_ERASED_FILE_AST_ID_MARKER: ErasedFileAstId =
+    ErasedFileAstId(pack_hash_index_and_kind(0, 0, ErasedFileAstIdKind::NoDownmap as u32));
+
 /// This is a type erased FileAstId.
 #[derive(Clone, Copy, PartialEq, Eq, Hash)]
 pub struct ErasedFileAstId(u32);
@@ -95,6 +100,7 @@ impl fmt::Debug for ErasedFileAstId {
             BlockExpr,
             AsmExpr,
             Fixup,
+            NoDownmap,
         );
         if f.alternate() {
             write!(f, "{kind}[{:04X}, {}]", self.hash_value(), self.index())
@@ -150,6 +156,9 @@ enum ErasedFileAstIdKind {
     // because incrementality is not a problem, they will always be the only item in the macro file,
     // and memory usage also not because they're rare.
     AsmExpr,
+    /// Represents a fake [`ErasedFileAstId`] that should not be mapped down to macro expansion
+    /// result.
+    NoDownmap,
     /// Keep this last.
     Root,
 }

crates/span/src/lib.rs

@@ -14,7 +14,7 @@ mod map;
 pub use self::{
     ast_id::{
         AstIdMap, AstIdNode, ErasedFileAstId, FIXUP_ERASED_FILE_AST_ID_MARKER, FileAstId,
-        ROOT_ERASED_FILE_AST_ID,
+        NO_DOWNMAP_ERASED_FILE_AST_ID_MARKER, ROOT_ERASED_FILE_AST_ID,
     },
     hygiene::{SyntaxContext, Transparency},
     map::{RealSpanMap, SpanMap},