From 724455bf5a6bfd5d7201215ec64263d79262aef4 Mon Sep 17 00:00:00 2001
From: Lukas Wirth <lukastw97@gmail.com>
Date: Sat, 25 Jan 2025 14:43:22 +0100
Subject: [PATCH] BodyCollector

---
 crates/hir-def/src/body.rs       | 112 +++++++++++++++++--------------
 crates/hir-def/src/body/lower.rs |  49 ++++++++------
 crates/hir-def/src/hir.rs        |   7 +-
 3 files changed, 91 insertions(+), 77 deletions(-)

diff --git a/crates/hir-def/src/body.rs b/crates/hir-def/src/body.rs
index 139e1ba896..334177751e 100644
--- a/crates/hir-def/src/body.rs
+++ b/crates/hir-def/src/body.rs
@@ -25,8 +25,8 @@ use crate::{
     db::DefDatabase,
     expander::Expander,
     hir::{
-        dummy_expr_id, Array, AsmOperand, Binding, BindingId, Expr, ExprId, ExprOrPatId, Label,
-        LabelId, Pat, PatId, RecordFieldPat, Statement,
+        Array, AsmOperand, Binding, BindingId, Expr, ExprId, ExprOrPatId, Label, LabelId, Pat,
+        PatId, RecordFieldPat, Statement,
     },
     item_tree::AttrOwner,
     nameres::DefMap,
@@ -81,7 +81,7 @@ pub struct Body {
     pub body_expr: ExprId,
     pub types: TypesMap,
     /// Block expressions in this body that may contain inner items.
-    block_scopes: Vec<BlockId>,
+    block_scopes: Box<[BlockId]>,
 
     /// A map from binding to its hygiene ID.
     ///
@@ -98,6 +98,64 @@ pub struct Body {
     ident_hygiene: FxHashMap<ExprOrPatId, HygieneId>,
 }
 
+/// The body of an item (function, const etc.).
+#[derive(Debug, Eq, PartialEq, Default)]
+pub struct BodyCollector {
+    pub exprs: Arena<Expr>,
+    pub pats: Arena<Pat>,
+    pub bindings: Arena<Binding>,
+    pub labels: Arena<Label>,
+    pub binding_owners: FxHashMap<BindingId, ExprId>,
+    pub types: TypesMap,
+    block_scopes: Vec<BlockId>,
+    binding_hygiene: FxHashMap<BindingId, HygieneId>,
+    ident_hygiene: FxHashMap<ExprOrPatId, HygieneId>,
+}
+
+impl BodyCollector {
+    fn finish(
+        self,
+        body_expr: ExprId,
+        self_param: Option<BindingId>,
+        params: Box<[PatId]>,
+    ) -> Body {
+        let Self {
+            block_scopes,
+            mut exprs,
+            mut labels,
+            mut pats,
+            mut bindings,
+            mut binding_owners,
+            mut binding_hygiene,
+            mut ident_hygiene,
+            mut types,
+        } = self;
+        exprs.shrink_to_fit();
+        labels.shrink_to_fit();
+        pats.shrink_to_fit();
+        bindings.shrink_to_fit();
+        binding_owners.shrink_to_fit();
+        binding_hygiene.shrink_to_fit();
+        ident_hygiene.shrink_to_fit();
+        types.shrink_to_fit();
+
+        Body {
+            exprs,
+            pats,
+            bindings,
+            labels,
+            binding_owners,
+            params,
+            self_param,
+            body_expr,
+            types,
+            block_scopes: block_scopes.into_boxed_slice(),
+            binding_hygiene,
+            ident_hygiene,
+        }
+    }
+}
+
 pub type ExprPtr = AstPtr<ast::Expr>;
 pub type ExprSource = InFile<ExprPtr>;
 
@@ -242,9 +300,8 @@ impl Body {
         };
         let module = def.module(db);
         let expander = Expander::new(db, file_id, module);
-        let (mut body, mut source_map) =
+        let (body, mut source_map) =
             Body::new(db, def, expander, params, body, module.krate, is_async_fn);
-        body.shrink_to_fit();
         source_map.shrink_to_fit();
 
         (Arc::new(body), Arc::new(source_map))
@@ -304,32 +361,6 @@ impl Body {
         lower::lower(db, owner, expander, params, body, krate, is_async_fn)
     }
 
-    fn shrink_to_fit(&mut self) {
-        let Self {
-            body_expr: _,
-            params: _,
-            self_param: _,
-            block_scopes,
-            exprs,
-            labels,
-            pats,
-            bindings,
-            binding_owners,
-            binding_hygiene,
-            ident_hygiene,
-            types,
-        } = self;
-        block_scopes.shrink_to_fit();
-        exprs.shrink_to_fit();
-        labels.shrink_to_fit();
-        pats.shrink_to_fit();
-        bindings.shrink_to_fit();
-        binding_owners.shrink_to_fit();
-        binding_hygiene.shrink_to_fit();
-        ident_hygiene.shrink_to_fit();
-        types.shrink_to_fit();
-    }
-
     pub fn walk_bindings_in_pat(&self, pat_id: PatId, mut f: impl FnMut(BindingId)) {
         self.walk_pats(pat_id, &mut |pat| {
             if let Pat::Bind { id, .. } = &self[pat] {
@@ -670,25 +701,6 @@ impl Body {
     }
 }
 
-impl Default for Body {
-    fn default() -> Self {
-        Self {
-            body_expr: dummy_expr_id(),
-            exprs: Default::default(),
-            pats: Default::default(),
-            bindings: Default::default(),
-            labels: Default::default(),
-            params: Default::default(),
-            block_scopes: Default::default(),
-            binding_owners: Default::default(),
-            self_param: Default::default(),
-            binding_hygiene: Default::default(),
-            ident_hygiene: Default::default(),
-            types: Default::default(),
-        }
-    }
-}
-
 impl Index<ExprId> for Body {
     type Output = Expr;
 
diff --git a/crates/hir-def/src/body/lower.rs b/crates/hir-def/src/body/lower.rs
index 583c6ac5e8..2b25dbdd7f 100644
--- a/crates/hir-def/src/body/lower.rs
+++ b/crates/hir-def/src/body/lower.rs
@@ -29,7 +29,9 @@ use triomphe::Arc;
 
 use crate::{
     attr::Attrs,
-    body::{Body, BodyDiagnostic, BodySourceMap, ExprPtr, HygieneId, LabelPtr, PatPtr},
+    body::{
+        Body, BodyCollector, BodyDiagnostic, BodySourceMap, ExprPtr, HygieneId, LabelPtr, PatPtr,
+    },
     builtin_type::BuiltinUint,
     data::adt::StructKind,
     db::DefDatabase,
@@ -82,7 +84,7 @@ pub(super) fn lower(
         def_map: expander.module.def_map(db),
         source_map: BodySourceMap::default(),
         ast_id_map: db.ast_id_map(expander.current_file_id()),
-        body: Body::default(),
+        body: BodyCollector::default(),
         expander,
         current_try_block_label: None,
         is_lowering_coroutine: false,
@@ -102,7 +104,7 @@ struct ExprCollector<'a> {
     def_map: Arc<DefMap>,
     ast_id_map: Arc<AstIdMap>,
     krate: CrateId,
-    body: Body,
+    body: BodyCollector,
     source_map: BodySourceMap,
 
     is_lowering_coroutine: bool,
@@ -214,6 +216,9 @@ impl ExprCollector<'_> {
         body: Option<ast::Expr>,
         is_async_fn: bool,
     ) -> (Body, BodySourceMap) {
+        let mut self_param = None;
+        let mut params = vec![];
+
         let skip_body = match self.owner {
             DefWithBodyId::FunctionId(it) => self.db.attrs(it.into()),
             DefWithBodyId::StaticId(it) => self.db.attrs(it.into()),
@@ -226,29 +231,32 @@ impl ExprCollector<'_> {
         // If #[rust_analyzer::skip] annotated, only construct enough information for the signature
         // and skip the body.
         if skip_body {
-            self.body.body_expr = self.missing_expr();
             if let Some((param_list, mut attr_enabled)) = param_list {
-                if let Some(self_param) =
+                if let Some(self_param_syn) =
                     param_list.self_param().filter(|_| attr_enabled.next().unwrap_or(false))
                 {
-                    let is_mutable =
-                        self_param.mut_token().is_some() && self_param.amp_token().is_none();
+                    let is_mutable = self_param_syn.mut_token().is_some()
+                        && self_param_syn.amp_token().is_none();
                     let binding_id: la_arena::Idx<Binding> = self.alloc_binding(
                         Name::new_symbol_root(sym::self_.clone()),
                         BindingAnnotation::new(is_mutable, false),
                     );
-                    self.body.self_param = Some(binding_id);
+                    self_param = Some(binding_id);
                     self.source_map.self_param =
-                        Some(self.expander.in_file(AstPtr::new(&self_param)));
+                        Some(self.expander.in_file(AstPtr::new(&self_param_syn)));
                 }
-                self.body.params = param_list
+                params = param_list
                     .params()
                     .zip(attr_enabled)
                     .filter(|(_, enabled)| *enabled)
                     .map(|_| self.missing_pat())
                     .collect();
             };
-            return (self.body, self.source_map);
+            let body_expr = self.missing_expr();
+            return (
+                self.body.finish(body_expr, self_param, params.into_boxed_slice()),
+                self.source_map,
+            );
         }
 
         self.awaitable_context.replace(if is_async_fn {
@@ -264,25 +272,25 @@ impl ExprCollector<'_> {
             }
         });
         if let Some((param_list, mut attr_enabled)) = param_list {
-            let mut params = vec![];
-            if let Some(self_param) =
+            if let Some(self_param_syn) =
                 param_list.self_param().filter(|_| attr_enabled.next().unwrap_or(false))
             {
                 let is_mutable =
-                    self_param.mut_token().is_some() && self_param.amp_token().is_none();
+                    self_param_syn.mut_token().is_some() && self_param_syn.amp_token().is_none();
                 let binding_id: la_arena::Idx<Binding> = self.alloc_binding(
                     Name::new_symbol_root(sym::self_.clone()),
                     BindingAnnotation::new(is_mutable, false),
                 );
-                let hygiene = self_param
+                let hygiene = self_param_syn
                     .name()
                     .map(|name| self.hygiene_id_for(name.syntax().text_range().start()))
                     .unwrap_or(HygieneId::ROOT);
                 if !hygiene.is_root() {
                     self.body.binding_hygiene.insert(binding_id, hygiene);
                 }
-                self.body.self_param = Some(binding_id);
-                self.source_map.self_param = Some(self.expander.in_file(AstPtr::new(&self_param)));
+                self_param = Some(binding_id);
+                self.source_map.self_param =
+                    Some(self.expander.in_file(AstPtr::new(&self_param_syn)));
             }
 
             for (param, _) in param_list.params().zip(attr_enabled).filter(|(_, enabled)| *enabled)
@@ -290,9 +298,8 @@ impl ExprCollector<'_> {
                 let param_pat = self.collect_pat_top(param.pat());
                 params.push(param_pat);
             }
-            self.body.params = params.into_boxed_slice();
         };
-        self.body.body_expr = self.with_label_rib(RibKind::Closure, |this| {
+        let body_expr = self.with_label_rib(RibKind::Closure, |this| {
             if is_async_fn {
                 match body {
                     Some(e) => {
@@ -310,7 +317,7 @@ impl ExprCollector<'_> {
             }
         });
 
-        (self.body, self.source_map)
+        (self.body.finish(body_expr, self_param, params.into_boxed_slice()), self.source_map)
     }
 
     fn ctx(&mut self) -> LowerCtx<'_> {
@@ -1934,7 +1941,7 @@ impl ExprCollector<'_> {
         f: impl FnOnce(&mut Self) -> T,
     ) -> T {
         self.label_ribs.push(LabelRib::new(RibKind::Normal(
-            self.body[label].name.clone(),
+            self.body.labels[label].name.clone(),
             label,
             hygiene,
         )));
diff --git a/crates/hir-def/src/hir.rs b/crates/hir-def/src/hir.rs
index 9392e8d12d..0dcddf162b 100644
--- a/crates/hir-def/src/hir.rs
+++ b/crates/hir-def/src/hir.rs
@@ -19,7 +19,7 @@ use std::fmt;
 
 use hir_expand::{name::Name, MacroDefId};
 use intern::Symbol;
-use la_arena::{Idx, RawIdx};
+use la_arena::Idx;
 use rustc_apfloat::ieee::{Half as f16, Quad as f128};
 use syntax::ast;
 use type_ref::TypeRefId;
@@ -37,11 +37,6 @@ pub type BindingId = Idx<Binding>;
 
 pub type ExprId = Idx<Expr>;
 
-/// FIXME: this is a hacky function which should be removed
-pub(crate) fn dummy_expr_id() -> ExprId {
-    ExprId::from_raw(RawIdx::from(u32::MAX))
-}
-
 pub type PatId = Idx<Pat>;
 
 // FIXME: Encode this as a single u32, we won't ever reach all 32 bits especially given these counts