diff --git a/llvm/include/llvm/Support/raw_ostream_proxy.h b/llvm/include/llvm/Support/raw_ostream_proxy.h new file mode 100644 --- /dev/null +++ b/llvm/include/llvm/Support/raw_ostream_proxy.h @@ -0,0 +1,158 @@ +//===- 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 { + +/// Common bits for \a raw_ostream_proxy_adaptor<>, split out to dedup in +/// template instantions. +class raw_ostream_proxy_adaptor_base { +protected: + raw_ostream_proxy_adaptor_base() = delete; + raw_ostream_proxy_adaptor_base(const raw_ostream_proxy_adaptor_base &) = + delete; + + explicit raw_ostream_proxy_adaptor_base(raw_ostream &OS) + : 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(); + } + + ~raw_ostream_proxy_adaptor_base() { + assert(!OS && "Derived objects should call resetProxiedOS()"); + } + + /// Stop proxying the stream, taking the derived object by reference as \p + /// ThisProxyOS. Updates \p ThisProxyOS to stop buffering before setting \a + /// OS to \c nullptr, ensuring that future writes crash immediately. + void resetProxiedOS(raw_ostream &ThisProxyOS) { + ThisProxyOS.SetUnbuffered(); + OS = nullptr; + } + + 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; } + +private: + raw_ostream *OS; + + /// Caches the value of OS->GetBufferSize() at construction time. + size_t PreferredBufferSize; +}; + +/// 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 raw_ostream_proxy_adaptor : public RawOstreamT, + public raw_ostream_proxy_adaptor_base { + 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); + } + + ~raw_ostream_proxy_adaptor() override { resetProxiedOS(); } + +protected: + template + explicit raw_ostream_proxy_adaptor(raw_ostream &OS, ArgsT &&...Args) + : RawOstreamT(std::forward(Args)...), + raw_ostream_proxy_adaptor_base(OS) {} + + /// 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() { + raw_ostream_proxy_adaptor_base::resetProxiedOS(*this); + } + void resetProxiedOS(raw_ostream &) = delete; +}; + +/// Adaptor for creating a stream that proxies a \a raw_pwrite_stream. +template +class raw_pwrite_stream_proxy_adaptor + : public raw_ostream_proxy_adaptor { + using RawOstreamAdaptorT = raw_ostream_proxy_adaptor; + + 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 + explicit raw_pwrite_stream_proxy_adaptor(raw_pwrite_stream &OS, + ArgsT &&...Args) + : RawOstreamAdaptorT(OS, std::forward(Args)...) {} + + raw_pwrite_stream &getProxiedOS() const { + return static_cast(RawOstreamAdaptorT::getProxiedOS()); + } +}; + +/// Non-owning proxy for a \a raw_ostream. Enables passing a stream into of 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 of 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 --- a/llvm/lib/Support/CMakeLists.txt +++ b/llvm/lib/Support/CMakeLists.txt @@ -247,6 +247,7 @@ YAMLTraits.cpp raw_os_ostream.cpp raw_ostream.cpp + raw_ostream_proxy.cpp regcomp.c regerror.c regexec.c diff --git a/llvm/lib/Support/raw_ostream_proxy.cpp b/llvm/lib/Support/raw_ostream_proxy.cpp new file mode 100644 --- /dev/null +++ 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 --- a/llvm/unittests/Support/CMakeLists.txt +++ b/llvm/unittests/Support/CMakeLists.txt @@ -101,6 +101,7 @@ 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 xxhashTest.cpp diff --git a/llvm/unittests/Support/raw_ostream_proxy_test.cpp b/llvm/unittests/Support/raw_ostream_proxy_test.cpp new file mode 100644 --- /dev/null +++ b/llvm/unittests/Support/raw_ostream_proxy_test.cpp @@ -0,0 +1,219 @@ +//===- 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. + constexpr static const 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 &Vector) + : Vector(Vector) {} + ~BufferedNoPwriteSmallVectorStream() override { flush(); } + + SmallVectorImpl &Vector; + bool IsDisplayed = false; +}; + +constexpr const 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 + 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); + raw_ostream_proxy ProxyOS(DestOS); + ProxyOS.enable_colors(true); + + WithColor(ProxyOS, HighlightColor::Error, ColorMode::Disable) << "test"; + EXPECT_EQ("", Dest); + ProxyOS.flush(); + EXPECT_EQ("test", Dest); + } + + { + SmallString<128> Dest; + BufferedNoPwriteSmallVectorStream DestOS(Dest); + raw_ostream_proxy ProxyOS(DestOS); + ProxyOS.enable_colors(true); + + WithColor(ProxyOS, HighlightColor::Error, 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, HighlightColor::Error, 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