[PECOFF] Remove COFFDefinedFileAtom::originalOffset().

The aim of this patch is to reduce the dependency from COFFDefinedAtom
to COFF structs defined in llvm/Object/COFF.h. Currently many attributes
of the atom are computed in the atom. That provide a simple interface but
does not work well in some cases.

There are some cases that the same type atom is created from different
parts of a COFF file. One example is the BSS atom, which can be created
from the defined symbol in the .bss section or from the undefined symbol.
Computing attributes from different sources in the atom complicates the
code. We should compute it outside the atom.

In the next patch, I'll move more code from Atoms.h to ReaderCOFF.cpp.

llvm-svn: 187681

2 files changed, 25 insertions, 15 deletions

diff --git a/lld/lib/ReaderWriter/PECOFF/Atoms.h b/lld/lib/ReaderWriter/PECOFF/Atoms.h
index badcff2b4a2..49d77887891 100644
--- a/lld/lib/ReaderWriter/PECOFF/Atoms.h
+++ b/lld/lib/ReaderWriter/PECOFF/Atoms.h
@@ -164,7 +164,6 @@ public:
         _ordinal(ordinal) {}
 
   virtual uint64_t ordinal() const { return _ordinal; }
-  uint64_t originalOffset() const { return _symbol->Value; }
   virtual StringRef getSectionName() const { return _sectionName; }
 
   static bool classof(const COFFBaseDefinedAtom *atom) {
diff --git a/lld/lib/ReaderWriter/PECOFF/ReaderCOFF.cpp b/lld/lib/ReaderWriter/PECOFF/ReaderCOFF.cpp
index 4212857e7db..c9cd43dd08a 100644
--- a/lld/lib/ReaderWriter/PECOFF/ReaderCOFF.cpp
+++ b/lld/lib/ReaderWriter/PECOFF/ReaderCOFF.cpp
@@ -268,8 +268,10 @@ private:
     // Create an atom for the entire section.
     if (symbols.empty()) {
       ArrayRef<uint8_t> Data(secData.data(), secData.size());
-      atoms.push_back(new (_alloc) COFFDefinedAtom(
-          *this, "", nullptr, section, Data, sectionName, 0));
+      auto *atom = new (_alloc) COFFDefinedAtom(*this, "", nullptr, section,
+                                                Data, sectionName, 0);
+      atoms.push_back(atom);
+      _definedAtomLocations[section][0] = atom;
       return error_code::success();
     }
 
@@ -278,8 +280,10 @@ private:
     if (symbols[0]->Value != 0) {
       uint64_t size = symbols[0]->Value;
       ArrayRef<uint8_t> data(secData.data(), size);
-      atoms.push_back(new (_alloc) COFFDefinedAtom(
-          *this, "", nullptr, section, data, sectionName, ++ordinal));
+      auto *atom = new (_alloc) COFFDefinedAtom(
+          *this, "", nullptr, section, data, sectionName, ++ordinal);
+      atoms.push_back(atom);
+      _definedAtomLocations[section][0] = atom;
     }
 
     for (auto si = symbols.begin(), se = symbols.end(); si != se; ++si) {
@@ -293,6 +297,7 @@ private:
           *this, _symbolName[*si], *si, section, data, sectionName, ++ordinal);
       atoms.push_back(atom);
       _symbolAtom[*si] = atom;
+      _definedAtomLocations[section][(*si)->Value] = atom;
     }
     return error_code::success();
   }
@@ -320,14 +325,16 @@ private:
     return error_code::success();
   }
 
-  /// Find the atom that is at \p targetOffset in \p section. It is assumed
-  /// that \p atoms are sorted by position in the section.
-  error_code findAtomAt(uint32_t targetOffset,
-                        const vector<COFFDefinedFileAtom *> &atoms,
-                        COFFDefinedFileAtom *&result) const {
-    for (COFFDefinedFileAtom *atom : atoms) {
-      if (targetOffset < atom->originalOffset() + atom->size()) {
+  /// Find the atom that is at \p targetAddress in \p section.
+  error_code findAtomAt(const coff_section *section, uint32_t targetAddress,
+                        COFFDefinedFileAtom *&result, uint32_t &offsetInAtom) {
+    for (auto i : _definedAtomLocations[section]) {
+      uint32_t atomAddress = i.first;
+      COFFDefinedAtom *atom = i.second;
+      if (atomAddress <= targetAddress &&
+          targetAddress < atomAddress + atom->size()) {
         result = atom;
+        offsetInAtom = targetAddress - atomAddress;
         return error_code::success();
       }
     }
@@ -364,10 +371,9 @@ private:
       return ec;
 
     COFFDefinedFileAtom *atom;
-    if (error_code ec = findAtomAt(rel->VirtualAddress, atoms, atom))
+    uint32_t offsetInAtom;
+    if (error_code ec = findAtomAt(section, itemAddress, atom, offsetInAtom))
       return ec;
-    uint32_t offsetInAtom = itemAddress - atom->originalOffset();
-    assert(offsetInAtom < atom->size());
     atom->addReference(std::unique_ptr<COFFReference>(
         new COFFReference(targetAtom, offsetInAtom, rel->Type)));
     return error_code::success();
@@ -470,6 +476,11 @@ private:
   // A map from section to its atoms.
   std::map<const coff_section *, vector<COFFDefinedFileAtom *>> _sectionAtoms;
 
+  // A sorted map to find an atom from a section and an offset within
+  // the section.
+  std::map<const coff_section *,
+           std::map<uint32_t, COFFDefinedAtom *>> _definedAtomLocations;
+
   mutable llvm::BumpPtrAllocator _alloc;
   const TargetInfo &_targetInfo;
 };