1 files changed, 101 insertions, 8 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();
 }