diff --git a/llvm/include/llvm/Support/raw_ostream.h b/llvm/include/llvm/Support/raw_ostream.h index 6aba60d5e8c8..ed79292bd5ed 100644 --- a/llvm/include/llvm/Support/raw_ostream.h +++ b/llvm/include/llvm/Support/raw_ostream.h @@ -313,13 +313,22 @@ private: void copy_to_buffer(const char *Ptr, size_t Size); }; +/// An abstract base class for streams implementations that also support a +/// pwrite operation. This is usefull for code that can mostly stream out data, +/// but needs to patch in a header that needs to know the output size. +class raw_pwrite_stream : public raw_ostream { +public: + using raw_ostream::raw_ostream; + virtual void pwrite(const char *Ptr, size_t Size, uint64_t Offset) = 0; +}; + //===----------------------------------------------------------------------===// // File Output Streams //===----------------------------------------------------------------------===// /// A raw_ostream that writes to a file descriptor. /// -class raw_fd_ostream : public raw_ostream { +class raw_fd_ostream : public raw_pwrite_stream { int FD; bool ShouldClose; @@ -378,6 +387,8 @@ public: /// to the offset specified from the beginning of the file. uint64_t seek(uint64_t off); + void pwrite(const char *Ptr, size_t Size, uint64_t Offset) override; + /// Set the stream to attempt to use atomic writes for individual output /// routines where possible. /// @@ -460,7 +471,7 @@ public: /// A raw_ostream that writes to an SmallVector or SmallString. This is a /// simple adaptor class. This class does not encounter output errors. -class raw_svector_ostream : public raw_ostream { +class raw_svector_ostream : public raw_pwrite_stream { SmallVectorImpl &OS; /// See raw_ostream::write_impl. @@ -469,6 +480,12 @@ class raw_svector_ostream : public raw_ostream { /// Return the current position within the stream, not counting the bytes /// currently in the buffer. uint64_t current_pos() const override; + +protected: + // Like the regular constructor, but doesn't call init. + explicit raw_svector_ostream(SmallVectorImpl &O, unsigned); + void init(); + public: /// Construct a new raw_svector_ostream. /// @@ -477,6 +494,8 @@ public: explicit raw_svector_ostream(SmallVectorImpl &O); ~raw_svector_ostream() override; + void pwrite(const char *Ptr, size_t Size, uint64_t Offset) override; + /// This is called when the SmallVector we're appending to is changed outside /// of the raw_svector_ostream's control. It is only safe to do this if the /// raw_svector_ostream has previously been flushed. @@ -488,7 +507,7 @@ public: }; /// A raw_ostream that discards all output. -class raw_null_ostream : public raw_ostream { +class raw_null_ostream : public raw_pwrite_stream { /// See raw_ostream::write_impl. void write_impl(const char *Ptr, size_t size) override; @@ -499,6 +518,18 @@ class raw_null_ostream : public raw_ostream { public: explicit raw_null_ostream() {} ~raw_null_ostream() override; + void pwrite(const char *Ptr, size_t Size, uint64_t Offset) override; +}; + +class buffer_ostream : public raw_svector_ostream { + raw_ostream &OS; + SmallVector Buffer; + +public: + buffer_ostream(raw_ostream &OS) : raw_svector_ostream(Buffer, 0), OS(OS) { + init(); + } + ~buffer_ostream() { OS << str(); } }; } // end llvm namespace diff --git a/llvm/lib/Support/raw_ostream.cpp b/llvm/lib/Support/raw_ostream.cpp index 16c52c90e090..044f8e0ca864 100644 --- a/llvm/lib/Support/raw_ostream.cpp +++ b/llvm/lib/Support/raw_ostream.cpp @@ -516,8 +516,8 @@ raw_fd_ostream::raw_fd_ostream(StringRef Filename, std::error_code &EC, /// FD is the file descriptor that this writes to. If ShouldClose is true, this /// closes the file when the stream is destroyed. raw_fd_ostream::raw_fd_ostream(int fd, bool shouldClose, bool unbuffered) - : raw_ostream(unbuffered), FD(fd), - ShouldClose(shouldClose), Error(false), UseAtomicWrites(false) { + : raw_pwrite_stream(unbuffered), FD(fd), ShouldClose(shouldClose), + Error(false), UseAtomicWrites(false) { if (FD < 0 ) { ShouldClose = false; return; @@ -630,6 +630,13 @@ uint64_t raw_fd_ostream::seek(uint64_t off) { return pos; } +void raw_fd_ostream::pwrite(const char *Ptr, size_t Size, uint64_t Offset) { + uint64_t Pos = tell(); + seek(Offset); + write(Ptr, Size); + seek(Pos); +} + size_t raw_fd_ostream::preferred_buffer_size() const { #if !defined(_MSC_VER) && !defined(__MINGW32__) && !defined(__minix) // Windows and Minix have no st_blksize. @@ -753,7 +760,14 @@ void raw_string_ostream::write_impl(const char *Ptr, size_t Size) { // capacity. This allows raw_ostream to write directly into the correct place, // and we only need to set the vector size when the data is flushed. +raw_svector_ostream::raw_svector_ostream(SmallVectorImpl &O, unsigned) + : OS(O) {} + raw_svector_ostream::raw_svector_ostream(SmallVectorImpl &O) : OS(O) { + init(); +} + +void raw_svector_ostream::init() { // Set up the initial external buffer. We make sure that the buffer has at // least 128 bytes free; raw_ostream itself only requires 64, but we want to // make sure that we don't grow the buffer unnecessarily on destruction (when @@ -767,6 +781,17 @@ raw_svector_ostream::~raw_svector_ostream() { flush(); } +void raw_svector_ostream::pwrite(const char *Ptr, size_t Size, + uint64_t Offset) { + flush(); + + uint64_t End = Offset + Size; + if (End > OS.size()) + OS.resize(End); + + memcpy(OS.begin() + Offset, Ptr, Size); +} + /// resync - This is called when the SmallVector we're appending to is changed /// outside of the raw_svector_ostream's control. It is only safe to do this /// if the raw_svector_ostream has previously been flushed. @@ -821,3 +846,5 @@ void raw_null_ostream::write_impl(const char *Ptr, size_t Size) { uint64_t raw_null_ostream::current_pos() const { return 0; } + +void raw_null_ostream::pwrite(const char *Ptr, size_t Size, uint64_t Offset) {} diff --git a/llvm/unittests/Support/CMakeLists.txt b/llvm/unittests/Support/CMakeLists.txt index f3b55c3a4172..3c8a090c2bac 100644 --- a/llvm/unittests/Support/CMakeLists.txt +++ b/llvm/unittests/Support/CMakeLists.txt @@ -44,6 +44,7 @@ add_llvm_unittest(SupportTests YAMLParserTest.cpp formatted_raw_ostream_test.cpp raw_ostream_test.cpp + raw_pwrite_stream_test.cpp ) # ManagedStatic.cpp uses . diff --git a/llvm/unittests/Support/raw_pwrite_stream_test.cpp b/llvm/unittests/Support/raw_pwrite_stream_test.cpp new file mode 100644 index 000000000000..bcbb29b86bde --- /dev/null +++ b/llvm/unittests/Support/raw_pwrite_stream_test.cpp @@ -0,0 +1,25 @@ +//===- raw_pwrite_stream_test.cpp - raw_pwrite_stream tests ---------------===// +// +// The LLVM Compiler Infrastructure +// +// This file is distributed under the University of Illinois Open Source +// License. See LICENSE.TXT for details. +// +//===----------------------------------------------------------------------===// + +#include "gtest/gtest.h" +#include "llvm/ADT/SmallString.h" +#include "llvm/Support/raw_ostream.h" + +using namespace llvm; + +namespace { + +TEST(raw_pwrite_ostreamTest, TestSVector) { + SmallString<64> Buffer; + raw_svector_ostream OS(Buffer); + StringRef Test = "test"; + OS.pwrite(Test.data(), Test.size(), 0); + EXPECT_EQ(Test, OS.str()); +} +}