Support: Skip buffering buffer_unique_ostream's owned stream

[dexonsmith]

Summary

Change buffer_unique_ostream's constructor to call
raw_ostream::SetUnbuffered() on its owned stream. Otherwise,
buffer_unique_ostream's destructor could cause the owned stream to
temporarily allocate a buffer only to be immediately flushed.

Also add some tests for buffer_ostream and buffer_unique_ostream. Use
the same naming scheme as other raw_ostream-related tests (e.g.,
raw_ostreamTest for the fixture, raw_ostream_test.cpp for the
filename).

(I considered changing buffer_ostream in the same way (calling
SetUnbuffered on the referenced stream), but that seemed like overreach
since the client may have more things to write.)

(I considered merging buffer_ostream and buffer_unique_ostream into a
single class (with a raw_ostream& and a std::unique_ptr that is only
sometimes used), but that makes the class bigger and the small amount of
code deduplication seems uncompelling.)

[dexonsmith]

(Forgot to run clang-format; done now)

[steven_wu]

LGTM.

llvm/include/llvm/Support/raw_ostream.h

 class buffer_unique_ostream : public raw_svector_ostream {
   std::unique_ptr<raw_ostream> OS;
   SmallVector<char, 0> Buffer;
 
   virtual void anchor() override;
 
 public:
   buffer_unique_ostream(std::unique_ptr<raw_ostream> OS)
-      : raw_svector_ostream(Buffer), OS(std::move(OS)) {}
+      : raw_svector_ostream(Buffer), OS(std::move(OS)) {
+    // Turn off buffering on OS, which we now own, to avoid allocating a buffer
+    // when the destructor writes only to be immediately flushed again.
+    this->OS->SetUnbuffered();
+  }
   ~buffer_unique_ostream() override { *OS << str(); }
 };
 
 class Error;
 
 /// This helper creates an output stream and then passes it to \p Write.
 /// The stream created is based on the specified \p OutputFileName:
 /// llvm::outs for "-", raw_null_ostream for "/dev/null", and raw_fd_ostream

llvm/unittests/Support/CMakeLists.txt

add_llvm_unittest(SupportTests
   TrailingObjectsTest.cpp
   TrigramIndexTest.cpp
   UnicodeTest.cpp
   VersionTupleTest.cpp
   VirtualFileSystemTest.cpp
   WithColorTest.cpp
   YAMLIOTest.cpp
   YAMLParserTest.cpp
+  buffer_ostream_test.cpp
   formatted_raw_ostream_test.cpp
   raw_fd_stream_test.cpp
   raw_ostream_test.cpp
   raw_pwrite_stream_test.cpp
   raw_sha1_ostream_test.cpp
   xxhashTest.cpp
   )
 

llvm/unittests/Support/buffer_ostream_test.cpp

+//===- buffer_ostream_test.cpp - buffer_ostream tests ---------------------===//
+//
+// 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/raw_ostream.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 NaiveSmallVectorStream : public raw_ostream {
+public:
+  uint64_t current_pos() const override { return Vector.size(); }
+  void write_impl(const char *Ptr, size_t Size) override {
+    Vector.append(Ptr, Ptr + Size);
+  }
+
+  explicit NaiveSmallVectorStream(SmallVectorImpl<char> &Vector)
+      : Vector(Vector) {}
+  ~NaiveSmallVectorStream() override { flush(); }
+
+  SmallVectorImpl<char> &Vector;
+};
+
+TEST(buffer_ostreamTest, Reference) {
+  SmallString<128> Dest;
+  {
+    NaiveSmallVectorStream DestOS(Dest);
+    buffer_ostream BufferOS(DestOS);
+
+    // Writing and flushing should have no effect on Dest.
+    BufferOS << "abcd";
+    static_cast<raw_ostream &>(BufferOS).flush();
+    EXPECT_EQ("", Dest);
+    DestOS.flush();
+    EXPECT_EQ("", Dest);
+  }
+
+  // Write should land when constructor is called.
+  EXPECT_EQ("abcd", Dest);
+}
+
+TEST(buffer_ostreamTest, Owned) {
+  SmallString<128> Dest;
+  {
+    auto DestOS = std::make_unique<NaiveSmallVectorStream>(Dest);
+
+    // Confirm that NaiveSmallVectorStream is buffered by default.
+    EXPECT_NE(0u, DestOS->GetBufferSize());
+
+    // Confirm that passing ownership to buffer_unique_ostream sets it to
+    // unbuffered. Also steal a reference to DestOS.
+    NaiveSmallVectorStream &DestOSRef = *DestOS;
+    buffer_unique_ostream BufferOS(std::move(DestOS));
+    EXPECT_EQ(0u, DestOSRef.GetBufferSize());
+
+    // Writing and flushing should have no effect on Dest.
+    BufferOS << "abcd";
+    static_cast<raw_ostream &>(BufferOS).flush();
+    EXPECT_EQ("", Dest);
+    DestOSRef.flush();
+    EXPECT_EQ("", Dest);
+  }
+
+  // Write should land when constructor is called.
+  EXPECT_EQ("abcd", Dest);
+}
+
+} // end namespace