fix: Fix SyntaxContext::roots technically overlapping valid interneds

crates/proc-macro-srv/src/server_impl/rust_analyzer_span.rs

@@ -183,31 +183,10 @@ impl server::Server for RaSpanServer<'_> {
         Range { start: span.range.start().into(), end: span.range.end().into() }
     }
     fn span_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);
-            } else if second.ctx.is_root() {
-                return Some(first);
-            }
-        }
-        Some(Span {
-            range: first.range.cover(second.range),
-            anchor: second.anchor,
-            ctx: second.ctx,
+        first.join(second, |_, _| {
+            // 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.
+            None
         })
     }
     fn span_subspan(

crates/span/src/hygiene.rs

@@ -19,7 +19,9 @@
 //! # The Call-site Hierarchy
 //!
 //! `ExpnData::call_site` in rustc, `MacroCallLoc::call_site` in rust-analyzer.
+#[cfg(feature = "salsa")]
 use crate::Edition;
+
 use std::fmt;
 
 /// A syntax context describes a hierarchy tracking order of macro definitions.
@@ -281,7 +283,7 @@ const _: () = {
                     let fields = SyntaxContext::ingredient(zalsa).data(zalsa, id);
                     fields.edition
                 }
-                None => Edition::from_u32(SyntaxContext::MAX_ID - self.into_u32()),
+                None => Edition::from_u32(SyntaxContext::MAX_ROOT_ID - self.into_u32()),
             }
         }
 
@@ -331,32 +333,9 @@ const _: () = {
     }
 };
 
-impl SyntaxContext {
-    #[inline]
-    pub fn is_root(self) -> bool {
-        (SyntaxContext::MAX_ID - Edition::LATEST as u32) <= self.into_u32()
-            && self.into_u32() <= (SyntaxContext::MAX_ID - Edition::Edition2015 as u32)
-    }
-
-    #[inline]
-    pub fn remove_root_edition(&mut self) {
-        if self.is_root() {
-            *self = Self::root(Edition::Edition2015);
-        }
-    }
-
-    /// The root context, which is the parent of all other contexts. All `FileId`s have this context.
-    #[inline]
-    pub const fn root(edition: Edition) -> Self {
-        let edition = edition as u32;
-        // SAFETY: Roots are valid `SyntaxContext`s
-        unsafe { SyntaxContext::from_u32(SyntaxContext::MAX_ID - edition) }
-    }
-}
-
 #[cfg(feature = "salsa")]
 impl<'db> SyntaxContext {
-    const MAX_ID: u32 = salsa::Id::MAX_U32 - 1;
+    const MAX_ROOT_ID: u32 = salsa::Id::MAX_U32 + Edition::LATEST as u32;
 
     #[inline]
     pub const fn into_u32(self) -> u32 {
@@ -368,6 +347,7 @@ impl<'db> SyntaxContext {
     /// The ID must be a valid `SyntaxContext`.
     #[inline]
     pub const unsafe fn from_u32(u32: u32) -> Self {
+        debug_assert!(u32 != 0);
         // INVARIANT: Our precondition.
         Self(u32, std::marker::PhantomData)
     }
@@ -378,7 +358,8 @@ impl<'db> SyntaxContext {
             None
         } else {
             // SAFETY: By our invariant, this is either a root (which we verified it's not) or a valid `salsa::Id`.
-            unsafe { Some(salsa::Id::from_index(self.0)) }
+            // Note we use `from_bits_unchecked` here to avoid the debug assert ..., the generation is still 0 with this.
+            unsafe { Some(salsa::Id::from_bits_unchecked(self.0 as u64)) }
         }
     }
 
@@ -388,6 +369,27 @@ impl<'db> SyntaxContext {
         unsafe { Self::from_u32(id.index()) }
     }
 
