[Support] Add WritableMemoryBuffer class

Summary:
The motivation here is LLDB, where we need to fixup relocations in
mmapped files before their contents can be read correctly.  The
MemoryBuffer class does exactly what we need, *except* that it maps the
file in read-only mode.

WritableMemoryBuffer reuses the existing machinery for opening and
mmapping a file. The only difference is in the argument to the
mapped_file_region constructor -- we create a private copy-on-write
mapping, so that we can make changes to the mapped data, but the changes
aren't carried over to the underlying file.

This patch is based on an initial version by Zachary Turner.

Reviewers: mehdi_amini, rnk, rafael, dblaikie, zturner

Subscribers: llvm-commits

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

llvm-svn: 321071

llvm/include/llvm/Support/MemoryBuffer.h

@@ -15,6 +15,7 @@
 #define LLVM_SUPPORT_MEMORYBUFFER_H
 
 #include "llvm-c/Types.h"
+#include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/Twine.h"
 #include "llvm/Support/CBindingWrapping.h"
@@ -47,6 +48,9 @@ protected:
 
   void init(const char *BufStart, const char *BufEnd,
             bool RequiresNullTerminator);
+
+  static constexpr bool Writable = false;
+
 public:
   MemoryBuffer(const MemoryBuffer &) = delete;
   MemoryBuffer &operator=(const MemoryBuffer &) = delete;
@@ -122,6 +126,9 @@ public:
   /// Allocate a new MemoryBuffer of the specified size that is not initialized.
   /// Note that the caller should initialize the memory allocated by this
   /// method. The memory is owned by the MemoryBuffer object.
+  //
+  // TODO: Remove this and migrate callers to
+  //   WritableMemoryBuffer::getNewUninitMemBuffer
   static std::unique_ptr<MemoryBuffer>
   getNewUninitMemBuffer(size_t Size, const Twine &BufferName = "");
 
@@ -156,6 +163,59 @@ public:
   MemoryBufferRef getMemBufferRef() const;
 };
 
