Track labels in scopes

4 files changed, 57 insertions, 11 deletions

diff --git a/crates/hir/src/lib.rs b/crates/hir/src/lib.rs
index 30e577671..e34be7e42 100644
--- a/crates/hir/src/lib.rs
+++ b/crates/hir/src/lib.rs
@@ -2199,6 +2199,7 @@ pub enum ScopeDef {
     ImplSelfType(Impl),
     AdtSelfType(Adt),
     Local(Local),
+    Label(Label),
     Unknown,
 }
 
diff --git a/crates/hir/src/semantics.rs b/crates/hir/src/semantics.rs
index 15651bb22..1198e3f0b 100644
--- a/crates/hir/src/semantics.rs
+++ b/crates/hir/src/semantics.rs
@@ -839,6 +839,10 @@ impl<'a> SemanticsScope<'a> {
                     let parent = resolver.body_owner().unwrap();
                     ScopeDef::Local(Local { parent, pat_id })
                 }
+                resolver::ScopeDef::Label(label_id) => {
+                    let parent = resolver.body_owner().unwrap();
+                    ScopeDef::Label(Label { parent, label_id })
+                }
             };
             f(name, def)
         })
diff --git a/crates/hir_def/src/body/scope.rs b/crates/hir_def/src/body/scope.rs
index 1bbb54fc6..7f0d8f915 100644
--- a/crates/hir_def/src/body/scope.rs
+++ b/crates/hir_def/src/body/scope.rs
@@ -8,7 +8,7 @@ use rustc_hash::FxHashMap;
 use crate::{
     body::Body,
     db::DefDatabase,
-    expr::{Expr, ExprId, Pat, PatId, Statement},
+    expr::{Expr, ExprId, LabelId, Pat, PatId, Statement},
     BlockId, DefWithBodyId,
 };
 
@@ -40,6 +40,7 @@ impl ScopeEntry {
 pub struct ScopeData {
     parent: Option<ScopeId>,
     block: Option<BlockId>,
+    label: Option<(LabelId, Name)>,
     entries: Vec<ScopeEntry>,
 }
 
@@ -67,6 +68,11 @@ impl ExprScopes {
         self.scopes[scope].block
     }
 
+    /// If `scope` refers to a labeled expression scope, returns the corresponding `Label`.
+    pub fn label(&self, scope: ScopeId) -> Option<(LabelId, Name)> {
+        self.scopes[scope].label.clone()
+    }
+
     pub fn scope_chain(&self, scope: Option<ScopeId>) -> impl Iterator<Item = ScopeId> + '_ {
         std::iter::successors(scope, move |&scope| self.scopes[scope].parent)
     }
@@ -85,15 +91,34 @@ impl ExprScopes {
     }
 
     fn root_scope(&mut self) -> ScopeId {
-        self.scopes.alloc(ScopeData { parent: None, block: None, entries: vec![] })
+        self.scopes.alloc(ScopeData { parent: None, block: None, label: None, entries: vec![] })
     }
 
     fn new_scope(&mut self, parent: ScopeId) -> ScopeId {
-        self.scopes.alloc(ScopeData { parent: Some(parent), block: None, entries: vec![] })
+        self.scopes.alloc(ScopeData {
+            parent: Some(parent),
+            block: None,
+            label: None,
+            entries: vec![],
+        })
     }
 
-    fn new_block_scope(&mut self, parent: ScopeId, block: BlockId) -> ScopeId {
-        self.scopes.alloc(ScopeData { parent: Some(parent), block: Some(block), entries: vec![] })
+    fn new_labeled_scope(&mut self, parent: ScopeId, label: Option<(LabelId, Name)>) -> ScopeId {
+        self.scopes.alloc(ScopeData { parent: Some(parent), block: None, label, entries: vec![] })
+    }
+
+    fn new_block_scope(
+        &mut self,
+        parent: ScopeId,
+        block: BlockId,
+        label: Option<(LabelId, Name)>,
+    ) -> ScopeId {
+        self.scopes.alloc(ScopeData {
+            parent: Some(parent),
+            block: Some(block),
+            label,
+            entries: vec![],
+        })
     }
 
     fn add_bindings(&mut self, body: &Body, scope: ScopeId, pat: PatId) {
@@ -144,21 +169,33 @@ fn compute_block_scopes(
 }
 
 fn compute_expr_scopes(expr: ExprId, body: &Body, scopes: &mut ExprScopes, scope: ScopeId) {
+    let make_label =
+        |label: &Option<_>| label.map(|label| (label, body.labels[label].name.clone()));
+
     scopes.set_scope(expr, scope);
     match &body[expr] {
-        Expr::Block { statements, tail, id, .. } => {
-            let scope = scopes.new_block_scope(scope, *id);
+        Expr::Block { statements, tail, id, label } => {
+            let scope = scopes.new_block_scope(scope, *id, make_label(label));
             // Overwrite the old scope for the block expr, so that every block scope can be found
             // via the block itself (important for blocks that only contain items, no expressions).
             scopes.set_scope(expr, scope);
-            compute_block_scopes(&statements, *tail, body, scopes, scope);
+            compute_block_scopes(statements, *tail, body, scopes, scope);
         }
-        Expr::For { iterable, pat, body: body_expr, .. } => {
+        Expr::For { iterable, pat, body: body_expr, label } => {
             compute_expr_scopes(*iterable, body, scopes, scope);
-            let scope = scopes.new_scope(scope);
+            let scope = scopes.new_labeled_scope(scope, make_label(label));
             scopes.add_bindings(body, scope, *pat);
             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(*body_expr, body, scopes, scope);
+        }
+        Expr::Loop { body: body_expr, label } => {
+            let scope = scopes.new_labeled_scope(scope, make_label(label));
+            compute_expr_scopes(*body_expr, body, scopes, scope);
+        }
         Expr::Lambda { args, body: body_expr, .. } => {
             let scope = scopes.new_scope(scope);
             scopes.add_params_bindings(body, scope, &args);
diff --git a/crates/hir_def/src/resolver.rs b/crates/hir_def/src/resolver.rs
index 42736171e..4a2d1c087 100644
--- a/crates/hir_def/src/resolver.rs
+++ b/crates/hir_def/src/resolver.rs
@@ -12,7 +12,7 @@ use crate::{
     body::scope::{ExprScopes, ScopeId},
     builtin_type::BuiltinType,
     db::DefDatabase,
-    expr::{ExprId, PatId},
+    expr::{ExprId, LabelId, PatId},
     generics::GenericParams,
     item_scope::{BuiltinShadowMode, BUILTIN_SCOPE},
     nameres::DefMap,
@@ -409,6 +409,7 @@ pub enum ScopeDef {
     AdtSelfType(AdtId),
     GenericParam(GenericParamId),
     Local(PatId),
+    Label(LabelId),
 }
 
 impl Scope {
@@ -470,6 +471,9 @@ impl Scope {
                 f(name![Self], ScopeDef::AdtSelfType(*i));
             }
             Scope::ExprScope(scope) => {
+                if let Some((label, name)) = scope.expr_scopes.label(scope.scope_id) {
+                    f(name.clone(), ScopeDef::Label(label))
+                }
                 scope.expr_scopes.entries(scope.scope_id).iter().for_each(|e| {
                     f(e.name().clone(), ScopeDef::Local(e.pat()));
                 });