Merge pull request #18852 from ChayimFriedman2/proc-macro-panic

fix: Fix a bug that was caused by fixup reversing

3 files changed, 82 insertions, 14 deletions

diff --git a/src/tools/rust-analyzer/crates/hir-expand/src/fixup.rs b/src/tools/rust-analyzer/crates/hir-expand/src/fixup.rs
index ad001f72deb..3d2d52a0afe 100644
--- a/src/tools/rust-analyzer/crates/hir-expand/src/fixup.rs
+++ b/src/tools/rust-analyzer/crates/hir-expand/src/fixup.rs
@@ -388,6 +388,7 @@ pub(crate) fn reverse_fixups(tt: &mut TopSubtree, undo_info: &SyntaxFixupUndoInf
     reverse_fixups_(tt, undo_info);
 }
 
+#[derive(Debug)]
 enum TransformTtAction<'a> {
     Keep,
     ReplaceWith(tt::TokenTreesView<'a>),
@@ -399,27 +400,40 @@ impl TransformTtAction<'_> {
     }
 }
 
+/// This function takes a token tree, and calls `callback` with each token tree in it.
+/// Then it does what the callback says: keeps the tt or replaces it with a (possibly empty)
+/// tts view.
 fn transform_tt<'a, 'b>(
     tt: &'a mut Vec<tt::TokenTree>,
     mut callback: impl FnMut(&mut tt::TokenTree) -> TransformTtAction<'b>,
 ) {
+    // We need to keep a stack of the currently open subtrees, because we need to update
+    // them if we change the number of items in them.
     let mut subtrees_stack = Vec::new();
     let mut i = 0;
     while i < tt.len() {
-        while let Some(&subtree_idx) = subtrees_stack.last() {
+        'pop_finished_subtrees: while let Some(&subtree_idx) = subtrees_stack.last() {
             let tt::TokenTree::Subtree(subtree) = &tt[subtree_idx] else {
                 unreachable!("non-subtree on subtrees stack");
             };
-            if subtree_idx + 1 + subtree.usize_len() == i {
+            if i >= subtree_idx + 1 + subtree.usize_len() {
                 subtrees_stack.pop();
             } else {
-                break;
+                break 'pop_finished_subtrees;
             }
         }
 
         let action = callback(&mut tt[i]);
         match action {
-            TransformTtAction::Keep => {}
+            TransformTtAction::Keep => {
+                // This cannot be shared with the replaced case, because then we may push the same subtree
+                // twice, and will update it twice which will lead to errors.
+                if let tt::TokenTree::Subtree(_) = &tt[i] {
+                    subtrees_stack.push(i);
+                }
+
+                i += 1;
+            }
             TransformTtAction::ReplaceWith(replacement) => {
                 let old_len = 1 + match &tt[i] {
                     tt::TokenTree::Leaf(_) => 0,
@@ -427,23 +441,17 @@ fn transform_tt<'a, 'b>(
                 };
                 let len_diff = replacement.len() as i64 - old_len as i64;
                 tt.splice(i..i + old_len, replacement.flat_tokens().iter().cloned());
-                i = i.checked_add_signed(len_diff as isize).unwrap();
+                // `+1` for the loop.
+                i = i.checked_add_signed(len_diff as isize + 1).unwrap();
 
                 for &subtree_idx in &subtrees_stack {
                     let tt::TokenTree::Subtree(subtree) = &mut tt[subtree_idx] else {
                         unreachable!("non-subtree on subtrees stack");
                     };
-                    subtree.len = (subtree.len as i64 + len_diff).try_into().unwrap();
+                    subtree.len = (i64::from(subtree.len) + len_diff).try_into().unwrap();
                 }
             }
         }
-
-        // `tt[i]` might have been removed.
-        if let Some(tt::TokenTree::Subtree(_)) = tt.get(i) {
-            subtrees_stack.push(i);
-        }
-
-        i += 1;
     }
 }
 
diff --git a/src/tools/rust-analyzer/crates/ide/src/inlay_hints.rs b/src/tools/rust-analyzer/crates/ide/src/inlay_hints.rs
index aa99ba49bc8..faa65019eea 100644
--- a/src/tools/rust-analyzer/crates/ide/src/inlay_hints.rs
+++ b/src/tools/rust-analyzer/crates/ide/src/inlay_hints.rs
@@ -856,4 +856,18 @@ fn main() {
 }"#,
         );
     }
+
+    #[test]
+    fn regression_18840() {
+        check(
+            r#"
+//- proc_macros: issue_18840
+#[proc_macros::issue_18840]
+fn foo() {
+    let
+    loop {}
+}
+"#,
+        );
+    }
 }
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 1f9a14cc903..b40b7757c6e 100644
--- a/src/tools/rust-analyzer/crates/test-fixture/src/lib.rs
+++ b/src/tools/rust-analyzer/crates/test-fixture/src/lib.rs
@@ -376,7 +376,7 @@ impl ChangeFixture {
     }
 }
 
-fn default_test_proc_macros() -> [(String, ProcMacro); 6] {
+fn default_test_proc_macros() -> [(String, ProcMacro); 7] {
     [
         (
             r#"
@@ -468,6 +468,21 @@ pub fn issue_18089(_attr: TokenStream, _item: TokenStream) -> TokenStream {
                 disabled: false,
             },
         ),
+        (
+            r#"
+#[proc_macro_attribute]
+pub fn issue_18840(_attr: TokenStream, _item: TokenStream) -> TokenStream {
+    loop {}
+}
+"#
+            .into(),
+            ProcMacro {
+                name: Symbol::intern("issue_18840"),
+                kind: ProcMacroKind::Attr,
+                expander: sync::Arc::new(Issue18840ProcMacroExpander),
+                disabled: false,
+            },
+        ),
     ]
 }
 
@@ -646,6 +661,37 @@ impl ProcMacroExpander for AttributeInputReplaceProcMacroExpander {
 }
 
 #[derive(Debug)]
+struct Issue18840ProcMacroExpander;
+impl ProcMacroExpander for Issue18840ProcMacroExpander {
+    fn expand(
+        &self,
+        fn_: &TopSubtree,
+        _: Option<&TopSubtree>,
+        _: &Env,
+        def_site: Span,
+        _: Span,
+        _: Span,
+        _: Option<String>,
+    ) -> Result<TopSubtree, ProcMacroExpansionError> {
+        // Input:
+        // ```
+        // #[issue_18840]
+        // fn foo() { let loop {} }
+        // ```
+
+        // The span that was created by the fixup infra.
+        let fixed_up_span = fn_.token_trees().flat_tokens()[5].first_span();
+        let mut result =
+            quote! {fixed_up_span => ::core::compile_error! { "my cool compile_error!" } };
+        // Make it so we won't remove the top subtree when reversing fixups.
+        let top_subtree_delimiter_mut = result.top_subtree_delimiter_mut();
+        top_subtree_delimiter_mut.open = def_site;
+        top_subtree_delimiter_mut.close = def_site;
+        Ok(result)
+    }
+}
+
+#[derive(Debug)]
 struct MirrorProcMacroExpander;
 impl ProcMacroExpander for MirrorProcMacroExpander {
     fn expand(