[WebAssembly] Optimise relocation processing. NFC.

This is a rebased version https://reviews.llvm.org/D42176 which is patch
by Nicolas Wilson.

Addresses issue:
https://github.com/WebAssembly/tool-conventions/issues/32, and
https://bugs.llvm.org/show_bug.cgi?id=38650

Previously, for each function/segment we iterated over every relocation
to find the relevant ones, which is an n^2 operation. Now, we just make
a single pass.

Differential Revision: https://reviews.llvm.org/D51063

llvm-svn: 340428

lld/wasm/InputChunks.cpp

@@ -43,16 +43,6 @@ StringRef InputChunk::getComdatName() const {
   return File->getWasmObj()->linkingData().Comdats[Index];
 }
 
-void InputChunk::copyRelocations(const WasmSection &Section) {
-  if (Section.Relocations.empty())
-    return;
-  size_t Start = getInputSectionOffset();
-  size_t Size = getInputSize();
-  for (const WasmRelocation &R : Section.Relocations)
-    if (R.Offset >= Start && R.Offset < Start + Size)
-      Relocations.push_back(R);
-}
-
 void InputChunk::verifyRelocTargets() const {
   for (const WasmRelocation &Rel : Relocations) {
     uint32_t ExistingValue;
@@ -242,7 +232,7 @@ void InputFunction::calculateSize() {
   uint32_t End = Start + Function->Size;
 
   uint32_t LastRelocEnd = Start + FunctionSizeLength;
-  for (WasmRelocation &Rel : Relocations) {
+  for (const WasmRelocation &Rel : Relocations) {
     LLVM_DEBUG(dbgs() << "  region: " << (Rel.Offset - LastRelocEnd) << "\n");
     CompressedFuncSize += Rel.Offset - LastRelocEnd;
     CompressedFuncSize += getRelocWidth(Rel, File->calcNewValue(Rel));

lld/wasm/InputChunks.h

@@ -49,17 +49,18 @@ public:
   Kind kind() const { return SectionKind; }
 
   virtual uint32_t getSize() const { return data().size(); }
-
-  void copyRelocations(const WasmSection &Section);
+  virtual uint32_t getInputSize() const { return getSize(); };
 
   virtual void writeTo(uint8_t *SectionStart) const;
 
   ArrayRef<WasmRelocation> getRelocations() const { return Relocations; }
+  void setRelocations(ArrayRef<WasmRelocation> Rs) { Relocations = Rs; }
 
   virtual StringRef getName() const = 0;
   virtual StringRef getDebugName() const = 0;
   virtual uint32_t getComdat() const = 0;
   StringRef getComdatName() const;
+  virtual uint32_t getInputSectionOffset() const = 0;
 
   size_t NumRelocations() const { return Relocations.size(); }
   void writeRelocations(llvm::raw_ostream &OS) const;
@@ -77,14 +78,12 @@ protected:
       : File(F), Live(!Config->GcSections), SectionKind(K) {}
   virtual ~InputChunk() = default;
   virtual ArrayRef<uint8_t> data() const = 0;
-  virtual uint32_t getInputSectionOffset() const = 0;
-  virtual uint32_t getInputSize() const { return getSize(); };
 
   // Verifies the existing data at relocation targets matches our expectations.
   // This is performed only debug builds as an extra sanity check.
   void verifyRelocTargets() const;
 
-  std::vector<WasmRelocation> Relocations;
+  ArrayRef<WasmRelocation> Relocations;
   Kind SectionKind;
 };
 
@@ -107,15 +106,15 @@ public:
   StringRef getName() const override { return Segment.Data.Name; }
   StringRef getDebugName() const override { return StringRef(); }
   uint32_t getComdat() const override { return Segment.Data.Comdat; }
+  uint32_t getInputSectionOffset() const override {
+    return Segment.SectionOffset;
+  }
 
   const OutputSegment *OutputSeg = nullptr;
   int32_t OutputSegmentOffset = 0;
 
 protected:
   ArrayRef<uint8_t> data() const override { return Segment.Data.Content; }
-  uint32_t getInputSectionOffset() const override {
-    return Segment.SectionOffset;
-  }
 
   const WasmSegment &Segment;
 };
@@ -145,9 +144,13 @@ public:
     }
     return data().size();
   }
+  uint32_t getInputSize() const override { return Function->Size; }
   uint32_t getFunctionIndex() const { return FunctionIndex.getValue(); }
   bool hasFunctionIndex() const { return FunctionIndex.hasValue(); }
   void setFunctionIndex(uint32_t Index);
+  uint32_t getInputSectionOffset() const override {
+    return Function->CodeSectionOffset;
+  }
   uint32_t getTableIndex() const { return TableIndex.getValue(); }
   bool hasTableIndex() const { return TableIndex.hasValue(); }
   void setTableIndex(uint32_t Index);
@@ -167,12 +170,6 @@ protected:
                                             Function->Size);
   }
 
