[llvm] Add null-termination capability to SmallVectorMemoryBuffer
authorJan Svoboda <jan_svoboda@apple.com>
Wed, 8 Dec 2021 12:24:45 +0000 (13:24 +0100)
committerJan Svoboda <jan_svoboda@apple.com>
Thu, 9 Dec 2021 10:32:13 +0000 (11:32 +0100)
Most of `MemoryBuffer` interfaces expose a `RequiresNullTerminator` parameter that's being used to:
* determine how to open a file (`mmap` vs `open`),
* assert newly initialized buffer indeed has an implicit null terminator.

This patch adds the paramater to the `SmallVectorMemoryBuffer` constructors, meaning:
* null terminator can now be added to `SmallVector`s that didn't have one before,
* `SmallVectors` that had a null terminator before keep it even after the move.

In line with existing code, the new parameter is defaulted to `true`. This patch makes sure all calls to the `SmallVectorMemoryBuffer` constructor set it to `false` to preserve the current semantics.

Reviewed By: dexonsmith

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

clang/unittests/AST/ASTImporterTest.cpp
clang/unittests/Frontend/TextDiagnosticTest.cpp
llvm/include/llvm/Support/SmallVectorMemoryBuffer.h
llvm/lib/ExecutionEngine/MCJIT/MCJIT.cpp
llvm/lib/ExecutionEngine/Orc/CompileUtils.cpp
llvm/lib/LTO/ThinLTOCodeGenerator.cpp
llvm/lib/Object/ArchiveWriter.cpp
llvm/tools/llvm-objcopy/MachO/MachOObjcopy.cpp
llvm/tools/llvm-objcopy/llvm-objcopy.cpp
llvm/unittests/Support/MemoryBufferTest.cpp

index 83ccdb5..9e8bee9 100644 (file)
@@ -6452,8 +6452,8 @@ TEST_P(ImportSourceLocations, OverwrittenFileBuffer) {
 
     llvm::SmallVector<char, 64> Buffer;
     Buffer.append(Contents.begin(), Contents.end());
-    auto FileContents =
-        std::make_unique<llvm::SmallVectorMemoryBuffer>(std::move(Buffer), Path);
+    auto FileContents = std::make_unique<llvm::SmallVectorMemoryBuffer>(
+        std::move(Buffer), Path, /*RequiresNullTerminator=*/false);
     FromSM.overrideFileContents(&FE, std::move(FileContents));
 
     // Import the VarDecl to trigger the importing of the FileID.
index 220cff6..3acd89a 100644 (file)
@@ -54,7 +54,7 @@ TEST(TextDiagnostic, ShowLine) {
   llvm::SmallVector<char, 64> buffer;
   buffer.append(main_file_contents.begin(), main_file_contents.end());
   auto file_contents = std::make_unique<llvm::SmallVectorMemoryBuffer>(
-      std::move(buffer), file_path);
+      std::move(buffer), file_path, /*RequiresNullTerminator=*/false);
   SrcMgr.overrideFileContents(fe, std::move(file_contents));
 
   // Create the actual file id and use it as the main file.
index af748cf..f7f2d4e 100644 (file)
@@ -28,17 +28,21 @@ namespace llvm {
 /// MemoryBuffer).
 class SmallVectorMemoryBuffer : public MemoryBuffer {
 public:
-  /// Construct an SmallVectorMemoryBuffer from the given SmallVector
-  /// r-value.
-  SmallVectorMemoryBuffer(SmallVectorImpl<char> &&SV)
-      : SV(std::move(SV)), BufferName("<in-memory object>") {
-    init(this->SV.begin(), this->SV.end(), false);
-  }
-
-  /// Construct a named SmallVectorMemoryBuffer from the given
-  /// SmallVector r-value and StringRef.
-  SmallVectorMemoryBuffer(SmallVectorImpl<char> &&SV, StringRef Name)
+  /// Construct a SmallVectorMemoryBuffer from the given SmallVector r-value.
+  SmallVectorMemoryBuffer(SmallVectorImpl<char> &&SV,
+                          bool RequiresNullTerminator = true)
+      : SmallVectorMemoryBuffer(std::move(SV), "<in-memory object>",
+                                RequiresNullTerminator) {}
+
+  /// Construct a named SmallVectorMemoryBuffer from the given SmallVector
+  /// r-value and StringRef.
+  SmallVectorMemoryBuffer(SmallVectorImpl<char> &&SV, StringRef Name,
+                          bool RequiresNullTerminator = true)
       : SV(std::move(SV)), BufferName(std::string(Name)) {
+    if (RequiresNullTerminator) {
+      this->SV.push_back('\0');
+      this->SV.pop_back();
+    }
     init(this->SV.begin(), this->SV.end(), false);
   }
 
index 200f42a..ed91228 100644 (file)
@@ -170,8 +170,8 @@ std::unique_ptr<MemoryBuffer> MCJIT::emitObject(Module *M) {
   PM.run(*M);
   // Flush the output buffer to get the generated code into memory
 
-  std::unique_ptr<MemoryBuffer> CompiledObjBuffer(
-      new SmallVectorMemoryBuffer(std::move(ObjBufferSV)));
+  auto CompiledObjBuffer = std::make_unique<SmallVectorMemoryBuffer>(
+      std::move(ObjBufferSV), /*RequiresNullTerminator=*/false);
 
   // If we have an object cache, tell it about the new object.
   // Note that we're using the compiled image, not the loaded image (as below).
index f8efed1..f342470 100644 (file)
@@ -53,7 +53,8 @@ Expected<SimpleCompiler::CompileResult> SimpleCompiler::operator()(Module &M) {
   }
 
   auto ObjBuffer = std::make_unique<SmallVectorMemoryBuffer>(
-      std::move(ObjBufferSV), M.getModuleIdentifier() + "-jitted-objectbuffer");
+      std::move(ObjBufferSV), M.getModuleIdentifier() + "-jitted-objectbuffer",
+      /*RequiresNullTerminator=*/false);
 
   auto Obj = object::ObjectFile::createObjectFile(ObjBuffer->getMemBufferRef());
 
index 9474d8c..9aea27f 100644 (file)
@@ -378,7 +378,8 @@ std::unique_ptr<MemoryBuffer> codegenModule(Module &TheModule,
     // Run codegen now. resulting binary is in OutputBuffer.
     PM.run(TheModule);
   }
-  return std::make_unique<SmallVectorMemoryBuffer>(std::move(OutputBuffer));
+  return std::make_unique<SmallVectorMemoryBuffer>(
+      std::move(OutputBuffer), /*RequiresNullTerminator=*/false);
 }
 
 /// Manage caching for a single Module.
@@ -541,7 +542,8 @@ ProcessThinLTOModule(Module &TheModule, ModuleSummaryIndex &Index,
       auto Index = buildModuleSummaryIndex(TheModule, nullptr, &PSI);
       WriteBitcodeToFile(TheModule, OS, true, &Index);
     }
-    return std::make_unique<SmallVectorMemoryBuffer>(std::move(OutputBuffer));
+    return std::make_unique<SmallVectorMemoryBuffer>(
+        std::move(OutputBuffer), /*RequiresNullTerminator=*/false);
   }
 
   return codegenModule(TheModule, TM);
index ce99746..da8bcec 100644 (file)
@@ -696,7 +696,7 @@ writeArchiveToBuffer(ArrayRef<NewArchiveMember> NewMembers, bool WriteSymtab,
     return std::move(E);
 
   return std::make_unique<SmallVectorMemoryBuffer>(
-      std::move(ArchiveBufferVector));
+      std::move(ArchiveBufferVector), /*RequiresNullTerminator=*/false);
 }
 
 } // namespace llvm
