[MemoryBuffer] Allow optionally specifying desired buffer alignment

Underlying data may have requirements/expectations/etc. about
the run-time alignment. WritableMemoryBuffer currently uses
a 16 byte alignment, which works for many situations but not all.
Allowing a desired alignment makes it easier to reuse WritableMemoryBuffer
in situations of special alignment, and also removes a problem when
opening files with special alignment constraints. Large files generally
get mmaped, which has ~page alignment, but small files go through
WritableMemoryBuffer which has the much smaller alignment guarantee.

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

3 files changed, 86 insertions, 36 deletions

diff --git a/llvm/include/llvm/Support/MemoryBuffer.h b/llvm/include/llvm/Support/MemoryBuffer.h
index 6385805eba1d..ed975c86c125 100644
--- a/llvm/include/llvm/Support/MemoryBuffer.h
+++ b/llvm/include/llvm/Support/MemoryBuffer.h
@@ -17,6 +17,7 @@
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/Twine.h"
+#include "llvm/Support/Alignment.h"
 #include "llvm/Support/CBindingWrapping.h"
 #include "llvm/Support/ErrorOr.h"
 #include "llvm/Support/MemoryBufferRef.h"
@@ -90,9 +91,13 @@ public:
   /// \param IsVolatile Set to true to indicate that the contents of the file
   /// can change outside the user's control, e.g. when libclang tries to parse
   /// while the user is editing/updating the file or if the file is on an NFS.
+  ///
+  /// \param Alignment Set to indicate that the buffer should be aligned to at
+  /// least the specified alignment.
   static ErrorOr<std::unique_ptr<MemoryBuffer>>
   getFile(const Twine &Filename, bool IsText = false,
-          bool RequiresNullTerminator = true, bool IsVolatile = false);
+          bool RequiresNullTerminator = true, bool IsVolatile = false,
+          Optional<Align> Alignment = None);
 
   /// Read all of the specified file into a MemoryBuffer as a stream
   /// (i.e. until EOF reached). This is useful for special files that
@@ -105,7 +110,8 @@ public:
   /// Since this is in the middle of a file, the buffer is not null terminated.
   static ErrorOr<std::unique_ptr<MemoryBuffer>>
   getOpenFileSlice(sys::fs::file_t FD, const Twine &Filename, uint64_t MapSize,
-                   int64_t Offset, bool IsVolatile = false);
+                   int64_t Offset, bool IsVolatile = false,
+                   Optional<Align> Alignment = None);
 
   /// Given an already-open file descriptor, read the file and return a
   /// MemoryBuffer.
@@ -113,9 +119,13 @@ public:
   /// \param IsVolatile Set to true to indicate that the contents of the file
   /// can change outside the user's control, e.g. when libclang tries to parse
   /// while the user is editing/updating the file or if the file is on an NFS.
+  ///
+  /// \param Alignment Set to indicate that the buffer should be aligned to at
+  /// least the specified alignment.
   static ErrorOr<std::unique_ptr<MemoryBuffer>>
   getOpenFile(sys::fs::file_t FD, const Twine &Filename, uint64_t FileSize,
-              bool RequiresNullTerminator = true, bool IsVolatile = false);
+              bool RequiresNullTerminator = true, bool IsVolatile = false,
+              Optional<Align> Alignment = None);
 
   /// Open the specified memory range as a MemoryBuffer. Note that InputData
   /// must be null terminated if RequiresNullTerminator is true.
@@ -138,12 +148,13 @@ public:
   /// is "-".
   static ErrorOr<std::unique_ptr<MemoryBuffer>>
   getFileOrSTDIN(const Twine &Filename, bool IsText = false,
-                 bool RequiresNullTerminator = true);
+                 bool RequiresNullTerminator = true,
+                 Optional<Align> Alignment = None);
 
   /// Map a subrange of the specified file as a MemoryBuffer.
   static ErrorOr<std::unique_ptr<MemoryBuffer>>
   getFileSlice(const Twine &Filename, uint64_t MapSize, uint64_t Offset,
-               bool IsVolatile = false);
+               bool IsVolatile = false, Optional<Align> Alignment = None);
 
   //===--------------------------------------------------------------------===//
   // Provided for performance analysis.
