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,10 @@ // Configuration Interface //===--------------------------------------------------------------------===// + /// If possible, pre-allocate \p Capacity bytes for stream data. + /// reserve() does not change size/data of the stream. + virtual void reserve(size_t Capacity) {} + /// Set the stream to be buffered, with an automatically determined buffer /// size. void SetBuffered(); @@ -626,6 +630,8 @@ flush(); return OS; } + + void reserve(size_t Capacity) override { OS.reserve(Capacity); } }; /// A raw_ostream that writes to an SmallVector or SmallString. This is a @@ -659,6 +665,8 @@ /// Return a StringRef for the vector contents. StringRef str() const { return StringRef(OS.data(), OS.size()); } + + void reserve(size_t Capacity) override { OS.reserve(Capacity); } }; /// 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; + size_t OrigCapacity = Str.capacity(); + raw_string_ostream OS(Str); + size_t NewCapacity = OrigCapacity + 1000; + OS.reserve(NewCapacity); + EXPECT_EQ(Str.capacity(), NewCapacity); + OS << "hello"; + OS << 1; + OS << 'w' << 'o' << 'r' << 'l' << 'd'; + OS.flush(); + EXPECT_EQ("hello1world", Str); +} }