2 files changed, 37 insertions, 18 deletions

diff --git a/lld/ELF/InputFiles.cpp b/lld/ELF/InputFiles.cpp
index 7292f4a998c..308ff381316 100644
--- a/lld/ELF/InputFiles.cpp
+++ b/lld/ELF/InputFiles.cpp
@@ -202,31 +202,30 @@ void elf::ObjectFile<ELFT>::initializeSections(
       break;
     case SHT_RELA:
     case SHT_REL: {
-      uint32_t RelocatedSectionIndex = Sec.sh_info;
-      if (RelocatedSectionIndex >= Size)
-        fatal("invalid relocated section index");
-      InputSectionBase<ELFT> *RelocatedSection =
-          Sections[RelocatedSectionIndex];
-      // Strictly speaking, a relocation section must be included in the
-      // group of the section it relocates. However, LLVM 3.3 and earlier
-      // would fail to do so, so we gracefully handle that case.
-      if (RelocatedSection == InputSection<ELFT>::Discarded)
-        continue;
-      if (!RelocatedSection)
-        fatal("unsupported relocation reference");
+      // This section contains relocation information.
+      // If -r is given, we do not interpret or apply relocation
+      // but just copy relocation sections to output.
       if (Config->Relocatable) {
-        // For -r, relocation sections are handled as regular input sections.
         Sections[I] = new (Alloc) InputSection<ELFT>(this, &Sec);
-      } else if (auto *S = dyn_cast<InputSection<ELFT>>(RelocatedSection)) {
+        break;
+      }
+
+      // Find the relocation target section and associate this
+      // section with it.
+      InputSectionBase<ELFT> *Target = getRelocTarget(Sec);
+      if (!Target)
+        break;
+      if (auto *S = dyn_cast<InputSection<ELFT>>(Target)) {
         S->RelocSections.push_back(&Sec);
-      } else if (auto *S = dyn_cast<EHInputSection<ELFT>>(RelocatedSection)) {
+        break;
+      }
+      if (auto *S = dyn_cast<EHInputSection<ELFT>>(Target)) {
         if (S->RelocSection)
           fatal("multiple relocation sections to .eh_frame are not supported");
         S->RelocSection = &Sec;
-      } else {
-        fatal("relocations pointing to SHF_MERGE are not supported");
+        break;
       }
-      break;
+      fatal("relocations pointing to SHF_MERGE are not supported");
     }
     default:
       Sections[I] = createInputSection(Sec);
@@ -236,6 +235,25 @@ void elf::ObjectFile<ELFT>::initializeSections(
 
 template <class ELFT>
 InputSectionBase<ELFT> *
+elf::ObjectFile<ELFT>::getRelocTarget(const Elf_Shdr &Sec) {
+  uint32_t Idx = Sec.sh_info;
+  if (Idx >= Sections.size())
+    fatal("invalid relocated section index");
+  InputSectionBase<ELFT> *Target = Sections[Idx];
+
+  // Strictly speaking, a relocation section must be included in the
+  // group of the section it relocates. However, LLVM 3.3 and earlier
+  // would fail to do so, so we gracefully handle that case.
+  if (Target == InputSection<ELFT>::Discarded)
+    return nullptr;
+
+  if (!Target)
+    fatal("unsupported relocation reference");
+  return Target;
+}
+
+template <class ELFT>
+InputSectionBase<ELFT> *
 elf::ObjectFile<ELFT>::createInputSection(const Elf_Shdr &Sec) {
   StringRef Name = check(this->ELFObj.getSectionName(&Sec));
 
diff --git a/lld/ELF/InputFiles.h b/lld/ELF/InputFiles.h
index 5cd9df16ccc..82c06b81ce5 100644
--- a/lld/ELF/InputFiles.h
+++ b/lld/ELF/InputFiles.h
@@ -139,6 +139,7 @@ public:
 private:
   void initializeSections(llvm::DenseSet<StringRef> &ComdatGroups);
   void initializeSymbols();
+  InputSectionBase<ELFT> *getRelocTarget(const Elf_Shdr &Sec);
   InputSectionBase<ELFT> *createInputSection(const Elf_Shdr &Sec);
 
   SymbolBody *createSymbolBody(const Elf_Sym *Sym);