[JITLink] Factor basic common GOT and stub creation code into its own class.

llvm-svn: 358838

1 files changed, 82 insertions, 0 deletions

diff --git a/llvm/lib/ExecutionEngine/JITLink/BasicGOTAndStubsBuilder.h b/llvm/lib/ExecutionEngine/JITLink/BasicGOTAndStubsBuilder.h
new file mode 100644
index 00000000000..1271ad962b3
--- /dev/null
+++ b/llvm/lib/ExecutionEngine/JITLink/BasicGOTAndStubsBuilder.h
@@ -0,0 +1,82 @@
+//===--- BasicGOTAndStubsBuilder.h - Generic GOT/Stub creation --*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// A base for simple GOT and stub creation.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_LIB_EXECUTIONENGINE_JITLINK_BASICGOTANDSTUBSBUILDER_H
+#define LLVM_LIB_EXECUTIONENGINE_JITLINK_BASICGOTANDSTUBSBUILDER_H
+
+#include "llvm/ExecutionEngine/JITLink/JITLink.h"
+
+namespace llvm {
+namespace jitlink {
+
+template <typename BuilderImpl> class BasicGOTAndStubsBuilder {
+public:
+  BasicGOTAndStubsBuilder(AtomGraph &G) : G(G) {}
+
+  void run() {
+    // We're going to be adding new atoms, but we don't want to iterate over
+    // the newly added ones, so just copy the existing atoms out.
+    std::vector<DefinedAtom *> DAs(G.defined_atoms().begin(),
+                                   G.defined_atoms().end());
+
+    for (auto *DA : DAs)
+      for (auto &E : DA->edges())
+        if (impl().isGOTEdge(E))
+          impl().fixGOTEdge(E, getGOTEntryAtom(E.getTarget()));
+        else if (impl().isExternalBranchEdge(E))
+          impl().fixExternalBranchEdge(E, getStubAtom(E.getTarget()));
+  }
+
+protected:
+  Atom &getGOTEntryAtom(Atom &Target) {
+    assert(Target.hasName() && "GOT edge cannot point to anonymous target");
+
+    auto GOTEntryI = GOTEntries.find(Target.getName());
+
+    // Build the entry if it doesn't exist.
+    if (GOTEntryI == GOTEntries.end()) {
+      auto &GOTEntry = impl().createGOTEntry(Target);
+      GOTEntryI =
+          GOTEntries.insert(std::make_pair(Target.getName(), &GOTEntry)).first;
+    }
+
+    assert(GOTEntryI != GOTEntries.end() && "Could not get GOT entry atom");
+    return *GOTEntryI->second;
+  }
+
+  Atom &getStubAtom(Atom &Target) {
+    assert(Target.hasName() &&
+           "External branch edge can not point to an anonymous target");
+    auto StubI = Stubs.find(Target.getName());
+
+    if (StubI == Stubs.end()) {
+      auto &StubAtom = impl().createStub(Target);
+      StubI = Stubs.insert(std::make_pair(Target.getName(), &StubAtom)).first;
+    }
+
+    assert(StubI != Stubs.end() && "Count not get stub atom");
+    return *StubI->second;
+  }
+
+  AtomGraph &G;
+
+private:
+  BuilderImpl &impl() { return static_cast<BuilderImpl &>(*this); }
+
+  DenseMap<StringRef, DefinedAtom *> GOTEntries;
+  DenseMap<StringRef, DefinedAtom *> Stubs;
+};
+
+} // end namespace jitlink
+} // end namespace llvm
+
+#endif // LLVM_LIB_EXECUTIONENGINE_JITLINK_BASICGOTANDSTUBSBUILDER_H