Add raw_pwrite_stream type.

This is a raw_ostream that also supports pwrite.
I will be used in a sec.

llvm-svn: 234895

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<char> &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<char> &O, unsigned);
+  void init();
+
 public:
   /// Construct a new raw_svector_ostream.
   ///
@@ -477,6 +494,8 @@ public:
   explicit raw_svector_ostream(SmallVectorImpl<char> &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<char, 0> Buffer;
+
+public:
+  buffer_ostream(raw_ostream &OS) : raw_svector_ostream(Buffer, 0), OS(OS) {
+    init();
+  }
+  ~buffer_ostream() { OS << str(); }
 };
 
 } // end llvm namespace

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<char> &O, unsigned)
+    : OS(O) {}
+
 raw_svector_ostream::raw_svector_ostream(SmallVectorImpl<char> &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) {}

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

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());
+}
+}