From d073a852205a0b1a711b36d6dc2d87b67dde60e5 Mon Sep 17 00:00:00 2001
From: Chayim Refael Friedman <chayimfr@gmail.com>
Date: Sun, 29 Sep 2024 23:00:27 +0300
Subject: [PATCH] When resolving labels in `break` and `continue` for the IDE,
 do not resolve them textually, instead reuse the results of HIR lowering

This fixes a bug where labels inside macros were not resolved, but more importantly this prepares us to a future where we have hygiene, and textual equivalence isn't enough to resolve identifiers.
---
 .../rust-analyzer/crates/hir/src/semantics.rs | 30 +++++--------------
 .../crates/hir/src/semantics/source_to_def.rs | 16 +++++++++-
 .../crates/ide/src/goto_definition.rs         | 18 +++++++++++
 3 files changed, 40 insertions(+), 24 deletions(-)

diff --git a/src/tools/rust-analyzer/crates/hir/src/semantics.rs b/src/tools/rust-analyzer/crates/hir/src/semantics.rs
index fa14b53dbc3d..fa18182fbb68 100644
--- a/src/tools/rust-analyzer/crates/hir/src/semantics.rs
+++ b/src/tools/rust-analyzer/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> {
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 c1e4e1d1e271..fd6d52d6c9df 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
@@ -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()
diff --git a/src/tools/rust-analyzer/crates/ide/src/goto_definition.rs b/src/tools/rust-analyzer/crates/ide/src/goto_definition.rs
index 8836166d969d..f8a07200bccd 100644
--- a/src/tools/rust-analyzer/crates/ide/src/goto_definition.rs
+++ b/src/tools/rust-analyzer/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(