From 9e8e1242690a3a33ef1acda4a028c863d1b9fbcb Mon Sep 17 00:00:00 2001
From: Lukas Wirth <lukastw97@gmail.com>
Date: Fri, 22 Dec 2023 09:39:55 +0100
Subject: [PATCH] Special case fixup spans in server::Span impl, they are
 immutable

---
 crates/hir-expand/src/builtin_fn_macro.rs     |  2 +-
 crates/hir-expand/src/fixup.rs                | 44 +++++++++----------
 .../src/server/rust_analyzer_span.rs          | 38 +++++++++++++---
 crates/proc-macro-srv/src/server/token_id.rs  |  3 ++
 crates/proc-macro-srv/src/tests/utils.rs      | 11 ++---
 crates/span/src/lib.rs                        |  7 +++
 6 files changed, 67 insertions(+), 38 deletions(-)

diff --git a/crates/hir-expand/src/builtin_fn_macro.rs b/crates/hir-expand/src/builtin_fn_macro.rs
index 2d8202b8c653..f99a89176233 100644
--- a/crates/hir-expand/src/builtin_fn_macro.rs
+++ b/crates/hir-expand/src/builtin_fn_macro.rs
@@ -776,7 +776,7 @@ fn quote_expand(
     _db: &dyn ExpandDatabase,
     _arg_id: MacroCallId,
     _tt: &tt::Subtree,
-    span: SpanData,
+    span: Span,
 ) -> ExpandResult<tt::Subtree> {
     ExpandResult::new(
         tt::Subtree::empty(tt::DelimSpan { open: span, close: span }),
diff --git a/crates/hir-expand/src/fixup.rs b/crates/hir-expand/src/fixup.rs
index 53e87f819772..d241d94b8c40 100644
--- a/crates/hir-expand/src/fixup.rs
+++ b/crates/hir-expand/src/fixup.rs
@@ -1,10 +1,9 @@
 //! To make attribute macros work reliably when typing, we need to take care to
 //! fix up syntax errors in the code we're passing to them.
 
-use la_arena::RawIdx;
 use rustc_hash::{FxHashMap, FxHashSet};
 use smallvec::SmallVec;
-use span::{ErasedFileAstId, FileId, Span, SpanAnchor, SpanData};
+use span::{ErasedFileAstId, Span, SpanAnchor, SpanData, FIXUP_ERASED_FILE_AST_ID_MARKER};
 use stdx::never;
 use syntax::{
     ast::{self, AstNode, HasLoopBody},
@@ -39,13 +38,11 @@ impl SyntaxFixupUndoInfo {
     pub(crate) const NONE: Self = SyntaxFixupUndoInfo { original: None };
 }
 
-// censoring -> just don't convert the node
-// replacement -> censor + append
-// append -> insert a fake node, here we need to assemble some dummy span that we can figure out how
-// to remove later
-const FIXUP_DUMMY_FILE: FileId = FileId::from_raw(FileId::MAX_FILE_ID);
-const FIXUP_DUMMY_AST_ID: ErasedFileAstId = ErasedFileAstId::from_raw(RawIdx::from_u32(!0));
+// We mark spans with `FIXUP_DUMMY_AST_ID` to indicate that they are fake.
+const FIXUP_DUMMY_AST_ID: ErasedFileAstId = FIXUP_ERASED_FILE_AST_ID_MARKER;
 const FIXUP_DUMMY_RANGE: TextRange = TextRange::empty(TextSize::new(0));
+// If the fake span has this range end, that means that the range start is an index into the
+// `original` list in `SyntaxFixupUndoInfo`.
 const FIXUP_DUMMY_RANGE_END: TextSize = TextSize::new(!0);
 
 pub(crate) fn fixup_syntax(
@@ -58,13 +55,13 @@ pub(crate) fn fixup_syntax(
     let mut preorder = node.preorder();
     let mut original = Vec::new();
     let dummy_range = FIXUP_DUMMY_RANGE;
-    // we use a file id of `FileId(!0)` to signal a fake node, and the text range's start offset as
-    // the index into the replacement vec but only if the end points to !0
-    let dummy_anchor = SpanAnchor { file_id: FIXUP_DUMMY_FILE, ast_id: FIXUP_DUMMY_AST_ID };
-    let fake_span = |range| SpanData {
-        range: dummy_range,
-        anchor: dummy_anchor,
-        ctx: span_map.span_for_range(range).ctx,
+    let fake_span = |range| {
+        let span = span_map.span_for_range(range);
+        SpanData {
+            range: dummy_range,
+            anchor: SpanAnchor { ast_id: FIXUP_DUMMY_AST_ID, ..span.anchor },
+            ctx: span.ctx,
+        }
     };
     while let Some(event) = preorder.next() {
         let syntax::WalkEvent::Enter(node) = event else { continue };
@@ -76,12 +73,13 @@ pub(crate) fn fixup_syntax(
             let original_tree = mbe::syntax_node_to_token_tree(&node, span_map, call_site);
             let idx = original.len() as u32;
             original.push(original_tree);
+            let span = span_map.span_for_range(node_range);
             let replacement = Leaf::Ident(Ident {
                 text: "__ra_fixup".into(),
                 span: SpanData {
                     range: TextRange::new(TextSize::new(idx), FIXUP_DUMMY_RANGE_END),
-                    anchor: dummy_anchor,
-                    ctx: span_map.span_for_range(node_range).ctx,
+                    anchor: SpanAnchor { ast_id: FIXUP_DUMMY_AST_ID, ..span.anchor },
+                    ctx: span.ctx,
                 },
             });
             append.insert(node.clone().into(), vec![replacement]);
@@ -304,8 +302,8 @@ pub(crate) fn reverse_fixups(tt: &mut Subtree, undo_info: &SyntaxFixupUndoInfo)
     let undo_info = &**undo_info;
     #[allow(deprecated)]
     if never!(
-        tt.delimiter.close.anchor.file_id == FIXUP_DUMMY_FILE
-            || tt.delimiter.open.anchor.file_id == FIXUP_DUMMY_FILE
+        tt.delimiter.close.anchor.ast_id == FIXUP_DUMMY_AST_ID
+            || tt.delimiter.open.anchor.ast_id == FIXUP_DUMMY_AST_ID
     ) {
         tt.delimiter.close = SpanData::DUMMY;
         tt.delimiter.open = SpanData::DUMMY;
@@ -321,7 +319,7 @@ fn reverse_fixups_(tt: &mut Subtree, undo_info: &[Subtree]) {
         .filter(|tt| match tt {
             tt::TokenTree::Leaf(leaf) => {
                 let span = leaf.span();
-                let is_real_leaf = span.anchor.file_id != FIXUP_DUMMY_FILE;
+                let is_real_leaf = span.anchor.ast_id != FIXUP_DUMMY_AST_ID;
                 let is_replaced_node = span.range.end() == FIXUP_DUMMY_RANGE_END;
                 is_real_leaf || is_replaced_node
             }
@@ -329,8 +327,8 @@ fn reverse_fixups_(tt: &mut Subtree, undo_info: &[Subtree]) {
         })
         .flat_map(|tt| match tt {
             tt::TokenTree::Subtree(mut tt) => {
-                if tt.delimiter.close.anchor.file_id == FIXUP_DUMMY_FILE
-                    || tt.delimiter.open.anchor.file_id == FIXUP_DUMMY_FILE
+                if tt.delimiter.close.anchor.ast_id == FIXUP_DUMMY_AST_ID
+                    || tt.delimiter.open.anchor.ast_id == FIXUP_DUMMY_AST_ID
                 {
                     // Even though fixup never creates subtrees with fixup spans, the old proc-macro server
                     // might copy them if the proc-macro asks for it, so we need to filter those out
@@ -341,7 +339,7 @@ fn reverse_fixups_(tt: &mut Subtree, undo_info: &[Subtree]) {
                 SmallVec::from_const([tt.into()])
             }
             tt::TokenTree::Leaf(leaf) => {
-                if leaf.span().anchor.file_id == FIXUP_DUMMY_FILE {
+                if leaf.span().anchor.ast_id == FIXUP_DUMMY_AST_ID {
                     // we have a fake node here, we need to replace it again with the original
                     let original = undo_info[u32::from(leaf.span().range.start()) as usize].clone();
                     if original.delimiter.kind == tt::DelimiterKind::Invisible {
diff --git a/crates/proc-macro-srv/src/server/rust_analyzer_span.rs b/crates/proc-macro-srv/src/server/rust_analyzer_span.rs
index 160981697110..bcf3600d2736 100644
--- a/crates/proc-macro-srv/src/server/rust_analyzer_span.rs
+++ b/crates/proc-macro-srv/src/server/rust_analyzer_span.rs
@@ -12,7 +12,7 @@ use std::{
 
 use ::tt::{TextRange, TextSize};
 use proc_macro::bridge::{self, server};
-use span::Span;
+use span::{Span, FIXUP_ERASED_FILE_AST_ID_MARKER};
 
 use crate::server::{
     delim_to_external, delim_to_internal, token_stream::TokenStreamBuilder, LiteralFormatter,
@@ -55,6 +55,10 @@ impl server::Types for RaSpanServer {
 }
 
 impl server::FreeFunctions for RaSpanServer {
+    fn injected_env_var(&mut self, _: &str) -> Option<std::string::String> {
+        None
+    }
+
     fn track_env_var(&mut self, var: &str, value: Option<&str>) {
         self.tracked_env_vars.insert(var.into(), value.map(Into::into));
     }
@@ -124,9 +128,7 @@ impl server::TokenStream for RaSpanServer {
                 });
 
                 let literal = tt::Literal { text, span: literal.0.span };
-                let leaf: ::tt::Leaf<
-                    ::tt::SpanData<base_db::span::SpanAnchor, base_db::span::SyntaxContextId>,
-                > = tt::Leaf::from(literal);
+                let leaf: tt::Leaf = tt::Leaf::from(literal);
                 let tree = tt::TokenTree::from(leaf);
                 Self::TokenStream::from_iter(iter::once(tree))
             }
@@ -246,6 +248,7 @@ impl server::Span for RaSpanServer {
         format!("{:?}", span)
     }
     fn source_file(&mut self, _span: Self::Span) -> Self::SourceFile {
+        // FIXME stub, requires db
         SourceFile {}
     }
     fn save_span(&mut self, _span: Self::Span) -> usize {
@@ -261,7 +264,7 @@ impl server::Span for RaSpanServer {
     /// See PR:
     /// https://github.com/rust-lang/rust/pull/55780
     fn source_text(&mut self, _span: Self::Span) -> Option<String> {
-        // FIXME requires db
+        // FIXME requires db, needs special handling wrt fixup spans
         None
     }
 
@@ -278,9 +281,20 @@ impl server::Span for RaSpanServer {
         Range { start: span.range.start().into(), end: span.range.end().into() }
     }
     fn join(&mut self, first: Self::Span, second: Self::Span) -> Option<Self::Span> {
+        // We can't modify the span range for fixup spans, those are meaningful to fixup, so just
+        // prefer the non-fixup span.
+        if first.anchor.ast_id == FIXUP_ERASED_FILE_AST_ID_MARKER {
+            return Some(second);
+        }
+        if second.anchor.ast_id == FIXUP_ERASED_FILE_AST_ID_MARKER {
+            return Some(first);
+        }
+        // FIXME: Once we can talk back to the client, implement a "long join" request for anchors
+        // that differ in [AstId]s as joining those spans requires resolving the AstIds.
         if first.anchor != second.anchor {
             return None;
         }
+        // Differing context, we can't merge these so prefer the one that's root
         if first.ctx != second.ctx {
             if first.ctx.is_root() {
                 return Some(second);
@@ -300,8 +314,10 @@ impl server::Span for RaSpanServer {
         start: Bound<usize>,
         end: Bound<usize>,
     ) -> Option<Self::Span> {
-        // FIXME requires db to resolve the ast id, THIS IS NOT INCREMENTAL as it works on absolute
-        // ranges
+        // We can't modify the span range for fixup spans, those are meaningful to fixup.
+        if span.anchor.ast_id == FIXUP_ERASED_FILE_AST_ID_MARKER {
+            return Some(span);
+        }
         let length = span.range.len().into();
 
         let start: u32 = match start {
@@ -341,10 +357,18 @@ impl server::Span for RaSpanServer {
     }
 
     fn end(&mut self, span: Self::Span) -> Self::Span {
+        // We can't modify the span range for fixup spans, those are meaningful to fixup.
+        if span.anchor.ast_id == FIXUP_ERASED_FILE_AST_ID_MARKER {
+            return span;
+        }
         Span { range: TextRange::empty(span.range.end()), ..span }
     }
 
     fn start(&mut self, span: Self::Span) -> Self::Span {
+        // We can't modify the span range for fixup spans, those are meaningful to fixup.
+        if span.anchor.ast_id == FIXUP_ERASED_FILE_AST_ID_MARKER {
+            return span;
+        }
         Span { range: TextRange::empty(span.range.start()), ..span }
     }
 
diff --git a/crates/proc-macro-srv/src/server/token_id.rs b/crates/proc-macro-srv/src/server/token_id.rs
index d07bceb1c8a1..12526ad4f3ae 100644
--- a/crates/proc-macro-srv/src/server/token_id.rs
+++ b/crates/proc-macro-srv/src/server/token_id.rs
@@ -53,6 +53,9 @@ impl server::Types for TokenIdServer {
 }
 
 impl server::FreeFunctions for TokenIdServer {
+    fn injected_env_var(&mut self, _: &str) -> Option<std::string::String> {
+        None
+    }
     fn track_env_var(&mut self, _var: &str, _value: Option<&str>) {}
     fn track_path(&mut self, _path: &str) {}
     fn literal_from_str(
diff --git a/crates/proc-macro-srv/src/tests/utils.rs b/crates/proc-macro-srv/src/tests/utils.rs
index 8755e5b3ff6e..9a1311d9550a 100644
--- a/crates/proc-macro-srv/src/tests/utils.rs
+++ b/crates/proc-macro-srv/src/tests/utils.rs
@@ -1,11 +1,8 @@
 //! utils used in proc-macro tests
 
-use base_db::{
-    span::{ErasedFileAstId, SpanAnchor, SpanData, SyntaxContextId},
-    FileId,
-};
 use expect_test::Expect;
 use proc_macro_api::msg::TokenId;
+use span::{ErasedFileAstId, FileId, Span, SpanAnchor, SyntaxContextId};
 use tt::TextRange;
 
 use crate::{dylib, proc_macro_test_dylib_path, ProcMacroSrv};
@@ -20,7 +17,7 @@ fn parse_string_spanned(
     anchor: SpanAnchor,
     call_site: SyntaxContextId,
     src: &str,
-) -> crate::server::TokenStream<SpanData> {
+) -> crate::server::TokenStream<Span> {
     crate::server::TokenStream::with_subtree(
         mbe::parse_to_token_tree(anchor, call_site, src).unwrap(),
     )
@@ -68,7 +65,7 @@ fn assert_expand_impl(
         .unwrap();
     expect.assert_eq(&format!("{res:?}"));
 
-    let def_site = SpanData {
+    let def_site = Span {
         range: TextRange::new(0.into(), 150.into()),
         anchor: SpanAnchor {
             file_id: FileId::from_raw(41),
@@ -76,7 +73,7 @@ fn assert_expand_impl(
         },
         ctx: SyntaxContextId::ROOT,
     };
-    let call_site = SpanData {
+    let call_site = Span {
         range: TextRange::new(0.into(), 100.into()),
         anchor: SpanAnchor {
             file_id: FileId::from_raw(42),
diff --git a/crates/span/src/lib.rs b/crates/span/src/lib.rs
index 7341a529a656..7617acde64a2 100644
--- a/crates/span/src/lib.rs
+++ b/crates/span/src/lib.rs
@@ -30,6 +30,13 @@ pub type ErasedFileAstId = la_arena::Idx<syntax::SyntaxNodePtr>;
 pub const ROOT_ERASED_FILE_AST_ID: ErasedFileAstId =
     la_arena::Idx::from_raw(la_arena::RawIdx::from_u32(0));
 
+/// FileId used as the span for syntax node fixups. Any Span containing this file id is to be
+/// considered fake.
+pub const FIXUP_ERASED_FILE_AST_ID_MARKER: ErasedFileAstId =
+    // we pick the second to last for this in case we every consider making this a NonMaxU32, this
+    // is required to be stable for the proc-macro-server
+    la_arena::Idx::from_raw(la_arena::RawIdx::from_u32(!0 - 1));
+
 #[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
 pub struct SpanData<Ctx> {
     /// The text range of this span, relative to the anchor.