1 files changed, 88 insertions, 6 deletions

diff --git a/llvm/tools/llvm-objcopy/MachO/Object.cpp b/llvm/tools/llvm-objcopy/MachO/Object.cpp
index d3b4fdc2f6338..de8cb0af108d7 100644
--- a/llvm/tools/llvm-objcopy/MachO/Object.cpp
+++ b/llvm/tools/llvm-objcopy/MachO/Object.cpp
@@ -1,5 +1,15 @@
+//===- Object.cpp - Mach-O object file model --------------------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
 #include "Object.h"
 #include "../llvm-objcopy.h"
+#include "llvm/ADT/SmallPtrSet.h"
+#include <unordered_set>
 
 namespace llvm {
 namespace objcopy {
@@ -22,11 +32,83 @@ void SymbolTable::removeSymbols(
       std::end(Symbols));
 }
 
-void Object::removeSections(function_ref<bool(const Section &)> ToRemove) {
-  for (LoadCommand &LC : LoadCommands)
-    LC.Sections.erase(std::remove_if(std::begin(LC.Sections),
-                                     std::end(LC.Sections), ToRemove),
-                      std::end(LC.Sections));
+void Object::updateLoadCommandIndexes() {
+  // Update indices of special load commands
+  for (size_t Index = 0, Size = LoadCommands.size(); Index < Size; ++Index) {
+    LoadCommand &LC = LoadCommands[Index];
+    switch (LC.MachOLoadCommand.load_command_data.cmd) {
+    case MachO::LC_SYMTAB:
+      SymTabCommandIndex = Index;
+      break;
+    case MachO::LC_DYSYMTAB:
+      DySymTabCommandIndex = Index;
+      break;
+    case MachO::LC_DYLD_INFO:
+    case MachO::LC_DYLD_INFO_ONLY:
+      DyLdInfoCommandIndex = Index;
+      break;
+    case MachO::LC_DATA_IN_CODE:
+      DataInCodeCommandIndex = Index;
+      break;
+    case MachO::LC_FUNCTION_STARTS:
+      FunctionStartsCommandIndex = Index;
+      break;
+    }
+  }
+}
+
+Error Object::removeLoadCommands(
+    function_ref<bool(const LoadCommand &)> ToRemove) {
+  auto It = std::stable_partition(
+      LoadCommands.begin(), LoadCommands.end(),
+      [&](const LoadCommand &LC) { return !ToRemove(LC); });
+  LoadCommands.erase(It, LoadCommands.end());
+
+  updateLoadCommandIndexes();
+  return Error::success();
+}
+
+Error Object::removeSections(
+    function_ref<bool(const std::unique_ptr<Section> &)> ToRemove) {
+  DenseMap<uint32_t, const Section *> OldIndexToSection;
+  uint32_t NextSectionIndex = 1;
+  for (LoadCommand &LC : LoadCommands) {
+    auto It = std::stable_partition(
+        std::begin(LC.Sections), std::end(LC.Sections),
+        [&](const std::unique_ptr<Section> &Sec) { return !ToRemove(Sec); });
+    for (auto I = LC.Sections.begin(), End = It; I != End; ++I) {
+      OldIndexToSection[(*I)->Index] = I->get();
+      (*I)->Index = NextSectionIndex++;
+    }
+    LC.Sections.erase(It, LC.Sections.end());
+  }
+
+  auto IsDead = [&](const std::unique_ptr<SymbolEntry> &S) -> bool {
+    Optional<uint32_t> Section = S->section();
+    return (Section && !OldIndexToSection.count(*Section));
+  };
+
+  SmallPtrSet<const SymbolEntry *, 2> DeadSymbols;
+  for (const std::unique_ptr<SymbolEntry> &Sym : SymTable.Symbols)
+    if (IsDead(Sym))
+      DeadSymbols.insert(Sym.get());
+
+  for (const LoadCommand &LC : LoadCommands)
+    for (const std::unique_ptr<Section> &Sec : LC.Sections)
+      for (const RelocationInfo &R : Sec->Relocations)
+        if (R.Symbol && *R.Symbol && DeadSymbols.count(*R.Symbol))
+          return createStringError(std::errc::invalid_argument,
+                                   "symbol '%s' defined in section with index "
+                                   "'%u' cannot be removed because it is "
+                                   "referenced by a relocation in section '%s'",
+                                   (*R.Symbol)->Name.c_str(),
+                                   *((*R.Symbol)->section()),
+                                   Sec->CanonicalName.c_str());
+  SymTable.removeSymbols(IsDead);
+  for (std::unique_ptr<SymbolEntry> &S : SymTable.Symbols)
+    if (S->section())
+      S->n_sect = OldIndexToSection[S->n_sect]->Index;
+  return Error::success();
 }
 
 void Object::addLoadCommand(LoadCommand LC) {
@@ -52,7 +134,7 @@ LoadCommand &Object::addSegment(StringRef SegName) {
     constructSegment(LC.MachOLoadCommand.segment_command_data,
                      MachO::LC_SEGMENT, SegName);
 
-  LoadCommands.push_back(LC);
+  LoadCommands.push_back(std::move(LC));
   return LoadCommands.back();
 }