@@ -188,18 +199,23 @@ public:
   }
 
   static ErrorOr<std::unique_ptr<WritableMemoryBuffer>>
-  getFile(const Twine &Filename, bool IsVolatile = false);
+  getFile(const Twine &Filename, bool IsVolatile = false,
+          Optional<Align> Alignment = None);
 
   /// 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);
+               bool IsVolatile = false, Optional<Align> Alignment = None);
 
   /// 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.
+  ///
+  /// \param Alignment Set to indicate that the buffer should be aligned to at
+  /// least the specified alignment.
   static std::unique_ptr<WritableMemoryBuffer>
-  getNewUninitMemBuffer(size_t Size, const Twine &BufferName = "");
+  getNewUninitMemBuffer(size_t Size, const Twine &BufferName = "",
+                        Optional<Align> Alignment = None);
 
   /// Allocate a new zero-initialized MemoryBuffer of the specified size. Note
   /// that the caller need not initialize the memory allocated by this method.
diff --git a/llvm/lib/Support/MemoryBuffer.cpp b/llvm/lib/Support/MemoryBuffer.cpp
index 9872dfa78b26..6bb046e9b3db 100644
--- a/llvm/lib/Support/MemoryBuffer.cpp
+++ b/llvm/lib/Support/MemoryBuffer.cpp
@@ -13,6 +13,7 @@
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/Config/config.h"
+#include "llvm/Support/Alignment.h"
 #include "llvm/Support/Errc.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/ErrorHandling.h"
@@ -109,7 +110,8 @@ public:
 template <typename MB>
 static ErrorOr<std::unique_ptr<MB>>
 getFileAux(const Twine &Filename, uint64_t MapSize, uint64_t Offset,
-           bool IsText, bool RequiresNullTerminator, bool IsVolatile);
+           bool IsText, bool RequiresNullTerminator, bool IsVolatile,
+           Optional<Align> Alignment);
 
 std::unique_ptr<MemoryBuffer>
 MemoryBuffer::getMemBuffer(StringRef InputData, StringRef BufferName,
@@ -144,21 +146,24 @@ MemoryBuffer::getMemBufferCopy(StringRef InputData, const Twine &BufferName) {
 
 ErrorOr<std::unique_ptr<MemoryBuffer>>
 MemoryBuffer::getFileOrSTDIN(const Twine &Filename, bool IsText,
-                             bool RequiresNullTerminator) {
+                             bool RequiresNullTerminator,
+                             Optional<Align> Alignment) {
   SmallString<256> NameBuf;
   StringRef NameRef = Filename.toStringRef(NameBuf);
 
   if (NameRef == "-")
     return getSTDIN();
   return getFile(Filename, IsText, RequiresNullTerminator,
-                 /*IsVolatile=*/false);
+                 /*IsVolatile=*/false, Alignment);
 }
 
 ErrorOr<std::unique_ptr<MemoryBuffer>>
 MemoryBuffer::getFileSlice(const Twine &FilePath, uint64_t MapSize,
-                           uint64_t Offset, bool IsVolatile) {
+                           uint64_t Offset, bool IsVolatile,
+                           Optional<Align> Alignment) {
   return getFileAux<MemoryBuffer>(FilePath, MapSize, Offset, /*IsText=*/false,
-                                  /*RequiresNullTerminator=*/false, IsVolatile);
+                                  /*RequiresNullTerminator=*/false, IsVolatile,
+                                  Alignment);
 }
 
 //===----------------------------------------------------------------------===//
