Reapply r360194 "[JITLink] Add support for MachO .alt_entry atoms." with fixes.

This patch modifies MachOAtomGraphBuilder to use setLayoutNext rather than
addEdge, and fixes a bug in the section layout algorithm that could result in
atoms appearing more than once in the section ordering (which resulted in those
atoms being assigned invalid addresses during layout).

llvm-svn: 360205

4 files changed, 128 insertions, 29 deletions

diff --git a/llvm/lib/ExecutionEngine/JITLink/JITLinkGeneric.cpp b/llvm/lib/ExecutionEngine/JITLink/JITLinkGeneric.cpp
index 38ad1ec7c59..8e4cfa05363 100644
--- a/llvm/lib/ExecutionEngine/JITLink/JITLinkGeneric.cpp
+++ b/llvm/lib/ExecutionEngine/JITLink/JITLinkGeneric.cpp
@@ -170,31 +170,31 @@ void JITLinkerBase::layOutAtoms() {
     auto &SL = KV.second;
     for (auto *SIList : {&SL.ContentSections, &SL.ZeroFillSections}) {
       for (auto &SI : *SIList) {
-        std::vector<DefinedAtom *> LayoutHeads;
-        LayoutHeads.reserve(SI.S->atoms_size());
-
-        // First build the list of layout-heads (i.e. "heads" of layout-next
-        // chains).
-        DenseSet<DefinedAtom *> AlreadyLayedOut;
-        for (auto *DA : SI.S->atoms()) {
-          if (AlreadyLayedOut.count(DA))
-            continue;
-          LayoutHeads.push_back(DA);
-          while (DA->hasLayoutNext()) {
-            auto &Next = DA->getLayoutNext();
-            AlreadyLayedOut.insert(&Next);
-            DA = &Next;
-          }
-        }
+        // First build the set of layout-heads (i.e. "heads" of layout-next
+        // chains) by copying the section atoms, then eliminating any that
+        // appear as layout-next targets.
+        DenseSet<DefinedAtom *> LayoutHeads;
+        for (auto *DA : SI.S->atoms())
+          LayoutHeads.insert(DA);
+
+        for (auto *DA : SI.S->atoms())
+          if (DA->hasLayoutNext())
+            LayoutHeads.erase(&DA->getLayoutNext());
+
+        // Next, sort the layout heads by address order.
+        std::vector<DefinedAtom *> OrderedLayoutHeads;
+        OrderedLayoutHeads.reserve(LayoutHeads.size());
+        for (auto *DA : LayoutHeads)
+          OrderedLayoutHeads.push_back(DA);
 
         // Now sort the list of layout heads by address.
-        std::sort(LayoutHeads.begin(), LayoutHeads.end(),
+        std::sort(OrderedLayoutHeads.begin(), OrderedLayoutHeads.end(),
                   [](const DefinedAtom *LHS, const DefinedAtom *RHS) {
                     return LHS->getAddress() < RHS->getAddress();
                   });
 
         // Now populate the SI.Atoms field by appending each of the chains.
-        for (auto *DA : LayoutHeads) {
+        for (auto *DA : OrderedLayoutHeads) {
           SI.Atoms.push_back(DA);
           while (DA->hasLayoutNext()) {
             auto &Next = DA->getLayoutNext();
diff --git a/llvm/lib/ExecutionEngine/JITLink/MachOAtomGraphBuilder.cpp b/llvm/lib/ExecutionEngine/JITLink/MachOAtomGraphBuilder.cpp
index 670fc3ecf7a..5feff393ebd 100644
--- a/llvm/lib/ExecutionEngine/JITLink/MachOAtomGraphBuilder.cpp
+++ b/llvm/lib/ExecutionEngine/JITLink/MachOAtomGraphBuilder.cpp
@@ -44,6 +44,33 @@ void MachOAtomGraphBuilder::addCustomAtomizer(StringRef SectionName,
   CustomAtomizeFunctions[SectionName] = std::move(Atomizer);
 }
 
+bool MachOAtomGraphBuilder::areLayoutLocked(const Atom &A, const Atom &B) {
+  // If these atoms are the same then they're trivially "locked".
+  if (&A == &B)
+    return true;
+
+  // If A and B are different, check whether either is undefined. (in which
+  // case they are not locked).
+  if (!A.isDefined() || !B.isDefined())
+    return false;
+
+  // A and B are different, but they're both defined atoms. We need to check
+  // whether they're part of the same alt_entry chain.
+  auto &DA = static_cast<const DefinedAtom &>(A);
+  auto &DB = static_cast<const DefinedAtom &>(B);
+
+  auto AStartItr = AltEntryStarts.find(&DA);
+  if (AStartItr == AltEntryStarts.end()) // If A is not in a chain bail out.
+    return false;
+
+  auto BStartItr = AltEntryStarts.find(&DB);
+  if (BStartItr == AltEntryStarts.end()) // If B is not in a chain bail out.
+    return false;
+
+  // A and B are layout locked if they're in the same chain.
+  return AStartItr->second == BStartItr->second;
+}
+
 unsigned
 MachOAtomGraphBuilder::getPointerSize(const object::MachOObjectFile &Obj) {
   return Obj.is64Bit() ? 8 : 4;
@@ -126,6 +153,9 @@ Error MachOAtomGraphBuilder::addNonCustomAtoms() {
   DenseMap<MachOSection *, AddrToAtomMap> SecToAtoms;
 
   DenseMap<MachOSection *, unsigned> FirstOrdinal;
+  std::vector<DefinedAtom *> AltEntryAtoms;
+
+  DenseSet<StringRef> ProcessedSymbols; // Used to check for duplicate defs.
 
   for (auto SymI = Obj.symbol_begin(), SymE = Obj.symbol_end(); SymI != SymE;
        ++SymI) {
@@ -135,6 +165,14 @@ Error MachOAtomGraphBuilder::addNonCustomAtoms() {
     if (!Name)
       return Name.takeError();
 
+    // Bail out on duplicate definitions: There should never be more than one
+    // definition for a symbol in a given object file.
+    if (ProcessedSymbols.count(*Name))
+      return make_error<JITLinkError>("Duplicate definition within object: " +
+                                      *Name);
+    else
+      ProcessedSymbols.insert(*Name);
+
     auto Addr = Sym.getAddress();
     if (!Addr)
       return Addr.takeError();
@@ -189,24 +227,35 @@ Error MachOAtomGraphBuilder::addNonCustomAtoms() {
 
     auto &Sec = SecByIndexItr->second;
 
-    auto &A = G->addDefinedAtom(Sec.getGenericSection(), *Name, *Addr,
-                                std::max(Sym.getAlignment(), 1U));
+    auto &DA = G->addDefinedAtom(Sec.getGenericSection(), *Name, *Addr,
+                                 std::max(Sym.getAlignment(), 1U));
+
+    DA.setGlobal(Flags & object::SymbolRef::SF_Global);
+    DA.setExported(Flags & object::SymbolRef::SF_Exported);
+    DA.setWeak(Flags & object::SymbolRef::SF_Weak);
 
-    A.setGlobal(Flags & object::SymbolRef::SF_Global);
-    A.setExported(Flags & object::SymbolRef::SF_Exported);
-    A.setWeak(Flags & object::SymbolRef::SF_Weak);
+    DA.setCallable(*SymType & object::SymbolRef::ST_Function);
 
-    A.setCallable(*SymType & object::SymbolRef::ST_Function);
+    // Check alt-entry.
+    {
+      uint16_t NDesc = 0;
+      if (Obj.is64Bit())
+        NDesc = Obj.getSymbolTableEntry(SymI->getRawDataRefImpl()).n_desc;
+      else
+        NDesc = Obj.getSymbolTableEntry(SymI->getRawDataRefImpl()).n_desc;
+      if (NDesc & MachO::N_ALT_ENTRY)
+        AltEntryAtoms.push_back(&DA);
+    }
 
     LLVM_DEBUG({
       dbgs() << "  Added " << *Name
              << " addr: " << format("0x%016" PRIx64, *Addr)
-             << ", align: " << A.getAlignment()
+             << ", align: " << DA.getAlignment()
              << ", section: " << Sec.getGenericSection().getName() << "\n";
     });
 
     auto &SecAtoms = SecToAtoms[&Sec];
-    SecAtoms[A.getAddress() - Sec.getAddress()] = &A;
+    SecAtoms[DA.getAddress() - Sec.getAddress()] = &DA;
   }
 
   // Add anonymous atoms.
@@ -263,6 +312,50 @@ Error MachOAtomGraphBuilder::addNonCustomAtoms() {
     }
   }
 
+  LLVM_DEBUG(dbgs() << "Adding alt-entry starts\n");
+
+  // Sort alt-entry atoms by address in ascending order.
+  llvm::sort(AltEntryAtoms.begin(), AltEntryAtoms.end(),
+             [](const DefinedAtom *LHS, const DefinedAtom *RHS) {
+               return LHS->getAddress() < RHS->getAddress();
+             });
+
+  // Process alt-entry atoms in address order to build the table of alt-entry
+  // atoms to alt-entry chain starts.
+  for (auto *DA : AltEntryAtoms) {
+    assert(!AltEntryStarts.count(DA) && "Duplicate entry in AltEntryStarts");
+
+    // DA is an alt-entry atom. Look for the predecessor atom that it is locked
+    // to, bailing out if we do not find one.
+    auto AltEntryPred = G->findAtomByAddress(DA->getAddress() - 1);
+    if (!AltEntryPred)
+      return AltEntryPred.takeError();
+
+    // Add a LayoutNext edge from the predecessor to this atom.
+    AltEntryPred->setLayoutNext(*DA);
+
+    // Check to see whether the predecessor itself is an alt-entry atom.
+    auto AltEntryStartItr = AltEntryStarts.find(&*AltEntryPred);
+    if (AltEntryStartItr != AltEntryStarts.end()) {
+      // If the predecessor was an alt-entry atom then re-use its value.
+      AltEntryStarts[DA] = AltEntryStartItr->second;
+      LLVM_DEBUG({
+        dbgs() << "  " << *DA << " -> " << *AltEntryStartItr->second
+               << " (based on existing entry for " << *AltEntryPred << ")\n";
+      });
+    } else {
+      // If the predecessor does not have an entry then add an entry for this
+      // atom (i.e. the alt_entry atom) and a self-reference entry for the
+      /// predecessory atom that is the start of this chain.
+      AltEntryStarts[&*AltEntryPred] = &*AltEntryPred;
+      AltEntryStarts[DA] = &*AltEntryPred;
+      LLVM_DEBUG({
+        dbgs() << "  " << *AltEntryPred << " -> " << *AltEntryPred << "\n"
+               << "  " << *DA << " -> " << *AltEntryPred << "\n";
+      });
+    }
+  }
+
   return Error::success();
 }
 
diff --git a/llvm/lib/ExecutionEngine/JITLink/MachOAtomGraphBuilder.h b/llvm/lib/ExecutionEngine/JITLink/MachOAtomGraphBuilder.h
index 340a11dfc0e..540e2c366f0 100644
--- a/llvm/lib/ExecutionEngine/JITLink/MachOAtomGraphBuilder.h
+++ b/llvm/lib/ExecutionEngine/JITLink/MachOAtomGraphBuilder.h
@@ -96,6 +96,10 @@ protected:
 
   virtual Error addRelocations() = 0;
 
+  /// Returns true if Atom A and Atom B are at a fixed offset from one another
+  /// (i.e. if they're part of the same alt-entry chain).
+  bool areLayoutLocked(const Atom &A, const Atom &B);
+
 private:
   static unsigned getPointerSize(const object::MachOObjectFile &Obj);
   static support::endianness getEndianness(const object::MachOObjectFile &Obj);
@@ -108,6 +112,7 @@ private:
 
   const object::MachOObjectFile &Obj;
   std::unique_ptr<AtomGraph> G;
+  DenseMap<const DefinedAtom *, const DefinedAtom *> AltEntryStarts;
   DenseMap<unsigned, MachOSection> Sections;
   StringMap<CustomAtomizeFunction> CustomAtomizeFunctions;
   Optional<MachOSection> CommonSymbolsSection;
diff --git a/llvm/lib/ExecutionEngine/JITLink/MachO_x86_64.cpp b/llvm/lib/ExecutionEngine/JITLink/MachO_x86_64.cpp
index 75725dbc526..d9472e40ead 100644
--- a/llvm/lib/ExecutionEngine/JITLink/MachO_x86_64.cpp
+++ b/llvm/lib/ExecutionEngine/JITLink/MachO_x86_64.cpp
@@ -181,19 +181,20 @@ private:
     MachOX86RelocationKind DeltaKind;
     Atom *TargetAtom;
     uint64_t Addend;
-    if (&AtomToFix == &*FromAtom) {
+    if (areLayoutLocked(AtomToFix, *FromAtom)) {
       TargetAtom = ToAtom;
       DeltaKind = (SubRI.r_length == 3) ? Delta64 : Delta32;
       Addend = FixupValue + (FixupAddress - FromAtom->getAddress());
       // FIXME: handle extern 'from'.
-    } else if (&AtomToFix == ToAtom) {
+    } else if (areLayoutLocked(AtomToFix, *ToAtom)) {
       TargetAtom = &*FromAtom;
       DeltaKind = (SubRI.r_length == 3) ? NegDelta64 : NegDelta32;
       Addend = FixupValue - (FixupAddress - ToAtom->getAddress());
     } else {
       // AtomToFix was neither FromAtom nor ToAtom.
       return make_error<JITLinkError>("SUBTRACTOR relocation must fix up "
-                                      "either 'A' or 'B'");
+                                      "either 'A' or 'B' (or an atom in one "
+                                      "of their alt-entry groups)");
     }
 
     return PairRelocInfo(DeltaKind, TargetAtom, Addend);