[llvm] [llvm-jitlink] Refactor GOT and stubs registration (NFC) (PR #78365)

Stefan Gränitz via llvm-commits llvm-commits at lists.llvm.org
Tue Jan 16 15:51:19 PST 2024


https://github.com/weliveindetail created https://github.com/llvm/llvm-project/pull/78365

Add methods `registerGOTEntry()` and `registerStubEntry()` in `Session::FileInfo` to factor out generic code from the individual object type implementations.

>From 6e3b317fcf8a34c4fe36df51105172cda08281c4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Stefan=20Gr=C3=A4nitz?= <stefan.graenitz at gmail.com>
Date: Tue, 16 Jan 2024 22:33:02 +0100
Subject: [PATCH] [llvm-jitlink] Refactor GOT and stubs registration (NFC)

---
 llvm/tools/llvm-jitlink/llvm-jitlink-coff.cpp | 35 +++-----
 llvm/tools/llvm-jitlink/llvm-jitlink-elf.cpp  | 79 ++++++++-----------
 .../tools/llvm-jitlink/llvm-jitlink-macho.cpp | 30 ++-----
 llvm/tools/llvm-jitlink/llvm-jitlink.cpp      | 30 +++++++
 llvm/tools/llvm-jitlink/llvm-jitlink.h        | 10 +++
 5 files changed, 89 insertions(+), 95 deletions(-)

diff --git a/llvm/tools/llvm-jitlink/llvm-jitlink-coff.cpp b/llvm/tools/llvm-jitlink/llvm-jitlink-coff.cpp
index 283e655205d780..5271fdb5565904 100644
--- a/llvm/tools/llvm-jitlink/llvm-jitlink-coff.cpp
+++ b/llvm/tools/llvm-jitlink/llvm-jitlink-coff.cpp
@@ -108,34 +108,17 @@ Error registerCOFFGraphInfo(Session &S, LinkGraph &G) {
       if (Sym->getAddress() > LastSym->getAddress())
         LastSym = Sym;
 
-      if (isGOTSection) {
-        if (Sym->isSymbolZeroFill())
-          return make_error<StringError>("zero-fill atom in GOT section",
-                                         inconvertibleErrorCode());
-
-        // If this is a GOT symbol with size (i.e. not the GOT start symbol)
-        // then add it to the GOT entry info table.
-        if (Sym->getSize() != 0) {
-          if (auto TS = getCOFFGOTTarget(G, Sym->getBlock()))
-            FileInfo.GOTEntryInfos[TS->getName()] = {
-                Sym->getSymbolContent(), Sym->getAddress().getValue(),
-                Sym->getTargetFlags()};
-          else
-            return TS.takeError();
+      if (isGOTSection || isStubsSection) {
+        if (isGOTSection) {
+          // Skip the GOT start symbol
+          if (Sym->getSize() != 0)
+            if (Error E = FileInfo.registerGOTEntry(G, *Sym, getCOFFGOTTarget))
+              return E;
+        } else {
+          if (Error E = FileInfo.registerStubEntry(G, *Sym, getCOFFStubTarget))
+            return E;
         }
         SectionContainsContent = true;
-      } else if (isStubsSection) {
-        if (Sym->isSymbolZeroFill())
-          return make_error<StringError>("zero-fill atom in Stub section",
-                                         inconvertibleErrorCode());
-
-        if (auto TS = getCOFFStubTarget(G, Sym->getBlock()))
-          FileInfo.StubInfos[TS->getName()] = {Sym->getSymbolContent(),
-                                               Sym->getAddress().getValue(),
-                                               Sym->getTargetFlags()};
-        else
-          return TS.takeError();
-        SectionContainsContent = true;
       }
 
       if (Sym->hasName()) {
diff --git a/llvm/tools/llvm-jitlink/llvm-jitlink-elf.cpp b/llvm/tools/llvm-jitlink/llvm-jitlink-elf.cpp
index a02468758b33f6..c6b4218aad7af8 100644
--- a/llvm/tools/llvm-jitlink/llvm-jitlink-elf.cpp
+++ b/llvm/tools/llvm-jitlink/llvm-jitlink-elf.cpp
@@ -72,13 +72,29 @@ static Expected<Symbol &> getELFStubTarget(LinkGraph &G, Block &B) {
   return getELFGOTTarget(G, GOTSym.getBlock());
 }
 
-static Expected<std::string> getELFAArch32StubTargetName(LinkGraph &G,
-                                                         Block &B) {
+static Expected<Symbol &> getELFAArch32StubTarget(LinkGraph &G, Block &B) {
   auto E = getFirstRelocationEdge(G, B);
   if (!E)
     return E.takeError();
-  Symbol &StubTarget = E->getTarget();
-  return StubTarget.getName().str();
+  return E->getTarget();
+}
+
+enum SectionType { GOT, Stubs, AArch32Stubs, Other };
+
+static Error registerSymbol(LinkGraph &G, Symbol &Sym, Session::FileInfo &FI,
+                            SectionType SecType) {
+  switch (SecType) {
+  case GOT:
+    if (Sym.getSize() == 0)
+      return Error::success(); // Skip the GOT start symbol
+    return FI.registerGOTEntry(G, Sym, getELFGOTTarget);
+  case Stubs:
+    return FI.registerStubEntry(G, Sym, getELFStubTarget);
+  case AArch32Stubs:
+    return FI.registerStubEntry(G, Sym, getELFAArch32StubTarget);
+  case Other:
+    return Error::success();
+  }
 }
 
 namespace llvm {
@@ -113,9 +129,16 @@ Error registerELFGraphInfo(Session &S, LinkGraph &G) {
                                          "\"",
                                      inconvertibleErrorCode());
 
-    bool isGOTSection = isELFGOTSection(Sec);
-    bool isStubsSection = isELFStubsSection(Sec);
-    bool isAArch32StubsSection = isELFAArch32StubsSection(Sec);
+    SectionType SecType;
+    if (isELFGOTSection(Sec)) {
+      SecType = GOT;
+    } else if (isELFStubsSection(Sec)) {
+      SecType = Stubs;
+    } else if (isELFAArch32StubsSection(Sec)) {
+      SecType = AArch32Stubs;
+    } else {
+      SecType = Other;
+    }
 
     bool SectionContainsContent = false;
     bool SectionContainsZeroFill = false;
@@ -128,45 +151,9 @@ Error registerELFGraphInfo(Session &S, LinkGraph &G) {
       if (Sym->getAddress() > LastSym->getAddress())
         LastSym = Sym;
 
-      if (isGOTSection) {
-        if (Sym->isSymbolZeroFill())
-          return make_error<StringError>("zero-fill atom in GOT section",
-                                         inconvertibleErrorCode());
-
-        // If this is a GOT symbol with size (i.e. not the GOT start symbol)
-        // then add it to the GOT entry info table.
-        if (Sym->getSize() != 0) {
-          if (auto TS = getELFGOTTarget(G, Sym->getBlock()))
-            FileInfo.GOTEntryInfos[TS->getName()] = {
-                Sym->getSymbolContent(), Sym->getAddress().getValue(),
-                Sym->getTargetFlags()};
-          else
-            return TS.takeError();
-        }
-        SectionContainsContent = true;
-      } else if (isStubsSection) {
-        if (Sym->isSymbolZeroFill())
-          return make_error<StringError>("zero-fill atom in Stub section",
-                                         inconvertibleErrorCode());
-
-        if (auto TS = getELFStubTarget(G, Sym->getBlock()))
-          FileInfo.StubInfos[TS->getName()] = {Sym->getSymbolContent(),
-                                               Sym->getAddress().getValue(),
-                                               Sym->getTargetFlags()};
-        else
-          return TS.takeError();
-        SectionContainsContent = true;
-      } else if (isAArch32StubsSection) {
-        if (Sym->isSymbolZeroFill())
-          return make_error<StringError>("zero-fill atom in Stub section",
-                                         inconvertibleErrorCode());
-
-        if (auto Name = getELFAArch32StubTargetName(G, Sym->getBlock()))
-          FileInfo.StubInfos[*Name] = {Sym->getSymbolContent(),
-                                       Sym->getAddress().getValue(),
-                                       Sym->getTargetFlags()};
-        else
-          return Name.takeError();
+      if (SecType != Other) {
+        if (Error Err = registerSymbol(G, *Sym, FileInfo, SecType))
+          return Err;
         SectionContainsContent = true;
       }
 
diff --git a/llvm/tools/llvm-jitlink/llvm-jitlink-macho.cpp b/llvm/tools/llvm-jitlink/llvm-jitlink-macho.cpp
index 7dcadd94c2365a..2c60c802293a1e 100644
--- a/llvm/tools/llvm-jitlink/llvm-jitlink-macho.cpp
+++ b/llvm/tools/llvm-jitlink/llvm-jitlink-macho.cpp
@@ -111,29 +111,13 @@ Error registerMachOGraphInfo(Session &S, LinkGraph &G) {
         FirstSym = Sym;
       if (Sym->getAddress() > LastSym->getAddress())
         LastSym = Sym;
-      if (isGOTSection) {
-        if (Sym->isSymbolZeroFill())
-          return make_error<StringError>("zero-fill atom in GOT section",
-                                         inconvertibleErrorCode());
-
-        if (auto TS = getMachOGOTTarget(G, Sym->getBlock()))
-          FileInfo.GOTEntryInfos[TS->getName()] = {Sym->getSymbolContent(),
-                                                   Sym->getAddress().getValue(),
-                                                   Sym->getTargetFlags()};
-        else
-          return TS.takeError();
-        SectionContainsContent = true;
-      } else if (isStubsSection) {
-        if (Sym->isSymbolZeroFill())
-          return make_error<StringError>("zero-fill atom in Stub section",
-                                         inconvertibleErrorCode());
-
-        if (auto TS = getMachOStubTarget(G, Sym->getBlock()))
-          FileInfo.StubInfos[TS->getName()] = {Sym->getSymbolContent(),
-                                               Sym->getAddress().getValue(),
-                                               Sym->getTargetFlags()};
-        else
-          return TS.takeError();
+      if (isGOTSection || isStubsSection) {
+        Error Err =
+            isGOTSection
+                ? FileInfo.registerGOTEntry(G, *Sym, getMachOGOTTarget)
+                : FileInfo.registerStubEntry(G, *Sym, getMachOStubTarget);
+        if (Err)
+          return Err;
         SectionContainsContent = true;
       } else if (Sym->hasName()) {
         if (Sym->isSymbolZeroFill()) {
diff --git a/llvm/tools/llvm-jitlink/llvm-jitlink.cpp b/llvm/tools/llvm-jitlink/llvm-jitlink.cpp
index aa032c97485fbf..8c18610313ce8f 100644
--- a/llvm/tools/llvm-jitlink/llvm-jitlink.cpp
+++ b/llvm/tools/llvm-jitlink/llvm-jitlink.cpp
@@ -1183,6 +1183,36 @@ Error Session::loadAndLinkDynamicLibrary(JITDylib &JD, StringRef LibPath) {
   return Error::success();
 }
 
+Error Session::FileInfo::registerGOTEntry(
+    LinkGraph &G, Symbol &Sym, GetSymbolTargetFunction GetSymbolTarget) {
+  if (Sym.isSymbolZeroFill())
+    return make_error<StringError>("Unexpected zero-fill symbol in section " +
+                                       Sym.getBlock().getSection().getName(),
+                                   inconvertibleErrorCode());
+  auto TS = GetSymbolTarget(G, Sym.getBlock());
+  if (!TS)
+    return TS.takeError();
+  GOTEntryInfos[TS->getName()] = {Sym.getSymbolContent(),
+                                  Sym.getAddress().getValue(),
+                                  Sym.getTargetFlags()};
+  return Error::success();
+}
+
+Error Session::FileInfo::registerStubEntry(
+    LinkGraph &G, Symbol &Sym, GetSymbolTargetFunction GetSymbolTarget) {
+  if (Sym.isSymbolZeroFill())
+    return make_error<StringError>("Unexpected zero-fill symbol in section " +
+                                       Sym.getBlock().getSection().getName(),
+                                   inconvertibleErrorCode());
+  auto TS = GetSymbolTarget(G, Sym.getBlock());
+  if (!TS)
+    return TS.takeError();
+  StubInfos[TS->getName()] = {Sym.getSymbolContent(),
+                              Sym.getAddress().getValue(),
+                              Sym.getTargetFlags()};
+  return Error::success();
+}
+
 Expected<Session::FileInfo &> Session::findFileInfo(StringRef FileName) {
   auto FileInfoItr = FileInfos.find(FileName);
   if (FileInfoItr == FileInfos.end())
diff --git a/llvm/tools/llvm-jitlink/llvm-jitlink.h b/llvm/tools/llvm-jitlink/llvm-jitlink.h
index 3ff406b7b82dfd..93a00266b15043 100644
--- a/llvm/tools/llvm-jitlink/llvm-jitlink.h
+++ b/llvm/tools/llvm-jitlink/llvm-jitlink.h
@@ -51,6 +51,16 @@ struct Session {
     StringMap<MemoryRegionInfo> SectionInfos;
     StringMap<MemoryRegionInfo> StubInfos;
     StringMap<MemoryRegionInfo> GOTEntryInfos;
+
+    using Symbol = jitlink::Symbol;
+    using LinkGraph = jitlink::LinkGraph;
+    using GetSymbolTargetFunction =
+        unique_function<Expected<Symbol &>(LinkGraph &G, jitlink::Block &)>;
+
+    Error registerGOTEntry(LinkGraph &G, Symbol &Sym,
+                           GetSymbolTargetFunction GetSymbolTarget);
+    Error registerStubEntry(LinkGraph &G, Symbol &Sym,
+                            GetSymbolTargetFunction GetSymbolTarget);
   };
 
   using DynLibJDMap = std::map<std::string, orc::JITDylib *>;



More information about the llvm-commits mailing list