Implement compiletest `--new-output-capture`, in stable Rust

4 files changed, 76 insertions, 2 deletions

diff --git a/src/tools/compiletest/src/common.rs b/src/tools/compiletest/src/common.rs
index 89b1b4f84b6..62fdee98735 100644
--- a/src/tools/compiletest/src/common.rs
+++ b/src/tools/compiletest/src/common.rs
@@ -667,6 +667,10 @@ pub struct Config {
     /// to avoid `!nocapture` double-negatives.
     pub nocapture: bool,
 
+    /// True if the experimental new output-capture implementation should be
+    /// used, avoiding the need for `#![feature(internal_output_capture)]`.
+    pub new_output_capture: bool,
+
     /// Needed both to construct [`build_helper::git::GitConfig`].
     pub nightly_branch: String,
     pub git_merge_commit_email: String,
@@ -784,6 +788,7 @@ impl Config {
             builtin_cfg_names: Default::default(),
             supported_crate_types: Default::default(),
             nocapture: Default::default(),
+            new_output_capture: Default::default(),
             nightly_branch: Default::default(),
             git_merge_commit_email: Default::default(),
             profiler_runtime: Default::default(),
diff --git a/src/tools/compiletest/src/executor.rs b/src/tools/compiletest/src/executor.rs
index 818d1719885..b0dc24798b6 100644
--- a/src/tools/compiletest/src/executor.rs
+++ b/src/tools/compiletest/src/executor.rs
@@ -168,12 +168,17 @@ enum CaptureKind {
     /// Use the old output-capture implementation, which relies on the unstable
     /// library feature `#![feature(internal_output_capture)]`.
     Old { buf: Arc<Mutex<Vec<u8>>> },
+
+    /// Use the new output-capture implementation, which only uses stable Rust.
+    New { buf: output_capture::CaptureBuf },
 }
 
 impl CaptureKind {
     fn for_config(config: &Config) -> Self {
         if config.nocapture {
             Self::None
+        } else if config.new_output_capture {
+            Self::New { buf: output_capture::CaptureBuf::new() }
         } else {
             // Create a capure buffer for `io::set_output_capture`.
             Self::Old { buf: Default::default() }
@@ -184,21 +189,30 @@ impl CaptureKind {
         match self {
             Self::None => false,
             Self::Old { .. } => true,
+            Self::New { .. } => true,
         }
     }
 
     fn stdout(&self) -> &dyn ConsoleOut {
-        &output_capture::Stdout
+        self.capture_buf_or(&output_capture::Stdout)
     }
 
     fn stderr(&self) -> &dyn ConsoleOut {
-        &output_capture::Stderr
+        self.capture_buf_or(&output_capture::Stderr)
+    }
+
+    fn capture_buf_or<'a>(&'a self, fallback: &'a dyn ConsoleOut) -> &'a dyn ConsoleOut {
+        match self {
+            Self::None | Self::Old { .. } => fallback,
+            Self::New { buf } => buf,
+        }
     }
 
     fn into_inner(self) -> Option<Vec<u8>> {
         match self {
             Self::None => None,
             Self::Old { buf } => Some(buf.lock().unwrap_or_else(|e| e.into_inner()).to_vec()),
+            Self::New { buf } => Some(buf.into_inner().into()),
         }
     }
 }
diff --git a/src/tools/compiletest/src/lib.rs b/src/tools/compiletest/src/lib.rs
index 875d497af1a..f647f96a9bf 100644
--- a/src/tools/compiletest/src/lib.rs
+++ b/src/tools/compiletest/src/lib.rs
@@ -178,6 +178,12 @@ pub fn parse_config(args: Vec<String>) -> Config {
         // FIXME: Temporarily retained so we can point users to `--no-capture`
         .optflag("", "nocapture", "")
         .optflag("", "no-capture", "don't capture stdout/stderr of tests")
+        .optopt(
+            "N",
+            "new-output-capture",
+            "enables or disables the new output-capture implementation",
+            "off|on",
+        )
         .optflag("", "profiler-runtime", "is the profiler runtime enabled for this target")
         .optflag("h", "help", "show this message")
         .reqopt("", "channel", "current Rust channel", "CHANNEL")
@@ -462,6 +468,14 @@ pub fn parse_config(args: Vec<String>) -> Config {
         supported_crate_types: OnceLock::new(),
 
         nocapture: matches.opt_present("no-capture"),
+        new_output_capture: {
+            let value = matches
+                .opt_str("new-output-capture")
+                .or_else(|| env::var("COMPILETEST_NEW_OUTPUT_CAPTURE").ok())
+                .unwrap_or_else(|| "off".to_owned());
+            parse_bool_option(&value)
+                .unwrap_or_else(|| panic!("unknown `--new-output-capture` value `{value}` given"))
+        },
 
         nightly_branch: matches.opt_str("nightly-branch").unwrap(),
         git_merge_commit_email: matches.opt_str("git-merge-commit-email").unwrap(),
@@ -477,6 +491,19 @@ pub fn parse_config(args: Vec<String>) -> Config {
     }
 }
 
+/// Parses the same set of boolean values accepted by rustc command-line arguments.
+///
+/// Accepting all of these values is more complicated than just picking one
+/// pair, but has the advantage that contributors who are used to rustc
+/// shouldn't have to think about which values are legal.
+fn parse_bool_option(value: &str) -> Option<bool> {
+    match value {
+        "off" | "no" | "n" | "false" => Some(false),
+        "on" | "yes" | "y" | "true" => Some(true),
+        _ => None,
+    }
+}
+
 pub fn opt_str(maybestr: &Option<String>) -> &str {
     match *maybestr {
         None => "(none)",
diff --git a/src/tools/compiletest/src/output_capture.rs b/src/tools/compiletest/src/output_capture.rs
index e5e3e14f288..de1aea11ade 100644
--- a/src/tools/compiletest/src/output_capture.rs
+++ b/src/tools/compiletest/src/output_capture.rs
@@ -1,5 +1,6 @@
 use std::fmt;
 use std::panic::RefUnwindSafe;
+use std::sync::Mutex;
 
 pub trait ConsoleOut: fmt::Debug + RefUnwindSafe {
     fn write_fmt(&self, args: fmt::Arguments<'_>);
@@ -22,3 +23,30 @@ impl ConsoleOut for Stderr {
         eprint!("{args}");
     }
 }
+
+pub(crate) struct CaptureBuf {
+    inner: Mutex<String>,
+}
+
+impl CaptureBuf {
+    pub(crate) fn new() -> Self {
+        Self { inner: Mutex::new(String::new()) }
+    }
+
+    pub(crate) fn into_inner(self) -> String {
+        self.inner.into_inner().unwrap_or_else(|e| e.into_inner())
+    }
+}
+
+impl fmt::Debug for CaptureBuf {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("CaptureBuf").finish_non_exhaustive()
+    }
+}
+
+impl ConsoleOut for CaptureBuf {
+    fn write_fmt(&self, args: fmt::Arguments<'_>) {
+        let mut s = self.inner.lock().unwrap_or_else(|e| e.into_inner());
+        <String as fmt::Write>::write_fmt(&mut s, args).unwrap();
+    }
+}