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

aarch64::GOTTableManager and aarch64::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.

llvm/include/llvm/ExecutionEngine/JITLink/aarch64.h

@@ -786,6 +786,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;
     const char *BlockWorkingMem = B->getContent().data();
@@ -848,16 +853,21 @@ private:
     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() == aarch64::Branch26PCRel && !E.getTarget().isDefined()) {
       DEBUG_WITH_TYPE("jitlink", {
@@ -884,6 +894,8 @@ public:
     return *StubsSection;
   }
 
+  void registerExistingEntries();
+
   GOTTableManager &GOT;
   Section *StubsSection = nullptr;
 };

llvm/lib/ExecutionEngine/JITLink/ELF_aarch64.cpp

@@ -659,8 +659,8 @@ const uint8_t TLSDescTableManager_ELF_aarch64::TLSDescEntryContent[16] = {
 Error buildTables_ELF_aarch64(LinkGraph &G) {
   LLVM_DEBUG(dbgs() << "Visiting edges in graph:\n");
 
-  aarch64::GOTTableManager GOT;
-  aarch64::PLTTableManager PLT(GOT);
+  aarch64::GOTTableManager GOT(G);
+  aarch64::PLTTableManager PLT(G, GOT);
   TLSInfoTableManager_ELF_aarch64 TLSInfo;
   TLSDescTableManager_ELF_aarch64 TLSDesc(TLSInfo);
   visitExistingEdges(G, GOT, PLT, TLSDesc, TLSInfo);

llvm/lib/ExecutionEngine/JITLink/MachO_arm64.cpp

@@ -558,8 +558,8 @@ namespace jitlink {
 Error buildTables_MachO_arm64(LinkGraph &G) {
   LLVM_DEBUG(dbgs() << "Visiting edges in graph:\n");
 
-  aarch64::GOTTableManager GOT;
-  aarch64::PLTTableManager PLT(GOT);
+  aarch64::GOTTableManager GOT(G);
+  aarch64::PLTTableManager PLT(G, GOT);
   visitExistingEdges(G, GOT, PLT);
   return Error::success();
 }

llvm/lib/ExecutionEngine/JITLink/aarch64.cpp

@@ -202,6 +202,26 @@ static Error writeStoreRegSeq(AppendFtor &Append, unsigned DstLocReg,
   return Append(Instr);
 }
 
+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() == 2 &&
+           "PLT block edge count != 2");
+    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);
+  }
+}
+
 const char *getPointerSigningFunctionSectionName() { return "$__ptrauth_sign"; }
 
 /// Creates a pointer signing function section, block, and symbol to reserve

llvm/unittests/ExecutionEngine/JITLink/AArch64Tests.cpp

@@ -0,0 +1,90 @@
+//===------- AArch64Tests.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/aarch64.h>
+
+#include "gtest/gtest.h"
+
+using namespace llvm;
+using namespace llvm::jitlink;
+using namespace llvm::jitlink::aarch64;
+
+TEST(AArch64, EmptyLinkGraph) {
+  LinkGraph G("foo", std::make_shared<orc::SymbolStringPool>(),
+              Triple("arm64-apple-darwin"), SubtargetFeatures(),
+              getEdgeKindName);
+  EXPECT_EQ(G.getName(), "foo");
+  EXPECT_EQ(G.getTargetTriple().str(), "arm64-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(AArch64, GOTAndStubs) {
+  LinkGraph G("foo", std::make_shared<orc::SymbolStringPool>(),
+              Triple("arm64-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);
+  }
+}

llvm/unittests/ExecutionEngine/JITLink/CMakeLists.txt

@@ -10,6 +10,7 @@ set(LLVM_LINK_COMPONENTS
 add_llvm_unittest(JITLinkTests
     AArch32Tests.cpp
     AArch32ErrorTests.cpp
+    AArch64Tests.cpp
     EHFrameSupportTests.cpp
     JITLinkTestUtils.cpp
     LinkGraphTests.cpp