@@ -237,58 +242,67 @@ getMemoryBufferForStream(sys::fs::file_t FD, const Twine &BufferName) {
 
 ErrorOr<std::unique_ptr<MemoryBuffer>>
 MemoryBuffer::getFile(const Twine &Filename, bool IsText,
-                      bool RequiresNullTerminator, bool IsVolatile) {
+                      bool RequiresNullTerminator, bool IsVolatile,
+                      Optional<Align> Alignment) {
   return getFileAux<MemoryBuffer>(Filename, /*MapSize=*/-1, /*Offset=*/0,
-                                  IsText, RequiresNullTerminator, IsVolatile);
+                                  IsText, RequiresNullTerminator, IsVolatile,
+                                  Alignment);
 }
 
 template <typename MB>
 static ErrorOr<std::unique_ptr<MB>>
 getOpenFileImpl(sys::fs::file_t FD, const Twine &Filename, uint64_t FileSize,
                 uint64_t MapSize, int64_t Offset, bool RequiresNullTerminator,
-                bool IsVolatile);
+                bool IsVolatile, Optional<Align> Alignment);
 
 template <typename MB>
 static ErrorOr<std::unique_ptr<MB>>
 getFileAux(const Twine &Filename, uint64_t MapSize, uint64_t Offset,
-           bool IsText, bool RequiresNullTerminator, bool IsVolatile) {
+           bool IsText, bool RequiresNullTerminator, bool IsVolatile,
+           Optional<Align> Alignment) {
   Expected<sys::fs::file_t> FDOrErr = sys::fs::openNativeFileForRead(
       Filename, IsText ? sys::fs::OF_TextWithCRLF : sys::fs::OF_None);
   if (!FDOrErr)
     return errorToErrorCode(FDOrErr.takeError());
   sys::fs::file_t FD = *FDOrErr;
   auto Ret = getOpenFileImpl<MB>(FD, Filename, /*FileSize=*/-1, MapSize, Offset,
-                                 RequiresNullTerminator, IsVolatile);
+                                 RequiresNullTerminator, IsVolatile, Alignment);
   sys::fs::closeFile(FD);
   return Ret;
 }
 
 ErrorOr<std::unique_ptr<WritableMemoryBuffer>>
-WritableMemoryBuffer::getFile(const Twine &Filename, bool IsVolatile) {
+WritableMemoryBuffer::getFile(const Twine &Filename, bool IsVolatile,
+                              Optional<Align> Alignment) {
   return getFileAux<WritableMemoryBuffer>(
       Filename, /*MapSize=*/-1, /*Offset=*/0, /*IsText=*/false,
-      /*RequiresNullTerminator=*/false, IsVolatile);
+      /*RequiresNullTerminator=*/false, IsVolatile, Alignment);
 }
 
 ErrorOr<std::unique_ptr<WritableMemoryBuffer>>
 WritableMemoryBuffer::getFileSlice(const Twine &Filename, uint64_t MapSize,
-                                   uint64_t Offset, bool IsVolatile) {
+                                   uint64_t Offset, bool IsVolatile,
+                                   Optional<Align> Alignment) {
   return getFileAux<WritableMemoryBuffer>(
       Filename, MapSize, Offset, /*IsText=*/false,
-      /*RequiresNullTerminator=*/false, IsVolatile);
+      /*RequiresNullTerminator=*/false, IsVolatile, Alignment);
 }
 
 std::unique_ptr<WritableMemoryBuffer>
