From cbd325707bc44c08f60ee3312af815cfb96ee86a Mon Sep 17 00:00:00 2001
From: Lukas Wirth <lukastw97@gmail.com>
Date: Sun, 21 Mar 2021 00:59:45 +0100
Subject: Track labels in scopes

---
 crates/hir_def/src/body/scope.rs | 57 +++++++++++++++++++++++++++++++++-------
 crates/hir_def/src/resolver.rs   |  6 ++++-
 2 files changed, 52 insertions(+), 11 deletions(-)

(limited to 'crates/hir_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()));
                 });
-- 
cgit v1.2.3