diff --git a/llvm/include/llvm/Support/raw_ostream.h b/llvm/include/llvm/Support/raw_ostream.h --- a/llvm/include/llvm/Support/raw_ostream.h +++ b/llvm/include/llvm/Support/raw_ostream.h @@ -137,6 +137,13 @@ // Configuration Interface //===--------------------------------------------------------------------===// + /// If possible, pre-allocate \p ExtraSize bytes for stream data. + /// i.e. it extends internal buffers to keep additional ExtraSize bytes. + /// So that the stream could keep at least tell() + ExtraSize bytes + /// without re-allocations. reserveExtraSpace() does not change + /// the size/data of the stream. + virtual void reserveExtraSpace(uint64_t ExtraSize) {} + /// Set the stream to be buffered, with an automatically determined buffer /// size. void SetBuffered(); @@ -626,6 +633,10 @@ flush(); return OS; } + + void reserveExtraSpace(uint64_t ExtraSize) override { + OS.reserve(tell() + ExtraSize); + } }; /// A raw_ostream that writes to an SmallVector or SmallString. This is a @@ -659,6 +670,10 @@ /// Return a StringRef for the vector contents. StringRef str() const { return StringRef(OS.data(), OS.size()); } + + void reserveExtraSpace(uint64_t ExtraSize) override { + OS.reserve(tell() + ExtraSize); + } }; /// A raw_ostream that discards all output. diff --git a/llvm/unittests/Support/raw_ostream_test.cpp b/llvm/unittests/Support/raw_ostream_test.cpp --- a/llvm/unittests/Support/raw_ostream_test.cpp +++ b/llvm/unittests/Support/raw_ostream_test.cpp @@ -455,4 +455,18 @@ TiedStream << "0"; EXPECT_EQ("acegostuv", TiedToBuffer); } + +TEST(raw_ostreamTest, reserve_stream) { + std::string Str; + raw_string_ostream OS(Str); + OS << "11111111111111111111"; + uint64_t CurrentPos = OS.tell(); + OS.reserveExtraSpace(1000); + EXPECT_TRUE(Str.capacity() >= CurrentPos + 1000); + OS << "hello"; + OS << 1; + OS << 'w' << 'o' << 'r' << 'l' << 'd'; + OS.flush(); + EXPECT_EQ("11111111111111111111hello1world", Str); +} }