[llvm-branch-commits] [llvm] [llvm-objcopy] Stream ELF output to reduce peak memory usage (PR #217706)

[James Henderson via llvm-branch-commits]

================
@@ -33,10 +36,87 @@ using namespace llvm::objcopy::elf;
 using namespace llvm::object;
 using namespace llvm::support;
 
-template <class ELFT> void ELFWriter<ELFT>::writePhdr(const Segment &Seg) {
-  uint8_t *B = reinterpret_cast<uint8_t *>(Buf->getBufferStart()) +
-               Obj.ProgramHdrSegment.Offset + Seg.Index * sizeof(Elf_Phdr);
-  Elf_Phdr &Phdr = *reinterpret_cast<Elf_Phdr *>(B);
+namespace {
+
+template <class T> void writeObject(raw_ostream &Out, const T &Value) {
+  Out.write(reinterpret_cast<const char *>(&Value), sizeof(Value));
+}
+
+} // namespace
+
+Error ELFWriterOutput::reserve(uint64_t Size) {
+  auto *FD = dyn_cast<raw_fd_stream>(&Dest);
+  if (FD && FD->supportsSeeking() && FD->isRegularFile()) {
+    Start = FD->tell();
+    std::optional<uint64_t> Last = checkedAddUnsigned(Start, Size);
+    if (!Last)
+      return createStringError(errc::file_too_large,
+                               "output exceeds the addressable file range");
+    // Discard existing contents after the current position, then extend the
+    // file so unwritten ranges read as zero without materializing them.
+    if (Error E = FD->resize(Start))
+      return E;
+    if (Error E = FD->resize(*Last))
+      return E;
+    Stream = FD;
+    End = *Last;
+    return Error::success();
+  }
+
+  Buf = WritableMemoryBuffer::getNewMemBuffer(Size);
+  if (!Buf)
+    return createStringError(errc::not_enough_memory,
+                             "failed to allocate memory buffer of " +
+                                 Twine::utohexstr(Size) + " bytes");
+  BufStream.reset({Buf->getBufferStart(), Buf->getBufferSize()});
+  return Error::success();
+}
+
+raw_ostream &ELFWriterOutput::streamAt(uint64_t Offset) {
+  if (!Stream) {
+    BufStream.seek(Offset);
+    return BufStream;
+  }
+  if (!Stream->has_error())
+    Stream->seek(Start + Offset);
+  return *Stream;
+}
+
+void ELFWriterOutput::write(ArrayRef<uint8_t> Data, uint64_t Offset) {
+  if (!Data.empty())
+    streamAt(Offset).write(reinterpret_cast<const char *>(Data.data()),
+                           Data.size());
+}
+
+void ELFWriterOutput::writeZeros(uint64_t Offset, uint64_t Size) {
+  if (!Size)
+    return;
+  // raw_ostream::write_zeros writes in 16-byte units, which is too small for
+  // section-sized runs. Fill from a scratch block instead.
+  SmallVector<char, 0> Zeros(std::min<uint64_t>(Size, 64 * 1024));
+  raw_ostream &Out = streamAt(Offset);
+  while (Size) {
+    size_t Chunk = std::min<uint64_t>(Size, Zeros.size());
+    Out.write(Zeros.data(), Chunk);
+    Size -= Chunk;
+  }
+}
+
+Error ELFWriterOutput::finalize() {
+  if (!Stream) {
+    Dest.write(Buf->getBufferStart(), Buf->getBufferSize());
+    return Error::success();
+  }
+  // Leave the stream positioned past the object, as a writer that appended to
+  // it would.
+  if (!Stream->has_error())
+    Stream->seek(End);
+  return Stream->takeError();
+}
+
+template <class ELFT>
+void ELFWriter<ELFT>::writePhdr(raw_ostream &Out, const Segment &Seg) {
+  Elf_Phdr Phdr = {};
----------------
jh7370 wrote:

There's no need for the explicit instantiation of the Elf_Phdr: the struct members are all written to below.

https://github.com/llvm/llvm-project/pull/217706