Replace ID based TokenMap with proper relative text-ranges / spans

crates/base-db/src/fixture.rs

@@ -8,11 +8,12 @@ use test_utils::{
     ESCAPED_CURSOR_MARKER,
 };
 use triomphe::Arc;
-use tt::token_id::{Leaf, Subtree, TokenTree};
+use tt::{Leaf, Subtree, TokenTree};
 use vfs::{file_set::FileSet, VfsPath};
 
 use crate::{
     input::{CrateName, CrateOrigin, LangCrateOrigin},
+    span::SpanData,
     Change, CrateDisplayName, CrateGraph, CrateId, Dependency, DependencyKind, Edition, Env,
     FileId, FilePosition, FileRange, ProcMacro, ProcMacroExpander, ProcMacroExpansionError,
     ProcMacros, ReleaseChannel, SourceDatabaseExt, SourceRoot, SourceRootId,
@@ -539,10 +540,10 @@ struct IdentityProcMacroExpander;
 impl ProcMacroExpander for IdentityProcMacroExpander {
     fn expand(
         &self,
-        subtree: &Subtree,
-        _: Option<&Subtree>,
+        subtree: &Subtree<SpanData>,
+        _: Option<&Subtree<SpanData>>,
         _: &Env,
-    ) -> Result<Subtree, ProcMacroExpansionError> {
+    ) -> Result<Subtree<SpanData>, ProcMacroExpansionError> {
         Ok(subtree.clone())
     }
 }
@@ -553,10 +554,10 @@ struct AttributeInputReplaceProcMacroExpander;
 impl ProcMacroExpander for AttributeInputReplaceProcMacroExpander {
     fn expand(
         &self,
-        _: &Subtree,
-        attrs: Option<&Subtree>,
+        _: &Subtree<SpanData>,
+        attrs: Option<&Subtree<SpanData>>,
         _: &Env,
-    ) -> Result<Subtree, ProcMacroExpansionError> {
+    ) -> Result<Subtree<SpanData>, ProcMacroExpansionError> {
         attrs
             .cloned()
             .ok_or_else(|| ProcMacroExpansionError::Panic("Expected attribute input".into()))
@@ -568,11 +569,11 @@ struct MirrorProcMacroExpander;
 impl ProcMacroExpander for MirrorProcMacroExpander {
     fn expand(
         &self,
-        input: &Subtree,
-        _: Option<&Subtree>,
+        input: &Subtree<SpanData>,
+        _: Option<&Subtree<SpanData>>,
         _: &Env,
-    ) -> Result<Subtree, ProcMacroExpansionError> {
-        fn traverse(input: &Subtree) -> Subtree {
+    ) -> Result<Subtree<SpanData>, ProcMacroExpansionError> {
+        fn traverse(input: &Subtree<SpanData>) -> Subtree<SpanData> {
             let mut token_trees = vec![];
             for tt in input.token_trees.iter().rev() {
                 let tt = match tt {
@@ -595,13 +596,13 @@ struct ShortenProcMacroExpander;
 impl ProcMacroExpander for ShortenProcMacroExpander {
     fn expand(
         &self,
-        input: &Subtree,
-        _: Option<&Subtree>,
+        input: &Subtree<SpanData>,
+        _: Option<&Subtree<SpanData>>,
         _: &Env,
-    ) -> Result<Subtree, ProcMacroExpansionError> {
+    ) -> Result<Subtree<SpanData>, ProcMacroExpansionError> {
         return Ok(traverse(input));
 
-        fn traverse(input: &Subtree) -> Subtree {
+        fn traverse(input: &Subtree<SpanData>) -> Subtree<SpanData> {
             let token_trees = input
                 .token_trees
                 .iter()
@@ -613,7 +614,7 @@ impl ProcMacroExpander for ShortenProcMacroExpander {
             Subtree { delimiter: input.delimiter, token_trees }
         }
 
-        fn modify_leaf(leaf: &Leaf) -> Leaf {
+        fn modify_leaf(leaf: &Leaf<SpanData>) -> Leaf<SpanData> {
             let mut leaf = leaf.clone();
             match &mut leaf {
                 Leaf::Literal(it) => {

crates/base-db/src/input.rs

@@ -13,9 +13,10 @@ use la_arena::{Arena, Idx};
 use rustc_hash::{FxHashMap, FxHashSet};
 use syntax::SmolStr;
 use triomphe::Arc;
-use tt::token_id::Subtree;
 use vfs::{file_set::FileSet, AbsPathBuf, AnchoredPath, FileId, VfsPath};
 
+use crate::span::SpanData;
+
 // Map from crate id to the name of the crate and path of the proc-macro. If the value is `None`,
 // then the crate for the proc-macro hasn't been build yet as the build data is missing.
 pub type ProcMacroPaths = FxHashMap<CrateId, Result<(Option<String>, AbsPathBuf), String>>;
@@ -255,10 +256,10 @@ pub enum ProcMacroKind {
 pub trait ProcMacroExpander: fmt::Debug + Send + Sync + RefUnwindSafe {
     fn expand(
         &self,
-        subtree: &Subtree,
-        attrs: Option<&Subtree>,
+        subtree: &tt::Subtree<SpanData>,
+        attrs: Option<&tt::Subtree<SpanData>>,
         env: &Env,
-    ) -> Result<Subtree, ProcMacroExpansionError>;
+    ) -> Result<tt::Subtree<SpanData>, ProcMacroExpansionError>;
 }
 
 #[derive(Debug)]

crates/base-db/src/lib.rs

@@ -5,6 +5,7 @@
 mod input;
 mod change;
 pub mod fixture;
+pub mod span;
 
 use std::panic;
 

crates/base-db/src/span.rs

@@ -0,0 +1,166 @@
+use std::fmt;
+
+use salsa::InternId;
+use vfs::FileId;
+
+pub type ErasedFileAstId = la_arena::Idx<syntax::SyntaxNodePtr>;
+
+// The first inde is always the root node's AstId
+pub const ROOT_ERASED_FILE_AST_ID: ErasedFileAstId =
+    la_arena::Idx::from_raw(la_arena::RawIdx::from_u32(0));
+
+#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
+pub struct SyntaxContext;
+
+pub type SpanData = tt::SpanData<SpanAnchor>;
+
+#[derive(Copy, Clone, PartialEq, Eq, Hash)]
+pub struct SpanAnchor {
+    pub file_id: HirFileId,
+    pub ast_id: ErasedFileAstId,
+}
+
+impl fmt::Debug for SpanAnchor {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_tuple("SpanAnchor").field(&self.file_id).field(&self.ast_id.into_raw()).finish()
+    }
+}
+
+impl tt::Span for SpanAnchor {
+    const DUMMY: Self = SpanAnchor { file_id: HirFileId(0), ast_id: ROOT_ERASED_FILE_AST_ID };
+}
+
+/// Input to the analyzer is a set of files, where each file is identified by
+/// `FileId` and contains source code. However, another source of source code in
+/// Rust are macros: each macro can be thought of as producing a "temporary
+/// file". To assign an id to such a file, we use the id of the macro call that
+/// produced the file. So, a `HirFileId` is either a `FileId` (source code
+/// written by user), or a `MacroCallId` (source code produced by macro).
+///
+/// What is a `MacroCallId`? Simplifying, it's a `HirFileId` of a file
+/// containing the call plus the offset of the macro call in the file. Note that
+/// this is a recursive definition! However, the size_of of `HirFileId` is
+/// finite (because everything bottoms out at the real `FileId`) and small
+/// (`MacroCallId` uses the location interning. You can check details here:
+/// <https://en.wikipedia.org/wiki/String_interning>).
+///
+/// The two variants are encoded in a single u32 which are differentiated by the MSB.
+/// If the MSB is 0, the value represents a `FileId`, otherwise the remaining 31 bits represent a
+/// `MacroCallId`.
+#[derive(Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
+pub struct HirFileId(u32);
+
+impl From<HirFileId> for u32 {
+    fn from(value: HirFileId) -> Self {
+        value.0
+    }
+}
+
+impl From<u32> for HirFileId {
+    fn from(value: u32) -> Self {
+        HirFileId(value)
+    }
+}
+
+impl From<MacroCallId> for HirFileId {
+    fn from(value: MacroCallId) -> Self {
+        value.as_file()
+    }
+}
+
+impl fmt::Debug for HirFileId {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        self.repr().fmt(f)
+    }
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+pub struct MacroFile {
+    pub macro_call_id: MacroCallId,
+}
+
+/// `MacroCallId` identifies a particular macro invocation, like
+/// `println!("Hello, {}", world)`.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+pub struct MacroCallId(salsa::InternId);
+crate::impl_intern_key!(MacroCallId);
+
+impl MacroCallId {
+    pub fn as_file(self) -> HirFileId {
+        MacroFile { macro_call_id: self }.into()
+    }
+
+    pub fn as_macro_file(self) -> MacroFile {
+        MacroFile { macro_call_id: self }
+    }
+}
+
+#[derive(Clone, Copy, PartialEq, Eq, Hash)]
+pub enum HirFileIdRepr {
+    FileId(FileId),
+    MacroFile(MacroFile),
+}
+
+impl fmt::Debug for HirFileIdRepr {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        match self {
+            Self::FileId(arg0) => f.debug_tuple("FileId").field(&arg0.0).finish(),
+            Self::MacroFile(arg0) => {
+                f.debug_tuple("MacroFile").field(&arg0.macro_call_id.0).finish()
+            }
+        }
+    }
+}
+
+impl From<FileId> for HirFileId {
+    fn from(FileId(id): FileId) -> Self {
+        assert!(id < Self::MAX_FILE_ID);
+        HirFileId(id)
+    }
+}
+
+impl From<MacroFile> for HirFileId {
+    fn from(MacroFile { macro_call_id: MacroCallId(id) }: MacroFile) -> Self {
+        let id = id.as_u32();
+        assert!(id < Self::MAX_FILE_ID);
+        HirFileId(id | Self::MACRO_FILE_TAG_MASK)
+    }
+}
+
+impl HirFileId {
+    const MAX_FILE_ID: u32 = u32::MAX ^ Self::MACRO_FILE_TAG_MASK;
+    const MACRO_FILE_TAG_MASK: u32 = 1 << 31;
+
+    #[inline]
+    pub fn is_macro(self) -> bool {
+        self.0 & Self::MACRO_FILE_TAG_MASK != 0
+    }
+
+    #[inline]
+    pub fn macro_file(self) -> Option<MacroFile> {
+        match self.0 & Self::MACRO_FILE_TAG_MASK {
+            0 => None,
+            _ => Some(MacroFile {
+                macro_call_id: MacroCallId(InternId::from(self.0 ^ Self::MACRO_FILE_TAG_MASK)),
+            }),
+        }
+    }
+
+    #[inline]
+    pub fn file_id(self) -> Option<FileId> {
+        match self.0 & Self::MACRO_FILE_TAG_MASK {
+            0 => Some(FileId(self.0)),
+            _ => None,
+        }
+    }
+
+    #[inline]
+    pub fn repr(self) -> HirFileIdRepr {
+        match self.0 & Self::MACRO_FILE_TAG_MASK {
+            0 => HirFileIdRepr::FileId(FileId(self.0)),
+            _ => HirFileIdRepr::MacroFile(MacroFile {
+                macro_call_id: MacroCallId(InternId::from(self.0 ^ Self::MACRO_FILE_TAG_MASK)),
+            }),
+        }
+    }
+}