Merge set_panic and set_print into set_output_capture.

There were no use cases for setting them separately.
Merging them simplifies some things.

5 files changed, 38 insertions, 98 deletions

diff --git a/library/std/src/io/mod.rs b/library/std/src/io/mod.rs
index e6b9314fd88..d1e5942cba8 100644
--- a/library/std/src/io/mod.rs
+++ b/library/std/src/io/mod.rs
@@ -269,20 +269,18 @@ pub use self::buffered::{BufReader, BufWriter, LineWriter};
 pub use self::cursor::Cursor;
 #[stable(feature = "rust1", since = "1.0.0")]
 pub use self::error::{Error, ErrorKind, Result};
+#[unstable(feature = "internal_output_capture", issue = "none")]
+#[doc(no_inline, hidden)]
+pub use self::stdio::set_output_capture;
 #[stable(feature = "rust1", since = "1.0.0")]
 pub use self::stdio::{stderr, stdin, stdout, Stderr, Stdin, Stdout};
 #[stable(feature = "rust1", since = "1.0.0")]
 pub use self::stdio::{StderrLock, StdinLock, StdoutLock};
 #[unstable(feature = "print_internals", issue = "none")]
 pub use self::stdio::{_eprint, _print};
-#[unstable(feature = "libstd_io_internals", issue = "42788")]
-#[doc(no_inline, hidden)]
-pub use self::stdio::{set_panic, set_print};
 #[stable(feature = "rust1", since = "1.0.0")]
 pub use self::util::{copy, empty, repeat, sink, Empty, Repeat, Sink};
 
-pub(crate) use self::stdio::clone_io;
-
 mod buffered;
 mod cursor;
 mod error;
diff --git a/library/std/src/io/stdio.rs b/library/std/src/io/stdio.rs
index 2e71ad29584..f1dfb4c7730 100644
--- a/library/std/src/io/stdio.rs
+++ b/library/std/src/io/stdio.rs
@@ -14,37 +14,29 @@ use crate::sync::{Arc, Mutex, MutexGuard};
 use crate::sys::stdio;
 use crate::sys_common;
 use crate::sys_common::remutex::{ReentrantMutex, ReentrantMutexGuard};
-use crate::thread::LocalKey;
 
 type LocalStream = Arc<Mutex<Vec<u8>>>;
 
 thread_local! {
-    /// Used by the test crate to capture the output of the print! and println! macros.
-    static LOCAL_STDOUT: Cell<Option<LocalStream>> = {
+    /// Used by the test crate to capture the output of the print macros and panics.
+    static OUTPUT_CAPTURE: Cell<Option<LocalStream>> = {
         Cell::new(None)
     }
 }
 
-thread_local! {
-    /// Used by the test crate to capture the output of the eprint! and eprintln! macros, and panics.
-    static LOCAL_STDERR: Cell<Option<LocalStream>> = {
-        Cell::new(None)
-    }
-}
-
-/// Flag to indicate LOCAL_STDOUT and/or LOCAL_STDERR is used.
+/// Flag to indicate OUTPUT_CAPTURE is used.
 ///
-/// If both are None and were never set on any thread, this flag is set to
-/// false, and both LOCAL_STDOUT and LOCAL_STDOUT can be safely ignored on all
-/// threads, saving some time and memory registering an unused thread local.
+/// If it is None and was never set on any thread, this flag is set to false,
+/// and OUTPUT_CAPTURE can be safely ignored on all threads, saving some time
+/// and memory registering an unused thread local.
 ///
-/// Note about memory ordering: This contains information about whether two
-/// thread local variables might be in use. Although this is a global flag, the
+/// Note about memory ordering: This contains information about whether a
+/// thread local variable might be in use. Although this is a global flag, the
 /// memory ordering between threads does not matter: we only want this flag to
-/// have a consistent order between set_print/set_panic and print_to *within
+/// have a consistent order between set_output_capture and print_to *within
 /// the same thread*. Within the same thread, things always have a perfectly
 /// consistent order. So Ordering::Relaxed is fine.
-static LOCAL_STREAMS: AtomicBool = AtomicBool::new(false);
+static OUTPUT_CAPTURE_USED: AtomicBool = AtomicBool::new(false);
 
 /// A handle to a raw instance of the standard input stream of this process.
 ///
@@ -890,70 +882,24 @@ impl fmt::Debug for StderrLock<'_> {
     }
 }
 
