From baec3076dbc612433210b7ab3366736e445f8c30 Mon Sep 17 00:00:00 2001
From: Zalathar <Zalathar@users.noreply.github.com>
Date: Sat, 3 Feb 2024 11:15:05 +1100
Subject: [PATCH] Allow tests to specify a `//@ filecheck-flags:` header

Any flags specified here will be passed to LLVM's `filecheck` tool, in tests
that use that tool.
---
 src/tools/compiletest/src/header.rs             |  8 ++++++++
 src/tools/compiletest/src/runtest.rs            | 11 ++++++-----
 tests/codegen/meta-filecheck/filecheck-flags.rs |  8 ++++++++
 3 files changed, 22 insertions(+), 5 deletions(-)
 create mode 100644 tests/codegen/meta-filecheck/filecheck-flags.rs

diff --git a/src/tools/compiletest/src/header.rs b/src/tools/compiletest/src/header.rs
index c76cba824a3..6de445a5783 100644
--- a/src/tools/compiletest/src/header.rs
+++ b/src/tools/compiletest/src/header.rs
@@ -197,6 +197,8 @@ pub struct TestProps {
     /// 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>,
+    /// Extra flags to pass to LLVM's `filecheck` tool, in tests that use it.
+    pub filecheck_flags: Vec<String>,
 }
 
 mod directives {
@@ -236,6 +238,7 @@ mod directives {
     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";
+    pub const FILECHECK_FLAGS: &'static str = "filecheck-flags";
     // This isn't a real directive, just one that is probably mistyped often
     pub const INCORRECT_COMPILER_FLAGS: &'static str = "compiler-flags";
 }
@@ -286,6 +289,7 @@ impl TestProps {
             mir_unit_test: None,
             remap_src_base: false,
             llvm_cov_flags: vec![],
+            filecheck_flags: vec![],
         }
     }
 
@@ -542,6 +546,10 @@ impl TestProps {
                     if let Some(flags) = config.parse_name_value_directive(ln, LLVM_COV_FLAGS) {
                         self.llvm_cov_flags.extend(split_flags(&flags));
                     }
+
+                    if let Some(flags) = config.parse_name_value_directive(ln, FILECHECK_FLAGS) {
+                        self.filecheck_flags.extend(split_flags(&flags));
+                    }
                 },
             );
 
diff --git a/src/tools/compiletest/src/runtest.rs b/src/tools/compiletest/src/runtest.rs
index 3c417f3161b..6e0a7e6909a 100644
--- a/src/tools/compiletest/src/runtest.rs
+++ b/src/tools/compiletest/src/runtest.rs
@@ -2907,11 +2907,8 @@ impl<'test> TestCx<'test> {
         let mut filecheck = Command::new(self.config.llvm_filecheck.as_ref().unwrap());
         filecheck.arg("--input-file").arg(output).arg(&self.testpaths.file);
 
-        // It would be more appropriate to make most of the arguments configurable through
-        // a comment-attribute similar to `compile-flags`. For example, --check-prefixes is a very
-        // useful flag.
-        //
-        // For now, though…
+        // FIXME: Consider making some of these prefix flags opt-in per test,
+        // via `filecheck-flags` or by adding new header directives.
 
         // Because we use custom prefixes, we also have to register the default prefix.
         filecheck.arg("--check-prefix=CHECK");
@@ -2931,6 +2928,10 @@ impl<'test> TestCx<'test> {
 
         // Provide more context on failures.
         filecheck.args(&["--dump-input-context", "100"]);
+
+        // Add custom flags supplied by the `filecheck-flags:` test header.
+        filecheck.args(&self.props.filecheck_flags);
+
         self.compose_and_run(filecheck, "", None, None)
     }
 
diff --git a/tests/codegen/meta-filecheck/filecheck-flags.rs b/tests/codegen/meta-filecheck/filecheck-flags.rs
new file mode 100644
index 00000000000..8e451cf4fdc
--- /dev/null
+++ b/tests/codegen/meta-filecheck/filecheck-flags.rs
@@ -0,0 +1,8 @@
+// Arguments provided via `filecheck-flags` should be passed to `filecheck`.
+
+//@ revisions: good bad
+//@ [good] filecheck-flags: --check-prefix=CUSTOM
+//@ [bad] should-fail
+
+// CUSTOM: main
+fn main() {}