Always cache macro expansions' root node in Semantics

Previously some expansions were not cached, but were cached in the expansion cache, which caused panics when later queries tried to lookup the node from the expansion cache.

7 files changed, 166 insertions, 48 deletions

diff --git a/src/tools/rust-analyzer/crates/hir-expand/src/builtin.rs b/src/tools/rust-analyzer/crates/hir-expand/src/builtin.rs
index 252430e4e95..7b9b7f36e2c 100644
--- a/src/tools/rust-analyzer/crates/hir-expand/src/builtin.rs
+++ b/src/tools/rust-analyzer/crates/hir-expand/src/builtin.rs
@@ -1,6 +1,6 @@
 //! Builtin macros and attributes
 #[macro_use]
-mod quote;
+pub mod quote;
 
 mod attr_macro;
 mod derive_macro;
diff --git a/src/tools/rust-analyzer/crates/hir-expand/src/builtin/quote.rs b/src/tools/rust-analyzer/crates/hir-expand/src/builtin/quote.rs
index fc248c906da..418d8d9660b 100644
--- a/src/tools/rust-analyzer/crates/hir-expand/src/builtin/quote.rs
+++ b/src/tools/rust-analyzer/crates/hir-expand/src/builtin/quote.rs
@@ -18,6 +18,7 @@ pub(crate) fn dollar_crate(span: Span) -> tt::Ident<Span> {
 // 2. #()* pattern repetition not supported now
 //    * But we can do it manually, see `test_quote_derive_copy_hack`
 #[doc(hidden)]
+#[macro_export]
 macro_rules! quote_impl__ {
     ($span:ident) => {
         Vec::<$crate::tt::TokenTree>::new()
@@ -27,8 +28,8 @@ macro_rules! quote_impl__ {
         {
             let children = $crate::builtin::quote::__quote!($span $($tt)*);
             $crate::tt::Subtree {
-                delimiter: crate::tt::Delimiter {
-                    kind: crate::tt::DelimiterKind::$delim,
+                delimiter: $crate::tt::Delimiter {
+                    kind: $crate::tt::DelimiterKind::$delim,
                     open: $span,
                     close: $span,
                 },
@@ -40,9 +41,9 @@ macro_rules! quote_impl__ {
     ( @PUNCT($span:ident) $first:literal ) => {
         {
             vec![
-                crate::tt::Leaf::Punct(crate::tt::Punct {
+                $crate::tt::Leaf::Punct($crate::tt::Punct {
                     char: $first,
-                    spacing: crate::tt::Spacing::Alone,
+                    spacing: $crate::tt::Spacing::Alone,
                     span: $span,
                 }).into()
             ]
@@ -52,14 +53,14 @@ macro_rules! quote_impl__ {
     ( @PUNCT($span:ident) $first:literal, $sec:literal ) => {
         {
             vec![
-                crate::tt::Leaf::Punct(crate::tt::Punct {
+                $crate::tt::Leaf::Punct($crate::tt::Punct {
                     char: $first,
-                    spacing: crate::tt::Spacing::Joint,
+                    spacing: $crate::tt::Spacing::Joint,
                     span: $span,
                 }).into(),
-                crate::tt::Leaf::Punct(crate::tt::Punct {
+                $crate::tt::Leaf::Punct($crate::tt::Punct {
                     char: $sec,
-                    spacing: crate::tt::Spacing::Alone,
+                    spacing: $crate::tt::Spacing::Alone,
                     span: $span,
                 }).into()
             ]
@@ -98,7 +99,7 @@ macro_rules! quote_impl__ {
     // Ident
     ($span:ident $tt:ident ) => {
         vec![ {
-            crate::tt::Leaf::Ident(crate::tt::Ident {
+            $crate::tt::Leaf::Ident($crate::tt::Ident {
                 sym: intern::Symbol::intern(stringify!($tt)),
                 span: $span,
                 is_raw: tt::IdentIsRaw::No,
@@ -109,6 +110,7 @@ macro_rules! quote_impl__ {
     // Puncts
     // FIXME: Not all puncts are handled
     ($span:ident -> ) => {$crate::builtin::quote::__quote!(@PUNCT($span) '-', '>')};
+    ($span:ident => ) => {$crate::builtin::quote::__quote!(@PUNCT($span) '=', '>')};
     ($span:ident & ) => {$crate::builtin::quote::__quote!(@PUNCT($span) '&')};
     ($span:ident , ) => {$crate::builtin::quote::__quote!(@PUNCT($span) ',')};
     ($span:ident : ) => {$crate::builtin::quote::__quote!(@PUNCT($span) ':')};
@@ -118,6 +120,9 @@ macro_rules! quote_impl__ {
     ($span:ident < ) => {$crate::builtin::quote::__quote!(@PUNCT($span) '<')};
     ($span:ident > ) => {$crate::builtin::quote::__quote!(@PUNCT($span) '>')};
     ($span:ident ! ) => {$crate::builtin::quote::__quote!(@PUNCT($span) '!')};
+    ($span:ident # ) => {$crate::builtin::quote::__quote!(@PUNCT($span) '#')};
+    ($span:ident $ ) => {$crate::builtin::quote::__quote!(@PUNCT($span) '$')};
+    ($span:ident * ) => {$crate::builtin::quote::__quote!(@PUNCT($span) '*')};
 
     ($span:ident $first:tt $($tail:tt)+ ) => {
         {
@@ -129,18 +134,19 @@ macro_rules! quote_impl__ {
         }
     };
 }
-pub(super) use quote_impl__ as __quote;
+pub use quote_impl__ as __quote;
 
 /// FIXME:
 /// It probably should implement in proc-macro
-macro_rules! quote_impl {
+#[macro_export]
+macro_rules! quote {
     ($span:ident=> $($tt:tt)* ) => {
         $crate::builtin::quote::IntoTt::to_subtree($crate::builtin::quote::__quote!($span $($tt)*), $span)
     }
 }
-pub(super) use quote_impl as quote;
+pub(super) use quote;
 
-pub(crate) trait IntoTt {
+pub trait IntoTt {
     fn to_subtree(self, span: Span) -> crate::tt::Subtree;
     fn to_tokens(self) -> Vec<crate::tt::TokenTree>;
 }
@@ -168,7 +174,7 @@ impl IntoTt for crate::tt::Subtree {
     }
 }
 
-pub(crate) trait ToTokenTree {
+pub trait ToTokenTree {
     fn to_token(self, span: Span) -> crate::tt::TokenTree;
 }
 
diff --git a/src/tools/rust-analyzer/crates/hir/src/semantics.rs b/src/tools/rust-analyzer/crates/hir/src/semantics.rs
index 882ac229dcf..dcaff044420 100644
--- a/src/tools/rust-analyzer/crates/hir/src/semantics.rs
+++ b/src/tools/rust-analyzer/crates/hir/src/semantics.rs
@@ -449,7 +449,7 @@ impl<'db> SemanticsImpl<'db> {
             Some(
                 calls
                     .into_iter()
-                    .map(|call| macro_call_to_macro_id(ctx, call?).map(|id| Macro { id }))
+                    .map(|call| macro_call_to_macro_id(self, ctx, call?).map(|id| Macro { id }))
                     .collect(),
             )
         })
@@ -892,16 +892,7 @@ impl<'db> SemanticsImpl<'db> {
             let InMacroFile { file_id, value: mapped_tokens } = self.with_ctx(|ctx| {
                 Some(
                     ctx.cache
-                        .expansion_info_cache
-                        .entry(macro_file)
-                        .or_insert_with(|| {
-                            let exp_info = macro_file.expansion_info(self.db.upcast());
-
-                            let InMacroFile { file_id, value } = exp_info.expanded();
-                            self.cache(value, file_id.into());
-
-                            exp_info
-                        })
+                        .get_or_insert_expansion(self, macro_file)
                         .map_range_down(span)?
                         .map(SmallVec::<[_; 2]>::from_iter),
                 )
@@ -1187,11 +1178,7 @@ impl<'db> SemanticsImpl<'db> {
                     let macro_file = file_id.macro_file()?;
 
                     self.with_ctx(|ctx| {
-                        let expansion_info = ctx
-                            .cache
-                            .expansion_info_cache
-                            .entry(macro_file)
-                            .or_insert_with(|| macro_file.expansion_info(self.db.upcast()));
+                        let expansion_info = ctx.cache.get_or_insert_expansion(self, macro_file);
                         expansion_info.arg().map(|node| node?.parent()).transpose()
                     })
                 }
@@ -1407,7 +1394,7 @@ impl<'db> SemanticsImpl<'db> {
         let macro_call = self.find_file(macro_call.syntax()).with_value(macro_call);
         self.with_ctx(|ctx| {
             ctx.macro_call_to_macro_call(macro_call)
-                .and_then(|call| macro_call_to_macro_id(ctx, call))
+                .and_then(|call| macro_call_to_macro_id(self, ctx, call))
                 .map(Into::into)
         })
         .or_else(|| {
@@ -1449,7 +1436,7 @@ impl<'db> SemanticsImpl<'db> {
         let item_in_file = self.wrap_node_infile(item.clone());
         let id = self.with_ctx(|ctx| {
             let macro_call_id = ctx.item_to_macro_call(item_in_file.as_ref())?;
-            macro_call_to_macro_id(ctx, macro_call_id)
+            macro_call_to_macro_id(self, ctx, macro_call_id)
         })?;
         Some(Macro { id })
     }
@@ -1769,6 +1756,7 @@ impl<'db> SemanticsImpl<'db> {
 }
 
 fn macro_call_to_macro_id(
+    sema: &SemanticsImpl<'_>,
     ctx: &mut SourceToDefCtx<'_, '_>,
     macro_call_id: MacroCallId,
 ) -> Option<MacroId> {
@@ -1784,11 +1772,7 @@ fn macro_call_to_macro_id(
                     it.to_ptr(db).to_node(&db.parse(file_id).syntax_node())
                 }
                 HirFileIdRepr::MacroFile(macro_file) => {
-                    let expansion_info = ctx
-                        .cache
-                        .expansion_info_cache
-                        .entry(macro_file)
-                        .or_insert_with(|| macro_file.expansion_info(ctx.db.upcast()));
+                    let expansion_info = ctx.cache.get_or_insert_expansion(sema, macro_file);
                     it.to_ptr(db).to_node(&expansion_info.expanded().value)
                 }
             };
@@ -1800,11 +1784,7 @@ fn macro_call_to_macro_id(
                     it.to_ptr(db).to_node(&db.parse(file_id).syntax_node())
                 }
                 HirFileIdRepr::MacroFile(macro_file) => {
-                    let expansion_info = ctx
-                        .cache
-                        .expansion_info_cache
-                        .entry(macro_file)
-                        .or_insert_with(|| macro_file.expansion_info(ctx.db.upcast()));
+                    let expansion_info = ctx.cache.get_or_insert_expansion(sema, macro_file);
                     it.to_ptr(db).to_node(&expansion_info.expanded().value)
                 }
             };
diff --git a/src/tools/rust-analyzer/crates/hir/src/semantics/source_to_def.rs b/src/tools/rust-analyzer/crates/hir/src/semantics/source_to_def.rs
index fc629b9c6d0..c1e4e1d1e27 100644
--- a/src/tools/rust-analyzer/crates/hir/src/semantics/source_to_def.rs
+++ b/src/tools/rust-analyzer/crates/hir/src/semantics/source_to_def.rs
@@ -99,7 +99,8 @@ use hir_def::{
     VariantId,
 };
 use hir_expand::{
-    attrs::AttrId, name::AsName, ExpansionInfo, HirFileId, HirFileIdExt, MacroCallId,
+    attrs::AttrId, name::AsName, ExpansionInfo, HirFileId, HirFileIdExt, InMacroFile, MacroCallId,
+    MacroFileIdExt,
 };
 use rustc_hash::FxHashMap;
 use smallvec::SmallVec;
@@ -110,15 +111,32 @@ use syntax::{
     AstNode, AstPtr, SyntaxNode,
 };
 
-use crate::{db::HirDatabase, InFile, InlineAsmOperand};
+use crate::{db::HirDatabase, InFile, InlineAsmOperand, SemanticsImpl};
 
 #[derive(Default)]
 pub(super) struct SourceToDefCache {
     pub(super) dynmap_cache: FxHashMap<(ChildContainer, HirFileId), DynMap>,
-    pub(super) expansion_info_cache: FxHashMap<MacroFileId, ExpansionInfo>,
+    expansion_info_cache: FxHashMap<MacroFileId, ExpansionInfo>,
     pub(super) file_to_def_cache: FxHashMap<FileId, SmallVec<[ModuleId; 1]>>,
 }
 
+impl SourceToDefCache {
+    pub(super) fn get_or_insert_expansion(
+        &mut self,
+        sema: &SemanticsImpl<'_>,
+        macro_file: MacroFileId,
+    ) -> &ExpansionInfo {
+        self.expansion_info_cache.entry(macro_file).or_insert_with(|| {
+            let exp_info = macro_file.expansion_info(sema.db.upcast());
+
+            let InMacroFile { file_id, value } = exp_info.expanded();
+            sema.cache(value, file_id.into());
+
+            exp_info
+        })
+    }
+}
+
 pub(super) struct SourceToDefCtx<'db, 'cache> {
     pub(super) db: &'db dyn HirDatabase,
     pub(super) cache: &'cache mut SourceToDefCache,
diff --git a/src/tools/rust-analyzer/crates/ide/src/syntax_highlighting/test_data/highlight_issue_18089.html b/src/tools/rust-analyzer/crates/ide/src/syntax_highlighting/test_data/highlight_issue_18089.html
new file mode 100644
index 00000000000..d07ba74db24
--- /dev/null
+++ b/src/tools/rust-analyzer/crates/ide/src/syntax_highlighting/test_data/highlight_issue_18089.html
@@ -0,0 +1,53 @@
+
+<style>
+body                { margin: 0; }
+pre                 { color: #DCDCCC; background: #3F3F3F; font-size: 22px; padding: 0.4em; }
+
+.lifetime           { color: #DFAF8F; font-style: italic; }
+.label              { color: #DFAF8F; font-style: italic; }
+.comment            { color: #7F9F7F; }
+.documentation      { color: #629755; }
+.intra_doc_link     { font-style: italic; }
+.injected           { opacity: 0.65 ; }
+.struct, .enum      { color: #7CB8BB; }
+.enum_variant       { color: #BDE0F3; }
+.string_literal     { color: #CC9393; }
+.field              { color: #94BFF3; }
+.function           { color: #93E0E3; }
+.function.unsafe    { color: #BC8383; }
+.trait.unsafe       { color: #BC8383; }
+.operator.unsafe    { color: #BC8383; }
+.mutable.unsafe     { color: #BC8383; text-decoration: underline; }
+.keyword.unsafe     { color: #BC8383; font-weight: bold; }
+.macro.unsafe       { color: #BC8383; }
+.parameter          { color: #94BFF3; }
+.text               { color: #DCDCCC; }
+.type               { color: #7CB8BB; }
+.builtin_type       { color: #8CD0D3; }
+.type_param         { color: #DFAF8F; }
+.attribute          { color: #94BFF3; }
+.numeric_literal    { color: #BFEBBF; }
+.bool_literal       { color: #BFE6EB; }
+.macro              { color: #94BFF3; }
+.proc_macro         { color: #94BFF3; text-decoration: underline; }
+.derive             { color: #94BFF3; font-style: italic; }
+.module             { color: #AFD8AF; }
+.value_param        { color: #DCDCCC; }
+.variable           { color: #DCDCCC; }
+.format_specifier   { color: #CC696B; }
+.mutable            { text-decoration: underline; }
+.escape_sequence    { color: #94BFF3; }
+.keyword            { color: #F0DFAF; font-weight: bold; }
+.control            { font-style: italic; }
+.reference          { font-style: italic; font-weight: bold; }
+.const              { font-weight: bolder; }
+
+.invalid_escape_sequence { color: #FC5555; text-decoration: wavy underline; }
+.unresolved_reference    { color: #FC5555; text-decoration: wavy underline; }
+</style>
+<pre><code><span class="keyword">fn</span> <span class="function declaration">main</span><span class="parenthesis">(</span><span class="parenthesis">)</span> <span class="brace">{</span>
+    <span class="macro">template</span><span class="macro_bang">!</span><span class="parenthesis macro">(</span><span class="none macro">template</span><span class="parenthesis macro">)</span><span class="semicolon">;</span>
+<span class="brace">}</span>
+
+<span class="attribute_bracket attribute">#</span><span class="attribute_bracket attribute">[</span><span class="module attribute crate_root library">proc_macros</span><span class="operator attribute">::</span><span class="attribute attribute library">issue_18089</span><span class="attribute_bracket attribute">]</span>
+<span class="keyword">fn</span> <span class="function declaration">template</span><span class="parenthesis">(</span><span class="parenthesis">)</span> <span class="brace">{</span><span class="brace">}</span></code></pre>
\ No newline at end of file
diff --git a/src/tools/rust-analyzer/crates/ide/src/syntax_highlighting/tests.rs b/src/tools/rust-analyzer/crates/ide/src/syntax_highlighting/tests.rs
index 82833d716b5..94cee4ef43b 100644
--- a/src/tools/rust-analyzer/crates/ide/src/syntax_highlighting/tests.rs
+++ b/src/tools/rust-analyzer/crates/ide/src/syntax_highlighting/tests.rs
@@ -1358,3 +1358,20 @@ pub unsafe fn bootstrap(msp: *const u32, rv: *const u32) -> ! {
         false,
     );
 }
+
+#[test]
+fn issue_18089() {
+    check_highlighting(
+        r#"
+//- proc_macros: issue_18089
+fn main() {
+    template!(template);
+}
+
+#[proc_macros::issue_18089]
+fn template() {}
+"#,
+        expect_file!["./test_data/highlight_issue_18089.html"],
+        false,
+    );
+}
diff --git a/src/tools/rust-analyzer/crates/test-fixture/src/lib.rs b/src/tools/rust-analyzer/crates/test-fixture/src/lib.rs
index 3443f1f9003..593e31c2fbb 100644
--- a/src/tools/rust-analyzer/crates/test-fixture/src/lib.rs
+++ b/src/tools/rust-analyzer/crates/test-fixture/src/lib.rs
@@ -13,7 +13,7 @@ use hir_expand::{
     proc_macro::{
         ProcMacro, ProcMacroExpander, ProcMacroExpansionError, ProcMacroKind, ProcMacrosBuilder,
     },
-    FileRange,
+    quote, FileRange,
 };
 use intern::Symbol;
 use rustc_hash::FxHashMap;
@@ -374,7 +374,7 @@ impl ChangeFixture {
     }
 }
 
-fn default_test_proc_macros() -> [(String, ProcMacro); 5] {
+fn default_test_proc_macros() -> [(String, ProcMacro); 6] {
     [
         (
             r#"
@@ -451,6 +451,21 @@ pub fn shorten(input: TokenStream) -> TokenStream {
                 disabled: false,
             },
         ),
+        (
+            r#"
+#[proc_macro_attribute]
+pub fn issue_18089(_attr: TokenStream, _item: TokenStream) -> TokenStream {
+    loop {}
+}
+"#
+            .into(),
+            ProcMacro {
+                name: Symbol::intern("issue_18089"),
+                kind: ProcMacroKind::Attr,
+                expander: sync::Arc::new(Issue18089ProcMacroExpander),
+                disabled: false,
+            },
+        ),
     ]
 }
 
@@ -577,6 +592,35 @@ impl ProcMacroExpander for IdentityProcMacroExpander {
     }
 }
 
+// Expands to a macro_rules! macro, for issue #18089.
+#[derive(Debug)]
+struct Issue18089ProcMacroExpander;
+impl ProcMacroExpander for Issue18089ProcMacroExpander {
+    fn expand(
+        &self,
+        subtree: &Subtree<Span>,
+        _: Option<&Subtree<Span>>,
+        _: &Env,
+        _: Span,
+        call_site: Span,
+        _: Span,
+        _: Option<String>,
+    ) -> Result<Subtree<Span>, ProcMacroExpansionError> {
+        let macro_name = &subtree.token_trees[1];
+        Ok(quote! { call_site =>
+            #[macro_export]
+            macro_rules! my_macro___ {
+                ($($token:tt)*) => {{
+                }};
+            }
+
+            pub use my_macro___ as #macro_name;
+
+            #subtree
+        })
+    }
+}
+
 // Pastes the attribute input as its output
 #[derive(Debug)]
 struct AttributeInputReplaceProcMacroExpander;