Thinner TokenMap

13 files changed, 123 insertions, 259 deletions

diff --git a/Cargo.lock b/Cargo.lock
index 3618d69c749..775231f3ea2 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1752,7 +1752,7 @@ dependencies = [
  "always-assert",
  "backtrace",
  "crossbeam-channel",
- "itertools 0.12.0",
+ "itertools",
  "jod-thread",
  "libc",
  "miow",
diff --git a/crates/base-db/src/span.rs b/crates/base-db/src/span.rs
index 370c732813a..607b8027ca9 100644
--- a/crates/base-db/src/span.rs
+++ b/crates/base-db/src/span.rs
@@ -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)]
diff --git a/crates/hir-def/src/macro_expansion_tests/mbe.rs b/crates/hir-def/src/macro_expansion_tests/mbe.rs
index fc17dcde9a0..39079685a4d 100644
--- a/crates/hir-def/src/macro_expansion_tests/mbe.rs
+++ b/crates/hir-def/src/macro_expansion_tests/mbe.rs
@@ -36,7 +36,7 @@ macro_rules! f {
 }
 
 struct#FileId(0):1@58..64\2# MyTraitMap2#FileId(0):2@31..42\0# {#FileId(0):1@72..73\2#
-    map#FileId(0):1@86..89\2#:#FileId(0):1@89..90\2# #FileId(0):1@89..90\2#::#FileId(0):1@92..93\2#std#FileId(0):1@93..96\2#::#FileId(0):1@97..98\2#collections#FileId(0):1@98..109\2#::#FileId(0):1@110..111\2#HashSet#FileId(0):1@111..118\2#<#FileId(0):1@118..119\2#(#FileId(0):1@119..120\2#)#FileId(0):1@120..121\2#>#FileId(0):1@121..122\2#,#FileId(0):1@122..123\2#
+    map#FileId(0):1@86..89\2#:#FileId(0):1@89..90\2# #FileId(0):1@89..90\2#::#FileId(0):1@91..92\2#std#FileId(0):1@93..96\2#::#FileId(0):1@96..97\2#collections#FileId(0):1@98..109\2#::#FileId(0):1@109..110\2#HashSet#FileId(0):1@111..118\2#<#FileId(0):1@118..119\2#(#FileId(0):1@119..120\2#)#FileId(0):1@120..121\2#>#FileId(0):1@121..122\2#,#FileId(0):1@122..123\2#
 }#FileId(0):1@132..133\2#
 "#]],
     );
diff --git a/crates/hir-expand/src/attrs.rs b/crates/hir-expand/src/attrs.rs
index f1619db2420..23a8fffa8c2 100644
--- a/crates/hir-expand/src/attrs.rs
+++ b/crates/hir-expand/src/attrs.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,
diff --git a/crates/hir-expand/src/db.rs b/crates/hir-expand/src/db.rs
index fed1705fb7d..601a754abbd 100644
--- a/crates/hir-expand/src/db.rs
+++ b/crates/hir-expand/src/db.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>>> {
diff --git a/crates/hir-expand/src/lib.rs b/crates/hir-expand/src/lib.rs
index 4e5aa903126..9027ea1c27b 100644
--- a/crates/hir-expand/src/lib.rs
+++ b/crates/hir-expand/src/lib.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;
diff --git a/crates/hir-expand/src/span.rs b/crates/hir-expand/src/span.rs
index 589f415de56..c2399259fac 100644
--- a/crates/hir-expand/src/span.rs
+++ b/crates/hir-expand/src/span.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),
         }
     }
diff --git a/crates/hir/src/semantics.rs b/crates/hir/src/semantics.rs
index f6ee836c529..38c4f081b76 100644
--- a/crates/hir/src/semantics.rs
+++ b/crates/hir/src/semantics.rs
@@ -17,6 +17,7 @@ use hir_def::{
 };
 use hir_expand::{
     db::ExpandDatabase, files::InRealFile, name::AsName, ExpansionInfo, HirFileIdExt, MacroCallId,
+    MacroFileId, MacroFileIdExt,
 };
 use itertools::Itertools;
 use rustc_hash::{FxHashMap, FxHashSet};
