compiletest: add `needs-crate-type` directive

The `//@ needs-crate-type: $crate_types...` directive takes a
comma-separated list of crate types that the target platform must
support in order for the test to be run.

5 files changed, 117 insertions, 2 deletions

diff --git a/src/tools/compiletest/src/common.rs b/src/tools/compiletest/src/common.rs
index 36218b2f904..6750b5288f4 100644
--- a/src/tools/compiletest/src/common.rs
+++ b/src/tools/compiletest/src/common.rs
@@ -395,6 +395,7 @@ pub struct Config {
 
     pub target_cfgs: OnceLock<TargetCfgs>,
     pub builtin_cfg_names: OnceLock<HashSet<String>>,
+    pub supported_crate_types: OnceLock<HashSet<String>>,
 
     pub nocapture: bool,
 
@@ -472,6 +473,11 @@ impl Config {
         self.builtin_cfg_names.get_or_init(|| builtin_cfg_names(self))
     }
 
+    /// Get the list of crate types that the target platform supports.
+    pub fn supported_crate_types(&self) -> &HashSet<String> {
+        self.supported_crate_types.get_or_init(|| supported_crate_types(self))
+    }
+
     pub fn has_threads(&self) -> bool {
         // Wasm targets don't have threads unless `-threads` is in the target
         // name, such as `wasm32-wasip1-threads`.
@@ -745,6 +751,31 @@ fn builtin_cfg_names(config: &Config) -> HashSet<String> {
     .collect()
 }
 
+pub const KNOWN_CRATE_TYPES: &[&str] =
+    &["bin", "cdylib", "dylib", "lib", "proc-macro", "rlib", "staticlib"];
+
+fn supported_crate_types(config: &Config) -> HashSet<String> {
+    let crate_types: HashSet<_> = rustc_output(
+        config,
+        &["--target", &config.target, "--print=supported-crate-types", "-Zunstable-options"],
+        Default::default(),
+    )
+    .lines()
+    .map(|l| l.to_string())
+    .collect();
+
+    for crate_type in crate_types.iter() {
+        assert!(
+            KNOWN_CRATE_TYPES.contains(&crate_type.as_str()),
+            "unexpected crate type `{}`: known crate types are {:?}",
+            crate_type,
+            KNOWN_CRATE_TYPES
+        );
+    }
+
+    crate_types
+}
+
 fn rustc_output(config: &Config, args: &[&str], envs: HashMap<String, String>) -> String {
     let mut command = Command::new(&config.rustc_path);
     add_dylib_path(&mut command, iter::once(&config.compile_lib_path));
diff --git a/src/tools/compiletest/src/directive-list.rs b/src/tools/compiletest/src/directive-list.rs
index 44d9c0330f7..086a8a67456 100644
--- a/src/tools/compiletest/src/directive-list.rs
+++ b/src/tools/compiletest/src/directive-list.rs
@@ -133,6 +133,7 @@ const KNOWN_DIRECTIVE_NAMES: &[&str] = &[
     "min-llvm-version",
     "min-system-llvm-version",
     "needs-asm-support",
+    "needs-crate-type",
     "needs-deterministic-layouts",
     "needs-dlltool",
     "needs-dynamic-linking",
diff --git a/src/tools/compiletest/src/header/needs.rs b/src/tools/compiletest/src/header/needs.rs
index 12f0790fb10..2ace40c490b 100644
--- a/src/tools/compiletest/src/header/needs.rs
+++ b/src/tools/compiletest/src/header/needs.rs
@@ -1,4 +1,4 @@
-use crate::common::{Config, KNOWN_TARGET_HAS_ATOMIC_WIDTHS, Sanitizer};
+use crate::common::{Config, KNOWN_CRATE_TYPES, KNOWN_TARGET_HAS_ATOMIC_WIDTHS, Sanitizer};
 use crate::header::{IgnoreDecision, llvm_has_libzstd};
 
 pub(super) fn handle_needs(
@@ -6,7 +6,7 @@ pub(super) fn handle_needs(
     config: &Config,
     ln: &str,
 ) -> IgnoreDecision {
-    // Note thet we intentionally still put the needs- prefix here to make the file show up when
+    // Note that we intentionally still put the needs- prefix here to make the file show up when
     // grepping for a directive name, even though we could technically strip that.
     let needs = &[
         Need {
@@ -224,6 +224,50 @@ pub(super) fn handle_needs(
         }
     }
 
+    // FIXME(jieyouxu): share multi-value directive logic with `needs-target-has-atomic` above.
+    if name == "needs-crate-type" {
+        let Some(rest) = rest else {
+            return IgnoreDecision::Error {
+                message:
+                    "expected `needs-crate-type` to have a comma-separated list of crate types"
+                        .to_string(),
+            };
+        };
+
+        // Expect directive value to be a list of comma-separated crate-types.
+        let specified_crate_types = rest
+            .split(',')
+            .map(|crate_type| crate_type.trim())
+            .map(ToString::to_string)
+            .collect::<Vec<String>>();
+
+        for crate_type in &specified_crate_types {
+            if !KNOWN_CRATE_TYPES.contains(&crate_type.as_str()) {
+                return IgnoreDecision::Error {
+                    message: format!(
+                        "unknown crate type specified in `needs-crate-type`: `{crate_type}` is not \
+                        a known crate type, known values are `{:?}`",
+                        KNOWN_CRATE_TYPES
+                    ),
+                };
+            }
+        }
+
+        let satisfies_all_crate_types = specified_crate_types
+            .iter()
+            .all(|specified| config.supported_crate_types().contains(specified));
+        if satisfies_all_crate_types {
+            return IgnoreDecision::Continue;
+        } else {
+            return IgnoreDecision::Ignore {
+                reason: format!(
+                    "skipping test as target does not support all of the crate types `{:?}`",
+                    specified_crate_types
+                ),
+            };
+        }
+    }
+
     if !name.starts_with("needs-") {
         return IgnoreDecision::Continue;
     }
diff --git a/src/tools/compiletest/src/header/tests.rs b/src/tools/compiletest/src/header/tests.rs
index 1a796cdd7cd..f3461f3c244 100644
--- a/src/tools/compiletest/src/header/tests.rs
+++ b/src/tools/compiletest/src/header/tests.rs
@@ -902,3 +902,41 @@ fn test_rustc_abi() {
     assert!(!check_ignore(&config, "//@ ignore-rustc_abi-x86-sse2"));
     assert!(check_ignore(&config, "//@ only-rustc_abi-x86-sse2"));
 }
+
+#[test]
+fn test_supported_crate_types() {
+    // Basic assumptions check on under-test compiler's `--print=supported-crate-types` output based
+    // on knowledge about the cherry-picked `x86_64-unknown-linux-gnu` and `wasm32-unknown-unknown`
+    // targets. Also smoke tests the `needs-crate-type` directive itself.
+
+    use std::collections::HashSet;
+
+    let config = cfg().target("x86_64-unknown-linux-gnu").build();
+    assert_eq!(
+        config.supported_crate_types().iter().map(String::as_str).collect::<HashSet<_>>(),
+        HashSet::from(["bin", "cdylib", "dylib", "lib", "proc-macro", "rlib", "staticlib"]),
+    );
+    assert!(!check_ignore(&config, "//@ needs-crate-type: rlib"));
+    assert!(!check_ignore(&config, "//@ needs-crate-type: dylib"));
+    assert!(!check_ignore(
+        &config,
+        "//@ needs-crate-type: bin, cdylib, dylib, lib, proc-macro, rlib, staticlib"
+    ));
+
+    let config = cfg().target("wasm32-unknown-unknown").build();
+    assert_eq!(
+        config.supported_crate_types().iter().map(String::as_str).collect::<HashSet<_>>(),
+        HashSet::from(["bin", "cdylib", "lib", "rlib", "staticlib"]),
+    );
+
+    // rlib is supported
+    assert!(!check_ignore(&config, "//@ needs-crate-type: rlib"));
+    // dylib is not
+    assert!(check_ignore(&config, "//@ needs-crate-type: dylib"));
+    // If multiple crate types are specified, then all specified crate types need to be supported.
+    assert!(check_ignore(&config, "//@ needs-crate-type: cdylib, dylib"));
+    assert!(check_ignore(
+        &config,
+        "//@ needs-crate-type: bin, cdylib, dylib, lib, proc-macro, rlib, staticlib"
+    ));
+}
diff --git a/src/tools/compiletest/src/lib.rs b/src/tools/compiletest/src/lib.rs
index 782f6e0f2d8..8fa11c81c6a 100644
--- a/src/tools/compiletest/src/lib.rs
+++ b/src/tools/compiletest/src/lib.rs
@@ -431,6 +431,7 @@ pub fn parse_config(args: Vec<String>) -> Config {
 
         target_cfgs: OnceLock::new(),
         builtin_cfg_names: OnceLock::new(),
+        supported_crate_types: OnceLock::new(),
 
         nocapture: matches.opt_present("no-capture"),