[Object][MachO] Refactor MachOUniversalWriter

[alexander-shaposhnikov]

Summary

This diff refactors writeUniversalBinary and adds writeUniversalBinaryToBuffer.
This is a preparation for adding support for universal binaries to llvm-objcopy.
The idea behind this diff is similar to D84858.

Test plan: make check-all

[MaskRay]

LGTM. It may be good to have the dependent patch posted to demonstrate the refactoring.

[smeenai]

LGTM

llvm/include/llvm/Object/MachOUniversalWriter.h

if (Rhs.CPUType == MachO::CPU_TYPE_ARM64)
       return true;
     // Sort by alignment to minimize file size
     return Lhs.P2Alignment < Rhs.P2Alignment;
   }
 };
 
 Error writeUniversalBinary(ArrayRef<Slice> Slices, StringRef OutputFileName);
 
+Expected<std::unique_ptr<MemoryBuffer>>
+writeUniversalBinaryToBuffer(ArrayRef<Slice> Slices);
+
 } // end namespace object
 
 } // end namespace llvm
 
 #endif // LLVM_OBJECT_MACHOUNIVERSALWRITER_H

llvm/lib/Object/MachOUniversalWriter.cpp

 #include "llvm/Object/MachOUniversalWriter.h"
 #include "llvm/ADT/Triple.h"
 #include "llvm/Object/Archive.h"
 #include "llvm/Object/Binary.h"
 #include "llvm/Object/Error.h"
 #include "llvm/Object/IRObjectFile.h"
 #include "llvm/Object/MachO.h"
 #include "llvm/Object/MachOUniversal.h"
-#include "llvm/Support/FileOutputBuffer.h"
+#include "llvm/Support/SmallVectorMemoryBuffer.h"
 
 using namespace llvm;
 using namespace object;
 
 // For compatibility with cctools lipo, a file's alignment is calculated as the
 // minimum aligment of all segments. For object files, the file's alignment is
 // the maximum alignment of its sections.
 static uint32_t calculateFileAlignment(const MachOObjectFile &O) {
for (const auto &S : Slices) {
     FatArch.size = S.getBinary()->getMemoryBufferRef().getBufferSize();
     FatArch.align = S.getP2Alignment();
     Offset += FatArch.size;
     FatArchList.push_back(FatArch);
   }
   return FatArchList;
 }
 
-Error object::writeUniversalBinary(ArrayRef<Slice> Slices,
-                                   StringRef OutputFileName) {
+static Error writeUniversalBinaryToStream(ArrayRef<Slice> Slices,
+                                          raw_ostream &Out) {
   MachO::fat_header FatHeader;
   FatHeader.magic = MachO::FAT_MAGIC;
   FatHeader.nfat_arch = Slices.size();
 
   Expected<SmallVector<MachO::fat_arch, 2>> FatArchListOrErr =
       buildFatArchList(Slices);
   if (!FatArchListOrErr)
     return FatArchListOrErr.takeError();
   SmallVector<MachO::fat_arch, 2> FatArchList = *FatArchListOrErr;
 
-  const bool IsExecutable = any_of(Slices, [](Slice S) {
-    return sys::fs::can_execute(S.getBinary()->getFileName());
-  });
-  const uint64_t OutputFileSize =
-      static_cast<uint64_t>(FatArchList.back().offset) +
-      FatArchList.back().size;
-  Expected<std::unique_ptr<FileOutputBuffer>> OutFileOrError =
-      FileOutputBuffer::create(OutputFileName, OutputFileSize,
-                               IsExecutable ? FileOutputBuffer::F_executable
-                                            : 0);
-  if (!OutFileOrError)
-    return createFileError(OutputFileName, OutFileOrError.takeError());
-  std::unique_ptr<FileOutputBuffer> OutFile = std::move(OutFileOrError.get());
-  std::memset(OutFile->getBufferStart(), 0, OutputFileSize);
-
   if (sys::IsLittleEndianHost)
     MachO::swapStruct(FatHeader);
-  std::memcpy(OutFile->getBufferStart(), &FatHeader, sizeof(MachO::fat_header));
+  Out.write(reinterpret_cast<const char *>(&FatHeader),
+            sizeof(MachO::fat_header));
 
-  for (size_t Index = 0, Size = Slices.size(); Index < Size; ++Index) {
-    MemoryBufferRef BufferRef = Slices[Index].getBinary()->getMemoryBufferRef();
-    std::copy(BufferRef.getBufferStart(), BufferRef.getBufferEnd(),
-              OutFile->getBufferStart() + FatArchList[Index].offset);
-  }
-
-  // FatArchs written after Slices in order to reduce the number of swaps for
-  // the LittleEndian case
   if (sys::IsLittleEndianHost)
     for (MachO::fat_arch &FA : FatArchList)
       MachO::swapStruct(FA);
-  std::memcpy(OutFile->getBufferStart() + sizeof(MachO::fat_header),
-              FatArchList.begin(),
-              sizeof(MachO::fat_arch) * FatArchList.size());
+  Out.write(reinterpret_cast<const char *>(FatArchList.data()),
 
-  if (Error E = OutFile->commit())
-    return createFileError(OutputFileName, std::move(E));
+  if (sys::IsLittleEndianHost)
+    for (MachO::fat_arch &FA : FatArchList)
+      MachO::swapStruct(FA);
+
+  size_t Offset =
+      sizeof(MachO::fat_header) + sizeof(MachO::fat_arch) * FatArchList.size();
+  for (size_t Index = 0, Size = Slices.size(); Index < Size; ++Index) {
+    MemoryBufferRef BufferRef = Slices[Index].getBinary()->getMemoryBufferRef();
+    assert((Offset <= FatArchList[Index].offset) && "Incorrect slice offset");
+    Out.write_zeros(FatArchList[Index].offset - Offset);
+    Out.write(BufferRef.getBufferStart(), BufferRef.getBufferSize());
+    Offset = FatArchList[Index].offset + BufferRef.getBufferSize();
+  }
 
+  Out.flush();
   return Error::success();
 }
+
+Error object::writeUniversalBinary(ArrayRef<Slice> Slices,
+                                   StringRef OutputFileName) {
+  const bool IsExecutable = any_of(Slices, [](Slice S) {
+    return sys::fs::can_execute(S.getBinary()->getFileName());
+  });
+  unsigned Mode = sys::fs::all_read | sys::fs::all_write;
+  if (IsExecutable)
+    Mode |= sys::fs::all_exe;
+  Expected<sys::fs::TempFile> Temp = sys::fs::TempFile::create(
+      OutputFileName + ".temp-universal-%%%%%%", Mode);
+  if (!Temp)
+    return Temp.takeError();
+  raw_fd_ostream Out(Temp->FD, false);
+  if (Error E = writeUniversalBinaryToStream(Slices, Out)) {
+    if (Error DiscardError = Temp->discard())
+      return joinErrors(std::move(E), std::move(DiscardError));
+    return E;
+  }
+  return Temp->keep(OutputFileName);
+}
+
+Expected<std::unique_ptr<MemoryBuffer>>
+object::writeUniversalBinaryToBuffer(ArrayRef<Slice> Slices) {
+  SmallVector<char, 0> Buffer;
+  raw_svector_ostream Out(Buffer);
+
+  if (Error E = writeUniversalBinaryToStream(Slices, Out))
+    return std::move(E);
+
+  return std::make_unique<SmallVectorMemoryBuffer>(std::move(Buffer));
+}