+    #[inline]
+    pub fn is_root(self) -> bool {
+        (SyntaxContext::MAX_ROOT_ID - Edition::LATEST as u32) <= self.into_u32()
+            && self.into_u32() <= (SyntaxContext::MAX_ROOT_ID - Edition::Edition2015 as u32)
+    }
+
+    #[inline]
+    pub fn remove_root_edition(&mut self) {
+        if self.is_root() {
+            *self = Self::root(Edition::Edition2015);
+        }
+    }
+
+    /// The root context, which is the parent of all other contexts. All `FileId`s have this context.
+    #[inline]
+    pub const fn root(edition: Edition) -> Self {
+        let edition = edition as u32;
+        // SAFETY: Roots are valid `SyntaxContext`s
+        unsafe { SyntaxContext::from_u32(SyntaxContext::MAX_ROOT_ID - edition) }
+    }
+
     #[inline]
     pub fn outer_mark(
         self,
@@ -446,15 +448,8 @@ impl<'db> SyntaxContext {
 #[derive(Copy, Clone, PartialEq, PartialOrd, Eq, Ord, Hash)]
 pub struct SyntaxContext(u32);
 
-#[allow(dead_code)]
-const SALSA_MAX_ID_MIRROR: u32 = u32::MAX - 0xFF;
-#[cfg(feature = "salsa")]
-const _: () = assert!(salsa::Id::MAX_U32 == SALSA_MAX_ID_MIRROR);
-
 #[cfg(not(feature = "salsa"))]
 impl SyntaxContext {
-    const MAX_ID: u32 = SALSA_MAX_ID_MIRROR - 1;
-
     pub const fn into_u32(self) -> u32 {
         self.0
     }
@@ -463,6 +458,7 @@ impl SyntaxContext {
     ///
     /// None. This is always safe to call without the `salsa` feature.
     pub const unsafe fn from_u32(u32: u32) -> Self {
+        debug_assert!(u32 != 0);
         Self(u32)
     }
 }
@@ -495,16 +491,28 @@ impl Transparency {
     }
 }
 
+#[cfg(feature = "salsa")]
 impl fmt::Display for SyntaxContext {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         if self.is_root() {
-            write!(f, "ROOT{}", Edition::from_u32(SyntaxContext::MAX_ID - self.into_u32()).number())
+            write!(
+                f,
+                "ROOT{}",
+                Edition::from_u32(SyntaxContext::MAX_ROOT_ID - self.into_u32()).number()
+            )
         } else {
             write!(f, "{}", self.into_u32())
         }
     }
 }
 
