From f4c52b40bd867c9aacd84f4890b4a12bba814f69 Mon Sep 17 00:00:00 2001
From: hkalbasi <hamidrezakalbasi@protonmail.com>
Date: Thu, 1 Jun 2023 17:51:53 +0330
Subject: [PATCH] fix bug in labeled for loop desugaring

---
 crates/hir-def/src/body/lower.rs              | 17 +++++++++++--
 .../src/handlers/undeclared_label.rs          | 25 +++++++++++++++++++
 2 files changed, 40 insertions(+), 2 deletions(-)

diff --git a/crates/hir-def/src/body/lower.rs b/crates/hir-def/src/body/lower.rs
index acc9943481..76301782af 100644
--- a/crates/hir-def/src/body/lower.rs
+++ b/crates/hir-def/src/body/lower.rs
@@ -744,10 +744,13 @@ impl ExprCollector<'_> {
             args: Box::new([self.collect_pat_top(e.pat())]),
             ellipsis: None,
         };
+        let label = e.label().map(|label| self.collect_label(label));
         let some_arm = MatchArm {
             pat: self.alloc_pat_desugared(some_pat),
             guard: None,
-            expr: self.collect_expr_opt(e.loop_body().map(|x| x.into())),
+            expr: self.with_opt_labeled_rib(label, |this| {
+                this.collect_expr_opt(e.loop_body().map(|x| x.into()))
+            }),
         };
         let iter_name = Name::generate_new_name();
         let iter_binding = self.alloc_binding(iter_name.clone(), BindingAnnotation::Mutable);
@@ -769,7 +772,6 @@ impl ExprCollector<'_> {
             Expr::Match { expr: iter_next_expr, arms: Box::new([none_arm, some_arm]) },
             syntax_ptr.clone(),
         );
-        let label = e.label().map(|label| self.collect_label(label));
         let loop_outer =
             self.alloc_expr(Expr::Loop { body: loop_inner, label }, syntax_ptr.clone());
         let iter_pat = self.alloc_pat_desugared(Pat::Bind { id: iter_binding, subpat: None });
@@ -1426,6 +1428,17 @@ impl ExprCollector<'_> {
         self.label_ribs.pop();
         res
     }
+
+    fn with_opt_labeled_rib<T>(
+        &mut self,
+        label: Option<LabelId>,
+        f: impl FnOnce(&mut Self) -> T,
+    ) -> T {
+        match label {
+            None => f(self),
+            Some(label) => self.with_labeled_rib(label, f),
+        }
+    }
     // endregion: labels
 }
 
diff --git a/crates/ide-diagnostics/src/handlers/undeclared_label.rs b/crates/ide-diagnostics/src/handlers/undeclared_label.rs
index 768efecb08..034e4fcfb8 100644
--- a/crates/ide-diagnostics/src/handlers/undeclared_label.rs
+++ b/crates/ide-diagnostics/src/handlers/undeclared_label.rs
@@ -33,6 +33,31 @@ fn foo() {
         );
     }
 
+    #[test]
+    fn for_loop() {
+        check_diagnostics(
+            r#"
+//- minicore: iterator
+fn foo() {
+    'xxx: for _ in unknown {
+        'yyy: for _ in unknown {
+            break 'xxx;
+            continue 'yyy;
+            break 'zzz;
+                //^^^^ error: use of undeclared label `'zzz`
+        }
+        continue 'xxx;
+        continue 'yyy;
+               //^^^^ error: use of undeclared label `'yyy`
+        break 'xxx;
+        break 'yyy;
+            //^^^^ error: use of undeclared label `'yyy`
+    }
+}
+"#,
+        );
+    }
+
     #[test]
     fn try_operator_desugar_works() {
         check_diagnostics(