[WebAssembly] Optimise relocation iteration to remove n^2 loop. NFC.

[sbc100]

Summary

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.

[sbc100]

• cleanup

[sbc100]

This change depends on relocations being alwasy stored in offset order, and I'm not convinced that is always that case. I created this change to ensure that llvm always create reloction sections like this but I'm not under what circumstance is might currently fail to do so: https://reviews.llvm.org/D51065

[ncw]

Thanks for rebasing, looks fine to me. I'm glad it's of use to someone.

wasm/InputChunks.h

 
 class InputChunk {
 public:
   enum Kind { DataSegment, Function, SyntheticFunction, Section };
 
   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;
 
   ObjFile *File;
   int32_t OutputOffset = 0;
 
   // Signals that the section is part of the output.  The garbage collector,
   // and COMDAT handling can set a sections' Live bit.
   // If GC is disabled, all sections start out as live by default.
   unsigned Live : 1;
 
 protected:
   InputChunk(ObjFile *F, Kind K)
       : 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;
 };
 
 // Represents a WebAssembly data segment which can be included as part of
 // an output data segments.  Note that in WebAssembly, unlike ELF and other
 // formats, used the term "data segment" to refer to the continous regions of
 // memory that make on the data section. See:
 // https://webassembly.github.io/spec/syntax/modules.html#syntax-data
 //
 // For example, by default, clang will produce a separate data section for
 // each global variable.
 class InputSegment : public InputChunk {
 public:
   InputSegment(const WasmSegment &Seg, ObjFile *F)
       : InputChunk(F, InputChunk::DataSegment), Segment(Seg) {}
 
   static bool classof(const InputChunk *C) { return C->kind() == DataSegment; }
 
   uint32_t getAlignment() const { return Segment.Data.Alignment; }
   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;
 };
 
 // Represents a single wasm function within and input file.  These are
 // combined to create the final output CODE section.
 class InputFunction : public InputChunk {
 public:
public:
   uint32_t getFunctionCodeOffset() const { return Function->CodeOffset; }
   uint32_t getSize() const override {
     if (Config->CompressRelocTargets && File) {
       assert(CompressedSize);
       return CompressedSize;
     }
     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);
 
   // The size of a given input function can depend on the values of the
   // LEB relocations within it.  This finalizeContents method is called after
   // all the symbol values have be calcualted but before getSize() is ever
   // called.
   void calculateSize();
 
   const WasmSignature &Signature;
 
 protected:
   ArrayRef<uint8_t> data() const override {
     assert(!Config->CompressRelocTargets);
     return File->CodeSection->Content.slice(getInputSectionOffset(),
                                             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;
   uint32_t CompressedFuncSize = 0;
   uint32_t CompressedSize = 0;
 };
 
 class SyntheticFunction : public InputFunction {

wasm/InputChunks.cpp

 
 StringRef InputChunk::getComdatName() const {
   uint32_t Index = getComdat();
   if (Index == UINT32_MAX)
     return StringRef();
   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;
     unsigned BytesRead = 0;
     uint32_t Offset = Rel.Offset - getInputSectionOffset();
     const uint8_t *Loc = data().data() + Offset;
     switch (Rel.Type) {
     case R_WEBASSEMBLY_TYPE_INDEX_LEB:
void InputFunction::calculateSize() {
   const uint8_t *FuncStart = SecStart + getInputSectionOffset();
   uint32_t FunctionSizeLength;
   decodeULEB128(FuncStart, &FunctionSizeLength);
 
   uint32_t Start = getInputSectionOffset();
   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));
     LastRelocEnd = Rel.Offset + getRelocWidthPadded(Rel);
   }
   LLVM_DEBUG(dbgs() << "  final region: " << (End - LastRelocEnd) << "\n");
   CompressedFuncSize += End - LastRelocEnd;
 

wasm/InputFiles.cpp

case R_WEBASSEMBLY_FUNCTION_OFFSET_I32:
     return 0;
   case R_WEBASSEMBLY_SECTION_OFFSET_I32:
     return getSectionSymbol(Reloc.Index)->Section->OutputOffset + Reloc.Addend;
   default:
     llvm_unreachable("unknown relocation type");
   }
 }
 
+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");
   std::unique_ptr<Binary> Bin = CHECK(createBinary(MB), toString(this));
 
   auto *Obj = dyn_cast<WasmObjectFile>(Bin.get());
   if (!Obj)
     fatal(toString(this) + ": not a wasm file");
void ObjFile::parse() {
   for (const SectionRef &Sec : WasmObj->sections()) {
     const WasmSection &Section = WasmObj->getWasmSection(Sec);
     if (Section.Type == WASM_SEC_CODE) {
       CodeSection = &Section;
     } else if (Section.Type == WASM_SEC_DATA) {
       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++;
   }
 
   TypeMap.resize(getWasmObj()->types().size());
   TypeIsUsed.resize(getWasmObj()->types().size(), false);
 
   ArrayRef<StringRef> Comdats = WasmObj->linkingData().Comdats;
   UsedComdats.resize(Comdats.size());
   for (unsigned I = 0; I < Comdats.size(); ++I)
     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();
   ArrayRef<uint32_t> FuncTypes = WasmObj->functionTypes();
   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())
     Globals.emplace_back(make<InputGlobal>(G, this));
 
   // Populate `Symbols` based on the WasmSymbols in the object.
   Symbols.reserve(WasmObj->getNumberOfSymbols());
   for (const SymbolRef &Sym : WasmObj->symbols()) {