[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

1 files changed, 96 insertions, 61 deletions

diff --git a/llvm/lib/Support/MemoryBuffer.cpp b/llvm/lib/Support/MemoryBuffer.cpp
index 85e782b2c04..8a93f97930c 100644
--- a/llvm/lib/Support/MemoryBuffer.cpp
+++ b/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,21 +97,22 @@ public:
     return StringRef(reinterpret_cast<const char *>(this + 1));
   }
 
-  BufferKind getBufferKind() const override {
-    return MemoryBuffer_Malloc;
+  MemoryBuffer::BufferKind getBufferKind() const override {
+    return MemoryBuffer::MemoryBuffer_Malloc;
   }
 };
 }
 
-static ErrorOr<std::unique_ptr<MemoryBuffer>>
-getFileAux(const Twine &Filename, int64_t FileSize, uint64_t MapSize, 
+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);
 
 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 =
-      getNewUninitMemBuffer(InputData.size(), BufferName);
-  if (!Buf)
-    return nullptr;
-  memcpy(const_cast<char*>(Buf->getBufferStart()), InputData.data(),
-         InputData.size());
-  return Buf;
+  auto Buf = getMemBufferCopyImpl(InputData, BufferName);
+  if (Buf)
+    return std::move(*Buf);
+  return nullptr;
 }
 
 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
-  // 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);
+  return WritableMemoryBuffer::getNewUninitMemBuffer(Size, BufferName);
 }
 
 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 {
-    return MemoryBuffer_MMap;
+  MemoryBuffer::BufferKind getBufferKind() const override {
+    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,
-                    RequiresNullTerminator, IsVolatile);
+  return getFileAux<MemoryBuffer>(Filename, FileSize, FileSize, 0,
+                                  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 =
-      getOpenFileImpl(FD, Filename, FileSize, MapSize, Offset,
-                      RequiresNullTerminator, IsVolatile);
+  auto Ret = getOpenFileImpl<MB>(FD, Filename, FileSize, MapSize, Offset,
+                                 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(
-        new (NamedBufferAlloc(Filename))
-        MemoryBufferMMapFile(RequiresNullTerminator, FD, MapSize, Offset, EC));
+    std::unique_ptr<MB> Result(
+        new (NamedBufferAlloc(Filename)) MemoryBufferMMapFile<MB>(
+            RequiresNullTerminator, FD, MapSize, Offset, EC));
     if (!EC)
       return std::move(Result);
   }
 
-  std::unique_ptr<MemoryBuffer> Buf =
-      MemoryBuffer::getNewUninitMemBuffer(MapSize, Filename);
+  auto Buf = WritableMemoryBuffer::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() {