-  uint32_t getInputSize() const override { return Function->Size; }
-
-  uint32_t getInputSectionOffset() const override {
-    return Function->CodeSectionOffset;
-  }
-
   const WasmFunction *Function;
   llvm::Optional<uint32_t> FunctionIndex;
   llvm::Optional<uint32_t> TableIndex;

lld/wasm/InputFiles.cpp

@@ -159,6 +159,39 @@ uint32_t ObjFile::calcNewValue(const WasmRelocation &Reloc) const {
   }
 }
 
+template <class T>
+static void setRelocs(const std::vector<T *> &Chunks,
+                      const WasmSection *Section) {
+  if (!Section)
+    return;
+
+  ArrayRef<WasmRelocation> Relocs = Section->Relocations;
+  assert(std::is_sorted(
+      Relocs.begin(), Relocs.end(),
+      [](const WasmRelocation &R1, const WasmRelocation &R2) {
+        return R1.Offset < R2.Offset;
+      }));
+  assert(std::is_sorted(
+      Chunks.begin(), Chunks.end(),
+      [](InputChunk *C1, InputChunk *C2) {
+        return C1->getInputSectionOffset() < C2->getInputSectionOffset();
+      }));
+
+  auto RelocsNext = Relocs.begin();
+  auto RelocsEnd = Relocs.end();
+  auto RelocLess = [](const WasmRelocation &R, uint32_t Val) {
+    return R.Offset < Val;
+  };
+  for (InputChunk *C : Chunks) {
+    auto RelocsStart = std::lower_bound(RelocsNext, RelocsEnd,
+                                        C->getInputSectionOffset(), RelocLess);
+    RelocsNext = std::lower_bound(RelocsStart, RelocsEnd,
+                                  C->getInputSectionOffset() + C->getInputSize(),
+                                  RelocLess);
+    C->setRelocations(ArrayRef<WasmRelocation>(RelocsStart, RelocsNext));
+  }
+}
+
 void ObjFile::parse() {
   // Parse a memory buffer as a wasm file.
   LLVM_DEBUG(dbgs() << "Parsing object: " << toString(this) << "\n");
@@ -200,7 +233,7 @@ void ObjFile::parse() {
       DataSection = &Section;
     } else if (Section.Type == WASM_SEC_CUSTOM) {
       CustomSections.emplace_back(make<InputSection>(Section, this));
-      CustomSections.back()->copyRelocations(Section);
+      CustomSections.back()->setRelocations(Section.Relocations);
       CustomSectionsByIndex[SectionIndex] = CustomSections.back();
     }
     SectionIndex++;
@@ -215,11 +248,9 @@ void ObjFile::parse() {
     UsedComdats[I] = Symtab->addComdat(Comdats[I]);
 
   // Populate `Segments`.
-  for (const WasmSegment &S : WasmObj->dataSegments()) {
-    InputSegment *Seg = make<InputSegment>(S, this);
-    Seg->copyRelocations(*DataSection);
-    Segments.emplace_back(Seg);
-  }
+  for (const WasmSegment &S : WasmObj->dataSegments())
+    Segments.emplace_back(make<InputSegment>(S, this));
+  setRelocs(Segments, DataSection);
 
   // Populate `Functions`.
   ArrayRef<WasmFunction> Funcs = WasmObj->functions();
@@ -227,12 +258,9 @@ void ObjFile::parse() {
   ArrayRef<WasmSignature> Types = WasmObj->types();
   Functions.reserve(Funcs.size());
 
-  for (size_t I = 0, E = Funcs.size(); I != E; ++I) {
-    InputFunction *F =
-        make<InputFunction>(Types[FuncTypes[I]], &Funcs[I], this);
-    F->copyRelocations(*CodeSection);
-    Functions.emplace_back(F);
-  }
+  for (size_t I = 0, E = Funcs.size(); I != E; ++I)
+    Functions.emplace_back(make<InputFunction>(Types[FuncTypes[I]], &Funcs[I], this));
+  setRelocs(Functions, CodeSection);
 
   // Populate `Globals`.
   for (const WasmGlobal &G : WasmObj->globals())