-WritableMemoryBuffer::getNewUninitMemBuffer(size_t Size, const Twine &BufferName) {
+WritableMemoryBuffer::getNewUninitMemBuffer(size_t Size,
+                                            const Twine &BufferName,
+                                            Optional<Align> Alignment) {
   using MemBuffer = MemoryBufferMem<WritableMemoryBuffer>;
+
+  // Use 16-byte alignment if no alignment is specified.
+  Align BufAlign = Alignment.value_or(Align(16));
+
   // 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;
+  size_t StringLen = sizeof(MemBuffer) + NameRef.size() + 1;
+  size_t RealLen = StringLen + Size + 1 + BufAlign.value();
   if (RealLen <= Size) // Check for rollover.
     return nullptr;
   char *Mem = static_cast<char*>(operator new(RealLen, std::nothrow));
@@ -299,7 +313,7 @@ WritableMemoryBuffer::getNewUninitMemBuffer(size_t Size, const Twine &BufferName
   CopyStringRef(Mem + sizeof(MemBuffer), NameRef);
 
   // The buffer begins after the name and must be aligned.
-  char *Buf = Mem + AlignedStringLen;
+  char *Buf = (char *)alignAddr(Mem + StringLen, BufAlign);
   Buf[Size] = 0; // Null terminate buffer.
 
   auto *Ret = new (Mem) MemBuffer(StringRef(Buf, Size), true);
@@ -427,7 +441,7 @@ template <typename MB>
 static ErrorOr<std::unique_ptr<MB>>
 getOpenFileImpl(sys::fs::file_t FD, const Twine &Filename, uint64_t FileSize,
                 uint64_t MapSize, int64_t Offset, bool RequiresNullTerminator,
-                bool IsVolatile) {
+                bool IsVolatile, Optional<Align> Alignment) {
   static int PageSize = sys::Process::getPageSizeEstimate();
 
   // Default is to map the full file.
@@ -469,7 +483,8 @@ getOpenFileImpl(sys::fs::file_t FD, const Twine &Filename, uint64_t FileSize,
     return EC;
 #endif
 
-  auto Buf = WritableMemoryBuffer::getNewUninitMemBuffer(MapSize, Filename);
+  auto Buf =
+      WritableMemoryBuffer::getNewUninitMemBuffer(MapSize, Filename, Alignment);
   if (!Buf) {
     // Failed to create a buffer. The only way it can fail is if
     // new(std::nothrow) returns 0.
@@ -495,18 +510,21 @@ getOpenFileImpl(sys::fs::file_t FD, const Twine &Filename, uint64_t FileSize,
 }
 
 ErrorOr<std::unique_ptr<MemoryBuffer>>
-MemoryBuffer::getOpenFile(sys::fs::file_t FD, const Twine &Filename, uint64_t FileSize,
-                          bool RequiresNullTerminator, bool IsVolatile) {
+MemoryBuffer::getOpenFile(sys::fs::file_t FD, const Twine &Filename,
+                          uint64_t FileSize, bool RequiresNullTerminator,
+                          bool IsVolatile, Optional<Align> Alignment) {
   return getOpenFileImpl<MemoryBuffer>(FD, Filename, FileSize, FileSize, 0,
-                         RequiresNullTerminator, IsVolatile);
+                                       RequiresNullTerminator, IsVolatile,
+                                       Alignment);
 }
 
 ErrorOr<std::unique_ptr<MemoryBuffer>>
-MemoryBuffer::getOpenFileSlice(sys::fs::file_t FD, const Twine &Filename, uint64_t MapSize,
-                               int64_t Offset, bool IsVolatile) {
+MemoryBuffer::getOpenFileSlice(sys::fs::file_t FD, const Twine &Filename,
+                               uint64_t MapSize, int64_t Offset,
+                               bool IsVolatile, Optional<Align> Alignment) {
   assert(MapSize != uint64_t(-1));
   return getOpenFileImpl<MemoryBuffer>(FD, Filename, -1, MapSize, Offset, false,
-                                       IsVolatile);
+                                       IsVolatile, Alignment);
 }
 
 ErrorOr<std::unique_ptr<MemoryBuffer>> MemoryBuffer::getSTDIN() {
diff --git a/llvm/unittests/Support/MemoryBufferTest.cpp b/llvm/unittests/Support/MemoryBufferTest.cpp
index 423d8f761181..d3969fb106d2 100644
--- a/llvm/unittests/Support/MemoryBufferTest.cpp
+++ b/llvm/unittests/Support/MemoryBufferTest.cpp
@@ -226,6 +226,22 @@ TEST_F(MemoryBufferTest, make_new) {
   EXPECT_EQ(nullptr, Five.get());
 }
 
+TEST_F(MemoryBufferTest, getNewAligned) {
+  auto CheckAlignment = [](size_t AlignmentValue) {
+    Align Alignment(AlignmentValue);
+    OwningBuffer AlignedBuffer =
+        WritableMemoryBuffer::getNewUninitMemBuffer(0, "", Alignment);
+    EXPECT_TRUE(isAddrAligned(Alignment, AlignedBuffer->getBufferStart()));
+  };
+
+  // Test allocation with different alignments.
+  CheckAlignment(16);
+  CheckAlignment(32);
+  CheckAlignment(64);
+  CheckAlignment(128);
+  CheckAlignment(256);
+}
+
 void MemoryBufferTest::testGetOpenFileSlice(bool Reopen) {
   // Test that MemoryBuffer::getOpenFile works properly when no null
   // terminator is requested and the size is large enough to trigger