From 6739c7e95e3ba7e187cb4bdbb9bf5e12c5729774 Mon Sep 17 00:00:00 2001
From: hi-rustin <rustin.liu@gmail.com>
Date: Wed, 27 Dec 2023 10:27:23 +0800
Subject: [PATCH] fix: set OUT_DIR for all units with build scripts

Signed-off-by: hi-rustin <rustin.liu@gmail.com>
---
 src/cargo/core/compiler/context/mod.rs | 47 +++++++++++++++-----------
 tests/testsuite/test.rs                |  2 --
 2 files changed, 27 insertions(+), 22 deletions(-)

diff --git a/src/cargo/core/compiler/context/mod.rs b/src/cargo/core/compiler/context/mod.rs
index cfbfccb30..3aedf515c 100644
--- a/src/cargo/core/compiler/context/mod.rs
+++ b/src/cargo/core/compiler/context/mod.rs
@@ -12,6 +12,7 @@ use crate::util::errors::CargoResult;
 use crate::util::profile;
 use anyhow::{bail, Context as _};
 use filetime::FileTime;
+use itertools::Itertools;
 use jobserver::Client;
 
 use super::build_plan::BuildPlan;
@@ -185,6 +186,32 @@ impl<'a, 'cfg> Context<'a, 'cfg> {
             plan.output_plan(self.bcx.config);
         }
 
+        // Add `OUT_DIR` to env vars if unit has a build script.
+        let units_with_build_script = &self
+            .bcx
+            .roots
+            .iter()
+            .filter(|unit| self.build_scripts.contains_key(unit))
+            .dedup_by(|x, y| x.pkg.package_id() == y.pkg.package_id())
+            .collect::<Vec<_>>();
+        for unit in units_with_build_script {
+            for dep in &self.bcx.unit_graph[unit] {
+                if dep.unit.mode.is_run_custom_build() {
+                    let out_dir = self
+                        .files()
+                        .build_script_out_dir(&dep.unit)
+                        .display()
+                        .to_string();
+                    let script_meta = self.get_run_build_script_metadata(&dep.unit);
+                    self.compilation
+                        .extra_env
+                        .entry(script_meta)
+                        .or_insert_with(Vec::new)
+                        .push(("OUT_DIR".to_string(), out_dir));
+                }
+            }
+        }
+
         // Collect the result of the build into `self.compilation`.
         for unit in &self.bcx.roots {
             // Collect tests and executables.
@@ -213,26 +240,6 @@ impl<'a, 'cfg> Context<'a, 'cfg> {
                 }
             }
 
-            // If the unit has a build script, add `OUT_DIR` to the
-            // environment variables.
-            if unit.target.is_lib() {
-                for dep in &self.bcx.unit_graph[unit] {
-                    if dep.unit.mode.is_run_custom_build() {
-                        let out_dir = self
-                            .files()
-                            .build_script_out_dir(&dep.unit)
-                            .display()
-                            .to_string();
-                        let script_meta = self.get_run_build_script_metadata(&dep.unit);
-                        self.compilation
-                            .extra_env
-                            .entry(script_meta)
-                            .or_insert_with(Vec::new)
-                            .push(("OUT_DIR".to_string(), out_dir));
-                    }
-                }
-            }
-
             // Collect information for `rustdoc --test`.
             if unit.mode.is_doc_test() {
                 let mut unstable_opts = false;
diff --git a/tests/testsuite/test.rs b/tests/testsuite/test.rs
index 8a9e4ab36..8a1d42326 100644
--- a/tests/testsuite/test.rs
+++ b/tests/testsuite/test.rs
@@ -4941,7 +4941,5 @@ fn cargo_test_set_out_dir_env_var() {
 
     p.cargo("test").run();
     p.cargo("test --package foo --test case -- tests::test_add --exact --nocapture")
-        .with_stdout_contains("test tests::test_add ... FAILED")
-        .with_status(101)
         .run();
 }