Auto merge of #18210 - ChayimFriedman2:label-macro, r=Veykril

fix: Fix resolution of label inside macro

When working on Something Else (TM) (I left a hint in the commits :P), I noticed to my surprise that labels inside macros are not resolved. This led to a discovery of *two* unrelated bugs, which are hereby fixed in two commits.

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

@@ -36,7 +36,7 @@ macro_rules! f {
 }
 
 struct#0:1@58..64#1# MyTraitMap2#0:2@31..42#0# {#0:1@72..73#1#
-    map#0:1@86..89#1#:#0:1@89..90#1# #0:1@89..90#1#::#0:1@91..92#1#std#0:1@93..96#1#::#0:1@96..97#1#collections#0:1@98..109#1#::#0:1@109..110#1#HashSet#0:1@111..118#1#<#0:1@118..119#1#(#0:1@119..120#1#)#0:1@120..121#1#>#0:1@121..122#1#,#0:1@122..123#1#
+    map#0:1@86..89#1#:#0:1@89..90#1# #0:1@89..90#1#::#0:1@91..93#1#std#0:1@93..96#1#::#0:1@96..98#1#collections#0:1@98..109#1#::#0:1@109..111#1#HashSet#0:1@111..118#1#<#0:1@118..119#1#(#0:1@119..120#1#)#0:1@120..121#1#>#0:1@121..122#1#,#0:1@122..123#1#
 }#0:1@132..133#1#
 "#]],
     );

crates/hir/src/semantics.rs

