From 70e21dc30b5cbaecad463adb9ec28e32ae5c35c8 Mon Sep 17 00:00:00 2001
From: Lukas Wirth <lukastw97@gmail.com>
Date: Fri, 1 Sep 2023 19:17:57 +0200
Subject: [PATCH] Remove some allocations in borrowck

---
 crates/hir-ty/src/mir/borrowck.rs | 77 ++++++++++++++++---------------
 1 file changed, 40 insertions(+), 37 deletions(-)

diff --git a/crates/hir-ty/src/mir/borrowck.rs b/crates/hir-ty/src/mir/borrowck.rs
index c70d7f63fd..c651ddc2cf 100644
--- a/crates/hir-ty/src/mir/borrowck.rs
+++ b/crates/hir-ty/src/mir/borrowck.rs
@@ -42,30 +42,27 @@ pub struct BorrowckResult {
 fn all_mir_bodies(
     db: &dyn HirDatabase,
     def: DefWithBodyId,
-) -> Box<dyn Iterator<Item = Result<Arc<MirBody>, MirLowerError>> + '_> {
+    mut cb: impl FnMut(Arc<MirBody>),
+) -> Result<(), MirLowerError> {
     fn for_closure(
         db: &dyn HirDatabase,
         c: ClosureId,
-    ) -> Box<dyn Iterator<Item = Result<Arc<MirBody>, MirLowerError>> + '_> {
+        cb: &mut impl FnMut(Arc<MirBody>),
+    ) -> Result<(), MirLowerError> {
         match db.mir_body_for_closure(c) {
             Ok(body) => {
-                let closures = body.closures.clone();
-                Box::new(
-                    iter::once(Ok(body))
-                        .chain(closures.into_iter().flat_map(|it| for_closure(db, it))),
-                )
+                cb(body.clone());
+                body.closures.iter().map(|&it| for_closure(db, it, cb)).collect()
             }
-            Err(e) => Box::new(iter::once(Err(e))),
+            Err(e) => Err(e),
         }
     }
     match db.mir_body(def) {
         Ok(body) => {
-            let closures = body.closures.clone();
-            Box::new(
-                iter::once(Ok(body)).chain(closures.into_iter().flat_map(|it| for_closure(db, it))),
-            )
+            cb(body.clone());
+            body.closures.iter().map(|&it| for_closure(db, it, &mut cb)).collect()
         }
-        Err(e) => Box::new(iter::once(Err(e))),
+        Err(e) => Err(e),
     }
 }
 
@@ -74,17 +71,15 @@ pub fn borrowck_query(
     def: DefWithBodyId,
 ) -> Result<Arc<[BorrowckResult]>, MirLowerError> {
     let _p = profile::span("borrowck_query");
-    let r = all_mir_bodies(db, def)
-        .map(|body| {
-            let body = body?;
-            Ok(BorrowckResult {
-                mutability_of_locals: mutability_of_locals(db, &body),
-                moved_out_of_ref: moved_out_of_ref(db, &body),
-                mir_body: body,
-            })
-        })
-        .collect::<Result<Vec<_>, MirLowerError>>()?;
-    Ok(r.into())
+    let mut res = vec![];
+    all_mir_bodies(db, def, |body| {
+        res.push(BorrowckResult {
+            mutability_of_locals: mutability_of_locals(db, &body),
+            moved_out_of_ref: moved_out_of_ref(db, &body),
+            mir_body: body,
+        });
+    })?;
+    Ok(res.into())
 }
 
 fn moved_out_of_ref(db: &dyn HirDatabase, body: &MirBody) -> Vec<MovedOutOfRef> {
@@ -277,21 +272,35 @@ fn ever_initialized_map(
             );
             return;
         };
-        let targets = match &terminator.kind {
-            TerminatorKind::Goto { target } => vec![*target],
-            TerminatorKind::SwitchInt { targets, .. } => targets.all_targets().to_vec(),
+        let mut process = |target, is_ever_initialized| {
+            if !result[target].contains_idx(l) || !result[target][l] && is_ever_initialized {
+                result[target].insert(l, is_ever_initialized);
+                dfs(db, body, target, l, result);
+            }
+        };
+        match &terminator.kind {
+            TerminatorKind::Goto { target } => process(*target, is_ever_initialized),
+            TerminatorKind::SwitchInt { targets, .. } => {
+                targets.all_targets().iter().for_each(|&it| process(it, is_ever_initialized));
+            }
             TerminatorKind::UnwindResume
             | TerminatorKind::Abort
             | TerminatorKind::Return
-            | TerminatorKind::Unreachable => vec![],
+            | TerminatorKind::Unreachable => (),
             TerminatorKind::Call { target, cleanup, destination, .. } => {
                 if destination.projection.len() == 0 && destination.local == l {
                     is_ever_initialized = true;
                 }
-                target.into_iter().chain(cleanup.into_iter()).copied().collect()
+                target
+                    .into_iter()
+                    .chain(cleanup.into_iter())
+                    .for_each(|&it| process(it, is_ever_initialized));
             }
             TerminatorKind::Drop { target, unwind, place: _ } => {
-                Some(target).into_iter().chain(unwind.into_iter()).copied().collect()
+                iter::once(target)
+                    .into_iter()
+                    .chain(unwind.into_iter())
+                    .for_each(|&it| process(it, is_ever_initialized));
             }
             TerminatorKind::DropAndReplace { .. }
             | TerminatorKind::Assert { .. }
@@ -300,13 +309,7 @@ fn ever_initialized_map(
             | TerminatorKind::FalseEdge { .. }
             | TerminatorKind::FalseUnwind { .. } => {
                 never!("We don't emit these MIR terminators yet");
-                vec![]
-            }
-        };
-        for target in targets {
-            if !result[target].contains_idx(l) || !result[target][l] && is_ever_initialized {
-                result[target].insert(l, is_ever_initialized);
-                dfs(db, body, target, l, result);
+                ()
             }
         }
     }