index 9e7b91d..915394b 100644 (file)
@@ -476,9 +476,8 @@ Error objcopy::macho::executeObjcopyOnMachOUniversalBinary(
                                          **ObjOrErr, MemStream))
       return E;
 
-    std::unique_ptr<MemoryBuffer> MB =
-        std::make_unique<SmallVectorMemoryBuffer>(std::move(Buffer),
-                                                  ArchFlagName);
+    auto MB = std::make_unique<SmallVectorMemoryBuffer>(
+        std::move(Buffer), ArchFlagName, /*RequiresNullTerminator=*/false);
     Expected<std::unique_ptr<Binary>> BinaryOrErr = object::createBinary(*MB);
     if (!BinaryOrErr)
       return BinaryOrErr.takeError();
index ad16648..a596398 100644 (file)
@@ -236,7 +236,8 @@ createNewArchiveMembers(const MultiFormatConfig &Config, const Archive &Ar) {
       return createFileError(Ar.getFileName(), Member.takeError());
 
     Member->Buf = std::make_unique<SmallVectorMemoryBuffer>(
-        std::move(Buffer), ChildNameOrErr.get());
+        std::move(Buffer), ChildNameOrErr.get(),
+        /*RequiresNullTerminator=*/false);
     Member->MemberName = Member->Buf->getBufferIdentifier();
     NewArchiveMembers.push_back(std::move(*Member));
   }
index 97bb62d..c3e7b3c 100644 (file)
@@ -11,6 +11,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Support/MemoryBuffer.h"
+#include "llvm/Support/SmallVectorMemoryBuffer.h"
 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/FileUtilities.h"
@@ -408,4 +409,27 @@ TEST_F(MemoryBufferTest, mmapVolatileNoNull) {
   EXPECT_EQ(MB->getBufferSize(), std::size_t(FileWrites * 8));
   EXPECT_TRUE(MB->getBuffer().startswith("01234567"));
 }
+
+// Test that SmallVector without a null terminator gets one.
+TEST(SmallVectorMemoryBufferTest, WithoutNullTerminatorRequiresNullTerminator) {
+  SmallString<0> Data("some data");
+
+  SmallVectorMemoryBuffer MB(std::move(Data),
+                             /*RequiresNullTerminator=*/true);
+  EXPECT_EQ(MB.getBufferSize(), 9u);
+  EXPECT_EQ(MB.getBufferEnd()[0], '\0');
+}
+
+// Test that SmallVector with a null terminator keeps it.
+TEST(SmallVectorMemoryBufferTest, WithNullTerminatorRequiresNullTerminator) {
+  SmallString<0> Data("some data");
+  Data.push_back('\0');
+  Data.pop_back();
+
+  SmallVectorMemoryBuffer MB(std::move(Data),
+                             /*RequiresNullTerminator=*/true);
+  EXPECT_EQ(MB.getBufferSize(), 9u);
+  EXPECT_EQ(MB.getBufferEnd()[0], '\0');
 }
+
+} // namespace