From 99c232223f8330debc199f616308bc559b38c361 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E8=AE=B8=E6=9D=B0=E5=8F=8B=20Jieyou=20Xu=20=28Joe=29?=
 <39484203+jieyouxu@users.noreply.github.com>
Date: Mon, 2 Dec 2024 14:48:01 +0800
Subject: [PATCH] Add `needs-target-has-atomic` directive

Before this commit, the test writer has to specify platforms and
architectures by hand for targets that have differing atomic width
support. `#[cfg(target_has_atomic)]` is not quite the same because (1)
you may have to specify additional matchers manually which has to be
maintained individually, and (2) the `#[cfg]` blocks does not
communicate to compiletest that a test would be ignored for a given
target.

This commit implements a `//@ needs-target-has-atomic` directive which
admits a comma-separated list of required atomic widths that the target
must satisfy in order for the test to run.

```
//@ needs-target-has-atomic: 8, 16, ptr
```

See <https://github.com/rust-lang/rust/issues/87377>.

Co-authored-by: kei519 <masaki.keigo.q00@kyoto-u.jp>
---
 src/tools/compiletest/src/common.rs         | 22 ++++++++-
 src/tools/compiletest/src/directive-list.rs |  1 +
 src/tools/compiletest/src/header/needs.rs   | 51 +++++++++++++++++++--
 src/tools/compiletest/src/header/tests.rs   | 37 +++++++++++++++
 4 files changed, 106 insertions(+), 5 deletions(-)

diff --git a/src/tools/compiletest/src/common.rs b/src/tools/compiletest/src/common.rs
index e6fbba943a48c..36f876bcdc604 100644
--- a/src/tools/compiletest/src/common.rs
+++ b/src/tools/compiletest/src/common.rs
@@ -1,4 +1,4 @@
-use std::collections::{HashMap, HashSet};
+use std::collections::{BTreeSet, HashMap, HashSet};
 use std::ffi::OsString;
 use std::path::{Path, PathBuf};
 use std::process::Command;
@@ -486,6 +486,9 @@ impl Config {
     }
 }
 
+/// Known widths of `target_has_atomic`.
+pub const KNOWN_TARGET_HAS_ATOMIC_WIDTHS: &[&str] = &["8", "16", "32", "64", "128", "ptr"];
+
 #[derive(Debug, Clone)]
 pub struct TargetCfgs {
     pub current: TargetCfg,
@@ -611,6 +614,17 @@ impl TargetCfgs {
                 ("panic", Some("abort")) => cfg.panic = PanicStrategy::Abort,
                 ("panic", Some("unwind")) => cfg.panic = PanicStrategy::Unwind,
                 ("panic", other) => panic!("unexpected value for panic cfg: {other:?}"),
+
+                ("target_has_atomic", Some(width))
+                    if KNOWN_TARGET_HAS_ATOMIC_WIDTHS.contains(&width) =>
+                {
+                    cfg.target_has_atomic.insert(width.to_string());
+                }
+                ("target_has_atomic", Some(other)) => {
+                    panic!("unexpected value for `target_has_atomic` cfg: {other:?}")
+                }
+                // Nightly-only std-internal impl detail.
+                ("target_has_atomic", None) => {}
                 _ => {}
             }
         }
