From 731ba80a6b53e3397eadeda37b18bd8fb3016aad Mon Sep 17 00:00:00 2001
From: Zalathar <Zalathar@users.noreply.github.com>
Date: Sat, 16 Dec 2023 13:30:47 +1100
Subject: [PATCH] Allow coverage tests to enable `llvm-cov --use-color`

---
 src/tools/compiletest/src/header.rs  |  9 +++++++++
 src/tools/compiletest/src/runtest.rs |  2 ++
 tests/coverage/color.coverage        | 13 +++++++++++++
 tests/coverage/color.rs              | 11 +++++++++++
 4 files changed, 35 insertions(+)
 create mode 100644 tests/coverage/color.coverage
 create mode 100644 tests/coverage/color.rs

diff --git a/src/tools/compiletest/src/header.rs b/src/tools/compiletest/src/header.rs
index a9f022664cfd9..e70e01e8757e0 100644
--- a/src/tools/compiletest/src/header.rs
+++ b/src/tools/compiletest/src/header.rs
@@ -178,6 +178,9 @@ pub struct TestProps {
     // Whether to tell `rustc` to remap the "src base" directory to a fake
     // directory.
     pub remap_src_base: bool,
+    /// Extra flags to pass to `llvm-cov` when producing coverage reports.
+    /// Only used by the "coverage-run" test mode.
+    pub llvm_cov_flags: Vec<String>,
 }
 
 mod directives {
@@ -216,6 +219,7 @@ mod directives {
     pub const MIR_UNIT_TEST: &'static str = "unit-test";
     pub const REMAP_SRC_BASE: &'static str = "remap-src-base";
     pub const COMPARE_OUTPUT_LINES_BY_SUBSET: &'static str = "compare-output-lines-by-subset";
+    pub const LLVM_COV_FLAGS: &'static str = "llvm-cov-flags";
     // This isn't a real directive, just one that is probably mistyped often
     pub const INCORRECT_COMPILER_FLAGS: &'static str = "compiler-flags";
 }
@@ -265,6 +269,7 @@ impl TestProps {
             stderr_per_bitwidth: false,
             mir_unit_test: None,
             remap_src_base: false,
+            llvm_cov_flags: vec![],
         }
     }
 
@@ -495,6 +500,10 @@ impl TestProps {
                     COMPARE_OUTPUT_LINES_BY_SUBSET,
                     &mut self.compare_output_lines_by_subset,
                 );
+
+                if let Some(flags) = config.parse_name_value_directive(ln, LLVM_COV_FLAGS) {
+                    self.llvm_cov_flags.extend(split_flags(&flags));
+                }
             });
         }
 
diff --git a/src/tools/compiletest/src/runtest.rs b/src/tools/compiletest/src/runtest.rs
index 1f5f77839de41..b258b748ca876 100644
--- a/src/tools/compiletest/src/runtest.rs
+++ b/src/tools/compiletest/src/runtest.rs
@@ -575,6 +575,8 @@ impl<'test> TestCx<'test> {
                 cmd.arg("--object");
                 cmd.arg(bin);
             }
+
+            cmd.args(&self.props.llvm_cov_flags);
         });
         if !proc_res.status.success() {
             self.fatal_proc_rec("llvm-cov show failed!", &proc_res);
diff --git a/tests/coverage/color.coverage b/tests/coverage/color.coverage
new file mode 100644
index 0000000000000..bc49fff9cb75d
--- /dev/null
+++ b/tests/coverage/color.coverage
@@ -0,0 +1,13 @@
+   LL|       |// edition: 2021
+   LL|       |// ignore-mode-coverage-map
+   LL|       |// ignore-windows
+   LL|       |// llvm-cov-flags: --use-color
+   LL|       |
+   LL|       |// Verify that telling `llvm-cov` to use colored output actually works.
+   LL|       |// Ignored on Windows because we can't tell the tool to use ANSI escapes.
+   LL|       |
+   LL|      1|fn main() {
+   LL|      1|    for _i in 0..0 {}
+                      ^0         ^0
+   LL|      1|}
+
diff --git a/tests/coverage/color.rs b/tests/coverage/color.rs
new file mode 100644
index 0000000000000..bd727946c781b
--- /dev/null
+++ b/tests/coverage/color.rs
@@ -0,0 +1,11 @@
+// edition: 2021
+// ignore-mode-coverage-map
+// ignore-windows
+// llvm-cov-flags: --use-color
+
+// Verify that telling `llvm-cov` to use colored output actually works.
+// Ignored on Windows because we can't tell the tool to use ANSI escapes.
+
+fn main() {
+    for _i in 0..0 {}
+}