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

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

Author: weliveindetail

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()) {

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;
       }
 

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()) {

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())

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 *>;