From c3186202a2bde98ad8d388e68cc3a86bc5236734 Mon Sep 17 00:00:00 2001
From: Lukas Wirth <lukastw97@gmail.com>
Date: Sun, 18 Jun 2023 12:03:04 +0200
Subject: [PATCH] Shrink size of hir::Binding

---
 crates/hir-def/src/body.rs         | 29 +++++++++++++++++++++++++++--
 crates/hir-def/src/body/lower.rs   | 19 +++++++++++--------
 crates/hir-def/src/hir.rs          | 16 ----------------
 crates/hir-ty/src/infer/closure.rs |  3 +--
 crates/hir-ty/src/mir/lower.rs     |  7 ++++---
 5 files changed, 43 insertions(+), 31 deletions(-)

diff --git a/crates/hir-def/src/body.rs b/crates/hir-def/src/body.rs
index a78bcd6c67fa..94dc39b11755 100644
--- a/crates/hir-def/src/body.rs
+++ b/crates/hir-def/src/body.rs
@@ -37,6 +37,9 @@ pub struct Body {
     pub pats: Arena<Pat>,
     pub bindings: Arena<Binding>,
     pub labels: Arena<Label>,
+    /// Id of the closure/generator that owns the corresponding binding. If a binding is owned by the
+    /// top level expression, it will not be listed in here.
+    pub binding_owners: FxHashMap<BindingId, ExprId>,
     /// 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).
@@ -206,14 +209,24 @@ impl Body {
     }
 
     fn shrink_to_fit(&mut self) {
-        let Self { _c: _, body_expr: _, block_scopes, exprs, labels, params, pats, bindings } =
-            self;
+        let Self {
+            _c: _,
+            body_expr: _,
+            block_scopes,
+            exprs,
+            labels,
+            params,
+            pats,
+            bindings,
+            binding_owners,
+        } = self;
         block_scopes.shrink_to_fit();
         exprs.shrink_to_fit();
         labels.shrink_to_fit();
         params.shrink_to_fit();
         pats.shrink_to_fit();
         bindings.shrink_to_fit();
+        binding_owners.shrink_to_fit();
     }
 
     pub fn walk_bindings_in_pat(&self, pat_id: PatId, mut f: impl FnMut(BindingId)) {
@@ -257,6 +270,17 @@ impl Body {
         f(pat_id);
         self.walk_pats_shallow(pat_id, |p| self.walk_pats(p, f));
     }
+
+    pub fn is_binding_upvar(&self, binding: BindingId, relative_to: ExprId) -> bool {
+        match self.binding_owners.get(&binding) {
+            Some(x) => {
+                // We assign expression ids in a way that outer closures will receive
+                // a lower id
+                x.into_raw() < relative_to.into_raw()
+            }
+            None => true,
+        }
+    }
 }
 
 impl Default for Body {
@@ -269,6 +293,7 @@ impl Default for Body {
             labels: Default::default(),
             params: Default::default(),
             block_scopes: Default::default(),
+            binding_owners: Default::default(),
             _c: Default::default(),
         }
     }
diff --git a/crates/hir-def/src/body/lower.rs b/crates/hir-def/src/body/lower.rs
index 53b80a17df1b..b375ec63a67b 100644
--- a/crates/hir-def/src/body/lower.rs
+++ b/crates/hir-def/src/body/lower.rs
@@ -11,7 +11,6 @@ use hir_expand::{
     AstId, ExpandError, InFile,
 };
 use intern::Interned;
-use la_arena::Arena;
 use profile::Count;
 use rustc_hash::FxHashMap;
 use smallvec::SmallVec;
@@ -60,10 +59,11 @@ pub(super) fn lower(
         source_map: BodySourceMap::default(),
         ast_id_map: db.ast_id_map(expander.current_file_id),
         body: Body {
-            exprs: Arena::default(),
-            pats: Arena::default(),
-            bindings: Arena::default(),
-            labels: Arena::default(),
+            exprs: Default::default(),
+            pats: Default::default(),
+            bindings: Default::default(),
+            binding_owners: Default::default(),
+            labels: Default::default(),
             params: Vec::new(),
             body_expr: dummy_expr_id(),
             block_scopes: Vec::new(),
@@ -1540,13 +1540,16 @@ impl ExprCollector<'_> {
     }
 
     fn alloc_binding(&mut self, name: Name, mode: BindingAnnotation) -> BindingId {
-        self.body.bindings.alloc(Binding {
+        let binding = self.body.bindings.alloc(Binding {
             name,
             mode,
             definitions: SmallVec::new(),
-            owner: self.current_binding_owner,
             problems: None,
-        })
+        });
+        if let Some(owner) = self.current_binding_owner {
+            self.body.binding_owners.insert(binding, owner);
+        }
+        binding
     }
 
     fn alloc_pat(&mut self, pat: Pat, ptr: PatPtr) -> PatId {
diff --git a/crates/hir-def/src/hir.rs b/crates/hir-def/src/hir.rs
index 77d879a77ba5..500e880061ac 100644
--- a/crates/hir-def/src/hir.rs
+++ b/crates/hir-def/src/hir.rs
@@ -501,25 +501,9 @@ pub struct Binding {
     pub name: Name,
     pub mode: BindingAnnotation,
     pub definitions: SmallVec<[PatId; 1]>,
-    /// Id of the closure/generator that owns this binding. If it is owned by the
-    /// top level expression, this field would be `None`.
-    pub owner: Option<ExprId>,
     pub problems: Option<BindingProblems>,
 }
 
-impl Binding {
-    pub fn is_upvar(&self, relative_to: ExprId) -> bool {
-        match self.owner {
-            Some(x) => {
-                // We assign expression ids in a way that outer closures will receive
-                // a lower id
-                x.into_raw() < relative_to.into_raw()
-            }
-            None => true,
-        }
-    }
-}
-
 #[derive(Debug, Clone, Eq, PartialEq)]
 pub struct RecordFieldPat {
     pub name: Name,
diff --git a/crates/hir-ty/src/infer/closure.rs b/crates/hir-ty/src/infer/closure.rs
index 23189f383e0e..ff64ae252bce 100644
--- a/crates/hir-ty/src/infer/closure.rs
+++ b/crates/hir-ty/src/infer/closure.rs
@@ -715,10 +715,9 @@ impl InferenceContext<'_> {
     }
 
     fn is_upvar(&self, place: &HirPlace) -> bool {
-        let b = &self.body[place.local];
         if let Some(c) = self.current_closure {
             let (_, root) = self.db.lookup_intern_closure(c.into());
-            return b.is_upvar(root);
+            return self.body.is_binding_upvar(place.local, root);
         }
         false
     }
diff --git a/crates/hir-ty/src/mir/lower.rs b/crates/hir-ty/src/mir/lower.rs
index e3401a36ceba..2cb29b4ab91d 100644
--- a/crates/hir-ty/src/mir/lower.rs
+++ b/crates/hir-ty/src/mir/lower.rs
@@ -1853,7 +1853,7 @@ pub fn mir_body_for_closure_query(
         .result
         .binding_locals
         .into_iter()
-        .filter(|x| ctx.body[x.0].owner == Some(expr))
+        .filter(|it| ctx.body.binding_owners.get(&it.0).copied() == Some(expr))
         .collect();
     if let Some(err) = err {
         return Err(MirLowerError::UnresolvedUpvar(err));
@@ -1909,10 +1909,11 @@ pub fn lower_to_mir(
     // 0 is return local
     ctx.result.locals.alloc(Local { ty: ctx.expr_ty_after_adjustments(root_expr) });
     let binding_picker = |b: BindingId| {
+        let owner = ctx.body.binding_owners.get(&b).copied();
         if root_expr == body.body_expr {
-            body[b].owner.is_none()
+            owner.is_none()
         } else {
-            body[b].owner == Some(root_expr)
+            owner == Some(root_expr)
         }
     };
     // 1 to param_len is for params