Thinner TokenMap

Author: Veykril

Cargo.lock

@@ -33,10 +33,11 @@ impl SyntaxContext for SyntaxContextId {
 impl SyntaxContextId {
     // TODO: This is very much UB, salsa exposes no way to create an InternId in a const context
     // currently (which kind of makes sense but we need it here!)
-    pub const ROOT: Self = SyntaxContextId(unsafe { core::mem::transmute(1) });
+    pub const ROOT: Self = SyntaxContextId(unsafe { InternId::new_unchecked(0) });
     // TODO: This is very much UB, salsa exposes no way to create an InternId in a const context
     // currently (which kind of makes sense but we need it here!)
-    pub const SELF_REF: Self = SyntaxContextId(unsafe { core::mem::transmute(!0u32) });
+    pub const SELF_REF: Self =
+        SyntaxContextId(unsafe { InternId::new_unchecked(InternId::MAX - 1) });
 
     pub fn is_root(self) -> bool {
         self == Self::ROOT
@@ -107,6 +108,7 @@ impl fmt::Debug for HirFileId {
 pub struct MacroFileId {
     pub macro_call_id: MacroCallId,
 }
+
 /// `MacroCallId` identifies a particular macro invocation, like
 /// `println!("Hello, {}", world)`.
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]

crates/base-db/src/span.rs

@@ -36,7 +36,7 @@ macro_rules! f {
 }
 
 struct#FileId(0):[email protected]\2# MyTraitMap2#FileId(0):[email protected]\0# {#FileId(0):[email protected]\2#
-    map#FileId(0):[email protected]\2#:#FileId(0):[email protected]\2# #FileId(0):[email protected]\2#::#FileId(0):1@92..93\2#std#FileId(0):[email protected]\2#::#FileId(0):1@97..98\2#collections#FileId(0):[email protected]\2#::#FileId(0):1@110..111\2#HashSet#FileId(0):[email protected]\2#<#FileId(0):[email protected]\2#(#FileId(0):[email protected]\2#)#FileId(0):[email protected]\2#>#FileId(0):[email protected]\2#,#FileId(0):[email protected]\2#
+    map#FileId(0):[email protected]\2#:#FileId(0):[email protected]\2# #FileId(0):[email protected]\2#::#FileId(0):1@91..92\2#std#FileId(0):[email protected]\2#::#FileId(0):1@96..97\2#collections#FileId(0):[email protected]\2#::#FileId(0):1@109..110\2#HashSet#FileId(0):[email protected]\2#<#FileId(0):[email protected]\2#(#FileId(0):[email protected]\2#)#FileId(0):[email protected]\2#>#FileId(0):[email protected]\2#,#FileId(0):[email protected]\2#
 }#FileId(0):[email protected]\2#
 "#]],
     );

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

@@ -308,7 +308,7 @@ impl Attr {
                     return None;
                 }
                 let path = meta.path()?;
-                let call_site = span_map.span_for_range(path.syntax().text_range()).ctx;
+                let call_site = span_map.span_at(path.syntax().text_range().start()).ctx;
                 Some((
                     ModPath::from_src(db, path, SpanMapRef::ExpansionSpanMap(&span_map))?,
                     call_site,

crates/hir-expand/src/attrs.rs

@@ -254,7 +254,7 @@ pub fn expand_speculative(
     };
 
     let expand_to = macro_expand_to(db, actual_macro_call);
-    let (node, rev_tmap) = token_tree_to_syntax_node(db, &speculative_expansion.value, expand_to);
+    let (node, rev_tmap) = token_tree_to_syntax_node(&speculative_expansion.value, expand_to);
 
     let syntax_node = node.syntax_node();
     let token = rev_tmap
@@ -312,7 +312,7 @@ fn parse_macro_expansion(
     tracing::debug!("expanded = {}", tt.as_debug_string());
     tracing::debug!("kind = {:?}", expand_to);
 
-    let (parse, rev_token_map) = token_tree_to_syntax_node(db, &tt, expand_to);
+    let (parse, rev_token_map) = token_tree_to_syntax_node(&tt, expand_to);
 
     ExpandResult { value: (parse, Arc::new(rev_token_map)), err }
 }
@@ -674,7 +674,6 @@ fn macro_expand_to(db: &dyn ExpandDatabase, id: MacroCallId) -> ExpandTo {
 }
 
 fn token_tree_to_syntax_node(
-    db: &dyn ExpandDatabase,
     tt: &tt::Subtree,
     expand_to: ExpandTo,
 ) -> (Parse<SyntaxNode>, ExpansionSpanMap) {
@@ -685,18 +684,7 @@ fn token_tree_to_syntax_node(
         ExpandTo::Type => mbe::TopEntryPoint::Type,
         ExpandTo::Expr => mbe::TopEntryPoint::Expr,
     };
-    let (parse, mut span_map) = mbe::token_tree_to_syntax_node(tt, entry_point);
-    // FIXME: now what the hell is going on here
-    span_map.span_map.sort_by(|(_, a), (_, b)| {
-        a.anchor.file_id.cmp(&b.anchor.file_id).then_with(|| {
-            let map = db.ast_id_map(a.anchor.file_id.into());
-            map.get_erased(a.anchor.ast_id)
-                .text_range()
-                .start()
-                .cmp(&map.get_erased(b.anchor.ast_id).text_range().start())
-        })
-    });
-    (parse, span_map)
+    mbe::token_tree_to_syntax_node(tt, entry_point)
 }
 
 fn check_tt_count(tt: &tt::Subtree) -> Result<(), ExpandResult<Arc<tt::Subtree>>> {

crates/hir-expand/src/db.rs

@@ -44,7 +44,7 @@ use crate::{
     db::TokenExpander,
     mod_path::ModPath,
     proc_macro::ProcMacroExpander,
-    span::ExpansionSpanMap,
+    span::{ExpansionSpanMap, SpanMap},
 };
 
 pub use crate::ast_id_map::{AstId, ErasedAstId, ErasedFileAstId};
@@ -172,7 +172,6 @@ pub trait HirFileIdExt {
     /// For macro-expansion files, returns the file original source file the
     /// expansion originated from.
     fn original_file(self, db: &dyn db::ExpandDatabase) -> FileId;
-    fn expansion_level(self, db: &dyn db::ExpandDatabase) -> u32;
 
     /// If this is a macro call, returns the syntax node of the call.
     fn call_node(self, db: &dyn db::ExpandDatabase) -> Option<InFile<SyntaxNode>>;
@@ -218,18 +217,6 @@ impl HirFileIdExt for HirFileId {
         }
     }
 
-    fn expansion_level(self, db: &dyn db::ExpandDatabase) -> u32 {
-        let mut level = 0;
-        let mut curr = self;
-        while let Some(macro_file) = curr.macro_file() {
-            let loc: MacroCallLoc = db.lookup_intern_macro_call(macro_file.macro_call_id);
-
-            level += 1;
-            curr = loc.kind.file_id();
-        }
-        level
-    }
-
     fn call_node(self, db: &dyn db::ExpandDatabase) -> Option<InFile<SyntaxNode>> {
         let macro_file = self.macro_file()?;
         let loc: MacroCallLoc = db.lookup_intern_macro_call(macro_file.macro_call_id);
@@ -330,6 +317,32 @@ impl HirFileIdExt for HirFileId {
     }
 }
 
+pub trait MacroFileIdExt {
+    fn expansion_level(self, db: &dyn db::ExpandDatabase) -> u32;
+    fn expansion_info(self, db: &dyn db::ExpandDatabase) -> ExpansionInfo;
+}
+
+impl MacroFileIdExt for MacroFileId {
+    fn expansion_level(self, db: &dyn db::ExpandDatabase) -> u32 {
+        let mut level = 0;
+        let mut macro_file = self;
+        loop {
+            let loc: MacroCallLoc = db.lookup_intern_macro_call(macro_file.macro_call_id);
+
+            level += 1;
+            macro_file = match loc.kind.file_id().repr() {
+                HirFileIdRepr::FileId(_) => break level,
+                HirFileIdRepr::MacroFile(it) => it,
+            };
+        }
+    }
+
+    /// Return expansion information if it is a macro-expansion file
+    fn expansion_info(self, db: &dyn db::ExpandDatabase) -> ExpansionInfo {
+        ExpansionInfo::new(db, self)
+    }
+}
+
 impl MacroDefId {
     pub fn as_lazy_macro(
         self,
@@ -398,7 +411,7 @@ impl MacroCallLoc {
         match file_id.repr() {
             HirFileIdRepr::FileId(file_id) => db.real_span_map(file_id).span_for_range(range),
             HirFileIdRepr::MacroFile(m) => {
-                db.parse_macro_expansion(m).value.1.span_for_range(range)
+                db.parse_macro_expansion(m).value.1.span_at(range.start())
             }
         }
     }
@@ -565,9 +578,8 @@ pub struct ExpansionInfo {
 
     macro_def: TokenExpander,
     macro_arg: Arc<tt::Subtree>,
-    exp_map: Arc<ExpansionSpanMap>,
-    /// [`None`] if the call is in a real file
-    arg_map: Option<Arc<ExpansionSpanMap>>,
+    pub exp_map: Arc<ExpansionSpanMap>,
+    arg_map: SpanMap,
 }
 
 impl ExpansionInfo {
@@ -582,38 +594,14 @@ impl ExpansionInfo {
     /// Maps the passed in file range down into a macro expansion if it is the input to a macro call.
     pub fn map_range_down<'a>(
         &'a self,
-        db: &'a dyn db::ExpandDatabase,
-        FileRange { file_id, range: absolute_range }: FileRange,
+        span: SpanData,
         // FIXME: use this for range mapping, so that we can resolve inline format args
         _relative_token_offset: Option<TextSize>,
         // FIXME: ret ty should be wrapped in InMacroFile
     ) -> Option<impl Iterator<Item = InFile<SyntaxToken>> + 'a> {
-        // search for all entries in the span map that have the given span and return the
-        // corresponding text ranges inside the expansion
-        // FIXME: Make this proper
-        let span_map = &self.exp_map.span_map;
-        let (start, end) = if span_map
-            .first()
-            .map_or(false, |(_, span)| span.anchor.file_id == file_id)
-        {
-            (0, span_map.partition_point(|a| a.1.anchor.file_id == file_id))
-        } else {
-            let start = span_map.partition_point(|a| a.1.anchor.file_id != file_id);
-            (start, start + span_map[start..].partition_point(|a| a.1.anchor.file_id == file_id))
-        };
-        let tokens = span_map[start..end]
-            .iter()
-            .filter_map(move |(range, span)| {
-                // we need to resolve the relative ranges here to make sure that we are in fact
-                // considering differently anchored spans (this might occur with proc-macros)
-                let offset = db
-                    .ast_id_map(span.anchor.file_id.into())
-                    .get_erased(span.anchor.ast_id)
-                    .text_range()
-                    .start();
-                let abs_range = span.range + offset;
-                absolute_range.eq(&abs_range).then_some(*range)
-            })
+        let tokens = self
+            .exp_map
+            .ranges_with_span(span)
             .flat_map(move |range| self.expanded.value.covering_element(range).into_token());
 
         Some(tokens.map(move |token| InFile::new(self.expanded.file_id.into(), token)))
@@ -626,7 +614,7 @@ impl ExpansionInfo {
         range: TextRange,
     ) -> (FileRange, SyntaxContextId) {
         debug_assert!(self.expanded.value.text_range().contains_range(range));
-        let span = self.exp_map.span_for_range(range);
+        let span = self.exp_map.span_at(range.start());
         let anchor_offset = db
             .ast_id_map(span.anchor.file_id.into())
             .get_erased(span.anchor.ast_id)
@@ -672,15 +660,15 @@ impl ExpansionInfo {
         token: TextRange,
     ) -> InFile<smallvec::SmallVec<[TextRange; 1]>> {
         debug_assert!(self.expanded.value.text_range().contains_range(token));
-        let span = self.exp_map.span_for_range(token);
+        let span = self.exp_map.span_at(token.start());
         match &self.arg_map {
-            None => {
+            SpanMap::RealSpanMap(_) => {
                 let file_id = span.anchor.file_id.into();
                 let anchor_offset =
                     db.ast_id_map(file_id).get_erased(span.anchor.ast_id).text_range().start();
                 InFile { file_id, value: smallvec::smallvec![span.range + anchor_offset] }
             }
-            Some(arg_map) => {
+            SpanMap::ExpansionSpanMap(arg_map) => {
                 let arg_range = self
                     .arg
                     .value
@@ -701,8 +689,7 @@ impl ExpansionInfo {
         let loc: MacroCallLoc = db.lookup_intern_macro_call(macro_file.macro_call_id);
 
         let arg_tt = loc.kind.arg(db);
-        let arg_map =
-            arg_tt.file_id.macro_file().map(|file| db.parse_macro_expansion(file).value.1);
+        let arg_map = db.span_map(arg_tt.file_id);
 
         let macro_def = db.macro_expander(loc.def);
         let (parse, exp_map) = db.parse_macro_expansion(macro_file).value;

crates/hir-expand/src/lib.rs

@@ -11,7 +11,7 @@ use crate::db::ExpandDatabase;
 pub type ExpansionSpanMap = TokenMap<SpanData>;
 
 /// Spanmap for a macro file or a real file
-#[derive(Clone, Debug)]
+#[derive(Clone, Debug, PartialEq, Eq)]
 pub enum SpanMap {
     /// Spanmap for a macro file
     ExpansionSpanMap(Arc<ExpansionSpanMap>),
@@ -46,7 +46,7 @@ impl mbe::SpanMapper<SpanData> for RealSpanMap {
 impl SpanMap {
     pub fn span_for_range(&self, range: TextRange) -> SpanData {
         match self {
-            Self::ExpansionSpanMap(span_map) => span_map.span_for_range(range),
+            Self::ExpansionSpanMap(span_map) => span_map.span_at(range.start()),
             Self::RealSpanMap(span_map) => span_map.span_for_range(range),
         }
     }
@@ -62,7 +62,7 @@ impl SpanMap {
 impl SpanMapRef<'_> {
     pub fn span_for_range(self, range: TextRange) -> SpanData {
         match self {
-            Self::ExpansionSpanMap(span_map) => span_map.span_for_range(range),
+            Self::ExpansionSpanMap(span_map) => span_map.span_at(range.start()),
             Self::RealSpanMap(span_map) => span_map.span_for_range(range),
         }
     }