From 64957acb5f359763395a54e314d1f5d5cfc6ccf3 Mon Sep 17 00:00:00 2001
From: Lukas Wirth <lukastw97@gmail.com>
Date: Sun, 21 Mar 2021 01:10:59 +0100
Subject: [PATCH] Fix incorrect scoping in while expressions

---
 crates/hir_def/src/body/scope.rs              |  2 +-
 .../src/completions/lifetime.rs               | 29 +++++++++++++++++++
 crates/ide_completion/src/context.rs          | 24 +++++++--------
 3 files changed, 41 insertions(+), 14 deletions(-)

diff --git a/crates/hir_def/src/body/scope.rs b/crates/hir_def/src/body/scope.rs
index 7f0d8f9159d..bd7005ca696 100644
--- a/crates/hir_def/src/body/scope.rs
+++ b/crates/hir_def/src/body/scope.rs
@@ -188,8 +188,8 @@ fn compute_expr_scopes(expr: ExprId, body: &Body, scopes: &mut ExprScopes, scope
             compute_expr_scopes(*body_expr, body, scopes, scope);
         }
         Expr::While { condition, body: body_expr, label } => {
-            compute_expr_scopes(*condition, body, scopes, scope);
             let scope = scopes.new_labeled_scope(scope, make_label(label));
+            compute_expr_scopes(*condition, body, scopes, scope);
             compute_expr_scopes(*body_expr, body, scopes, scope);
         }
         Expr::Loop { body: body_expr, label } => {
diff --git a/crates/ide_completion/src/completions/lifetime.rs b/crates/ide_completion/src/completions/lifetime.rs
index 07be28e9c41..628c1fb9b6c 100644
--- a/crates/ide_completion/src/completions/lifetime.rs
+++ b/crates/ide_completion/src/completions/lifetime.rs
@@ -253,4 +253,33 @@ fn foo() {
             "#]],
         );
     }
+
+    #[test]
+    fn complete_label_in_while_cond() {
+        check(
+            r#"
+fn foo() {
+    'outer: while { 'inner: loop { break '$0 } } {}
+}
+"#,
+            expect![[r#"
+                lb 'inner
+                lb 'outer
+            "#]],
+        );
+    }
+
+    #[test]
+    fn complete_label_in_for_iterable() {
+        check(
+            r#"
+fn foo() {
+    'outer: for _ in [{ 'inner: loop { break '$0 } }] {}
+}
+"#,
+            expect![[r#"
+                lb 'inner
+            "#]],
+        );
+    }
 }
diff --git a/crates/ide_completion/src/context.rs b/crates/ide_completion/src/context.rs
index 6cb7e526454..67e2d6f6c29 100644
--- a/crates/ide_completion/src/context.rs
+++ b/crates/ide_completion/src/context.rs
@@ -475,19 +475,17 @@ impl<'a> CompletionContext<'a> {
                 return;
             }
 
-            if parent.kind() != syntax::SyntaxKind::LABEL {
-                match_ast! {
-                    match parent {
-                        ast::LifetimeParam(_it) => {
-                            self.lifetime_allowed = true;
-                            self.lifetime_param_syntax =
-                                self.sema.find_node_at_offset_with_macros(original_file, offset);
-                        },
-                        ast::BreakExpr(_it) => self.is_label_ref = true,
-                        ast::ContinueExpr(_it) => self.is_label_ref = true,
-                        ast::Label(_it) => (),
-                        _ => self.lifetime_allowed = true,
-                    }
+            match_ast! {
+                match parent {
+                    ast::LifetimeParam(_it) => {
+                        self.lifetime_allowed = true;
+                        self.lifetime_param_syntax =
+                            self.sema.find_node_at_offset_with_macros(original_file, offset);
+                    },
+                    ast::BreakExpr(_it) => self.is_label_ref = true,
+                    ast::ContinueExpr(_it) => self.is_label_ref = true,
+                    ast::Label(_it) => (),
+                    _ => self.lifetime_allowed = true,
                 }
             }
         }