[llvm] Add thread-local overrides for `llvm::errs()` and `llvm::outs()`. (PR #90374)

[Chandler Carruth via llvm-commits]

https://github.com/chandlerc created https://github.com/llvm/llvm-project/pull/90374

This allows library users of LLVM to intercept these streams and control where they write. While some components of LLVM already accept stream parameters to control this, not all do and so it seems useful to add this as a fallback. The motivating case is the verbose output in Clang's Driver library.

>From 848c557c2cfc0b9b5b33df57f0f8378ab95c8942 Mon Sep 17 00:00:00 2001
From: Chandler Carruth <chandlerc at gmail.com>
Date: Sun, 28 Apr 2024 04:54:07 +0000
Subject: [PATCH] Add thread-local overrides for `llvm::errs()` and
 `llvm::outs()`.

This allows library users of LLVM to intercept these streams and control
where they write. While some components of LLVM already accept stream
parameters to control this, not all do and so it seems useful to add
this as a fallback. The motivating case is the verbose output in Clang's
Driver library.
---
 llvm/include/llvm/Support/raw_ostream.h       | 24 ++++++++++++++++
 llvm/lib/Support/raw_ostream.cpp              | 23 +++++++++++++++
 llvm/unittests/Support/raw_fd_stream_test.cpp | 28 +++++++++++++++++++
 3 files changed, 75 insertions(+)

diff --git a/llvm/include/llvm/Support/raw_ostream.h b/llvm/include/llvm/Support/raw_ostream.h
index 696290a5d99cf5..f32400d9a1e0c5 100644
--- a/llvm/include/llvm/Support/raw_ostream.h
+++ b/llvm/include/llvm/Support/raw_ostream.h
@@ -618,6 +618,30 @@ raw_fd_ostream &errs();
 /// This returns a reference to a raw_ostream which simply discards output.
 raw_ostream &nulls();
 
+/// Scoped thread-local override of `outs` and `errs`.
+///
+/// This can be used by libraries using LLVM that want to control where output
+/// is sent. While some systems explicitly accept a stream in their API, others
+/// directly use these and so we provide a scoped override as a fallback.
+///
+/// Note that these overrides rely on thread-local override, and so much be
+/// carefully instantiated on each thread where the override should be applied.
+class ScopedOutsAndErrsOverride {
+public:
+  // Install overrides for `outs` and `errs`. If the new stream is a null
+  // pointer it will cause that stream to use the system one as if no override
+  // is in place.
+  ScopedOutsAndErrsOverride(raw_fd_ostream *NewOuts, raw_fd_ostream *NewErrs);
+
+  // Restore any previous overrides when destroyed so that these overrides can
+  // be nested.
+  ~ScopedOutsAndErrsOverride();
+
+private:
+  raw_fd_ostream *PrevOuts;
+  raw_fd_ostream *PrevErrs;
+};
+
 //===----------------------------------------------------------------------===//
 // File Streams
 //===----------------------------------------------------------------------===//
diff --git a/llvm/lib/Support/raw_ostream.cpp b/llvm/lib/Support/raw_ostream.cpp
index 8cb7b5ac68ea7e..1f8ac00cbe9881 100644
--- a/llvm/lib/Support/raw_ostream.cpp
+++ b/llvm/lib/Support/raw_ostream.cpp
@@ -893,7 +893,14 @@ void raw_fd_ostream::anchor() {}
 //  outs(), errs(), nulls()
 //===----------------------------------------------------------------------===//
 
+static thread_local raw_fd_ostream *OutsOverride = nullptr;
+static thread_local raw_fd_ostream *ErrsOverride = nullptr;
+
 raw_fd_ostream &llvm::outs() {
+  if (auto *TLSOuts = OutsOverride; TLSOuts != nullptr) {
+    return *TLSOuts;
+  }
+  
   // Set buffer settings to model stdout behavior.
   std::error_code EC;
 #ifdef __MVS__
@@ -906,6 +913,10 @@ raw_fd_ostream &llvm::outs() {
 }
 
 raw_fd_ostream &llvm::errs() {
+  if (auto *TLSErrs = OutsOverride; TLSErrs != nullptr) {
+    return *TLSErrs;
+  }
+
   // Set standard error to be unbuffered and tied to outs() by default.
 #ifdef __MVS__
   std::error_code EC = enableAutoConversion(STDERR_FILENO);
@@ -921,6 +932,18 @@ raw_ostream &llvm::nulls() {
   return S;
 }
 
+ScopedOutsAndErrsOverride::ScopedOutsAndErrsOverride(raw_fd_ostream *NewOuts,
+                                                     raw_fd_ostream *NewErrs)
+    : PrevOuts(OutsOverride), PrevErrs(ErrsOverride) {
+  OutsOverride = NewOuts;
+  ErrsOverride = NewErrs;
+}
+
+ScopedOutsAndErrsOverride::~ScopedOutsAndErrsOverride() {
+  OutsOverride = PrevOuts;
+  ErrsOverride = PrevErrs;
+}
+
 //===----------------------------------------------------------------------===//
 // File Streams
 //===----------------------------------------------------------------------===//
diff --git a/llvm/unittests/Support/raw_fd_stream_test.cpp b/llvm/unittests/Support/raw_fd_stream_test.cpp
index 00d834da32101c..b7784f3b625ad4 100644
--- a/llvm/unittests/Support/raw_fd_stream_test.cpp
+++ b/llvm/unittests/Support/raw_fd_stream_test.cpp
@@ -64,4 +64,32 @@ TEST(raw_fd_streamTest, DynCast) {
   }
 }
 
+TEST(raw_fd_streamTest, OverrideOutsAndErrs) {
+  SmallString<64> Path;
+  int FD;
+  ASSERT_FALSE(sys::fs::createTemporaryFile("foo", "bar", FD, Path));
+  FileRemover Cleanup(Path);
+  std::error_code EC;
+  raw_fd_stream OS(Path, EC);
+  ASSERT_TRUE(!EC);
+
+  ScopedOutsAndErrsOverride Overrides(&OS, &OS);
+
+  // First test `outs`.
+  llvm::outs() << "outs";
+  llvm::outs().flush();
+  char Buffer[4];
+  OS.seek(0);
+  OS.read(Buffer, sizeof(Buffer));
+  EXPECT_EQ("outs", StringRef(Buffer, sizeof(Buffer)));
+
+  // Now test `errs`.
+  OS.seek(0);
+  llvm::errs() << "errs";
+  llvm::errs().flush();
+  OS.seek(0);
+  OS.read(Buffer, sizeof(Buffer));
+  EXPECT_EQ("errs", StringRef(Buffer, sizeof(Buffer)));
+}
+
 } // namespace