From 50b34279c3bb84ac5287626d8ee2bc4414e5d07a Mon Sep 17 00:00:00 2001
From: bjorn3 <17426603+bjorn3@users.noreply.github.com>
Date: Sat, 11 May 2024 17:39:51 +0000
Subject: [PATCH] Split cgus into todo and done before the main module codegen
 loop

---
 src/concurrency_limiter.rs | 13 --------
 src/driver/aot.rs          | 67 +++++++++++++++++++-------------------
 2 files changed, 34 insertions(+), 46 deletions(-)

diff --git a/src/concurrency_limiter.rs b/src/concurrency_limiter.rs
index 9678969134a..c7f543cf7cd 100644
--- a/src/concurrency_limiter.rs
+++ b/src/concurrency_limiter.rs
@@ -78,11 +78,6 @@ pub(super) fn acquire(&mut self, dcx: &rustc_errors::DiagCtxt) -> ConcurrencyLim
         }
     }
 
-    pub(super) fn job_already_done(&mut self) {
-        let mut state = self.state.lock().unwrap();
-        state.job_already_done();
-    }
-
     pub(crate) fn finished(mut self) {
         self.helper_thread.take();
 
@@ -190,14 +185,6 @@ pub(super) fn job_finished(&mut self) {
             self.assert_invariants();
         }
 
-        pub(super) fn job_already_done(&mut self) {
-            self.assert_invariants();
-            self.pending_jobs -= 1;
-            self.assert_invariants();
-            self.drop_excess_capacity();
-            self.assert_invariants();
-        }
-
         pub(super) fn poison(&mut self, error: String) {
             self.poisoned = true;
             self.stored_error = Some(error);
diff --git a/src/driver/aot.rs b/src/driver/aot.rs
index e8c96486041..2466cbed299 100644
--- a/src/driver/aot.rs
+++ b/src/driver/aot.rs
@@ -604,40 +604,41 @@ pub(crate) fn run_aot(
 
     let global_asm_config = Arc::new(crate::global_asm::GlobalAsmConfig::new(tcx));
 
-    let mut concurrency_limiter = ConcurrencyLimiter::new(tcx.sess, cgus.len());
+    let (todo_cgus, done_cgus) =
+        cgus.into_iter().enumerate().partition::<Vec<_>, _>(|&(i, _)| match cgu_reuse[i] {
+            _ if backend_config.disable_incr_cache => true,
+            CguReuse::No => true,
+            CguReuse::PreLto | CguReuse::PostLto => false,
+        });
 
-    let modules = tcx.sess.time("codegen mono items", || {
-        cgus.iter()
-            .enumerate()
-            .map(|(i, cgu)| {
-                let cgu_reuse =
-                    if backend_config.disable_incr_cache { CguReuse::No } else { cgu_reuse[i] };
-                match cgu_reuse {
-                    CguReuse::No => {
-                        let dep_node = cgu.codegen_dep_node(tcx);
-                        tcx.dep_graph
-                            .with_task(
-                                dep_node,
-                                tcx,
-                                (
-                                    backend_config.clone(),
-                                    global_asm_config.clone(),
-                                    cgu.name(),
-                                    concurrency_limiter.acquire(tcx.dcx()),
-                                ),
-                                module_codegen,
-                                Some(rustc_middle::dep_graph::hash_result),
-                            )
-                            .0
-                    }
-                    CguReuse::PreLto | CguReuse::PostLto => {
-                        concurrency_limiter.job_already_done();
-                        OngoingModuleCodegen::Sync(reuse_workproduct_for_cgu(tcx, cgu))
-                    }
-                }
-            })
-            .collect::<Vec<_>>()
-    });
+    let mut concurrency_limiter = ConcurrencyLimiter::new(tcx.sess, todo_cgus.len());
+
+    let modules =
+        tcx.sess.time("codegen mono items", || {
+            todo_cgus
+                .into_iter()
+                .map(|(_, cgu)| {
+                    let dep_node = cgu.codegen_dep_node(tcx);
+                    tcx.dep_graph
+                        .with_task(
+                            dep_node,
+                            tcx,
+                            (
+                                backend_config.clone(),
+                                global_asm_config.clone(),
+                                cgu.name(),
+                                concurrency_limiter.acquire(tcx.dcx()),
+                            ),
+                            module_codegen,
+                            Some(rustc_middle::dep_graph::hash_result),
+                        )
+                        .0
+                })
+                .chain(done_cgus.into_iter().map(|(_, cgu)| {
+                    OngoingModuleCodegen::Sync(reuse_workproduct_for_cgu(tcx, cgu))
+                }))
+                .collect::<Vec<_>>()
+        });
 
     let mut allocator_module = make_module(tcx.sess, &backend_config, "allocator_shim".to_string());
     let mut allocator_unwind_context = UnwindContext::new(allocator_module.isa(), true);