+#[cfg(not(feature = "salsa"))]
+impl fmt::Display for SyntaxContext {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "{}", self.into_u32())
+    }
+}
+
 impl std::fmt::Debug for SyntaxContext {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         if f.alternate() {
@@ -514,3 +522,49 @@ impl std::fmt::Debug for SyntaxContext {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_root_edition_is_root() {
+        for edition in Edition::iter() {
+            let ctx = SyntaxContext::root(edition);
+            assert!(ctx.is_root(), "{edition} root should be identified as root");
+        }
+    }
+
+    #[test]
+    fn test_root_edition_editions() {
+        let db = salsa::DatabaseImpl::new();
+        for edition in Edition::iter() {
+            let ctx = SyntaxContext::root(edition);
+            assert_eq!(edition, ctx.edition(&db), "{edition} root should have edition {edition}");
+        }
+    }
+
+    #[test]
+    fn test_roots_do_not_overlap_with_salsa_ids() {
+        for edition in Edition::iter() {
+            let root = SyntaxContext::root(edition);
+            let root_u32 = root.into_u32();
+            assert!(
+                root_u32 >= salsa::Id::MAX_U32,
+                "Root context for {:?} (value {}) must be >= salsa::Id::MAX_U32 ({}) to avoid collision",
+                edition,
+                root_u32,
+                salsa::Id::MAX_U32
+            );
+        }
+    }
+
+    #[test]
+    fn test_non_root_value_is_not_root() {
+        for edition in Edition::iter() {
+            // SAFETY: This is just for testing purposes
+            let ctx = unsafe { SyntaxContext::from_u32(edition as u32 + 1) };
+            assert!(!ctx.is_root(), "{edition} root should be identified as root");
+        }
+    }
+}

crates/span/src/lib.rs

@@ -51,13 +51,20 @@ impl Span {
         }
         // Differing context, we can't merge these so prefer the one that's root
         if self.ctx != other.ctx {
+            #[cfg(feature = "salsa")]
             if self.ctx.is_root() {
                 return Some(other);
             } else if other.ctx.is_root() {
                 return Some(self);
             }
+            None
+        } else {
+            Some(Span {
+                range: self.range.cover(other.range),
+                anchor: other.anchor,
+                ctx: other.ctx,
+            })
         }
-        Some(Span { range: self.range.cover(other.range), anchor: other.anchor, ctx: other.ctx })
     }
 
     pub fn eq_ignoring_ctx(self, other: Self) -> bool {

crates/span/src/map.rs

@@ -6,8 +6,8 @@ use std::{fmt, hash::Hash};
 use stdx::{always, itertools::Itertools};
 
 use crate::{
-    EditionedFileId, ErasedFileAstId, ROOT_ERASED_FILE_AST_ID, Span, SpanAnchor, SyntaxContext,
-    TextRange, TextSize,
+    EditionedFileId, ErasedFileAstId, ROOT_ERASED_FILE_AST_ID, Span, SyntaxContext, TextRange,
+    TextSize,
 };
 
 /// Maps absolute text ranges for the corresponding file to the relevant span data.
@@ -220,6 +220,7 @@ impl RealSpanMap {
         Self { file_id, pairs, end }
     }
 
+    #[cfg(feature = "salsa")]
     pub fn span_for_range(&self, range: TextRange) -> Span {
         assert!(
             range.end() <= self.end,
@@ -234,7 +235,7 @@ impl RealSpanMap {
         let (offset, ast_id) = self.pairs[idx - 1];
         Span {
             range: range - offset,
-            anchor: SpanAnchor { file_id: self.file_id, ast_id },
+            anchor: crate::SpanAnchor { file_id: self.file_id, ast_id },
             ctx: SyntaxContext::root(self.file_id.edition()),
         }
     }

crates/tt/src/lib.rs

@@ -765,27 +765,24 @@ impl Subtree {
 
 pub fn pretty(tkns: TokenTreesView<'_>) -> String {
     return dispatch_ref! {
-        match tkns.repr => tt => pretty_impl(tt)
+        match tkns.repr => tt => pretty_impl(tkns, tt)
     };
 
     use crate::storage::TokenTree;
 
-    fn tokentree_to_text<S: SpanStorage>(tkn: &TokenTree<S>, tkns: &mut &[TokenTree<S>]) -> String {
+    fn tokentree_to_text<S: SpanStorage>(
+        tkns_view: TokenTreesView<'_>,
+        tkn: &TokenTree<S>,
+        tkns: &mut &[TokenTree<S>],
+    ) -> String {
         match tkn {
             TokenTree::Ident { sym, is_raw, .. } => format!("{}{}", is_raw.as_str(), sym),
-            &TokenTree::Literal { ref text_and_suffix, kind, suffix_len, span: _ } => {
+            &TokenTree::Literal { ref text_and_suffix, kind, suffix_len, span } => {
                 format!(
                     "{}",
                     Literal {
                         text_and_suffix: text_and_suffix.clone(),
-                        span: Span {
-                            range: TextRange::empty(TextSize::new(0)),
-                            anchor: span::SpanAnchor {
-                                file_id: span::EditionedFileId::from_raw(0),
-                                ast_id: span::FIXUP_ERASED_FILE_AST_ID_MARKER
-                            },
-                            ctx: span::SyntaxContext::root(span::Edition::Edition2015)
-                        },
+                        span: span.span(tkns_view.span_parts),
                         kind,
                         suffix_len
                     }
@@ -794,7 +791,7 @@ pub fn pretty(tkns: TokenTreesView<'_>) -> String {
             TokenTree::Punct { char, .. } => format!("{}", char),
             TokenTree::Subtree { len, delim_kind, .. } => {
                 let (subtree_content, rest) = tkns.split_at(*len as usize);
-                let content = pretty_impl(subtree_content);
+                let content = pretty_impl(tkns_view, subtree_content);
                 *tkns = rest;
                 let (open, close) = match *delim_kind {
                     DelimiterKind::Brace => ("{", "}"),
@@ -807,13 +804,16 @@ pub fn pretty(tkns: TokenTreesView<'_>) -> String {
         }
     }
 
-    fn pretty_impl<S: SpanStorage>(mut tkns: &[TokenTree<S>]) -> String {
+    fn pretty_impl<S: SpanStorage>(
+        tkns_view: TokenTreesView<'_>,
+        mut tkns: &[TokenTree<S>],
+    ) -> String {
         let mut last = String::new();
         let mut last_to_joint = true;
 
         while let Some((tkn, rest)) = tkns.split_first() {
             tkns = rest;
-            last = [last, tokentree_to_text(tkn, &mut tkns)].join(if last_to_joint {
+            last = [last, tokentree_to_text(tkns_view, tkn, &mut tkns)].join(if last_to_joint {
                 ""
             } else {
                 " "