[llvm] e18ea6f - Support: Skip buffering buffer_unique_ostream's owned stream

[Duncan P. N. Exon Smith via llvm-commits]

Author: Duncan P. N. Exon Smith
Date: 2021-10-22T16:25:31-07:00
New Revision: e18ea6f2946a10042258b976e60166a926ee939a

URL: https://github.com/llvm/llvm-project/commit/e18ea6f2946a10042258b976e60166a926ee939a
DIFF: https://github.com/llvm/llvm-project/commit/e18ea6f2946a10042258b976e60166a926ee939a.diff

LOG: Support: Skip buffering buffer_unique_ostream's owned stream

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.)

Differential Revision: https://reviews.llvm.org/D110369

Added: 
    llvm/unittests/Support/buffer_ostream_test.cpp

Modified: 
    llvm/include/llvm/Support/raw_ostream.h
    llvm/unittests/Support/CMakeLists.txt

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/Support/raw_ostream.h b/llvm/include/llvm/Support/raw_ostream.h
index 00d9c7db33b60..98c26ef0b1e52 100644
--- a/llvm/include/llvm/Support/raw_ostream.h
+++ b/llvm/include/llvm/Support/raw_ostream.h
@@ -721,7 +721,11 @@ class buffer_unique_ostream : public raw_svector_ostream {
 
 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(); }
 };
 

diff  --git a/llvm/unittests/Support/CMakeLists.txt b/llvm/unittests/Support/CMakeLists.txt
index b44d9ee27188d..d331d1b4b187a 100644
--- a/llvm/unittests/Support/CMakeLists.txt
+++ b/llvm/unittests/Support/CMakeLists.txt
@@ -92,6 +92,7 @@ add_llvm_unittest(SupportTests
   WithColorTest.cpp
   YAMLIOTest.cpp
   YAMLParserTest.cpp
+  buffer_ostream_test.cpp
   formatted_raw_ostream_test.cpp
   raw_fd_stream_test.cpp
   raw_ostream_test.cpp

diff  --git a/llvm/unittests/Support/buffer_ostream_test.cpp b/llvm/unittests/Support/buffer_ostream_test.cpp
new file mode 100644
index 0000000000000..f29c1bb61afa0
--- /dev/null
+++ b/llvm/unittests/Support/buffer_ostream_test.cpp
@@ -0,0 +1,77 @@
+//===- 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