[llvm-objcopy][MachO] Fix --add-section

This diff fixes --add-section functionality and simplifies the tests organization.

Test plan: make check-all

Differential revision: https://reviews.llvm.org/D87497
diff --git a/llvm/tools/llvm-objcopy/MachO/MachOObjcopy.cpp b/llvm/tools/llvm-objcopy/MachO/MachOObjcopy.cpp
index fff3f01..e87487e 100644
--- a/llvm/tools/llvm-objcopy/MachO/MachOObjcopy.cpp
+++ b/llvm/tools/llvm-objcopy/MachO/MachOObjcopy.cpp
@@ -259,7 +259,11 @@
   for (LoadCommand &LC : Obj.LoadCommands) {
     Optional<StringRef> SegName = LC.getSegmentName();
     if (SegName && SegName == TargetSegName) {
+      uint64_t Addr = *LC.getSegmentVMAddr();
+      for (const std::unique_ptr<Section> &S : LC.Sections)
+        Addr = std::max(Addr, S->Addr + S->Size);
       LC.Sections.push_back(std::make_unique<Section>(Sec));
+      LC.Sections.back()->Addr = Addr;
       return Error::success();
     }
   }
@@ -268,6 +272,7 @@
   // Insert a new section into it.
   LoadCommand &NewSegment = Obj.addSegment(TargetSegName);
   NewSegment.Sections.push_back(std::make_unique<Section>(Sec));
+  NewSegment.Sections.back()->Addr = *NewSegment.getSegmentVMAddr();
   return Error::success();
 }
 
diff --git a/llvm/tools/llvm-objcopy/MachO/Object.cpp b/llvm/tools/llvm-objcopy/MachO/Object.cpp
index 0359e2a..6a89076 100644
--- a/llvm/tools/llvm-objcopy/MachO/Object.cpp
+++ b/llvm/tools/llvm-objcopy/MachO/Object.cpp
@@ -111,24 +111,53 @@
   return Error::success();
 }
 
+uint64_t Object::nextAvailableSegmentAddress() const {
+  uint64_t HeaderSize =
+      is64Bit() ? sizeof(MachO::mach_header_64) : sizeof(MachO::mach_header);
+  uint64_t Addr = HeaderSize + Header.SizeOfCmds;
+  for (const LoadCommand &LC : LoadCommands) {
+    const MachO::macho_load_command &MLC = LC.MachOLoadCommand;
+    switch (MLC.load_command_data.cmd) {
+    case MachO::LC_SEGMENT:
+      Addr = std::max(Addr,
+                      static_cast<uint64_t>(MLC.segment_command_data.vmaddr) +
+                          MLC.segment_command_data.vmsize);
+      break;
+    case MachO::LC_SEGMENT_64:
+      Addr = std::max(Addr, MLC.segment_command_64_data.vmaddr +
+                                MLC.segment_command_64_data.vmsize);
+      break;
+    default:
+      continue;
+    }
+  }
+  return Addr;
+}
+
 template <typename SegmentType>
 static void constructSegment(SegmentType &Seg,
                              llvm::MachO::LoadCommandType CmdType,
-                             StringRef SegName) {
+                             StringRef SegName, uint64_t SegVMAddr) {
   assert(SegName.size() <= sizeof(Seg.segname) && "too long segment name");
   memset(&Seg, 0, sizeof(SegmentType));
   Seg.cmd = CmdType;
   strncpy(Seg.segname, SegName.data(), SegName.size());
+  Seg.maxprot |=
+      (MachO::VM_PROT_READ | MachO::VM_PROT_WRITE | MachO::VM_PROT_EXECUTE);
+  Seg.initprot |=
+      (MachO::VM_PROT_READ | MachO::VM_PROT_WRITE | MachO::VM_PROT_EXECUTE);
+  Seg.vmaddr = SegVMAddr;
 }
 
 LoadCommand &Object::addSegment(StringRef SegName) {
   LoadCommand LC;
+  const uint64_t SegVMAddr = nextAvailableSegmentAddress();
   if (is64Bit())
     constructSegment(LC.MachOLoadCommand.segment_command_64_data,
-                     MachO::LC_SEGMENT_64, SegName);
+                     MachO::LC_SEGMENT_64, SegName, SegVMAddr);
   else
     constructSegment(LC.MachOLoadCommand.segment_command_data,
-                     MachO::LC_SEGMENT, SegName);
+                     MachO::LC_SEGMENT, SegName, SegVMAddr);
 
   LoadCommands.push_back(std::move(LC));
   return LoadCommands.back();
@@ -152,6 +181,18 @@
   }
 }
 
+Optional<uint64_t> LoadCommand::getSegmentVMAddr() const {
+  const MachO::macho_load_command &MLC = MachOLoadCommand;
+  switch (MLC.load_command_data.cmd) {
+  case MachO::LC_SEGMENT:
+    return MLC.segment_command_data.vmaddr;
+  case MachO::LC_SEGMENT_64:
+    return MLC.segment_command_64_data.vmaddr;
+  default:
+    return None;
+  }
+}
+
 } // end namespace macho
 } // end namespace objcopy
 } // end namespace llvm
diff --git a/llvm/tools/llvm-objcopy/MachO/Object.h b/llvm/tools/llvm-objcopy/MachO/Object.h
index 102aeed..0a930ba 100644
--- a/llvm/tools/llvm-objcopy/MachO/Object.h
+++ b/llvm/tools/llvm-objcopy/MachO/Object.h
@@ -94,6 +94,9 @@
 
   // Returns the segment name if the load command is a segment command.
   Optional<StringRef> getSegmentName() const;
+
+  // Returns the segment vm address if the load command is a segment command.
+  Optional<uint64_t> getSegmentVMAddr() const;
 };
 
 // A symbol information. Fields which starts with "n_" are same as them in the
@@ -340,6 +343,8 @@
     return Header.Magic == MachO::MH_MAGIC_64 ||
            Header.Magic == MachO::MH_CIGAM_64;
   }
+
+  uint64_t nextAvailableSegmentAddress() const;
 };
 
 } // end namespace macho