+/// This class is an extension of MemoryBuffer, which allows writing to the
+/// underlying contents.  It only supports creation methods that are guaranteed
+/// to produce a writable buffer.  For example, mapping a file read-only is not
+/// supported.
+class WritableMemoryBuffer : public MemoryBuffer {
+protected:
+  WritableMemoryBuffer() = default;
+
+  static constexpr bool Writable = true;
+
+public:
+  using MemoryBuffer::getBuffer;
+  using MemoryBuffer::getBufferEnd;
+  using MemoryBuffer::getBufferStart;
+
+  // const_cast is well-defined here, because the underlying buffer is
+  // guaranteed to have been initialized with a mutable buffer.
+  char *getBufferStart() {
+    return const_cast<char *>(MemoryBuffer::getBufferStart());
+  }
+  char *getBufferEnd() {
+    return const_cast<char *>(MemoryBuffer::getBufferEnd());
+  }
+  MutableArrayRef<char> getBuffer() {
+    return {getBufferStart(), getBufferEnd()};
+  }
+
+  static ErrorOr<std::unique_ptr<WritableMemoryBuffer>>
+  getFile(const Twine &Filename, int64_t FileSize = -1,
+          bool IsVolatile = false);
+
+  /// Map a subrange of the specified file as a WritableMemoryBuffer.
+  static ErrorOr<std::unique_ptr<WritableMemoryBuffer>>
+  getFileSlice(const Twine &Filename, uint64_t MapSize, uint64_t Offset,
+               bool IsVolatile = false);
+
+  static std::unique_ptr<WritableMemoryBuffer>
+  getNewUninitMemBuffer(size_t Size, const Twine &BufferName = "");
+
+private:
+  // Hide these base class factory function so one can't write
+  //   WritableMemoryBuffer::getXXX()
+  // and be surprised that he got a read-only Buffer.
+  using MemoryBuffer::getFileAsStream;
+  using MemoryBuffer::getFileOrSTDIN;
+  using MemoryBuffer::getMemBuffer;
+  using MemoryBuffer::getMemBufferCopy;
+  using MemoryBuffer::getNewMemBuffer;
+  using MemoryBuffer::getOpenFile;
+  using MemoryBuffer::getOpenFileSlice;
+  using MemoryBuffer::getSTDIN;
+};
+
 class MemoryBufferRef {
   StringRef Buffer;
   StringRef Identifier;

llvm/lib/Support/MemoryBuffer.cpp

@@ -80,10 +80,12 @@ void *operator new(size_t N, const NamedBufferAlloc &Alloc) {
 
 namespace {
 /// MemoryBufferMem - Named MemoryBuffer pointing to a block of memory.
-class MemoryBufferMem : public MemoryBuffer {
+template<typename MB>
+class MemoryBufferMem : public MB {
 public:
   MemoryBufferMem(StringRef InputData, bool RequiresNullTerminator) {
-    init(InputData.begin(), InputData.end(), RequiresNullTerminator);
+    MemoryBuffer::init(InputData.begin(), InputData.end(),
+                       RequiresNullTerminator);
   }
 
   /// Disable sized deallocation for MemoryBufferMem, because it has
@@ -95,13 +97,14 @@ public:
     return StringRef(reinterpret_cast<const char *>(this + 1));
   }
 
-  BufferKind getBufferKind() const override {
+  MemoryBuffer::BufferKind getBufferKind() const override {
-    return MemoryBuffer_Malloc;
+    return MemoryBuffer::MemoryBuffer_Malloc;
   }
 };
 }
 
-static ErrorOr<std::unique_ptr<MemoryBuffer>>
+template <typename MB>
+static ErrorOr<std::unique_ptr<MB>>
 getFileAux(const Twine &Filename, int64_t FileSize, uint64_t MapSize,
            uint64_t Offset, bool RequiresNullTerminator, bool IsVolatile);
 
@@ -109,7 +112,7 @@ std::unique_ptr<MemoryBuffer>
 MemoryBuffer::getMemBuffer(StringRef InputData, StringRef BufferName,
                            bool RequiresNullTerminator) {
   auto *Ret = new (NamedBufferAlloc(BufferName))
-      MemoryBufferMem(InputData, RequiresNullTerminator);
+      MemoryBufferMem<MemoryBuffer>(InputData, RequiresNullTerminator);
   return std::unique_ptr<MemoryBuffer>(Ret);
 }
 
@@ -119,41 +122,26 @@ MemoryBuffer::getMemBuffer(MemoryBufferRef Ref, bool RequiresNullTerminator) {
       Ref.getBuffer(), Ref.getBufferIdentifier(), RequiresNullTerminator));
 }
 
+static ErrorOr<std::unique_ptr<WritableMemoryBuffer>>
+getMemBufferCopyImpl(StringRef InputData, const Twine &BufferName) {
+  auto Buf = WritableMemoryBuffer::getNewUninitMemBuffer(InputData.size(), BufferName);
+  if (!Buf)
+    return make_error_code(errc::not_enough_memory);
+  memcpy(Buf->getBufferStart(), InputData.data(), InputData.size());
+  return std::move(Buf);
+}
+
 std::unique_ptr<MemoryBuffer>
 MemoryBuffer::getMemBufferCopy(StringRef InputData, const Twine &BufferName) {
-  std::unique_ptr<MemoryBuffer> Buf =
+  auto Buf = getMemBufferCopyImpl(InputData, BufferName);
-      getNewUninitMemBuffer(InputData.size(), BufferName);
+  if (Buf)
-  if (!Buf)
+    return std::move(*Buf);
-    return nullptr;
+  return nullptr;
-  memcpy(const_cast<char*>(Buf->getBufferStart()), InputData.data(),
-         InputData.size());
-  return Buf;
 }
 
 std::unique_ptr<MemoryBuffer>
 MemoryBuffer::getNewUninitMemBuffer(size_t Size, const Twine &BufferName) {
-  // Allocate space for the MemoryBuffer, the data and the name. It is important
+  return WritableMemoryBuffer::getNewUninitMemBuffer(Size, BufferName);
-  // that MemoryBuffer and data are aligned so PointerIntPair works with them.
-  // TODO: Is 16-byte alignment enough?  We copy small object files with large
-  // alignment expectations into this buffer.
-  SmallString<256> NameBuf;
-  StringRef NameRef = BufferName.toStringRef(NameBuf);
-  size_t AlignedStringLen =
-      alignTo(sizeof(MemoryBufferMem) + NameRef.size() + 1, 16);
-  size_t RealLen = AlignedStringLen + Size + 1;
-  char *Mem = static_cast<char*>(operator new(RealLen, std::nothrow));
-  if (!Mem)
-    return nullptr;
-
-  // The name is stored after the class itself.
-  CopyStringRef(Mem + sizeof(MemoryBufferMem), NameRef);
-
-  // The buffer begins after the name and must be aligned.
-  char *Buf = Mem + AlignedStringLen;
-  Buf[Size] = 0; // Null terminate buffer.
-
-  auto *Ret = new (Mem) MemoryBufferMem(StringRef(Buf, Size), true);
-  return std::unique_ptr<MemoryBuffer>(Ret);
 }
 
 std::unique_ptr<MemoryBuffer>
@@ -179,10 +167,10 @@ MemoryBuffer::getFileOrSTDIN(const Twine &Filename, int64_t FileSize,
 ErrorOr<std::unique_ptr<MemoryBuffer>>
 MemoryBuffer::getFileSlice(const Twine &FilePath, uint64_t MapSize, 
                            uint64_t Offset, bool IsVolatile) {
-  return getFileAux(FilePath, -1, MapSize, Offset, false, IsVolatile);
+  return getFileAux<MemoryBuffer>(FilePath, -1, MapSize, Offset, false,
+                                  IsVolatile);
 }
 
-
 //===----------------------------------------------------------------------===//
 // MemoryBuffer::getFile implementation.
 //===----------------------------------------------------------------------===//
@@ -191,7 +179,8 @@ namespace {
 /// \brief Memory maps a file descriptor using sys::fs::mapped_file_region.
 ///
 /// This handles converting the offset into a legal offset on the platform.
-class MemoryBufferMMapFile : public MemoryBuffer {
+template<typename MB>
+class MemoryBufferMMapFile : public MB {
   sys::fs::mapped_file_region MFR;
 
   static uint64_t getLegalMapOffset(uint64_t Offset) {
@@ -209,11 +198,13 @@ class MemoryBufferMMapFile : public MemoryBuffer {
 public:
   MemoryBufferMMapFile(bool RequiresNullTerminator, int FD, uint64_t Len,
                        uint64_t Offset, std::error_code &EC)
-      : MFR(FD, sys::fs::mapped_file_region::readonly,
+      : MFR(FD,
+            MB::Writable ? sys::fs::mapped_file_region::priv
+                         : sys::fs::mapped_file_region::readonly,
             getLegalMapSize(Len, Offset), getLegalMapOffset(Offset), EC) {
     if (!EC) {
       const char *Start = getStart(Len, Offset);
-      init(Start, Start + Len, RequiresNullTerminator);
+      MemoryBuffer::init(Start, Start + Len, RequiresNullTerminator);
     }
   }
 
@@ -226,13 +217,13 @@ public:
     return StringRef(reinterpret_cast<const char *>(this + 1));
   }
 
-  BufferKind getBufferKind() const override {
+  MemoryBuffer::BufferKind getBufferKind() const override {
-    return MemoryBuffer_MMap;
+    return MemoryBuffer::MemoryBuffer_MMap;
   }
 };
 }
 
-static ErrorOr<std::unique_ptr<MemoryBuffer>>
+static ErrorOr<std::unique_ptr<WritableMemoryBuffer>>
 getMemoryBufferForStream(int FD, const Twine &BufferName) {
   const ssize_t ChunkSize = 4096*4;
   SmallString<ChunkSize> Buffer;
@@ -246,37 +237,80 @@ getMemoryBufferForStream(int FD, const Twine &BufferName) {
     Buffer.set_size(Buffer.size() + ReadBytes);
   } while (ReadBytes != 0);
 
-  return MemoryBuffer::getMemBufferCopy(Buffer, BufferName);
+  return getMemBufferCopyImpl(Buffer, BufferName);
 }
 
 
 ErrorOr<std::unique_ptr<MemoryBuffer>>
 MemoryBuffer::getFile(const Twine &Filename, int64_t FileSize,
                       bool RequiresNullTerminator, bool IsVolatile) {
-  return getFileAux(Filename, FileSize, FileSize, 0,
+  return getFileAux<MemoryBuffer>(Filename, FileSize, FileSize, 0,
-                    RequiresNullTerminator, IsVolatile);
+                                  RequiresNullTerminator, IsVolatile);
 }
 
-static ErrorOr<std::unique_ptr<MemoryBuffer>>
+template <typename MB>
+static ErrorOr<std::unique_ptr<MB>>
 getOpenFileImpl(int FD, const Twine &Filename, uint64_t FileSize,
                 uint64_t MapSize, int64_t Offset, bool RequiresNullTerminator,
                 bool IsVolatile);
 
-static ErrorOr<std::unique_ptr<MemoryBuffer>>
+template <typename MB>
+static ErrorOr<std::unique_ptr<MB>>
 getFileAux(const Twine &Filename, int64_t FileSize, uint64_t MapSize,
            uint64_t Offset, bool RequiresNullTerminator, bool IsVolatile) {
   int FD;
   std::error_code EC = sys::fs::openFileForRead(Filename, FD);
+
   if (EC)
     return EC;
 
-  ErrorOr<std::unique_ptr<MemoryBuffer>> Ret =
+  auto Ret = getOpenFileImpl<MB>(FD, Filename, FileSize, MapSize, Offset,
-      getOpenFileImpl(FD, Filename, FileSize, MapSize, Offset,
+                                 RequiresNullTerminator, IsVolatile);
-                      RequiresNullTerminator, IsVolatile);
   close(FD);
   return Ret;
 }
 
+ErrorOr<std::unique_ptr<WritableMemoryBuffer>>
+WritableMemoryBuffer::getFile(const Twine &Filename, int64_t FileSize,
+                              bool IsVolatile) {
+  return getFileAux<WritableMemoryBuffer>(Filename, FileSize, FileSize, 0,
+                                          /*RequiresNullTerminator*/ false,
+                                          IsVolatile);
+}
+
+ErrorOr<std::unique_ptr<WritableMemoryBuffer>>
+WritableMemoryBuffer::getFileSlice(const Twine &Filename, uint64_t MapSize,
+                                   uint64_t Offset, bool IsVolatile) {
+  return getFileAux<WritableMemoryBuffer>(Filename, -1, MapSize, Offset, false,
+                                          IsVolatile);
+}
+
+std::unique_ptr<WritableMemoryBuffer>
+WritableMemoryBuffer::getNewUninitMemBuffer(size_t Size, const Twine &BufferName) {
+  using MemBuffer = MemoryBufferMem<WritableMemoryBuffer>;
+  // Allocate space for the MemoryBuffer, the data and the name. It is important
+  // that MemoryBuffer and data are aligned so PointerIntPair works with them.
+  // TODO: Is 16-byte alignment enough?  We copy small object files with large
+  // alignment expectations into this buffer.
+  SmallString<256> NameBuf;
+  StringRef NameRef = BufferName.toStringRef(NameBuf);
+  size_t AlignedStringLen = alignTo(sizeof(MemBuffer) + NameRef.size() + 1, 16);
+  size_t RealLen = AlignedStringLen + Size + 1;
+  char *Mem = static_cast<char*>(operator new(RealLen, std::nothrow));
+  if (!Mem)
+    return nullptr;
+
+  // The name is stored after the class itself.
+  CopyStringRef(Mem + sizeof(MemBuffer), NameRef);
+
+  // The buffer begins after the name and must be aligned.
+  char *Buf = Mem + AlignedStringLen;
+  Buf[Size] = 0; // Null terminate buffer.
+
+  auto *Ret = new (Mem) MemBuffer(StringRef(Buf, Size), true);
+  return std::unique_ptr<WritableMemoryBuffer>(Ret);
+}
+
 static bool shouldUseMmap(int FD,
                           size_t FileSize,
                           size_t MapSize,
@@ -332,7 +366,8 @@ static bool shouldUseMmap(int FD,
   return true;
 }
 
-static ErrorOr<std::unique_ptr<MemoryBuffer>>
+template <typename MB>
+static ErrorOr<std::unique_ptr<MB>>
 getOpenFileImpl(int FD, const Twine &Filename, uint64_t FileSize,
                 uint64_t MapSize, int64_t Offset, bool RequiresNullTerminator,
                 bool IsVolatile) {
@@ -364,22 +399,21 @@ getOpenFileImpl(int FD, const Twine &Filename, uint64_t FileSize,
   if (shouldUseMmap(FD, FileSize, MapSize, Offset, RequiresNullTerminator,
                     PageSize, IsVolatile)) {
     std::error_code EC;
-    std::unique_ptr<MemoryBuffer> Result(
+    std::unique_ptr<MB> Result(
-        new (NamedBufferAlloc(Filename))
+        new (NamedBufferAlloc(Filename)) MemoryBufferMMapFile<MB>(
-        MemoryBufferMMapFile(RequiresNullTerminator, FD, MapSize, Offset, EC));
+            RequiresNullTerminator, FD, MapSize, Offset, EC));
     if (!EC)
       return std::move(Result);
   }
 
-  std::unique_ptr<MemoryBuffer> Buf =
+  auto Buf = WritableMemoryBuffer::getNewUninitMemBuffer(MapSize, Filename);
-      MemoryBuffer::getNewUninitMemBuffer(MapSize, Filename);
   if (!Buf) {
     // Failed to create a buffer. The only way it can fail is if
     // new(std::nothrow) returns 0.
     return make_error_code(errc::not_enough_memory);
   }
 
-  char *BufPtr = const_cast<char *>(Buf->getBufferStart());
+  char *BufPtr = Buf.get()->getBufferStart();
 
   size_t BytesLeft = MapSize;
 #ifndef HAVE_PREAD
@@ -412,7 +446,7 @@ getOpenFileImpl(int FD, const Twine &Filename, uint64_t FileSize,
 ErrorOr<std::unique_ptr<MemoryBuffer>>
 MemoryBuffer::getOpenFile(int FD, const Twine &Filename, uint64_t FileSize,
                           bool RequiresNullTerminator, bool IsVolatile) {
-  return getOpenFileImpl(FD, Filename, FileSize, FileSize, 0,
+  return getOpenFileImpl<MemoryBuffer>(FD, Filename, FileSize, FileSize, 0,
                          RequiresNullTerminator, IsVolatile);
 }
 
@@ -420,7 +454,8 @@ ErrorOr<std::unique_ptr<MemoryBuffer>>
 MemoryBuffer::getOpenFileSlice(int FD, const Twine &Filename, uint64_t MapSize,
                                int64_t Offset, bool IsVolatile) {
   assert(MapSize != uint64_t(-1));
-  return getOpenFileImpl(FD, Filename, -1, MapSize, Offset, false, IsVolatile);
+  return getOpenFileImpl<MemoryBuffer>(FD, Filename, -1, MapSize, Offset, false,
+                                       IsVolatile);
 }
 
 ErrorOr<std::unique_ptr<MemoryBuffer>> MemoryBuffer::getSTDIN() {

llvm/unittests/Support/MemoryBufferTest.cpp

@@ -15,6 +15,7 @@
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/FileUtilities.h"
 #include "llvm/Support/raw_ostream.h"
+#include "llvm/Testing/Support/Error.h"
 #include "gtest/gtest.h"
 
 using namespace llvm;
@@ -226,4 +227,37 @@ TEST_F(MemoryBufferTest, slice) {
   EXPECT_TRUE(BufData2.substr(0x1800,8).equals("abcdefgh"));
   EXPECT_TRUE(BufData2.substr(0x2FF8,8).equals("abcdefgh"));
 }
+
+TEST_F(MemoryBufferTest, writableSlice) {
+  // Create a file initialized with some data
+  int FD;
+  SmallString<64> TestPath;
+  sys::fs::createTemporaryFile("MemoryBufferTest_WritableSlice", "temp", FD,
+                               TestPath);
+  FileRemover Cleanup(TestPath);
+  raw_fd_ostream OF(FD, true);
+  for (unsigned i = 0; i < 0x1000; ++i)
+    OF << "0123456789abcdef";
+  OF.close();
+
+  {
+    auto MBOrError =
+        WritableMemoryBuffer::getFileSlice(TestPath.str(), 0x6000, 0x2000);
+    ASSERT_FALSE(MBOrError.getError());
+    // Write some data.  It should be mapped private, so that upon completion
+    // the original file contents are not modified.
+    WritableMemoryBuffer &MB = **MBOrError;
+    ASSERT_EQ(0x6000u, MB.getBufferSize());
+    char *Start = MB.getBufferStart();
+    ASSERT_EQ(MB.getBufferEnd(), MB.getBufferStart() + MB.getBufferSize());
+    ::memset(Start, 'x', MB.getBufferSize());
+  }
+
+  auto MBOrError = MemoryBuffer::getFile(TestPath);
+  ASSERT_FALSE(MBOrError.getError());
+  auto &MB = **MBOrError;
+  ASSERT_EQ(0x10000u, MB.getBufferSize());
+  for (size_t i = 0; i < MB.getBufferSize(); i += 0x10)
+    EXPECT_EQ("0123456789abcdef", MB.getBuffer().substr(i, 0x10)) << "i: " << i;
+}
 }