[JITLink] Teach x86_64 GOT & PLT table managers to discover existing entries.

x86_64::GOTTableManager and x86_64::PLTTableManager will now look for existing
GOT and PLT sections and re-use existing entries if they're present.

This will be used for an upcoming MachO patch to enable compact unwind support.

This patch is the x86-64 counterpart 42595bd, which added the same
functionality to the GOT and PLT managers for aarch64.

Author: lhames

llvm/include/llvm/ExecutionEngine/JITLink/x86_64.h

@@ -671,6 +671,11 @@ class GOTTableManager : public TableManager<GOTTableManager> {
 public:
   static StringRef getSectionName() { return "$__GOT"; }
 
+  GOTTableManager(LinkGraph &G) {
+    if ((GOTSection = G.findSectionByName(getSectionName())))
+      registerExistingEntries();
+  }
+
   bool visitEdge(LinkGraph &G, Block *B, Edge &E) {
     Edge::Kind KindToSet = Edge::Invalid;
     switch (E.getKind()) {
@@ -721,16 +726,21 @@ class GOTTableManager : public TableManager<GOTTableManager> {
     return *GOTSection;
   }
 
+  void registerExistingEntries();
+
   Section *GOTSection = nullptr;
 };
 
 /// Procedure Linkage Table Builder.
 class PLTTableManager : public TableManager<PLTTableManager> {
 public:
-  PLTTableManager(GOTTableManager &GOT) : GOT(GOT) {}
-
   static StringRef getSectionName() { return "$__STUBS"; }
 
+  PLTTableManager(LinkGraph &G, GOTTableManager &GOT) : GOT(GOT) {
+    if ((StubsSection = G.findSectionByName(getSectionName())))
+      registerExistingEntries();
+  }
+
   bool visitEdge(LinkGraph &G, Block *B, Edge &E) {
     if (E.getKind() == x86_64::BranchPCRel32 && !E.getTarget().isDefined()) {
       DEBUG_WITH_TYPE("jitlink", {
@@ -754,14 +764,16 @@ class PLTTableManager : public TableManager<PLTTableManager> {
 
 public:
   Section &getStubsSection(LinkGraph &G) {
-    if (!PLTSection)
-      PLTSection = &G.createSection(getSectionName(),
-                                    orc::MemProt::Read | orc::MemProt::Exec);
-    return *PLTSection;
+    if (!StubsSection)
+      StubsSection = &G.createSection(getSectionName(),
+                                      orc::MemProt::Read | orc::MemProt::Exec);
+    return *StubsSection;
   }
 
+  void registerExistingEntries();
+
   GOTTableManager &GOT;
-  Section *PLTSection = nullptr;
+  Section *StubsSection = nullptr;
 };
 
 /// Optimize the GOT and Stub relocations if the edge target address is in range

llvm/lib/ExecutionEngine/JITLink/ELF_x86_64.cpp

@@ -88,8 +88,8 @@ const uint8_t TLSInfoTableManager_ELF_x86_64::TLSInfoEntryContent[16] = {
 Error buildTables_ELF_x86_64(LinkGraph &G) {
   LLVM_DEBUG(dbgs() << "Visiting edges in graph:\n");
 
-  x86_64::GOTTableManager GOT;
-  x86_64::PLTTableManager PLT(GOT);
+  x86_64::GOTTableManager GOT(G);
+  x86_64::PLTTableManager PLT(G, GOT);
   TLSInfoTableManager_ELF_x86_64 TLSInfo;
   visitExistingEdges(G, GOT, PLT, TLSInfo);
   return Error::success();

llvm/lib/ExecutionEngine/JITLink/MachO_x86_64.cpp

@@ -459,8 +459,8 @@ class MachOLinkGraphBuilder_x86_64 : public MachOLinkGraphBuilder {
 };
 
 Error buildGOTAndStubs_MachO_x86_64(LinkGraph &G) {
-  x86_64::GOTTableManager GOT;
-  x86_64::PLTTableManager PLT(GOT);
+  x86_64::GOTTableManager GOT(G);
+  x86_64::PLTTableManager PLT(G, GOT);
   visitExistingEdges(G, GOT, PLT);
   return Error::success();
 }

llvm/lib/ExecutionEngine/JITLink/x86_64.cpp

@@ -89,6 +89,26 @@ const char ReentryTrampolineContent[5] = {
   static_cast<char>(0xe8), 0x00, 0x00, 0x00, 0x00
 };
 
+void GOTTableManager::registerExistingEntries() {
+  for (auto *EntrySym : GOTSection->symbols()) {
+    assert(EntrySym->getBlock().edges_size() == 1 &&
+           "GOT block edge count != 1");
+    registerPreExistingEntry(EntrySym->getBlock().edges().begin()->getTarget(),
+                             *EntrySym);
+  }
+}
+
+void PLTTableManager::registerExistingEntries() {
+  for (auto *EntrySym : StubsSection->symbols()) {
+    assert(EntrySym->getBlock().edges_size() == 1 &&
+           "PLT block edge count != 1");
+    auto &GOTSym = EntrySym->getBlock().edges().begin()->getTarget();
+    assert(GOTSym.getBlock().edges_size() == 1 && "GOT block edge count != 1");
+    registerPreExistingEntry(GOTSym.getBlock().edges().begin()->getTarget(),
+                             *EntrySym);
+  }
+}
+
 Error optimizeGOTAndStubAccesses(LinkGraph &G) {
   LLVM_DEBUG(dbgs() << "Optimizing GOT entries and stubs:\n");
 

llvm/unittests/ExecutionEngine/JITLink/CMakeLists.txt

@@ -17,6 +17,7 @@ add_llvm_unittest(JITLinkTests
     MachOLinkGraphTests.cpp
     MemoryManagerErrorTests.cpp
     StubsTests.cpp
+    X86_64Tests.cpp
   )
 
 target_link_libraries(JITLinkTests PRIVATE LLVMTestingSupport)

llvm/unittests/ExecutionEngine/JITLink/X86_64Tests.cpp

@@ -0,0 +1,90 @@
+//===-------- X86_64Tests.cpp - Unit tests for the AArch64 backend --------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include <llvm/BinaryFormat/ELF.h>
+#include <llvm/ExecutionEngine/JITLink/x86_64.h>
+
+#include "gtest/gtest.h"
+
+using namespace llvm;
+using namespace llvm::jitlink;
+using namespace llvm::jitlink::x86_64;
+
+TEST(X86_64, EmptyLinkGraph) {
+  LinkGraph G("foo", std::make_shared<orc::SymbolStringPool>(),
+              Triple("x86_64-apple-darwin"), SubtargetFeatures(),
+              getEdgeKindName);
+  EXPECT_EQ(G.getName(), "foo");
+  EXPECT_EQ(G.getTargetTriple().str(), "x86_64-apple-darwin");
+  EXPECT_EQ(G.getPointerSize(), 8U);
+  EXPECT_EQ(G.getEndianness(), llvm::endianness::little);
+  EXPECT_TRUE(G.external_symbols().empty());
+  EXPECT_TRUE(G.absolute_symbols().empty());
+  EXPECT_TRUE(G.defined_symbols().empty());
+  EXPECT_TRUE(G.blocks().empty());
+}
+
+TEST(X86_64, GOTAndStubs) {
+  LinkGraph G("foo", std::make_shared<orc::SymbolStringPool>(),
+              Triple("x86_64-apple-darwin"), SubtargetFeatures(),
+              getEdgeKindName);
+
+  auto &External = G.addExternalSymbol("external", 0, false);
+
+  // First table accesses. We expect the graph to be empty:
+  EXPECT_EQ(G.findSectionByName(GOTTableManager::getSectionName()), nullptr);
+  EXPECT_EQ(G.findSectionByName(PLTTableManager::getSectionName()), nullptr);
+
+  {
+    // Create first GOT and PLT table managers and request a PLT stub. This
+    // should force creation of both a PLT stub and GOT entry.
+    GOTTableManager GOT(G);
+    PLTTableManager PLT(G, GOT);
+
+    PLT.getEntryForTarget(G, External);
+  }
+
+  auto *GOTSec = G.findSectionByName(GOTTableManager::getSectionName());
+  EXPECT_NE(GOTSec, nullptr);
+  if (GOTSec) {
+    // Expect one entry in the GOT now.
+    EXPECT_EQ(GOTSec->symbols_size(), 1U);
+    EXPECT_EQ(GOTSec->blocks_size(), 1U);
+  }
+
+  auto *PLTSec = G.findSectionByName(PLTTableManager::getSectionName());
+  EXPECT_NE(PLTSec, nullptr);
+  if (PLTSec) {
+    // Expect one entry in the PLT.
+    EXPECT_EQ(PLTSec->symbols_size(), 1U);
+    EXPECT_EQ(PLTSec->blocks_size(), 1U);
+  }
+
+  {
+    // Create second GOT and PLT table managers and request a PLT stub. This
+    // should force creation of both a PLT stub and GOT entry.
+    GOTTableManager GOT(G);
+    PLTTableManager PLT(G, GOT);
+
+    PLT.getEntryForTarget(G, External);
+  }
+
+  EXPECT_EQ(G.findSectionByName(GOTTableManager::getSectionName()), GOTSec);
+  if (GOTSec) {
+    // Expect the same one entry in the GOT.
+    EXPECT_EQ(GOTSec->symbols_size(), 1U);
+    EXPECT_EQ(GOTSec->blocks_size(), 1U);
+  }
+
+  EXPECT_EQ(G.findSectionByName(PLTTableManager::getSectionName()), PLTSec);
+  if (PLTSec) {
+    // Expect the same one entry in the GOT.
+    EXPECT_EQ(PLTSec->symbols_size(), 1U);
+    EXPECT_EQ(PLTSec->blocks_size(), 1U);
+  }
+}