[MemoryBuffer] Allow optionally specifying desired buffer alignment

[rriddle]

Summary

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.

Depends on D137764

[rriddle]

@dblaikie Would love thoughts if you think this is the right approach. I hit some interesting situations downstream where a small file with large alignment constraints(in the realm of 64/128/etc bytes) was seemingly miscompiling, but then I learned that MemoryBuffer only mmaps large files. This patch explicitly encodes alignment constraints, which felt safer than just changing the current alignment value hardcoded in WritableMemoryBuffer. The situation I hit was basically the same as the existing TODO in WritableMemoryBuffer::getNewUninitMemBuffer.

[dblaikie]

If you've got a use for it, this seems fine to me.

llvm/unittests/Support/MemoryBufferTest.cpp

232–233

Prefer = over () when both are valid - makes it easier to read because it limits the code to only implicit ctors (so, as a reader, I don't worry that the function is returning a raw pointer & this is taking ownership - and want to go check that the raw pointer is really conferring ownership, etc)

[lattner]

This LGTM and makes sense. Nice use of the alignment apis.

llvm/include/llvm/Support/MemoryBuffer.h

 
 #ifndef LLVM_SUPPORT_MEMORYBUFFER_H
 #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/Alignment.h"
 #include "llvm/Support/CBindingWrapping.h"
 #include "llvm/Support/ErrorOr.h"
 #include "llvm/Support/MemoryBufferRef.h"
 #include <cstddef>
 #include <cstdint>
 #include <memory>
 
 namespace llvm {
public:
   /// if successful, otherwise returning null.
   ///
   /// \param IsText Set to true to indicate that the file should be read in
   /// text mode.
   ///
   /// \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
   /// look like a regular file but have 0 size (e.g. /proc/cpuinfo on Linux).
   static ErrorOr<std::unique_ptr<MemoryBuffer>>
   getFileAsStream(const Twine &Filename);
 
   /// Given an already-open file descriptor, map some slice of it into a
   /// MemoryBuffer. The slice is specified by an \p Offset and \p MapSize.
   /// 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.
   ///
   /// \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.
   static std::unique_ptr<MemoryBuffer>
   getMemBuffer(StringRef InputData, StringRef BufferName = "",
                bool RequiresNullTerminator = true);
 
   static std::unique_ptr<MemoryBuffer>
   getMemBuffer(MemoryBufferRef Ref, bool RequiresNullTerminator = true);
 
   /// Open the specified memory range as a MemoryBuffer, copying the contents
   /// and taking ownership of it. InputData does not have to be null terminated.
   static std::unique_ptr<MemoryBuffer>
   getMemBufferCopy(StringRef InputData, const Twine &BufferName = "");
 
   /// Read all of stdin into a file buffer, and return it.
   static ErrorOr<std::unique_ptr<MemoryBuffer>> getSTDIN();
 
   /// Open the specified file as a MemoryBuffer, or open stdin if the Filename
   /// 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.
   //===--------------------------------------------------------------------===//
 
   /// The kind of memory backing used to support the MemoryBuffer.
   enum BufferKind {
     MemoryBuffer_Malloc,
public:
   char *getBufferEnd() {
     return const_cast<char *>(MemoryBuffer::getBufferEnd());
   }
   MutableArrayRef<char> getBuffer() {
     return {getBufferStart(), getBufferEnd()};
   }
 
   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.
   /// The memory is owned by the MemoryBuffer object.
   static std::unique_ptr<WritableMemoryBuffer>
   getNewMemBuffer(size_t Size, const Twine &BufferName = "");
 
 private:

llvm/lib/Support/MemoryBuffer.cpp

 //===--- MemoryBuffer.cpp - Memory Buffer implementation ------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 //  This file implements the MemoryBuffer interface.
 //
 //===----------------------------------------------------------------------===//
 
 #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"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/MathExtras.h"
 #include "llvm/Support/Process.h"
 #include "llvm/Support/Program.h"
 #include "llvm/Support/SmallVectorMemoryBuffer.h"
MemoryBuffer::BufferKind getBufferKind() const override {
     return MemoryBuffer::MemoryBuffer_Malloc;
   }
 };
 } // namespace
 
 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,
                            bool RequiresNullTerminator) {
   auto *Ret = new (NamedBufferAlloc(BufferName))
       MemoryBufferMem<MemoryBuffer>(InputData, RequiresNullTerminator);
   return std::unique_ptr<MemoryBuffer>(Ret);
 }
MemoryBuffer::getMemBufferCopy(StringRef InputData, const Twine &BufferName) {
   auto Buf = getMemBufferCopyImpl(InputData, BufferName);
   if (Buf)
     return std::move(*Buf);
   return nullptr;
 }
 
 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);
 }
 
 //===----------------------------------------------------------------------===//
 // MemoryBuffer::getFile implementation.
 //===----------------------------------------------------------------------===//
 
 namespace {
 
getMemoryBufferForStream(sys::fs::file_t FD, const Twine &BufferName) {
   SmallString<sys::fs::DefaultReadChunkSize> Buffer;
   if (Error E = sys::fs::readNativeFileToEOF(FD, Buffer))
     return errorToErrorCode(std::move(E));
   return getMemBufferCopyImpl(Buffer, 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));
   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;
+  char *Buf = (char *)alignAddr(Mem + StringLen, BufAlign);
   Buf[Size] = 0; // Null terminate buffer.
 
   auto *Ret = new (Mem) MemBuffer(StringRef(Buf, Size), true);
   return std::unique_ptr<WritableMemoryBuffer>(Ret);
 }
 
 std::unique_ptr<WritableMemoryBuffer>
 WritableMemoryBuffer::getNewMemBuffer(size_t Size, const Twine &BufferName) {
WriteThroughMemoryBuffer::getFileSlice(const Twine &Filename, uint64_t MapSize,
                                        uint64_t Offset) {
   return getReadWriteFile(Filename, -1, MapSize, Offset);
 }
 
 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.
   if (MapSize == uint64_t(-1)) {
     // If we don't know the file size, use fstat to find out.  fstat on an open
     // file descriptor is cheaper than stat on a random path.
     if (FileSize == uint64_t(-1)) {
       sys::fs::file_status Status;
getOpenFileImpl(sys::fs::file_t FD, const Twine &Filename, uint64_t FileSize,
   }
 
 #ifdef __MVS__
   // Set codepage auto-conversion for z/OS.
   if (auto EC = llvm::enableAutoConversion(FD))
     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.
     return make_error_code(errc::not_enough_memory);
   }
 
   // Read until EOF, zero-initialize the rest.
   MutableArrayRef<char> ToRead = Buf->getBuffer();
while (!ToRead.empty()) {
     ToRead = ToRead.drop_front(*ReadBytes);
     Offset += *ReadBytes;
   }
 
   return std::move(Buf);
 }
 
 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() {
   // Read in all of the data from stdin, we cannot mmap stdin.
   //
   // FIXME: That isn't necessarily true, we should try to mmap stdin and
   // fallback if it fails.
   sys::ChangeStdinMode(sys::fs::OF_Text);

llvm/unittests/Support/MemoryBufferTest.cpp

for (size_t i = 0; i < 123; ++i)
     EXPECT_EQ(0, Four->getBufferStart()[0]);
 
   // uninitialized buffer with rollover size
   OwningBuffer Five(
       WritableMemoryBuffer::getNewUninitMemBuffer(SIZE_MAX, "huge"));
   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
   // the usage of memory mapping.
   int TestFD;
   SmallString<64> TestPath;
   // Create a temporary file and write data into it.
   sys::fs::createTemporaryFile("prefix", "temp", TestFD, TestPath);