From 262b9c39824b58068d89d6c5cf53d8fea782b11c Mon Sep 17 00:00:00 2001
From: Lukas Wirth <lukastw97@gmail.com>
Date: Wed, 23 Dec 2020 16:34:30 +0100
Subject: [PATCH] Track labels in the HIR

---
 crates/hir/src/code_model.rs              | 30 +++++++-
 crates/hir/src/from_id.rs                 | 15 +++-
 crates/hir/src/lib.rs                     |  4 +-
 crates/hir/src/semantics.rs               | 33 +++++++-
 crates/hir/src/semantics/source_to_def.rs | 15 +++-
 crates/hir_def/src/body.rs                | 24 +++++-
 crates/hir_def/src/body/lower.rs          | 92 +++++++++++------------
 crates/hir_def/src/expr.rs                | 14 +++-
 crates/hir_ty/src/infer/expr.rs           |  8 +-
 9 files changed, 167 insertions(+), 68 deletions(-)

diff --git a/crates/hir/src/code_model.rs b/crates/hir/src/code_model.rs
index 1d7e5ddd7f6..1ddf68c08cc 100644
--- a/crates/hir/src/code_model.rs
+++ b/crates/hir/src/code_model.rs
@@ -9,7 +9,7 @@ use hir_def::{
     adt::StructKind,
     adt::VariantData,
     builtin_type::BuiltinType,
-    expr::{BindingAnnotation, Pat, PatId},
+    expr::{BindingAnnotation, LabelId, Pat, PatId},
     import_map,
     item_tree::ItemTreeNode,
     lang_item::LangItemTarget,
@@ -1205,6 +1205,34 @@ impl Local {
     }
 }
 
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
+pub struct Label {
+    pub(crate) parent: DefWithBodyId,
+    pub(crate) label_id: LabelId,
+}
+
+impl Label {
+    pub fn module(self, db: &dyn HirDatabase) -> Module {
+        self.parent(db).module(db)
+    }
+
+    pub fn parent(self, _db: &dyn HirDatabase) -> DefWithBody {
+        self.parent.into()
+    }
+
+    pub fn name(self, db: &dyn HirDatabase) -> Name {
+        let body = db.body(self.parent.into());
+        body[self.label_id].name.clone()
+    }
+
+    pub fn source(self, db: &dyn HirDatabase) -> InFile<ast::Label> {
+        let (_body, source_map) = db.body_with_source_map(self.parent.into());
+        let src = source_map.label_syntax(self.label_id);
+        let root = src.file_syntax(db.upcast());
+        src.map(|ast| ast.to_node(&root))
+    }
+}
+
 #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
 pub enum GenericParam {
     TypeParam(TypeParam),
diff --git a/crates/hir/src/from_id.rs b/crates/hir/src/from_id.rs
index 8e0c571b851..a0792b9a667 100644
--- a/crates/hir/src/from_id.rs
+++ b/crates/hir/src/from_id.rs
@@ -4,12 +4,15 @@
 //! are splitting the hir.
 
 use hir_def::{
-    expr::PatId, item_scope::ItemInNs, AdtId, AssocItemId, DefWithBodyId, EnumVariantId, FieldId,
-    GenericDefId, ModuleDefId, VariantId,
+    expr::{LabelId, PatId},
+    item_scope::ItemInNs,
+    AdtId, AssocItemId, DefWithBodyId, EnumVariantId, FieldId, GenericDefId, ModuleDefId,
+    VariantId,
 };
 
 use crate::{
-    Adt, AssocItem, DefWithBody, Field, GenericDef, Local, MacroDef, ModuleDef, Variant, VariantDef,
+    Adt, AssocItem, DefWithBody, Field, GenericDef, Label, Local, MacroDef, ModuleDef, Variant,
+    VariantDef,
 };
 
 macro_rules! from_id {
@@ -228,6 +231,12 @@ impl From<(DefWithBodyId, PatId)> for Local {
     }
 }
 
+impl From<(DefWithBodyId, LabelId)> for Label {
+    fn from((parent, label_id): (DefWithBodyId, LabelId)) -> Self {
+        Label { parent, label_id }
+    }
+}
+
 impl From<MacroDef> for ItemInNs {
     fn from(macro_def: MacroDef) -> Self {
         ItemInNs::Macros(macro_def.into())
diff --git a/crates/hir/src/lib.rs b/crates/hir/src/lib.rs
index bdd270c58d3..7ac9fd50799 100644
--- a/crates/hir/src/lib.rs
+++ b/crates/hir/src/lib.rs
@@ -35,8 +35,8 @@ pub use crate::{
     code_model::{
         Access, Adt, AsAssocItem, AssocItem, AssocItemContainer, Callable, CallableKind, Const,
         Crate, CrateDependency, DefWithBody, Enum, Field, FieldSource, Function, GenericDef,
-        HasVisibility, Impl, LifetimeParam, Local, MacroDef, Module, ModuleDef, ScopeDef, Static,
-        Struct, Trait, Type, TypeAlias, TypeParam, Union, Variant, VariantDef,
+        HasVisibility, Impl, Label, LifetimeParam, Local, MacroDef, Module, ModuleDef, ScopeDef,
+        Static, Struct, Trait, Type, TypeAlias, TypeParam, Union, Variant, VariantDef,
     },
     has_source::HasSource,
     semantics::{PathResolution, Semantics, SemanticsScope},
diff --git a/crates/hir/src/semantics.rs b/crates/hir/src/semantics.rs
index 25ebf73d88e..67cd16e311c 100644
--- a/crates/hir/src/semantics.rs
+++ b/crates/hir/src/semantics.rs
@@ -15,7 +15,7 @@ use itertools::Itertools;
 use rustc_hash::{FxHashMap, FxHashSet};
 use syntax::{
     algo::find_node_at_offset,
-    ast::{self, GenericParamsOwner},
+    ast::{self, GenericParamsOwner, LoopBodyOwner},
     match_ast, AstNode, SyntaxNode, SyntaxToken, TextSize,
 };
 
@@ -25,8 +25,8 @@ use crate::{
     diagnostics::Diagnostic,
     semantics::source_to_def::{ChildContainer, SourceToDefCache, SourceToDefCtx},
     source_analyzer::{resolve_hir_path, SourceAnalyzer},
-    AssocItem, Callable, Crate, Field, Function, HirFileId, Impl, InFile, LifetimeParam, Local,
-    MacroDef, Module, ModuleDef, Name, Path, ScopeDef, Trait, Type, TypeAlias, TypeParam,
+    AssocItem, Callable, Crate, Field, Function, HirFileId, Impl, InFile, Label, LifetimeParam,
+    Local, MacroDef, Module, ModuleDef, Name, Path, ScopeDef, Trait, Type, TypeAlias, TypeParam,
     VariantDef,
 };
 
@@ -182,6 +182,10 @@ impl<'db, DB: HirDatabase> Semantics<'db, DB> {
         self.imp.resolve_lifetime_param(lifetime)
     }
 
+    pub fn resolve_label(&self, lifetime: &ast::Lifetime) -> Option<Label> {
+        self.imp.resolve_label(lifetime)
+    }
+
     pub fn type_of_expr(&self, expr: &ast::Expr) -> Option<Type> {
         self.imp.type_of_expr(expr)
     }
@@ -425,6 +429,28 @@ impl<'db> SemanticsImpl<'db> {
         ToDef::to_def(self, src)
     }
 
+    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::EffectExpr(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.find_file(label.syntax().clone()).with_value(label);
+        ToDef::to_def(self, src)
+    }
+
     fn type_of_expr(&self, expr: &ast::Expr) -> Option<Type> {
         self.analyze(expr.syntax()).type_of_expr(self.db, expr)
     }
@@ -720,6 +746,7 @@ to_def_impls![
     (crate::LifetimeParam, ast::LifetimeParam, lifetime_param_to_def),
     (crate::MacroDef, ast::MacroRules, macro_rules_to_def),
     (crate::Local, ast::IdentPat, bind_pat_to_def),
+    (crate::Label, ast::Label, label_to_def),
 ];
 
 fn find_root(node: &SyntaxNode) -> SyntaxNode {
diff --git a/crates/hir/src/semantics/source_to_def.rs b/crates/hir/src/semantics/source_to_def.rs
index 3efca5baa5c..424e6e8a940 100644
--- a/crates/hir/src/semantics/source_to_def.rs
+++ b/crates/hir/src/semantics/source_to_def.rs
@@ -4,7 +4,7 @@ use base_db::FileId;
 use hir_def::{
     child_by_source::ChildBySource,
     dyn_map::DynMap,
-    expr::PatId,
+    expr::{LabelId, PatId},
     keys::{self, Key},
     ConstId, DefWithBodyId, EnumId, EnumVariantId, FieldId, FunctionId, GenericDefId, ImplId,
     LifetimeParamId, ModuleId, StaticId, StructId, TraitId, TypeAliasId, TypeParamId, UnionId,
@@ -108,12 +108,21 @@ impl SourceToDefCtx<'_, '_> {
         &mut self,
         src: InFile<ast::IdentPat>,
     ) -> Option<(DefWithBodyId, PatId)> {
-        let container = self.find_pat_container(src.as_ref().map(|it| it.syntax()))?;
+        let container = self.find_pat_or_label_container(src.as_ref().map(|it| it.syntax()))?;
         let (_body, source_map) = self.db.body_with_source_map(container);
         let src = src.map(ast::Pat::from);
         let pat_id = source_map.node_pat(src.as_ref())?;
         Some((container, pat_id))
     }
+    pub(super) fn label_to_def(
+        &mut self,
+        src: InFile<ast::Label>,
+    ) -> Option<(DefWithBodyId, LabelId)> {
+        let container = self.find_pat_or_label_container(src.as_ref().map(|it| it.syntax()))?;
+        let (_body, source_map) = self.db.body_with_source_map(container);
+        let label_id = source_map.node_label(src.as_ref())?;
+        Some((container, label_id))
+    }
 
     fn to_def<Ast: AstNode + 'static, ID: Copy + 'static>(
         &mut self,
@@ -237,7 +246,7 @@ impl SourceToDefCtx<'_, '_> {
         None
     }
 
-    fn find_pat_container(&mut self, src: InFile<&SyntaxNode>) -> Option<DefWithBodyId> {
+    fn find_pat_or_label_container(&mut self, src: InFile<&SyntaxNode>) -> Option<DefWithBodyId> {
         for container in src.cloned().ancestors_with_macros(self.db.upcast()).skip(1) {
             let res: DefWithBodyId = match_ast! {
                 match (container.value) {
diff --git a/crates/hir_def/src/body.rs b/crates/hir_def/src/body.rs
index 998b826010e..d07004b9d7c 100644
--- a/crates/hir_def/src/body.rs
+++ b/crates/hir_def/src/body.rs
@@ -26,7 +26,7 @@ pub(crate) use lower::LowerCtx;
 use crate::{
     attr::{Attrs, RawAttrs},
     db::DefDatabase,
-    expr::{Expr, ExprId, Pat, PatId},
+    expr::{Expr, ExprId, Label, LabelId, Pat, PatId},
     item_scope::BuiltinShadowMode,
     item_scope::ItemScope,
     nameres::CrateDefMap,
@@ -226,6 +226,7 @@ pub(crate) struct Mark {
 pub struct Body {
     pub exprs: Arena<Expr>,
     pub pats: Arena<Pat>,
+    pub labels: Arena<Label>,
     /// The patterns for the function's parameters. While the parameter types are
     /// part of the function signature, the patterns are not (they don't change
     /// the external type of the function).
@@ -244,6 +245,8 @@ pub type ExprSource = InFile<ExprPtr>;
 pub type PatPtr = Either<AstPtr<ast::Pat>, AstPtr<ast::SelfParam>>;
 pub type PatSource = InFile<PatPtr>;
 
+pub type LabelPtr = AstPtr<ast::Label>;
+pub type LabelSource = InFile<LabelPtr>;
 /// An item body together with the mapping from syntax nodes to HIR expression
 /// IDs. This is needed to go from e.g. a position in a file to the HIR
 /// expression containing it; but for type inference etc., we want to operate on
@@ -261,6 +264,8 @@ pub struct BodySourceMap {
     expr_map_back: ArenaMap<ExprId, Result<ExprSource, SyntheticSyntax>>,
     pat_map: FxHashMap<PatSource, PatId>,
     pat_map_back: ArenaMap<PatId, Result<PatSource, SyntheticSyntax>>,
+    label_map: FxHashMap<LabelSource, LabelId>,
+    label_map_back: ArenaMap<LabelId, LabelSource>,
     field_map: FxHashMap<(ExprId, usize), InFile<AstPtr<ast::RecordExprField>>>,
     expansions: FxHashMap<InFile<AstPtr<ast::MacroCall>>, HirFileId>,
 
@@ -334,6 +339,14 @@ impl Index<PatId> for Body {
     }
 }
 
+impl Index<LabelId> for Body {
+    type Output = Label;
+
+    fn index(&self, label: LabelId) -> &Label {
+        &self.labels[label]
+    }
+}
+
 impl BodySourceMap {
     pub fn expr_syntax(&self, expr: ExprId) -> Result<ExprSource, SyntheticSyntax> {
         self.expr_map_back[expr].clone()
@@ -363,6 +376,15 @@ impl BodySourceMap {
         self.pat_map.get(&src).cloned()
     }
 
+    pub fn label_syntax(&self, label: LabelId) -> LabelSource {
+        self.label_map_back[label].clone()
+    }
+
+    pub fn node_label(&self, node: InFile<&ast::Label>) -> Option<LabelId> {
+        let src = node.map(|it| AstPtr::new(it));
+        self.label_map.get(&src).cloned()
+    }
+
     pub fn field_syntax(&self, expr: ExprId, field: usize) -> InFile<AstPtr<ast::RecordExprField>> {
         self.field_map[&(expr, field)].clone()
     }
diff --git a/crates/hir_def/src/body/lower.rs b/crates/hir_def/src/body/lower.rs
index 978c3a32498..e1de39d0c23 100644
--- a/crates/hir_def/src/body/lower.rs
+++ b/crates/hir_def/src/body/lower.rs
@@ -22,13 +22,14 @@ use test_utils::mark;
 
 use crate::{
     adt::StructKind,
-    body::{Body, BodySourceMap, Expander, PatPtr, SyntheticSyntax},
+    body::{Body, BodySourceMap, Expander, LabelSource, PatPtr, SyntheticSyntax},
     builtin_type::{BuiltinFloat, BuiltinInt},
     db::DefDatabase,
     diagnostics::{InactiveCode, MacroError, UnresolvedProcMacro},
     expr::{
-        dummy_expr_id, ArithOp, Array, BinaryOp, BindingAnnotation, CmpOp, Expr, ExprId, Literal,
-        LogicOp, MatchArm, Ordering, Pat, PatId, RecordFieldPat, RecordLitField, Statement,
+        dummy_expr_id, ArithOp, Array, BinaryOp, BindingAnnotation, CmpOp, Expr, ExprId, Label,
+        LabelId, Literal, LogicOp, MatchArm, Ordering, Pat, PatId, RecordFieldPat, RecordLitField,
+        Statement,
     },
     item_scope::BuiltinShadowMode,
     item_tree::{ItemTree, ItemTreeId, ItemTreeNode},
@@ -72,6 +73,7 @@ pub(super) fn lower(
         body: Body {
             exprs: Arena::default(),
             pats: Arena::default(),
+            labels: Arena::default(),
             params: Vec::new(),
             body_expr: dummy_expr_id(),
             item_scope: Default::default(),
@@ -175,6 +177,18 @@ impl ExprCollector<'_> {
         id
     }
 
+    fn alloc_label(&mut self, label: Label, ptr: AstPtr<ast::Label>) -> LabelId {
+        let src = self.expander.to_source(ptr);
+        let id = self.make_label(label, src.clone());
+        self.source_map.label_map.insert(src, id);
+        id
+    }
+    fn make_label(&mut self, label: Label, src: LabelSource) -> LabelId {
+        let id = self.body.labels.alloc(label);
+        self.source_map.label_map_back.insert(id, src);
+        id
+    }
+
     fn collect_expr(&mut self, expr: ast::Expr) -> ExprId {
         let syntax_ptr = AstPtr::new(&expr);
         if self.check_cfg(&expr).is_none() {
@@ -228,19 +242,22 @@ impl ExprCollector<'_> {
                     self.alloc_expr(Expr::Unsafe { body }, syntax_ptr)
                 }
                 // FIXME: we need to record these effects somewhere...
-                ast::Effect::Label(label) => match e.block_expr() {
-                    Some(block) => {
-                        let res = self.collect_block(block);
-                        match &mut self.body.exprs[res] {
-                            Expr::Block { label: block_label, .. } => {
-                                *block_label = label.lifetime().map(|t| Name::new_lifetime(&t))
+                ast::Effect::Label(label) => {
+                    let label = self.collect_label(label);
+                    match e.block_expr() {
+                        Some(block) => {
+                            let res = self.collect_block(block);
+                            match &mut self.body.exprs[res] {
+                                Expr::Block { label: block_label, .. } => {
+                                    *block_label = Some(label);
+                                }
+                                _ => unreachable!(),
                             }
-                            _ => unreachable!(),
+                            res
                         }
-                        res
+                        None => self.missing_expr(),
                     }
-                    None => self.missing_expr(),
-                },
+                }
                 // FIXME: we need to record these effects somewhere...
                 ast::Effect::Async(_) => {
                     let body = self.collect_block_opt(e.block_expr());
@@ -249,16 +266,12 @@ impl ExprCollector<'_> {
             },
             ast::Expr::BlockExpr(e) => self.collect_block(e),
             ast::Expr::LoopExpr(e) => {
+                let label = e.label().map(|label| self.collect_label(label));
                 let body = self.collect_block_opt(e.loop_body());
-                self.alloc_expr(
-                    Expr::Loop {
-                        body,
-                        label: e.label().and_then(|l| l.lifetime()).map(|l| Name::new_lifetime(&l)),
-                    },
-                    syntax_ptr,
-                )
+                self.alloc_expr(Expr::Loop { body, label }, syntax_ptr)
             }
             ast::Expr::WhileExpr(e) => {
+                let label = e.label().map(|label| self.collect_label(label));
                 let body = self.collect_block_opt(e.loop_body());
 
                 let condition = match e.condition() {
@@ -279,42 +292,20 @@ impl ExprCollector<'_> {
                             ];
                             let match_expr =
                                 self.alloc_expr_desugared(Expr::Match { expr: match_expr, arms });
-                            return self.alloc_expr(
-                                Expr::Loop {
-                                    body: match_expr,
-                                    label: e
-                                        .label()
-                                        .and_then(|l| l.lifetime())
-                                        .map(|l| Name::new_lifetime(&l)),
-                                },
-                                syntax_ptr,
-                            );
+                            return self
+                                .alloc_expr(Expr::Loop { body: match_expr, label }, syntax_ptr);
                         }
                     },
                 };
 
-                self.alloc_expr(
-                    Expr::While {
-                        condition,
-                        body,
-                        label: e.label().and_then(|l| l.lifetime()).map(|l| Name::new_lifetime(&l)),
-                    },
-                    syntax_ptr,
-                )
+                self.alloc_expr(Expr::While { condition, body, label }, syntax_ptr)
             }
             ast::Expr::ForExpr(e) => {
+                let label = e.label().map(|label| self.collect_label(label));
                 let iterable = self.collect_expr_opt(e.iterable());
                 let pat = self.collect_pat_opt(e.pat());
                 let body = self.collect_block_opt(e.loop_body());
-                self.alloc_expr(
-                    Expr::For {
-                        iterable,
-                        pat,
-                        body,
-                        label: e.label().and_then(|l| l.lifetime()).map(|l| Name::new_lifetime(&l)),
-                    },
-                    syntax_ptr,
-                )
+                self.alloc_expr(Expr::For { iterable, pat, body, label }, syntax_ptr)
             }
             ast::Expr::CallExpr(e) => {
                 let callee = self.collect_expr_opt(e.expr());
@@ -814,6 +805,13 @@ impl ExprCollector<'_> {
         }
     }
 
+    fn collect_label(&mut self, ast_label: ast::Label) -> LabelId {
+        let label = Label {
+            name: ast_label.lifetime().as_ref().map_or_else(Name::missing, Name::new_lifetime),
+        };
+        self.alloc_label(label, AstPtr::new(&ast_label))
+    }
+
     fn collect_pat(&mut self, pat: ast::Pat) -> PatId {
         let pattern = match &pat {
             ast::Pat::IdentPat(bp) => {
diff --git a/crates/hir_def/src/expr.rs b/crates/hir_def/src/expr.rs
index e5d740a3653..b080f155385 100644
--- a/crates/hir_def/src/expr.rs
+++ b/crates/hir_def/src/expr.rs
@@ -29,6 +29,12 @@ pub(crate) fn dummy_expr_id() -> ExprId {
 
 pub type PatId = Idx<Pat>;
 
+#[derive(Debug, Clone, Eq, PartialEq)]
+pub struct Label {
+    pub name: Name,
+}
+pub type LabelId = Idx<Label>;
+
 #[derive(Debug, Clone, Eq, PartialEq)]
 pub enum Literal {
     String(String),
@@ -52,22 +58,22 @@ pub enum Expr {
     Block {
         statements: Vec<Statement>,
         tail: Option<ExprId>,
-        label: Option<Name>,
+        label: Option<LabelId>,
     },
     Loop {
         body: ExprId,
-        label: Option<Name>,
+        label: Option<LabelId>,
     },
     While {
         condition: ExprId,
         body: ExprId,
-        label: Option<Name>,
+        label: Option<LabelId>,
     },
     For {
         iterable: ExprId,
         pat: PatId,
         body: ExprId,
-        label: Option<Name>,
+        label: Option<LabelId>,
     },
     Call {
         callee: ExprId,
diff --git a/crates/hir_ty/src/infer/expr.rs b/crates/hir_ty/src/infer/expr.rs
index 2cdce2cefce..fb72e28b960 100644
--- a/crates/hir_ty/src/infer/expr.rs
+++ b/crates/hir_ty/src/infer/expr.rs
@@ -143,7 +143,7 @@ impl<'a> InferenceContext<'a> {
                     self.breakables.push(BreakableContext {
                         may_break: false,
                         break_ty: break_ty.clone(),
-                        label: label.clone(),
+                        label: label.map(|label| self.body[label].name.clone()),
                     });
                     let ty = self.infer_block(statements, *tail, &Expectation::has_type(break_ty));
                     let ctxt = self.breakables.pop().expect("breakable stack broken");
@@ -172,7 +172,7 @@ impl<'a> InferenceContext<'a> {
                 self.breakables.push(BreakableContext {
                     may_break: false,
                     break_ty: self.table.new_type_var(),
-                    label: label.clone(),
+                    label: label.map(|label| self.body[label].name.clone()),
                 });
                 self.infer_expr(*body, &Expectation::has_type(Ty::unit()));
 
@@ -191,7 +191,7 @@ impl<'a> InferenceContext<'a> {
                 self.breakables.push(BreakableContext {
                     may_break: false,
                     break_ty: Ty::Unknown,
-                    label: label.clone(),
+                    label: label.map(|label| self.body[label].name.clone()),
                 });
                 // while let is desugared to a match loop, so this is always simple while
                 self.infer_expr(*condition, &Expectation::has_type(Ty::simple(TypeCtor::Bool)));
@@ -207,7 +207,7 @@ impl<'a> InferenceContext<'a> {
                 self.breakables.push(BreakableContext {
                     may_break: false,
                     break_ty: Ty::Unknown,
-                    label: label.clone(),
+                    label: label.map(|label| self.body[label].name.clone()),
                 });
                 let pat_ty =
                     self.resolve_associated_type(iterable_ty, self.resolve_into_iter_item());