@@ -645,6 +659,12 @@ pub struct TargetCfg {
     pub(crate) xray: bool,
     #[serde(default = "default_reloc_model")]
     pub(crate) relocation_model: String,
+
+    // Not present in target cfg json output, additional derived information.
+    #[serde(skip)]
+    /// Supported target atomic widths: e.g. `8` to `128` or `ptr`. This is derived from the builtin
+    /// `target_has_atomic` `cfg`s e.g. `target_has_atomic="8"`.
+    pub(crate) target_has_atomic: BTreeSet<String>,
 }
 
 impl TargetCfg {
diff --git a/src/tools/compiletest/src/directive-list.rs b/src/tools/compiletest/src/directive-list.rs
index 952533e904c5a..c7e8d4b58a562 100644
--- a/src/tools/compiletest/src/directive-list.rs
+++ b/src/tools/compiletest/src/directive-list.rs
@@ -154,6 +154,7 @@ const KNOWN_DIRECTIVE_NAMES: &[&str] = &[
     "needs-sanitizer-thread",
     "needs-std-debug-assertions",
     "needs-symlink",
+    "needs-target-has-atomic",
     "needs-threads",
     "needs-unwind",
     "needs-wasmtime",
diff --git a/src/tools/compiletest/src/header/needs.rs b/src/tools/compiletest/src/header/needs.rs
index 77570c58db5cf..e19dcd992fc2e 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, Sanitizer};
+use crate::common::{Config, KNOWN_TARGET_HAS_ATOMIC_WIDTHS, Sanitizer};
 use crate::header::{IgnoreDecision, llvm_has_libzstd};
 
 pub(super) fn handle_needs(
@@ -171,11 +171,54 @@ pub(super) fn handle_needs(
         },
     ];
 
-    let (name, comment) = match ln.split_once([':', ' ']) {
-        Some((name, comment)) => (name, Some(comment)),
+    let (name, rest) = match ln.split_once([':', ' ']) {
+        Some((name, rest)) => (name, Some(rest)),
         None => (ln, None),
     };
 
+    // FIXME(jieyouxu): tighten up this parsing to reject using both `:` and ` ` as means to
+    // delineate value.
+    if name == "needs-target-has-atomic" {
+        let Some(rest) = rest else {
+            return IgnoreDecision::Error {
+                message: "expected `needs-target-has-atomic` to have a comma-separated list of atomic widths".to_string(),
+            };
+        };
+
+        // Expect directive value to be a list of comma-separated atomic widths.
+        let specified_widths = rest
+            .split(',')
+            .map(|width| width.trim())
+            .map(ToString::to_string)
+            .collect::<Vec<String>>();
+
+        for width in &specified_widths {
+            if !KNOWN_TARGET_HAS_ATOMIC_WIDTHS.contains(&width.as_str()) {
+                return IgnoreDecision::Error {
+                    message: format!(
+                        "unknown width specified in `needs-target-has-atomic`: `{width}` is not a \
+                        known `target_has_atomic_width`, known values are `{:?}`",
+                        KNOWN_TARGET_HAS_ATOMIC_WIDTHS
+                    ),
+                };
+            }
+        }
+
+        let satisfies_all_specified_widths = specified_widths
+            .iter()
+            .all(|specified| config.target_cfg().target_has_atomic.contains(specified));
+        if satisfies_all_specified_widths {
+            return IgnoreDecision::Continue;
+        } else {
+            return IgnoreDecision::Ignore {
+                reason: format!(
+                    "skipping test as target does not support all of the required `target_has_atomic` widths `{:?}`",
+                    specified_widths
+                ),
+            };
+        }
+    }
+
     if !name.starts_with("needs-") {
         return IgnoreDecision::Continue;
     }
@@ -193,7 +236,7 @@ pub(super) fn handle_needs(
                 break;
             } else {
                 return IgnoreDecision::Ignore {
-                    reason: if let Some(comment) = comment {
+                    reason: if let Some(comment) = rest {
                         format!("{} ({})", need.ignore_reason, comment.trim())
                     } else {
                         need.ignore_reason.into()
diff --git a/src/tools/compiletest/src/header/tests.rs b/src/tools/compiletest/src/header/tests.rs
index 4d75c38dd3206..cd7c6f8361ed3 100644
--- a/src/tools/compiletest/src/header/tests.rs
+++ b/src/tools/compiletest/src/header/tests.rs
@@ -787,3 +787,40 @@ fn test_not_trailing_directive() {
     run_path(&mut poisoned, Path::new("a.rs"), b"//@ revisions: incremental");
     assert!(!poisoned);
 }
+
+#[test]
+fn test_needs_target_has_atomic() {
+    use std::collections::BTreeSet;
+
+    // `x86_64-unknown-linux-gnu` supports `["8", "16", "32", "64", "ptr"]` but not `128`.
+
+    let config = cfg().target("x86_64-unknown-linux-gnu").build();
+    // Expectation sanity check.
+    assert_eq!(
+        config.target_cfg().target_has_atomic,
+        BTreeSet::from([
+            "8".to_string(),
+            "16".to_string(),
+            "32".to_string(),
+            "64".to_string(),
+            "ptr".to_string()
+        ]),
+        "expected `x86_64-unknown-linux-gnu` to not have 128-bit atomic support"
+    );
+
+    assert!(!check_ignore(&config, "//@ needs-target-has-atomic: 8"));
+    assert!(!check_ignore(&config, "//@ needs-target-has-atomic: 16"));
+    assert!(!check_ignore(&config, "//@ needs-target-has-atomic: 32"));
+    assert!(!check_ignore(&config, "//@ needs-target-has-atomic: 64"));
+    assert!(!check_ignore(&config, "//@ needs-target-has-atomic: ptr"));
+
+    assert!(check_ignore(&config, "//@ needs-target-has-atomic: 128"));
+
+    assert!(!check_ignore(&config, "//@ needs-target-has-atomic: 8,16,32,64,ptr"));
+
+    assert!(check_ignore(&config, "//@ needs-target-has-atomic: 8,16,32,64,ptr,128"));
+
+    // Check whitespace between widths is permitted.
+    assert!(!check_ignore(&config, "//@ needs-target-has-atomic: 8, ptr"));
+    assert!(check_ignore(&config, "//@ needs-target-has-atomic: 8, ptr, 128"));
+}