@@ -117,11 +118,11 @@ pub struct Semantics<'db, DB> {
 pub struct SemanticsImpl<'db> {
     pub db: &'db dyn HirDatabase,
     s2d_cache: RefCell<SourceToDefCache>,
-    expansion_info_cache: RefCell<FxHashMap<HirFileId, Option<ExpansionInfo>>>,
-    // Rootnode to HirFileId cache
+    expansion_info_cache: RefCell<FxHashMap<MacroFileId, ExpansionInfo>>,
+    /// Rootnode to HirFileId cache
     cache: RefCell<FxHashMap<SyntaxNode, HirFileId>>,
-    // MacroCall to its expansion's HirFileId cache
-    macro_call_cache: RefCell<FxHashMap<InFile<ast::MacroCall>, HirFileId>>,
+    /// MacroCall to its expansion's MacroFileId cache
+    macro_call_cache: RefCell<FxHashMap<InFile<ast::MacroCall>, MacroFileId>>,
 }
 
 impl<DB> fmt::Debug for Semantics<'_, DB> {
@@ -258,7 +259,7 @@ impl<'db> SemanticsImpl<'db> {
     pub fn expand(&self, macro_call: &ast::MacroCall) -> Option<SyntaxNode> {
         let sa = self.analyze_no_infer(macro_call.syntax())?;
         let file_id = sa.expand(self.db, InFile::new(sa.file_id, macro_call))?;
-        let node = self.parse_or_expand(file_id);
+        let node = self.parse_or_expand(file_id.into());
         Some(node)
     }
 
@@ -527,6 +528,7 @@ impl<'db> SemanticsImpl<'db> {
         res
     }
 
+    // FIXME: should only take real file inputs for simplicity
     fn descend_into_macros_impl(
         &self,
         token: SyntaxToken,
@@ -537,31 +539,22 @@ impl<'db> SemanticsImpl<'db> {
     ) {
         // FIXME: Clean this up
         let _p = profile::span("descend_into_macros");
-        let parent = match token.parent() {
+        let sa = match token.parent().and_then(|parent| self.analyze_no_infer(&parent)) {
             Some(it) => it,
             None => return,
         };
-        let sa = match self.analyze_no_infer(&parent) {
-            Some(it) => it,
-            None => return,
-        };
-        let def_map = sa.resolver.def_map();
-        let absolute_range = match sa.file_id.repr() {
+
+        let mut cache = self.expansion_info_cache.borrow_mut();
+        let mut mcache = self.macro_call_cache.borrow_mut();
+        let span = match sa.file_id.repr() {
             base_db::span::HirFileIdRepr::FileId(file_id) => {
-                FileRange { file_id, range: token.text_range() }
-            }
-            base_db::span::HirFileIdRepr::MacroFile(m) => {
-                let span =
-                    self.db.parse_macro_expansion(m).value.1.span_for_range(token.text_range());
-                let range = span.range
-                    + self
-                        .db
-                        .ast_id_map(span.anchor.file_id.into())
-                        .get_erased(span.anchor.ast_id)
-                        .text_range()
-                        .start();
-                FileRange { file_id: span.anchor.file_id, range }
+                self.db.real_span_map(file_id).span_for_range(token.text_range())
             }
+            base_db::span::HirFileIdRepr::MacroFile(macro_file) => cache
+                .entry(macro_file)
+                .or_insert_with(|| macro_file.expansion_info(self.db.upcast()))
+                .exp_map
+                .span_at(token.text_range().start()),
         };
 
         // fetch span information of token in real file, then use that look through expansions of
@@ -569,24 +562,21 @@ impl<'db> SemanticsImpl<'db> {
         // what about things where spans change? Due to being joined etc, that is we don't find the
         // exact span anymore?
 
+        let def_map = sa.resolver.def_map();
         let mut stack: SmallVec<[_; 4]> = smallvec![InFile::new(sa.file_id, token)];
-        let mut cache = self.expansion_info_cache.borrow_mut();
-        let mut mcache = self.macro_call_cache.borrow_mut();
 
         let mut process_expansion_for_token =
             |stack: &mut SmallVec<_>, macro_file, _token: InFile<&_>| {
                 let expansion_info = cache
                     .entry(macro_file)
-                    .or_insert_with(|| macro_file.expansion_info(self.db.upcast()))
-                    .as_ref()?;
+                    .or_insert_with(|| macro_file.expansion_info(self.db.upcast()));
 
                 {
                     let InFile { file_id, value } = expansion_info.expanded();
                     self.cache(value, file_id);
                 }
 
-                let mapped_tokens =
-                    expansion_info.map_range_down(self.db.upcast(), absolute_range, None)?;
+                let mapped_tokens = expansion_info.map_range_down(span, None)?;
                 let len = stack.len();
 
                 // requeue the tokens we got from mapping our current token down
@@ -599,9 +589,9 @@ impl<'db> SemanticsImpl<'db> {
         // either due to not being in a macro-call or because its unused push it into the result vec,
         // otherwise push the remapped tokens back into the queue as they can potentially be remapped again.
         while let Some(token) = stack.pop() {
-            self.db.unwind_if_cancelled();
             let was_not_remapped = (|| {
                 // First expand into attribute invocations
+
                 let containing_attribute_macro_call = self.with_ctx(|ctx| {
                     token.value.parent_ancestors().filter_map(ast::Item::cast).find_map(|item| {
                         if item.attrs().next().is_none() {
@@ -612,7 +602,7 @@ impl<'db> SemanticsImpl<'db> {
                     })
                 });
                 if let Some(call_id) = containing_attribute_macro_call {
-                    let file_id = call_id.as_file();
+                    let file_id = call_id.as_macro_file();
                     return process_expansion_for_token(&mut stack, file_id, token.as_ref());
                 }
 
@@ -629,7 +619,8 @@ impl<'db> SemanticsImpl<'db> {
                 }
 
                 if let Some(macro_call) = ast::MacroCall::cast(parent.clone()) {
-                    let mcall = token.with_value(macro_call);
+                    let mcall: hir_expand::files::InFileWrapper<HirFileId, ast::MacroCall> =
+                        token.with_value(macro_call);
                     let file_id = match mcache.get(&mcall) {
                         Some(&it) => it,
                         None => {
@@ -659,7 +650,7 @@ impl<'db> SemanticsImpl<'db> {
                         match derive_call {
                             Some(call_id) => {
                                 // resolved to a derive
-                                let file_id = call_id.as_file();
+                                let file_id = call_id.as_macro_file();
                                 return process_expansion_for_token(
                                     &mut stack,
                                     file_id,
@@ -698,7 +689,7 @@ impl<'db> SemanticsImpl<'db> {
                     for (.., derive) in helpers.iter().filter(|(helper, ..)| *helper == attr_name) {
                         res = res.or(process_expansion_for_token(
                             &mut stack,
-                            derive.as_file(),
+                            derive.as_macro_file(),
                             token.as_ref(),
                         ));
                     }
@@ -1052,7 +1043,7 @@ impl<'db> SemanticsImpl<'db> {
 
     fn with_ctx<F: FnOnce(&mut SourceToDefCtx<'_, '_>) -> T, T>(&self, f: F) -> T {
         let mut cache = self.s2d_cache.borrow_mut();
-        let mut ctx = SourceToDefCtx { db: self.db, cache: &mut cache };
+        let mut ctx = SourceToDefCtx { db: self.db, dynmap_cache: &mut cache };
         f(&mut ctx)
     }
 
diff --git a/crates/hir/src/semantics/source_to_def.rs b/crates/hir/src/semantics/source_to_def.rs
index 5b20c873157..df8c1e904fe 100644
--- a/crates/hir/src/semantics/source_to_def.rs
+++ b/crates/hir/src/semantics/source_to_def.rs
@@ -112,7 +112,7 @@ pub(super) type SourceToDefCache = FxHashMap<(ChildContainer, HirFileId), DynMap
 
 pub(super) struct SourceToDefCtx<'a, 'b> {
     pub(super) db: &'b dyn HirDatabase,
-    pub(super) cache: &'a mut SourceToDefCache,
+    pub(super) dynmap_cache: &'a mut SourceToDefCache,
 }
 
 impl SourceToDefCtx<'_, '_> {
@@ -300,7 +300,7 @@ impl SourceToDefCtx<'_, '_> {
 
     fn cache_for(&mut self, container: ChildContainer, file_id: HirFileId) -> &DynMap {
         let db = self.db;
-        self.cache
+        self.dynmap_cache
             .entry((container, file_id))
             .or_insert_with(|| container.child_by_source(db, file_id))
     }
diff --git a/crates/hir/src/source_analyzer.rs b/crates/hir/src/source_analyzer.rs
index 0f1093e6d14..8b114836896 100644
--- a/crates/hir/src/source_analyzer.rs
+++ b/crates/hir/src/source_analyzer.rs
@@ -29,7 +29,7 @@ use hir_expand::{
     mod_path::path,
     name,
     name::{AsName, Name},
-    HirFileId, HirFileIdExt, InFile,
+    HirFileId, HirFileIdExt, InFile, MacroFileId, MacroFileIdExt,
 };
 use hir_ty::{
     diagnostics::{
@@ -753,14 +753,15 @@ impl SourceAnalyzer {
         &self,
         db: &dyn HirDatabase,
         macro_call: InFile<&ast::MacroCall>,
-    ) -> Option<HirFileId> {
+    ) -> Option<MacroFileId> {
         let krate = self.resolver.krate();
         let macro_call_id = macro_call.as_call_id(db.upcast(), krate, |path| {
             self.resolver
                 .resolve_path_as_macro(db.upcast(), &path, Some(MacroSubNs::Bang))
                 .map(|(it, _)| macro_id_to_def_id(db.upcast(), it))
         })?;
-        Some(macro_call_id.as_file()).filter(|it| it.expansion_level(db.upcast()) < 64)
+        // why the 64?
+        Some(macro_call_id.as_macro_file()).filter(|it| it.expansion_level(db.upcast()) < 64)
     }
 
     pub(crate) fn resolve_variant(
diff --git a/crates/mbe/src/syntax_bridge.rs b/crates/mbe/src/syntax_bridge.rs
index bbf49670ce5..5d802ba86c0 100644
--- a/crates/mbe/src/syntax_bridge.rs
+++ b/crates/mbe/src/syntax_bridge.rs
@@ -23,7 +23,7 @@ pub trait SpanMapper<S: Span> {
 
 impl<S: Span> SpanMapper<S> for TokenMap<S> {
     fn span_for(&self, range: TextRange) -> S {
-        self.span_for_range(range)
+        self.span_at(range.start())
     }
 }
 
@@ -152,8 +152,8 @@ where
 {
     let mut map = TokenMap::empty();
     node.descendants_with_tokens().filter_map(NodeOrToken::into_token).for_each(|t| {
-        map.insert(
-            t.text_range(),
+        map.push(
+            t.text_range().start(),
             SpanData { range: t.text_range() - anchor_offset, anchor, ctx: Ctx::DUMMY },
         );
     });
@@ -730,15 +730,13 @@ where
                 self.inner.start_node(SyntaxKind::NAME_REF);
                 self.inner.token(SyntaxKind::INT_NUMBER, left);
                 self.inner.finish_node();
-                let range = TextRange::at(self.text_pos, TextSize::of(left));
-                self.token_map.insert(range, span);
+                self.token_map.push(self.text_pos + TextSize::of(left), span);
 
                 // here we move the exit up, the original exit has been deleted in process
                 self.inner.finish_node();
 
                 self.inner.token(SyntaxKind::DOT, ".");
-                let range = TextRange::at(range.end(), TextSize::of("."));
-                self.token_map.insert(range, span);
+                self.token_map.push(self.text_pos + TextSize::of(left) + TextSize::of("."), span);
 
                 if has_pseudo_dot {
                     assert!(right.is_empty(), "{left}.{right}");
@@ -746,8 +744,7 @@ where
                     assert!(!right.is_empty(), "{left}.{right}");
                     self.inner.start_node(SyntaxKind::NAME_REF);
                     self.inner.token(SyntaxKind::INT_NUMBER, right);
-                    let range = TextRange::at(range.end(), TextSize::of(right));
-                    self.token_map.insert(range, span);
+                    self.token_map.push(self.text_pos + TextSize::of(text), span);
                     self.inner.finish_node();
 
                     // the parser creates an unbalanced start node, we are required to close it here
@@ -772,7 +769,7 @@ where
                 break;
             }
             last = self.cursor;
-            let text: &str = loop {
+            let (text, span) = loop {
                 break match self.cursor.token_tree() {
                     Some(tt::buffer::TokenTreeRef::Leaf(leaf, _)) => {
                         // Mark the range if needed
@@ -788,19 +785,13 @@ where
                             }
                             tt::Leaf::Literal(lit) => (lit.text.as_str(), lit.span),
                         };
-                        let range = TextRange::at(self.text_pos, TextSize::of(text));
-                        self.token_map.insert(range, span);
                         self.cursor = self.cursor.bump();
-                        text
+                        (text, span)
                     }
                     Some(tt::buffer::TokenTreeRef::Subtree(subtree, _)) => {
                         self.cursor = self.cursor.subtree().unwrap();
                         match delim_to_str(subtree.delimiter.kind, false) {
-                            Some(it) => {
-                                let range = TextRange::at(self.text_pos, TextSize::of(it));
-                                self.token_map.insert(range, subtree.delimiter.open);
-                                it
-                            }
+                            Some(it) => (it, subtree.delimiter.open),
                             None => continue,
                         }
                     }
@@ -808,11 +799,7 @@ where
                         let parent = self.cursor.end().unwrap();
                         self.cursor = self.cursor.bump();
                         match delim_to_str(parent.delimiter.kind, true) {
-                            Some(it) => {
-                                let range = TextRange::at(self.text_pos, TextSize::of(it));
-                                self.token_map.insert(range, parent.delimiter.close);
-                                it
-                            }
+                            Some(it) => (it, parent.delimiter.close),
                             None => continue,
                         }
                     }
@@ -820,6 +807,7 @@ where
             };
             self.buf += text;
             self.text_pos += TextSize::of(text);
+            self.token_map.push(self.text_pos, span);
         }
 
         self.inner.token(kind, self.buf.as_str());
@@ -839,8 +827,8 @@ where
             // need to add whitespace either.
             if curr.spacing == tt::Spacing::Alone && curr.char != ';' && next.char != '\'' {
                 self.inner.token(WHITESPACE, " ");
-                self.token_map.insert(TextRange::at(self.text_pos, TextSize::of(' ')), curr.span);
                 self.text_pos += TextSize::of(' ');
+                self.token_map.push(self.text_pos, curr.span);
             }
         }
     }
diff --git a/crates/mbe/src/token_map.rs b/crates/mbe/src/token_map.rs
index 32f61af25e5..c2ec30ca72f 100644
--- a/crates/mbe/src/token_map.rs
+++ b/crates/mbe/src/token_map.rs
@@ -3,7 +3,7 @@
 use std::hash::Hash;
 
 use stdx::itertools::Itertools;
-use syntax::TextRange;
+use syntax::{TextRange, TextSize};
 use tt::Span;
 
 /// Maps absolute text ranges for the corresponding file to the relevant span data.
@@ -12,138 +12,46 @@ use tt::Span;
 pub struct TokenMap<S: Span> {
     // FIXME: This needs to be sorted by (FileId, AstId)
     // Then we can do a binary search on the file id,
-    // then a bin search on the ast id
-    pub span_map: Vec<(TextRange, S)>,
-    // span_map2: rustc_hash::FxHashMap<TextRange, usize>,
+    // then a bin search on the ast id?
+    spans: Vec<(TextSize, S)>,
 }
 
 impl<S: Span> TokenMap<S> {
-    pub fn empty() -> Self {
-        Self { span_map: Vec::new() }
+    pub(crate) fn empty() -> Self {
+        Self { spans: Vec::new() }
     }
 
-    pub fn finish(&mut self) {
-        debug_assert_eq!(
-            self.span_map
-                .iter()
-                .sorted_by_key(|it| (it.0.start(), it.0.end()))
-                .tuple_windows()
-                .find(|(range, next)| range.0.end() != next.0.start()),
-            None,
-            "span map has holes!"
-        );
-        self.span_map.shrink_to_fit();
+    pub(crate) fn finish(&mut self) {
+        assert!(self.spans.iter().tuple_windows().all(|(a, b)| a.0 < b.0));
+        self.spans.shrink_to_fit();
     }
 
-    pub(crate) fn insert(&mut self, range: TextRange, span: S) {
-        self.span_map.push((range, span));
+    pub(crate) fn push(&mut self, offset: TextSize, span: S) {
+        self.spans.push((offset, span));
     }
 
     pub fn ranges_with_span(&self, span: S) -> impl Iterator<Item = TextRange> + '_ {
         // FIXME: linear search
-        // FIXME: Disregards resolving spans to get more matches! See ExpansionInfo::map_token_down
-        self.span_map.iter().filter_map(
-            move |(range, s)| {
-                if s == &span {
-                    Some(*range)
-                } else {
-                    None
-                }
-            },
-        )
+        self.spans.iter().enumerate().filter_map(move |(idx, &(end, s))| {
+            if s != span {
+                return None;
+            }
+            let start = idx.checked_sub(1).map_or(TextSize::new(0), |prev| self.spans[prev].0);
+            Some(TextRange::new(start, end))
+        })
     }
 
     // FIXME: We need APIs for fetching the span of a token as well as for a whole node. The node
     // one *is* fallible though.
-    // Token span fetching technically only needs an offset really, as the entire file span is
-    // populated, where node fetching is more like fetching the spans at all source positions, and
-    // then we need to verify that all those positions have the same context, if not we fail! But
-    // how do we handle them having different span ranges?
-
-    pub fn span_for_range(&self, range: TextRange) -> S {
-        // TODO FIXME: make this proper
-        self.span_map
-            .iter()
-            .filter_map(|(r, s)| Some((r, s, r.intersect(range).filter(|it| !it.is_empty())?)))
-            .max_by_key(|(_, _, intersection)| intersection.len())
-            .map_or_else(
-                || panic!("no span for range {:?} in {:#?}", range, self.span_map),
-                |(_, &s, _)| s,
-            )
+    pub fn span_at(&self, offset: TextSize) -> S {
+        let entry = self.spans.partition_point(|&(it, _)| it <= offset);
+        self.spans[entry].1
     }
 
     pub fn spans_for_node_range(&self, range: TextRange) -> impl Iterator<Item = S> + '_ {
-        // TODO FIXME: make this proper
-        self.span_map
-            .iter()
-            .filter(move |(r, _)| r.intersect(range).filter(|it| !it.is_empty()).is_some())
-            .map(|&(_, s)| s)
+        let (start, end) = (range.start(), range.end());
+        let start_entry = self.spans.partition_point(|&(it, _)| it <= start);
+        let end_entry = self.spans[start_entry..].partition_point(|&(it, _)| it <= end); // FIXME: this might be wrong?
+        (&self.spans[start_entry..][..end_entry]).iter().map(|&(_, s)| s)
     }
-
-    // pub fn ranges_by_token(
-    //     &self,
-    //     token_id: tt::TokenId,
-    //     kind: SyntaxKind,
-    // ) -> impl Iterator<Item = TextRange> + '_ {
-    //     self.entries
-    //         .iter()
-    //         .filter(move |&&(tid, _)| tid == token_id)
-    //         .filter_map(move |(_, range)| range.by_kind(kind))
-    // }
-
-    // pub(crate) fn remove_delim(&mut self, idx: usize) {
-    //     // FIXME: This could be accidentally quadratic
-    //     self.entries.remove(idx);
-    // }
-
-    // pub fn entries(&self) -> impl Iterator<Item = (tt::TokenId, TextRange)> + '_ {
-    //     self.entries.iter().filter_map(|&(tid, tr)| match tr {
-    //         TokenTextRange::Token(range) => Some((tid, range)),
-    //         TokenTextRange::Delimiter(_) => None,
-    //     })
-    // }
-
-    // pub fn filter(&mut self, id: impl Fn(tt::TokenId) -> bool) {
-    //     self.entries.retain(|&(tid, _)| id(tid));
-    // }
-    // pub fn synthetic_token_id(&self, token_id: tt::TokenId) -> Option<SyntheticTokenId> {
-    //     self.synthetic_entries.iter().find(|(tid, _)| *tid == token_id).map(|(_, id)| *id)
-    // }
-
-    // pub fn first_range_by_token(
-    //     &self,
-    //     token_id: tt::TokenId,
-    //     kind: SyntaxKind,
-    // ) -> Option<TextRange> {
-    //     self.ranges_by_token(token_id, kind).next()
-    // }
-
-    // pub(crate) fn insert(&mut self, token_id: tt::TokenId, relative_range: TextRange) {
-    //     self.entries.push((token_id, TokenTextRange::Token(relative_range)));
-    // }
-
-    // pub(crate) fn insert_synthetic(&mut self, token_id: tt::TokenId, id: SyntheticTokenId) {
-    //     self.synthetic_entries.push((token_id, id));
-    // }
-
-    // pub(crate) fn insert_delim(
-    //     &mut self,
-    //     token_id: tt::TokenId,
-    //     open_relative_range: TextRange,
-    //     close_relative_range: TextRange,
-    // ) -> usize {
-    //     let res = self.entries.len();
-    //     let cover = open_relative_range.cover(close_relative_range);
-
-    //     self.entries.push((token_id, TokenTextRange::Delimiter(cover)));
-    //     res
-    // }
-
-    // pub(crate) fn update_close_delim(&mut self, idx: usize, close_relative_range: TextRange) {
-    //     let (_, token_text_range) = &mut self.entries[idx];
-    //     if let TokenTextRange::Delimiter(dim) = token_text_range {
-    //         let cover = dim.cover(close_relative_range);
-    //         *token_text_range = TokenTextRange::Delimiter(cover);
-    //     }
-    // }
 }
diff --git a/crates/rust-analyzer/src/integrated_benchmarks.rs b/crates/rust-analyzer/src/integrated_benchmarks.rs
index b7e839ac2a5..a29e18811d8 100644
--- a/crates/rust-analyzer/src/integrated_benchmarks.rs
+++ b/crates/rust-analyzer/src/integrated_benchmarks.rs
@@ -56,8 +56,7 @@ fn integrated_highlighting_benchmark() {
         analysis.highlight_as_html(file_id, false).unwrap();
     }
 
-    profile::init_from("*>100");
-    // let _s = profile::heartbeat_span();
+    profile::init_from("*>1");
 
     {
         let _it = stdx::timeit("change");