[JITLink] Add support for MachO .alt_entry atoms.

The MachO .alt_entry directive is applied to a symbol to indicate that it is
locked (in terms of address layout and liveness) to its predecessor atom. I.e.
it is an alternate entry point, at a fixed offset, for the previous atom.

This patch updates MachOAtomGraphBuilder to check for the .alt_entry flag on
symbols and add a corresponding LayoutNext edge to the atom-graph. It also
updates MachOAtomGraphBuilder_x86_64 to generalize handling of the
X86_64_RELOC_SUBTRACTOR relocation: previously either the minuend or
subtrahend of the subtraction had to be the same as the atom being fixed up,
now it is only necessary for the minuend or subtrahend to be locked (via any
chain of alt_entry directives) to the atom being fixed up.

llvm-svn: 360194

3 files changed, 110 insertions, 11 deletions

diff --git a/llvm/lib/ExecutionEngine/JITLink/MachOAtomGraphBuilder.cpp b/llvm/lib/ExecutionEngine/JITLink/MachOAtomGraphBuilder.cpp
index 670fc3ecf7a..4054f66b5c1 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->addEdge(Edge::LayoutNext, 0, *DA, 0);
+
+    // 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);