-/// Resets the thread-local stderr handle to the specified writer
-///
-/// This will replace the current thread's stderr handle, returning the old
-/// handle. All future calls to `panic!` and friends will emit their output to
-/// this specified handle.
-///
-/// Note that this does not need to be called for all new threads; the default
-/// output handle is to the process's stderr stream.
-#[unstable(
-    feature = "set_stdio",
-    reason = "this function may disappear completely or be replaced \
-                     with a more general mechanism",
-    issue = "none"
-)]
-#[doc(hidden)]
-pub fn set_panic(sink: Option<LocalStream>) -> Option<LocalStream> {
-    if sink.is_none() && !LOCAL_STREAMS.load(Ordering::Relaxed) {
-        // LOCAL_STDERR is definitely None since LOCAL_STREAMS is false.
-        return None;
-    }
-    LOCAL_STREAMS.store(true, Ordering::Relaxed);
-    LOCAL_STDERR.with(move |slot| slot.replace(sink))
-}
-
-/// Resets the thread-local stdout handle to the specified writer
-///
-/// This will replace the current thread's stdout handle, returning the old
-/// handle. All future calls to `print!` and friends will emit their output to
-/// this specified handle.
-///
-/// Note that this does not need to be called for all new threads; the default
-/// output handle is to the process's stdout stream.
+/// Sets the thread-local output capture buffer and returns the old one.
 #[unstable(
-    feature = "set_stdio",
-    reason = "this function may disappear completely or be replaced \
-                     with a more general mechanism",
+    feature = "internal_output_capture",
+    reason = "this function is meant for use in the test crate \
+        and may disappear in the future",
     issue = "none"
 )]
 #[doc(hidden)]