@@ -36,9 +36,9 @@ use span::{EditionedFileId, FileId, HirFileIdRepr};
 use stdx::TupleExt;
 use syntax::{
     algo::skip_trivia_token,
-    ast::{self, HasAttrs as _, HasGenericParams, HasLoopBody, IsString as _},
-    match_ast, AstNode, AstToken, Direction, SyntaxKind, SyntaxNode, SyntaxNodePtr, SyntaxToken,
-    TextRange, TextSize,
+    ast::{self, HasAttrs as _, HasGenericParams, IsString as _},
+    AstNode, AstToken, Direction, SyntaxKind, SyntaxNode, SyntaxNodePtr, SyntaxToken, TextRange,
+    TextSize,
 };
 
 use crate::{
@@ -1221,26 +1221,10 @@ impl<'db> SemanticsImpl<'db> {
         ToDef::to_def(self, src.as_ref())
     }
 
-    pub fn resolve_label(&self, lifetime: &ast::Lifetime) -> Option<Label> {
-        let text = lifetime.text();
-        let label = lifetime.syntax().ancestors().find_map(|syn| {
-            let label = match_ast! {
-                match syn {
-                    ast::ForExpr(it) => it.label(),
-                    ast::WhileExpr(it) => it.label(),
-                    ast::LoopExpr(it) => it.label(),
-                    ast::BlockExpr(it) => it.label(),
-                    _ => None,
-                }
-            };
-            label.filter(|l| {
-                l.lifetime()
-                    .and_then(|lt| lt.lifetime_ident_token())
-                    .map_or(false, |lt| lt.text() == text)
-            })
-        })?;
-        let src = self.wrap_node_infile(label);
-        ToDef::to_def(self, src.as_ref())
+    pub fn resolve_label(&self, label: &ast::Lifetime) -> Option<Label> {
+        let (parent, label_id) = self
+            .with_ctx(|ctx| ctx.label_ref_to_def(self.wrap_node_infile(label.clone()).as_ref()))?;
+        Some(Label { parent, label_id })
     }
 
     pub fn resolve_type(&self, ty: &ast::Type) -> Option<Type> {

crates/hir/src/semantics/source_to_def.rs

@@ -92,7 +92,7 @@ use hir_def::{
         keys::{self, Key},
         DynMap,
     },
-    hir::{BindingId, LabelId},
+    hir::{BindingId, Expr, LabelId},
     AdtId, BlockId, ConstId, ConstParamId, DefWithBodyId, EnumId, EnumVariantId, ExternCrateId,
     FieldId, FunctionId, GenericDefId, GenericParamId, ImplId, LifetimeParamId, Lookup, MacroId,
     ModuleId, StaticId, StructId, TraitAliasId, TraitId, TypeAliasId, TypeParamId, UnionId, UseId,
@@ -343,6 +343,20 @@ impl SourceToDefCtx<'_, '_> {
         Some((container, label_id))
     }
 
+    pub(super) fn label_ref_to_def(
+        &mut self,
+        src: InFile<&ast::Lifetime>,
+    ) -> Option<(DefWithBodyId, LabelId)> {
+        let break_or_continue = ast::Expr::cast(src.value.syntax().parent()?)?;
+        let container = self.find_pat_or_label_container(src.syntax_ref())?;
+        let (body, source_map) = self.db.body_with_source_map(container);
+        let break_or_continue = source_map.node_expr(src.with_value(&break_or_continue))?;
+        let (Expr::Break { label, .. } | Expr::Continue { label }) = body[break_or_continue] else {
+            return None;
+        };
+        Some((container, label?))
+    }
+
     pub(super) fn item_to_macro_call(&mut self, src: InFile<&ast::Item>) -> Option<MacroCallId> {
         let map = self.dyn_map(src)?;
         map[keys::ATTR_MACRO_CALL].get(&AstPtr::new(src.value)).copied()

crates/ide/src/goto_definition.rs

@@ -2660,6 +2660,24 @@ fn foo() {
         );
     }
 
+    #[test]
+    fn label_inside_macro() {
+        check(
+            r#"
+macro_rules! m {
+    ($s:stmt) => { $s };
+}
+
+fn foo() {
+    'label: loop {
+ // ^^^^^^
+        m!(continue 'label$0);
+    }
+}
+"#,
+        );
+    }
+
     #[test]
     fn goto_def_on_return_in_try() {
         check(

crates/syntax-bridge/src/lib.rs

@@ -148,6 +148,7 @@ pub fn token_tree_to_syntax_node<Ctx>(
 ) -> (Parse<SyntaxNode>, SpanMap<Ctx>)
 where
     SpanData<Ctx>: Copy + fmt::Debug,
+    Ctx: PartialEq,
 {
     let buffer = match tt {
         tt::Subtree {
@@ -892,6 +893,7 @@ fn delim_to_str(d: tt::DelimiterKind, closing: bool) -> Option<&'static str> {
 impl<Ctx> TtTreeSink<'_, Ctx>
 where
     SpanData<Ctx>: Copy + fmt::Debug,
+    Ctx: PartialEq,
 {
     /// Parses a float literal as if it was a one to two name ref nodes with a dot inbetween.
     /// This occurs when a float literal is used as a field access.
@@ -949,6 +951,7 @@ where
         }
 
         let mut last = self.cursor;
+        let mut combined_span = None;
         'tokens: for _ in 0..n_tokens {
             let tmp: u8;
             if self.cursor.eof() {
@@ -982,7 +985,10 @@ where
                             format_to!(self.buf, "{lit}");
                             debug_assert_ne!(self.buf.len() - buf_l, 0);
                             self.text_pos += TextSize::new((self.buf.len() - buf_l) as u32);
-                            self.token_map.push(self.text_pos, lit.span);
+                            combined_span = match combined_span {
+                                None => Some(lit.span),
+                                Some(prev_span) => Some(Self::merge_spans(prev_span, lit.span)),
+                            };
                             self.cursor = self.cursor.bump();
                             continue 'tokens;
                         }
@@ -1006,9 +1012,13 @@ where
             };
             self.buf += text;
             self.text_pos += TextSize::of(text);
-            self.token_map.push(self.text_pos, span);
+            combined_span = match combined_span {
+                None => Some(span),
+                Some(prev_span) => Some(Self::merge_spans(prev_span, span)),
+            }
         }
 
+        self.token_map.push(self.text_pos, combined_span.expect("expected at least one token"));
         self.inner.token(kind, self.buf.as_str());
         self.buf.clear();
         // FIXME: Emitting whitespace for this is really just a hack, we should get rid of it.
@@ -1043,4 +1053,22 @@ where
     fn error(&mut self, error: String) {
         self.inner.error(error, self.text_pos)
     }
+
+    fn merge_spans(a: SpanData<Ctx>, b: SpanData<Ctx>) -> SpanData<Ctx> {
+        // We don't do what rustc does exactly, rustc does something clever when the spans have different syntax contexts
+        // but this runs afoul of our separation between `span` and `hir-expand`.
+        SpanData {
+            range: if a.ctx == b.ctx {
+                TextRange::new(
+                    std::cmp::min(a.range.start(), b.range.start()),
+                    std::cmp::max(a.range.end(), b.range.end()),
+                )
+            } else {
+                // Combining ranges make no sense when they come from different syntax contexts.
+                a.range
+            },
+            anchor: a.anchor,
+            ctx: a.ctx,
+        }
+    }
 }