[llvm] Add null-termination capability to SmallVectorMemoryBuffer

[jansvoboda11]

Summary

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 SmallVectors 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.

[dexonsmith]

LGTM, except I'd prefer the #includes be changed separately since that cleanup seems unrelated to this change.

lldb/unittests/Expression/CppModuleConfigurationTest.cpp

14 ↗

(On Diff #392716)

Please clean up the includes in a prep commit (no need for review IMO) rather than a drive-by.

clang/unittests/AST/ASTImporterTest.cpp

TEST_P(ImportSourceLocations, OverwrittenFileBuffer) {
   {
     SourceManager &FromSM = FromTU->getASTContext().getSourceManager();
     clang::FileManager &FM = FromSM.getFileManager();
     const clang::FileEntry &FE =
         *FM.getVirtualFile(Path, static_cast<off_t>(Contents.size()), 0);
 
     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.
     auto Pattern = varDecl(hasName("X"));
     VarDecl *FromD = FirstDeclMatcher<VarDecl>().match(FromTU, Pattern);
     ImportedLoc = Import(FromD, Lang_C99)->getLocation();
   }
 

clang/unittests/Frontend/TextDiagnosticTest.cpp

TEST(TextDiagnostic, ShowLine) {
   const clang::FileEntryRef fe = FileMgr.getVirtualFileRef(
       file_path,
       /*Size=*/static_cast<off_t>(main_file_contents.size()),
       /*ModificationTime=*/0);
 
   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.
   clang::FileID fid =
       SrcMgr.createFileID(fe, SourceLocation(), clang::SrcMgr::C_User);
   SrcMgr.setMainFileID(fid);
 
   // Create the source location for the test diagnostic.

llvm/include/llvm/Support/SmallVectorMemoryBuffer.h

 /// SmallVector-backed MemoryBuffer instance.
 ///
 /// This class enables efficient construction of MemoryBuffers from SmallVector
 /// instances. This is useful for MCJIT and Orc, where object files are streamed
 /// into SmallVectors, then inspected using ObjectFile (which takes a
 /// 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);
   }
 
   // Key function.
   ~SmallVectorMemoryBuffer() override;
 
   StringRef getBufferIdentifier() const override { return BufferName; }
 

llvm/lib/ExecutionEngine/MCJIT/MCJIT.cpp

std::unique_ptr<MemoryBuffer> MCJIT::emitObject(Module *M) {
   // that may be executed.
   if (TM->addPassesToEmitMC(PM, Ctx, ObjStream, !getVerifyModules()))
     report_fatal_error("Target does not support MC emission!");
 
   // Initialize passes.
   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).
   if (ObjCache) {
     // MemoryBuffer is a thin wrapper around the actual memory, so it's OK
     // to create a temporary object here and delete it after the call.
     MemoryBufferRef MB = CompiledObjBuffer->getMemBufferRef();
     ObjCache->notifyObjectCompiled(M, MB);

llvm/lib/ExecutionEngine/Orc/CompileUtils.cpp

SmallVector<char, 0> ObjBufferSV;
     MCContext *Ctx;
     if (TM.addPassesToEmitMC(PM, Ctx, ObjStream))
       return make_error<StringError>("Target does not support MC emission",
                                      inconvertibleErrorCode());
     PM.run(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());
 
   if (!Obj)
     return Obj.takeError();
 
   notifyObjectCompiled(M, *ObjBuffer);
   return std::move(ObjBuffer);

llvm/lib/LTO/ThinLTOCodeGenerator.cpp

// CodeGen
     // Setup the codegen now.
     if (TM.addPassesToEmitFile(PM, OS, nullptr, CGFT_ObjectFile,
                                /* DisableVerify */ true))
       report_fatal_error("Failed to setup codegen");
 
     // 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.
 class ModuleCacheEntry {
   SmallString<128> EntryPath;
 
 public:
   // Create a cache entry. This compute a unique hash for the Module considering
if (DisableCodeGen) {
     // Configured to stop before CodeGen, serialize the bitcode and return.
     SmallVector<char, 128> OutputBuffer;
     {
       raw_svector_ostream OS(OutputBuffer);
       ProfileSummaryInfo PSI(TheModule);
       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);
 }
 
 /// Resolve prevailing symbols. Record resolutions in the \p ResolvedODR map
 /// for caching, and in the \p Index for application during the ThinLTO
 /// backends. This is needed for correctness for exported symbols (ensure

llvm/lib/Object/ArchiveWriter.cpp

writeArchiveToBuffer(ArrayRef<NewArchiveMember> NewMembers, bool WriteSymtab,
   SmallVector<char, 0> ArchiveBufferVector;
   raw_svector_ostream ArchiveStream(ArchiveBufferVector);
 
   if (Error E = writeArchiveToStream(ArchiveStream, NewMembers, WriteSymtab,
                                      Kind, Deterministic, Thin))
     return std::move(E);
 
   return std::make_unique<SmallVectorMemoryBuffer>(
-      std::move(ArchiveBufferVector));
+      std::move(ArchiveBufferVector), /*RequiresNullTerminator=*/false);
 }
 
 } // namespace llvm

llvm/tools/llvm-objcopy/MachO/MachOObjcopy.cpp

for (const auto &O : In.objects()) {
     Expected<const MachOConfig &> MachO = Config.getMachOConfig();
     if (!MachO)
       return MachO.takeError();
 
     if (Error E = executeObjcopyOnBinary(Config.getCommonConfig(), *MachO,
                                          **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();
     Binaries.emplace_back(std::move(*BinaryOrErr), std::move(MB));
     Slices.emplace_back(*cast<MachOObjectFile>(Binaries.back().getBinary()),
                         O.getAlign());
   }
 
   if (Error Err = writeUniversalBinaryToStream(Slices, Out))
     return Err;
 
   return Error::success();
 }

llvm/tools/llvm-objcopy/llvm-objcopy.cpp

if (Error E = executeObjcopyOnBinary(Config, *ChildOrErr->get(), MemStream))
       return std::move(E);
 
     Expected<NewArchiveMember> Member = NewArchiveMember::getOldMember(
         Child, Config.getCommonConfig().DeterministicArchives);
     if (!Member)
       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));
   }
   if (Err)
     return createFileError(Config.getCommonConfig().InputFilename,
                            std::move(Err));
   return std::move(NewArchiveMembers);
 }

llvm/unittests/Support/MemoryBufferTest.cpp

 //===- llvm/unittest/Support/MemoryBufferTest.cpp - MemoryBuffer tests ----===//
 //
 // 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 unit tests for the MemoryBuffer support class.
 //
 //===----------------------------------------------------------------------===//
 
 #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"
 #include "llvm/Support/Process.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Testing/Support/Error.h"
 #include "gtest/gtest.h"
 #if LLVM_ENABLE_THREADS
TEST_F(MemoryBufferTest, mmapVolatileNoNull) {
   auto MBOrError = MemoryBuffer::getOpenFile(*File, TestPath,
       /*FileSize=*/-1, /*RequiresNullTerminator=*/false, /*IsVolatile=*/true);
   ASSERT_NO_ERROR(MBOrError.getError())
   OwningBuffer MB = std::move(*MBOrError);
   EXPECT_EQ(MB->getBufferKind(), MemoryBuffer::MemoryBuffer_MMap);
   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