-pub fn set_print(sink: Option<LocalStream>) -> Option<LocalStream> {
-    if sink.is_none() && !LOCAL_STREAMS.load(Ordering::Relaxed) {
-        // LOCAL_STDOUT is definitely None since LOCAL_STREAMS is false.
+pub fn set_output_capture(sink: Option<LocalStream>) -> Option<LocalStream> {
+    if sink.is_none() && !OUTPUT_CAPTURE_USED.load(Ordering::Relaxed) {
+        // OUTPUT_CAPTURE is definitely None since OUTPUT_CAPTURE_USED is false.
         return None;
     }
-    LOCAL_STREAMS.store(true, Ordering::Relaxed);
-    LOCAL_STDOUT.with(move |slot| slot.replace(sink))
-}
-
-pub(crate) fn clone_io() -> (Option<LocalStream>, Option<LocalStream>) {
-    // Don't waste time when LOCAL_{STDOUT,STDERR} are definitely None.
-    if !LOCAL_STREAMS.load(Ordering::Relaxed) {
-        return (None, None);
-    }
-
-    let clone = |cell: &Cell<Option<LocalStream>>| {
-        let s = cell.take();
-        cell.set(s.clone());
-        s
-    };
-
-    (LOCAL_STDOUT.with(clone), LOCAL_STDERR.with(clone))
+    OUTPUT_CAPTURE_USED.store(true, Ordering::Relaxed);
+    OUTPUT_CAPTURE.with(move |slot| slot.replace(sink))
 }
 
-/// Write `args` to output stream `local_s` if possible, `global_s`
+/// Write `args` to the capture buffer if enabled and possible, or `global_s`
 /// otherwise. `label` identifies the stream in a panic message.
 ///
 /// This function is used to print error messages, so it takes extra
@@ -963,16 +909,12 @@ pub(crate) fn clone_io() -> (Option<LocalStream>, Option<LocalStream>) {
 /// thread, it will just fall back to the global stream.
 ///
 /// However, if the actual I/O causes an error, this function does panic.
-fn print_to<T>(
-    args: fmt::Arguments<'_>,
-    local_s: &'static LocalKey<Cell<Option<LocalStream>>>,
-    global_s: fn() -> T,
-    label: &str,
-) where
+fn print_to<T>(args: fmt::Arguments<'_>, global_s: fn() -> T, label: &str)
+where
     T: Write,
 {
-    if LOCAL_STREAMS.load(Ordering::Relaxed)
-        && local_s.try_with(|s| {
+    if OUTPUT_CAPTURE_USED.load(Ordering::Relaxed)
+        && OUTPUT_CAPTURE.try_with(|s| {
             // Note that we completely remove a local sink to write to in case
             // our printing recursively panics/prints, so the recursive
             // panic/print goes to the global sink instead of our local sink.
@@ -982,7 +924,7 @@ fn print_to<T>(
             })
         }) == Ok(Some(()))
     {
-        // Succesfully wrote to local stream.
+        // Succesfully wrote to capture buffer.
         return;
     }
 
@@ -999,7 +941,7 @@ fn print_to<T>(
 #[doc(hidden)]
 #[cfg(not(test))]
 pub fn _print(args: fmt::Arguments<'_>) {
-    print_to(args, &LOCAL_STDOUT, stdout, "stdout");
+    print_to(args, stdout, "stdout");
 }
 
 #[unstable(
@@ -1010,7 +952,7 @@ pub fn _print(args: fmt::Arguments<'_>) {
 #[doc(hidden)]
 #[cfg(not(test))]
 pub fn _eprint(args: fmt::Arguments<'_>) {
-    print_to(args, &LOCAL_STDERR, stderr, "stderr");
+    print_to(args, stderr, "stderr");
 }
 
 #[cfg(test)]
diff --git a/library/std/src/lib.rs b/library/std/src/lib.rs
index a0c2a27438a..334757aafd7 100644
--- a/library/std/src/lib.rs
+++ b/library/std/src/lib.rs
@@ -207,7 +207,7 @@
 // std may use features in a platform-specific way
 #![allow(unused_features)]
 #![cfg_attr(not(bootstrap), feature(rustc_allow_const_fn_unstable))]
-#![cfg_attr(test, feature(print_internals, set_stdio, update_panic_count))]
+#![cfg_attr(test, feature(internal_output_capture, print_internals, update_panic_count))]
 #![cfg_attr(
     all(target_vendor = "fortanix", target_env = "sgx"),
     feature(slice_index_methods, coerce_unsized, sgx_platform)
diff --git a/library/std/src/panicking.rs b/library/std/src/panicking.rs
index a5d4d72b00d..8ba3feccb6b 100644
--- a/library/std/src/panicking.rs
+++ b/library/std/src/panicking.rs
@@ -24,11 +24,11 @@ use crate::sys_common::{thread_info, util};
 use crate::thread;
 
 #[cfg(not(test))]
-use crate::io::set_panic;
+use crate::io::set_output_capture;
 // make sure to use the stderr output configured
 // by libtest in the real copy of std
 #[cfg(test)]
-use realstd::io::set_panic;
+use realstd::io::set_output_capture;
 
 // Binary interface to the panic runtime that the standard library depends on.
 //
@@ -218,9 +218,9 @@ fn default_hook(info: &PanicInfo<'_>) {
         }
     };
 
-    if let Some(local) = set_panic(None) {
+    if let Some(local) = set_output_capture(None) {
         write(&mut *local.lock().unwrap_or_else(|e| e.into_inner()));
-        set_panic(Some(local));
+        set_output_capture(Some(local));
     } else if let Some(mut out) = panic_output() {
         write(&mut out);
     }
diff --git a/library/std/src/thread/mod.rs b/library/std/src/thread/mod.rs
index fefaa77a2a1..5d65f960fcd 100644
--- a/library/std/src/thread/mod.rs
+++ b/library/std/src/thread/mod.rs
@@ -456,15 +456,15 @@ impl Builder {
         let my_packet: Arc<UnsafeCell<Option<Result<T>>>> = Arc::new(UnsafeCell::new(None));
         let their_packet = my_packet.clone();
 
-        let (stdout, stderr) = crate::io::clone_io();
+        let output_capture = crate::io::set_output_capture(None);
+        crate::io::set_output_capture(output_capture.clone());
 
         let main = move || {
             if let Some(name) = their_thread.cname() {
                 imp::Thread::set_name(name);
             }
 
-            crate::io::set_print(stdout);
-            crate::io::set_panic(stderr);
+            crate::io::set_output_capture(output_capture);
 
             // SAFETY: the stack guard passed is the one for the current thread.
             // This means the current thread's stack and the new thread's stack