Support: Add proxies for raw_ostream and raw_pwrite_stream (#113362)

cachemeifyoucan · web-flow · commit 90670b5ca890 · 2025-08-26T16:19:52.000-07:00
Add proxies classes for `raw_ostream` and `raw_pwrite_stream` called
`raw_ostream_proxy` and `raw_pwrite_stream_proxy`. Add adaptor classes,
`raw_ostream_proxy_adaptor&lt;&gt;` and `raw_pwrite_stream_proxy_adaptor&lt;&gt;`,
to allow subclasses to use a different parent class than `raw_ostream`
or `raw_pwrite_stream`.

The adaptors are used by a future patch to help a subclass of
`llvm::vfs::OutputFile`, an abstract subclass of `raw_pwrite_stream`, to
proxy a `raw_fd_ostream`.

Patched by dexonsmith.
diff --git a/llvm/include/llvm/Support/raw_ostream_proxy.h b/llvm/include/llvm/Support/raw_ostream_proxy.h
@@ -0,0 +1,132 @@
+//===- raw_ostream_proxy.h - Proxies for raw output streams -----*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_SUPPORT_RAW_OSTREAM_PROXY_H
+#define LLVM_SUPPORT_RAW_OSTREAM_PROXY_H
+
+#include "llvm/Support/raw_ostream.h"
+
+namespace llvm {
+
+/// Adaptor to create a stream class that proxies another \a raw_ostream.
+///
+/// Use \a raw_ostream_proxy_adaptor<> directly to implement an abstract
+/// derived class of \a raw_ostream as a proxy. Otherwise use \a
+/// raw_ostream_proxy.
+///
+/// Most operations are forwarded to the proxied stream.
+///
+/// If the proxied stream is buffered, the buffer is dropped and moved to this
+/// stream. This allows \a flush() to work correctly, flushing immediately from
+/// the proxy through to the final stream, and avoids any wasteful
+/// double-buffering.
+///
+/// \a enable_colors() changes both the proxied stream and the proxy itself.
+/// \a is_displayed() and \a has_colors() are forwarded to the proxy. \a
+/// changeColor(), resetColor(), and \a reverseColor() are not forwarded, since
+/// they need to call \a flush() and the buffer lives in the proxy.
+template <class RawOstreamT = raw_ostream>
+class raw_ostream_proxy_adaptor : public RawOstreamT {
+  void write_impl(const char *Ptr, size_t Size) override {
+    getProxiedOS().write(Ptr, Size);
+  }
+  uint64_t current_pos() const override { return getProxiedOS().tell(); }
+  size_t preferred_buffer_size() const override {
+    return getPreferredBufferSize();
+  }
+
+public:
+  void reserveExtraSpace(uint64_t ExtraSize) override {
+    getProxiedOS().reserveExtraSpace(ExtraSize);
+  }
+  bool is_displayed() const override { return getProxiedOS().is_displayed(); }
+  bool has_colors() const override { return getProxiedOS().has_colors(); }
+  void enable_colors(bool enable) override {
+    RawOstreamT::enable_colors(enable);
+    getProxiedOS().enable_colors(enable);
+  }
+  bool hasProxiedOS() const { return OS; }
+  raw_ostream &getProxiedOS() const {
+    assert(OS && "raw_ostream_proxy_adaptor use after reset");
+    return *OS;
+  }
+  size_t getPreferredBufferSize() const { return PreferredBufferSize; }
+
+  ~raw_ostream_proxy_adaptor() override { resetProxiedOS(); }
+
+protected:
+  template <class... ArgsT>
+  explicit raw_ostream_proxy_adaptor(raw_ostream &OS, ArgsT &&...Args)
+      : RawOstreamT(std::forward<ArgsT>(Args)...), OS(&OS),
+        PreferredBufferSize(OS.GetBufferSize()) {
+    // Drop OS's buffer to make this->flush() forward. This proxy will add a
+    // buffer in its place.
+    OS.SetUnbuffered();
+  }
+
+  /// Stop proxying the stream. Flush and set up a crash for future writes.
+  ///
+  /// For example, this can simplify logic when a subclass might have a longer
+  /// lifetime than the stream it proxies.
+  void resetProxiedOS() {
+    this->SetUnbuffered();
+    OS = nullptr;
+  }
+
+private:
+  raw_ostream *OS;
+
+  /// Caches the value of OS->GetBufferSize() at construction time.
+  size_t PreferredBufferSize;
+};
+
+/// Adaptor for creating a stream that proxies a \a raw_pwrite_stream.
+template <class RawPwriteStreamT = raw_pwrite_stream>
+class raw_pwrite_stream_proxy_adaptor
+    : public raw_ostream_proxy_adaptor<RawPwriteStreamT> {
+  using RawOstreamAdaptorT = raw_ostream_proxy_adaptor<RawPwriteStreamT>;
+
+  void pwrite_impl(const char *Ptr, size_t Size, uint64_t Offset) override {
+    this->flush();
+    getProxiedOS().pwrite(Ptr, Size, Offset);
+  }
+
+protected:
+  raw_pwrite_stream_proxy_adaptor() = default;
+  template <class... ArgsT>
+  explicit raw_pwrite_stream_proxy_adaptor(raw_pwrite_stream &OS,
+                                           ArgsT &&...Args)
+      : RawOstreamAdaptorT(OS, std::forward<ArgsT>(Args)...) {}
+
+  raw_pwrite_stream &getProxiedOS() const {
+    return static_cast<raw_pwrite_stream &>(RawOstreamAdaptorT::getProxiedOS());
+  }
+};
+
+/// Non-owning proxy for a \a raw_ostream. Enables passing a stream into an
+/// API that takes ownership.
+class raw_ostream_proxy : public raw_ostream_proxy_adaptor<> {
+  void anchor() override;
+
+public:
+  raw_ostream_proxy(raw_ostream &OS) : raw_ostream_proxy_adaptor<>(OS) {}
+};
+
+/// Non-owning proxy for a \a raw_pwrite_stream. Enables passing a stream
+/// into an API that takes ownership.
+class raw_pwrite_stream_proxy : public raw_pwrite_stream_proxy_adaptor<> {
+  void anchor() override;
+
+public:
+  raw_pwrite_stream_proxy(raw_pwrite_stream &OS)
+      : raw_pwrite_stream_proxy_adaptor<>(OS) {}
+};
+
+} // end namespace llvm
+
+#endif // LLVM_SUPPORT_RAW_OSTREAM_PROXY_H
diff --git a/llvm/lib/Support/CMakeLists.txt b/llvm/lib/Support/CMakeLists.txt
@@ -277,6 +277,7 @@ add_llvm_component_library(LLVMSupport
   YAMLTraits.cpp
   raw_os_ostream.cpp
   raw_ostream.cpp
+  raw_ostream_proxy.cpp
   raw_socket_stream.cpp
   regcomp.c
   regerror.c
diff --git a/llvm/lib/Support/raw_ostream_proxy.cpp b/llvm/lib/Support/raw_ostream_proxy.cpp
@@ -0,0 +1,15 @@
+//===- raw_ostream_proxy.cpp - Implement the raw_ostream proxies ----------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Support/raw_ostream_proxy.h"
+
+using namespace llvm;
+
+void raw_ostream_proxy::anchor() {}
+
+void raw_pwrite_stream_proxy::anchor() {}
diff --git a/llvm/unittests/Support/CMakeLists.txt b/llvm/unittests/Support/CMakeLists.txt
@@ -111,6 +111,7 @@ add_llvm_unittest(SupportTests
   formatted_raw_ostream_test.cpp
   raw_fd_stream_test.cpp
   raw_ostream_test.cpp
+  raw_ostream_proxy_test.cpp
   raw_pwrite_stream_test.cpp
   raw_sha1_ostream_test.cpp
   raw_socket_stream_test.cpp
diff --git a/llvm/unittests/Support/raw_ostream_proxy_test.cpp b/llvm/unittests/Support/raw_ostream_proxy_test.cpp
@@ -0,0 +1,225 @@
+//===- raw_ostream_proxy_test.cpp - Tests for raw ostream proxies ---------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/ADT/SmallString.h"
+#include "llvm/Support/WithColor.h"
+#include "llvm/Support/raw_ostream.h"
+#include "llvm/Support/raw_ostream_proxy.h"
+#include "gtest/gtest.h"
+
+using namespace llvm;
+
+namespace {
+
+/// Naive version of raw_svector_ostream that is buffered (by default) and
+/// doesn't support pwrite.
+class BufferedNoPwriteSmallVectorStream : public raw_ostream {
+public:
+  // Choose a strange buffer size to ensure it doesn't collide with the default
+  // on \a raw_ostream.
+  static constexpr size_t PreferredBufferSize = 63;
+
+  size_t preferred_buffer_size() const override { return PreferredBufferSize; }
+  uint64_t current_pos() const override { return Vector.size(); }
+  void write_impl(const char *Ptr, size_t Size) override {
+    Vector.append(Ptr, Ptr + Size);
+  }
+
+  bool is_displayed() const override { return IsDisplayed; }
+
+  explicit BufferedNoPwriteSmallVectorStream(SmallVectorImpl<char> &Vector)
+      : Vector(Vector) {}
+  ~BufferedNoPwriteSmallVectorStream() override { flush(); }
+
+  SmallVectorImpl<char> &Vector;
+  bool IsDisplayed = false;
+};
+
+constexpr size_t BufferedNoPwriteSmallVectorStream::PreferredBufferSize;
+
+TEST(raw_ostream_proxyTest, write) {
+  // Besides confirming that "write" works, this test confirms that the proxy
+  // takes on the buffer from the stream it's proxying, such that writes to the
+  // proxy are flushed to the underlying stream as if no proxy were present.
+  SmallString<128> Dest;
+  {
+    // Confirm that BufferedNoPwriteSmallVectorStream is buffered by default,
+    // and that setting up a proxy effectively transfers a buffer of the same
+    // size to the proxy.
+    BufferedNoPwriteSmallVectorStream DestOS(Dest);
+    EXPECT_EQ(BufferedNoPwriteSmallVectorStream::PreferredBufferSize,
+              DestOS.GetBufferSize());
+    raw_ostream_proxy ProxyOS(DestOS);
+    EXPECT_EQ(0u, DestOS.GetBufferSize());
+    EXPECT_EQ(BufferedNoPwriteSmallVectorStream::PreferredBufferSize,
+              ProxyOS.GetBufferSize());
+
+    // Flushing should send through to Dest.
+    ProxyOS << "abcd";
+    EXPECT_EQ("", Dest);
+    ProxyOS.flush();
+    EXPECT_EQ("abcd", Dest);
+
+    // Buffer should still work.
+    ProxyOS << "e";
+    EXPECT_EQ("abcd", Dest);
+  }
+
+  // Destructing ProxyOS should flush (and not crash).
+  EXPECT_EQ("abcde", Dest);
+
+  {
+    // Set up another stream, this time unbuffered.
+    BufferedNoPwriteSmallVectorStream DestOS(Dest);
+    DestOS.SetUnbuffered();
+    EXPECT_EQ(0u, DestOS.GetBufferSize());
+    raw_ostream_proxy ProxyOS(DestOS);
+    EXPECT_EQ(0u, DestOS.GetBufferSize());
+    EXPECT_EQ(0u, ProxyOS.GetBufferSize());
+
+    // Flushing should not be required.
+    ProxyOS << "f";
+    EXPECT_EQ("abcdef", Dest);
+  }
+  EXPECT_EQ("abcdef", Dest);
+}
+
+TEST(raw_ostream_proxyTest, pwrite) {
+  // This test confirms that the proxy takes on the buffer from the stream it's
+  // proxying, such that writes to the proxy are flushed to the underlying
+  // stream as if no proxy were present.
+  SmallString<128> Dest;
+  raw_svector_ostream DestOS(Dest);
+  raw_pwrite_stream_proxy ProxyOS(DestOS);
+  EXPECT_EQ(0u, ProxyOS.GetBufferSize());
+
+  // Get some initial data.
+  ProxyOS << "abcd";
+  EXPECT_EQ("abcd", Dest);
+
+  // Confirm that pwrite works.
+  ProxyOS.pwrite("BC", 2, 1);
+  EXPECT_EQ("aBCd", Dest);
+}
+
+TEST(raw_ostream_proxyTest, pwriteWithBuffer) {
+  // This test confirms that when a buffer is configured, pwrite still works.
+  SmallString<128> Dest;
+  raw_svector_ostream DestOS(Dest);
+  DestOS.SetBufferSize(256);
+  EXPECT_EQ(256u, DestOS.GetBufferSize());
+
+  // Confirm that the proxy steals the buffer.
+  raw_pwrite_stream_proxy ProxyOS(DestOS);
+  EXPECT_EQ(0u, DestOS.GetBufferSize());
+  EXPECT_EQ(256u, ProxyOS.GetBufferSize());
+
+  // Check that the buffer is working.
+  ProxyOS << "abcd";
+  EXPECT_EQ("", Dest);
+
+  // Confirm that pwrite flushes.
+  ProxyOS.pwrite("BC", 2, 1);
+  EXPECT_EQ("aBCd", Dest);
+}
+
+class ProxyWithReset : public raw_ostream_proxy_adaptor<> {
+public:
+  ProxyWithReset(raw_ostream &OS) : raw_ostream_proxy_adaptor<>(OS) {}
+
+  // Allow this to be called outside the class.
+  using raw_ostream_proxy_adaptor<>::hasProxiedOS;
+  using raw_ostream_proxy_adaptor<>::getProxiedOS;
+  using raw_ostream_proxy_adaptor<>::resetProxiedOS;
+};
+
+TEST(raw_ostream_proxyTest, resetProxiedOS) {
+  // Confirm that base classes can drop the proxied OS before destruction and
+  // get consistent crashes.
+  SmallString<128> Dest;
+  BufferedNoPwriteSmallVectorStream DestOS(Dest);
+  ProxyWithReset ProxyOS(DestOS);
+  EXPECT_TRUE(ProxyOS.hasProxiedOS());
+  EXPECT_EQ(&DestOS, &ProxyOS.getProxiedOS());
+
+  // Write some data.
+  ProxyOS << "abcd";
+  EXPECT_EQ("", Dest);
+
+  // Reset the underlying stream.
+  ProxyOS.resetProxiedOS();
+  EXPECT_EQ("abcd", Dest);
+  EXPECT_EQ(0u, ProxyOS.GetBufferSize());
+  EXPECT_FALSE(ProxyOS.hasProxiedOS());
+
+#if GTEST_HAS_DEATH_TEST && !defined(NDEBUG)
+  EXPECT_DEATH(ProxyOS << "e", "use after reset");
+  EXPECT_DEATH(ProxyOS.getProxiedOS(), "use after reset");
+#endif
+}
+
+TEST(raw_ostream_proxyTest, ColorMode) {
+  {
+    SmallString<128> Dest;
+    BufferedNoPwriteSmallVectorStream DestOS(Dest);
+    DestOS.IsDisplayed = true;
+    raw_ostream_proxy ProxyOS(DestOS);
+    ProxyOS.enable_colors(true);
+
+    WithColor(ProxyOS, raw_ostream::Colors::RED, /*Bold=*/true, /*BG=*/false,
+              ColorMode::Disable)
+        << "test";
+    EXPECT_EQ("", Dest);
+    ProxyOS.flush();
+    EXPECT_EQ("test", Dest);
+  }
+
+  {
+    SmallString<128> Dest;
+    BufferedNoPwriteSmallVectorStream DestOS(Dest);
+    raw_ostream_proxy ProxyOS(DestOS);
+
+    WithColor(ProxyOS, raw_ostream::Colors::RED, /*Bold=*/true, /*BG=*/false,
+              ColorMode::Auto)
+        << "test";
+    EXPECT_EQ("", Dest);
+    ProxyOS.flush();
+    EXPECT_EQ("test", Dest);
+  }
+
+#ifdef LLVM_ON_UNIX
+  {
+    SmallString<128> Dest;
+    BufferedNoPwriteSmallVectorStream DestOS(Dest);
+    raw_ostream_proxy ProxyOS(DestOS);
+    ProxyOS.enable_colors(true);
+
+    WithColor(ProxyOS, raw_ostream::Colors::RED, /*Bold=*/true, /*BG=*/false,
+              ColorMode::Enable)
+        << "test";
+    EXPECT_EQ("", Dest);
+    ProxyOS.flush();
+    EXPECT_EQ("\x1B[0;1;31mtest\x1B[0m", Dest);
+  }
+
+  {
+    SmallString<128> Dest;
+    BufferedNoPwriteSmallVectorStream DestOS(Dest);
+    DestOS.IsDisplayed = true;
+    raw_ostream_proxy ProxyOS(DestOS);
+    ProxyOS.enable_colors(true);
+
+    WithColor(ProxyOS, HighlightColor::Error, ColorMode::Auto) << "test";
+    EXPECT_EQ("", Dest);
+    ProxyOS.flush();
+    EXPECT_EQ("\x1B[0;1;31mtest\x1B[0m", Dest);
+  }
+#endif
+}
+
+} // end namespace
