3 files changed, 32 insertions, 16 deletions

diff --git a/llvm/tools/llvm-objcopy/ELF/ELFObjcopy.cpp b/llvm/tools/llvm-objcopy/ELF/ELFObjcopy.cpp
index 6f21db9db60..83f1f0e4743 100644
--- a/llvm/tools/llvm-objcopy/ELF/ELFObjcopy.cpp
+++ b/llvm/tools/llvm-objcopy/ELF/ELFObjcopy.cpp
@@ -232,27 +232,24 @@ static void replaceDebugSections(
     const CopyConfig &Config, Object &Obj, SectionPred &RemovePred,
     function_ref<bool(const SectionBase &)> shouldReplace,
     function_ref<SectionBase *(const SectionBase *)> addSection) {
+  // Build a list of the debug sections we are going to replace.
+  // We can't call `addSection` while iterating over sections,
+  // because it would mutate the sections array.
   SmallVector<SectionBase *, 13> ToReplace;
-  SmallVector<RelocationSection *, 13> RelocationSections;
-  for (auto &Sec : Obj.sections()) {
-    if (RelocationSection *R = dyn_cast<RelocationSection>(&Sec)) {
-      if (shouldReplace(*R->getSection()))
-        RelocationSections.push_back(R);
-      continue;
-    }
-
+  for (auto &Sec : Obj.sections())
     if (shouldReplace(Sec))
       ToReplace.push_back(&Sec);
-  }
 
-  for (SectionBase *S : ToReplace) {
-    SectionBase *NewSection = addSection(S);
+  // Build a mapping from original section to a new one.
+  DenseMap<SectionBase *, SectionBase *> FromTo;
+  for (SectionBase *S : ToReplace)
+    FromTo[S] = addSection(S);
 
-    for (RelocationSection *RS : RelocationSections) {
-      if (RS->getSection() == S)
-        RS->setSection(NewSection);
-    }
-  }
+  // Now we want to update the target sections of relocation
+  // sections. Also we will update the relocations themselves
+  // to update the symbol references.
+  for (auto &Sec : Obj.sections())
+    Sec.replaceSectionReferences(FromTo);
 
   RemovePred = [shouldReplace, RemovePred](const SectionBase &Sec) {
     return shouldReplace(Sec) || RemovePred(Sec);
diff --git a/llvm/tools/llvm-objcopy/ELF/Object.cpp b/llvm/tools/llvm-objcopy/ELF/Object.cpp
index ead2a889438..d859c7e47a6 100644
--- a/llvm/tools/llvm-objcopy/ELF/Object.cpp
+++ b/llvm/tools/llvm-objcopy/ELF/Object.cpp
@@ -62,6 +62,8 @@ Error SectionBase::removeSymbols(function_ref<bool(const Symbol &)> ToRemove) {
 void SectionBase::initialize(SectionTableRef SecTable) {}
 void SectionBase::finalize() {}
 void SectionBase::markSymbols() {}
+void SectionBase::replaceSectionReferences(
+    const DenseMap<SectionBase *, SectionBase *> &) {}
 
 template <class ELFT> void ELFWriter<ELFT>::writeShdr(const SectionBase &Sec) {
   uint8_t *B = Buf.getBufferStart();
@@ -634,6 +636,19 @@ void RelocationSection::markSymbols() {
     Reloc.RelocSymbol->Referenced = true;
 }
 
+void RelocationSection::replaceSectionReferences(
+    const DenseMap<SectionBase *, SectionBase *> &FromTo) {
+  // Update the target section if it was replaced.
+  if (SectionBase *To = FromTo.lookup(SecToApplyRel))
+    SecToApplyRel = To;
+
+  // Change the sections where symbols are defined in if their
+  // original sections were replaced.
+  for (const Relocation &R : Relocations)
+    if (SectionBase *To = FromTo.lookup(R.RelocSymbol->DefinedIn))
+      R.RelocSymbol->DefinedIn = To;
+}
+
 void SectionWriter::visit(const DynamicRelocationSection &Sec) {
   llvm::copy(Sec.Contents,
             Out.getBufferStart() + Sec.Offset);
diff --git a/llvm/tools/llvm-objcopy/ELF/Object.h b/llvm/tools/llvm-objcopy/ELF/Object.h
index 2c6c59065f2..dabec852321 100644
--- a/llvm/tools/llvm-objcopy/ELF/Object.h
+++ b/llvm/tools/llvm-objcopy/ELF/Object.h
@@ -280,6 +280,8 @@ public:
   virtual void accept(SectionVisitor &Visitor) const = 0;
   virtual void accept(MutableSectionVisitor &Visitor) = 0;
   virtual void markSymbols();
+  virtual void
+  replaceSectionReferences(const DenseMap<SectionBase *, SectionBase *> &);
 };
 
 class Segment {
@@ -596,6 +598,8 @@ public:
       function_ref<bool(const SectionBase *)> ToRemove) override;
   Error removeSymbols(function_ref<bool(const Symbol &)> ToRemove) override;
   void markSymbols() override;
+  void replaceSectionReferences(
+      const DenseMap<SectionBase *, SectionBase *> &FromTo) override;
 
   static bool classof(const SectionBase *S) {
     if (S->Flags & ELF::SHF_ALLOC)