[llvm-jitlink] Allow optional stub-kind filter in stub_addr() expressions (#78369)

We use `jitlink-check` lines in LIT tests as the primary tool for
testing JITLink backends. Parsing and evaluation of the expressions is
implemented in `RuntimeDyldChecker`. The `stub_addr(obj, name)`
expression allows to obtain the linker-generated stub for the external
symbol `name` in object file `obj`.

This patch adds support for a filter parameter to select one out of many
stubs. This is necessary for the AArch32 JITLink backend, which must be
able to emit two different kinds of stubs depending on the instruction
set state (Arm/Thumb) of the relocation site. Since the new parameter is
optional, we don't have to update existing tests.

Filters are regular expressions without brackets that match exactly one
existing stub. Given object file `armv7.o` with two stubs for external
function `ext` of kinds `armv7_abs_le` and `thumbv7_abs_le`, we get the
following filter results e.g.:
```
stub_addr(armv7.o, ext, thumb)        thumbv7_abs_le
stub_addr(armv7.o, ext, thumbv7)      thumbv7_abs_le
stub_addr(armv7.o, ext, armv7_abs_le) armv7_abs_le
stub_addr(armv7.o, ext, v7_.*_le)     Error: "ext" has 2 candidate stubs in file "armv7.o". Please refine stub-kind filter "v7_.*_le" for disambiguation (encountered kinds are "thumbv7_abs_le", "armv7_abs_le").
stub_addr(armv7.o, ext, v8)           Error: "ext" has 2 stubs in file "armv7.o", but none of them matches the stub-kind filter "v8" (all encountered kinds are "thumbv7_abs_le", "armv7_abs_le").
```

Author: weliveindetail

llvm/include/llvm/ExecutionEngine/RuntimeDyldChecker.h

@@ -154,7 +154,7 @@ class RuntimeDyldChecker {
   using GetSectionInfoFunction = std::function<Expected<MemoryRegionInfo>(
       StringRef FileName, StringRef SectionName)>;
   using GetStubInfoFunction = std::function<Expected<MemoryRegionInfo>(
-      StringRef StubContainer, StringRef TargetName)>;
+      StringRef StubContainer, StringRef TargetName, StringRef StubKindFilter)>;
   using GetGOTInfoFunction = std::function<Expected<MemoryRegionInfo>(
       StringRef GOTContainer, StringRef TargetName)>;
 

llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldChecker.cpp

@@ -400,15 +400,25 @@ class RuntimeDyldCheckerExprEval {
     StringRef Symbol;
     std::tie(Symbol, RemainingExpr) = parseSymbol(RemainingExpr);
 
+    // Parse optional parameter to filter by stub kind
+    StringRef KindNameFilter;
+    if (RemainingExpr.starts_with(",")) {
+      RemainingExpr = RemainingExpr.substr(1).ltrim();
+      size_t ClosingBracket = RemainingExpr.find(")");
+      KindNameFilter = RemainingExpr.substr(0, ClosingBracket);
+      RemainingExpr = RemainingExpr.substr(ClosingBracket);
+    }
+
     if (!RemainingExpr.starts_with(")"))
       return std::make_pair(
           unexpectedToken(RemainingExpr, Expr, "expected ')'"), "");
     RemainingExpr = RemainingExpr.substr(1).ltrim();
 
     uint64_t StubAddr;
     std::string ErrorMsg;
-    std::tie(StubAddr, ErrorMsg) = Checker.getStubOrGOTAddrFor(
-        StubContainerName, Symbol, PCtx.IsInsideLoad, IsStubAddr);
+    std::tie(StubAddr, ErrorMsg) =
+        Checker.getStubOrGOTAddrFor(StubContainerName, Symbol, KindNameFilter,
+                                    PCtx.IsInsideLoad, IsStubAddr);
 
     if (ErrorMsg != "")
       return std::make_pair(EvalResult(ErrorMsg), "");
@@ -985,11 +995,14 @@ std::pair<uint64_t, std::string> RuntimeDyldCheckerImpl::getSectionAddr(
 }
 
 std::pair<uint64_t, std::string> RuntimeDyldCheckerImpl::getStubOrGOTAddrFor(
-    StringRef StubContainerName, StringRef SymbolName, bool IsInsideLoad,
-    bool IsStubAddr) const {
-
-  auto StubInfo = IsStubAddr ? GetStubInfo(StubContainerName, SymbolName)
-                             : GetGOTInfo(StubContainerName, SymbolName);
+    StringRef StubContainerName, StringRef SymbolName, StringRef StubKindFilter,
+    bool IsInsideLoad, bool IsStubAddr) const {
+
+  assert((StubKindFilter.empty() || IsStubAddr) &&
+         "Kind name filter only supported for stubs");
+  auto StubInfo =
+      IsStubAddr ? GetStubInfo(StubContainerName, SymbolName, StubKindFilter)
+                 : GetGOTInfo(StubContainerName, SymbolName);
 
   if (!StubInfo) {
     std::string ErrMsg;

llvm/lib/ExecutionEngine/RuntimeDyld/RuntimeDyldCheckerImpl.h

@@ -64,7 +64,8 @@ class RuntimeDyldCheckerImpl {
 
   std::pair<uint64_t, std::string>
   getStubOrGOTAddrFor(StringRef StubContainerName, StringRef Symbol,
-                      bool IsInsideLoad, bool IsStubAddr) const;
+                      StringRef StubKindFilter, bool IsInsideLoad,
+                      bool IsStubAddr) const;
 
   std::optional<uint64_t> getSectionLoadAddress(void *LocalAddr) const;
 

llvm/tools/llvm-jitlink/llvm-jitlink-elf.cpp

@@ -91,7 +91,7 @@ static Error registerSymbol(LinkGraph &G, Symbol &Sym, Session::FileInfo &FI,
   case Stubs:
     return FI.registerStubEntry(G, Sym, getELFStubTarget);
   case AArch32Stubs:
-    return FI.registerStubEntry(G, Sym, getELFAArch32StubTarget);
+    return FI.registerMultiStubEntry(G, Sym, getELFAArch32StubTarget);
   case Other:
     return Error::success();
   }

llvm/tools/llvm-jitlink/llvm-jitlink.cpp

@@ -331,8 +331,12 @@ operator<<(raw_ostream &OS, const Session::FileInfo &FI) {
     OS << "  Section \"" << SIKV.first() << "\": " << SIKV.second << "\n";
   for (auto &GOTKV : FI.GOTEntryInfos)
     OS << "  GOT \"" << GOTKV.first() << "\": " << GOTKV.second << "\n";
-  for (auto &StubKV : FI.StubInfos)
-    OS << "  Stub \"" << StubKV.first() << "\": " << StubKV.second << "\n";
+  for (auto &StubKVs : FI.StubInfos) {
+    OS << "  Stubs \"" << StubKVs.first() << "\":";
+    for (auto MemRegion : StubKVs.second)
+      OS << " " << MemRegion;
+    OS << "\n";
+  }
   return OS;
 }
 
@@ -1207,9 +1211,35 @@ Error Session::FileInfo::registerStubEntry(
   auto TS = GetSymbolTarget(G, Sym.getBlock());
   if (!TS)
     return TS.takeError();
-  StubInfos[TS->getName()] = {Sym.getSymbolContent(),
-                              Sym.getAddress().getValue(),
-                              Sym.getTargetFlags()};
+
+  SmallVector<MemoryRegionInfo> &Entry = StubInfos[TS->getName()];
+  Entry.insert(Entry.begin(),
+               {Sym.getSymbolContent(), Sym.getAddress().getValue(),
+                Sym.getTargetFlags()});
+  return Error::success();
+}
+
+Error Session::FileInfo::registerMultiStubEntry(
+    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 Target = GetSymbolTarget(G, Sym.getBlock());
+  if (!Target)
+    return Target.takeError();
+
+  SmallVector<MemoryRegionInfo> &Entry = StubInfos[Target->getName()];
+  Entry.emplace_back(Sym.getSymbolContent(), Sym.getAddress().getValue(),
+                     Sym.getTargetFlags());
+
+  // Let's keep stubs ordered by ascending address.
+  std::sort(Entry.begin(), Entry.end(),
+            [](const MemoryRegionInfo &L, const MemoryRegionInfo &R) {
+              return L.getTargetAddress() < R.getTargetAddress();
+            });
+
   return Error::success();
 }
 
@@ -1235,8 +1265,14 @@ Session::findSectionInfo(StringRef FileName, StringRef SectionName) {
   return SecInfoItr->second;
 }
 
+static StringRef detectStubKind(const Session::MemoryRegionInfo &Stub) {
+  // Implement acutal stub kind detection
+  return "";
+}
+
 Expected<Session::MemoryRegionInfo &>
-Session::findStubInfo(StringRef FileName, StringRef TargetName) {
+Session::findStubInfo(StringRef FileName, StringRef TargetName,
+                      StringRef KindNameFilter) {
   auto FI = findFileInfo(FileName);
   if (!FI)
     return FI.takeError();
@@ -1246,7 +1282,38 @@ Session::findStubInfo(StringRef FileName, StringRef TargetName) {
                                        "\" registered for file \"" + FileName +
                                        "\"",
                                    inconvertibleErrorCode());
-  return StubInfoItr->second;
+  auto &StubsForTarget = StubInfoItr->second;
+  assert(!StubsForTarget.empty() && "At least 1 stub in each entry");
+  if (KindNameFilter.empty() && StubsForTarget.size() == 1)
+    return StubsForTarget[0]; // Regular single-stub match
+
+  std::string KindsStr;
+  SmallVector<MemoryRegionInfo *, 1> Matches;
+  Regex KindNameMatcher(KindNameFilter.empty() ? ".*" : KindNameFilter);
+  for (MemoryRegionInfo &Stub : StubsForTarget) {
+    StringRef Kind = detectStubKind(Stub);
+    if (KindNameMatcher.match(Kind))
+      Matches.push_back(&Stub);
+    KindsStr += "\"" + (Kind.empty() ? "<unknown>" : Kind.str()) + "\", ";
+  }
+  if (Matches.empty())
+    return make_error<StringError>(
+        "\"" + TargetName + "\" has " + Twine(StubsForTarget.size()) +
+            " stubs in file \"" + FileName +
+            "\", but none of them matches the stub-kind filter \"" +
+            KindNameFilter + "\" (all encountered kinds are " +
+            StringRef(KindsStr.data(), KindsStr.size() - 2) + ").",
+        inconvertibleErrorCode());
+  if (Matches.size() > 1)
+    return make_error<StringError>(
+        "\"" + TargetName + "\" has " + Twine(Matches.size()) +
+            " candidate stubs in file \"" + FileName +
+            "\". Please refine stub-kind filter \"" + KindNameFilter +
+            "\" for disambiguation (encountered kinds are " +
+            StringRef(KindsStr.data(), KindsStr.size() - 2) + ").",
+        inconvertibleErrorCode());
+
+  return *Matches[0];
 }
 
 Expected<Session::MemoryRegionInfo &>
@@ -2015,8 +2082,9 @@ static Error runChecks(Session &S, Triple TT, SubtargetFeatures Features) {
     return S.findSectionInfo(FileName, SectionName);
   };
 
-  auto GetStubInfo = [&S](StringRef FileName, StringRef SectionName) {
-    return S.findStubInfo(FileName, SectionName);
+  auto GetStubInfo = [&S](StringRef FileName, StringRef SectionName,
+                          StringRef KindNameFilter) {
+    return S.findStubInfo(FileName, SectionName, KindNameFilter);
   };
 
   auto GetGOTInfo = [&S](StringRef FileName, StringRef SectionName) {

llvm/tools/llvm-jitlink/llvm-jitlink.h

@@ -49,7 +49,7 @@ struct Session {
 
   struct FileInfo {
     StringMap<MemoryRegionInfo> SectionInfos;
-    StringMap<MemoryRegionInfo> StubInfos;
+    StringMap<SmallVector<MemoryRegionInfo, 1>> StubInfos;
     StringMap<MemoryRegionInfo> GOTEntryInfos;
 
     using Symbol = jitlink::Symbol;
@@ -61,6 +61,8 @@ struct Session {
                            GetSymbolTargetFunction GetSymbolTarget);
     Error registerStubEntry(LinkGraph &G, Symbol &Sym,
                             GetSymbolTargetFunction GetSymbolTarget);
+    Error registerMultiStubEntry(LinkGraph &G, Symbol &Sym,
+                                 GetSymbolTargetFunction GetSymbolTarget);
   };
 
   using DynLibJDMap = std::map<std::string, orc::JITDylib *>;
@@ -74,7 +76,8 @@ struct Session {
   Expected<MemoryRegionInfo &> findSectionInfo(StringRef FileName,
                                                StringRef SectionName);
   Expected<MemoryRegionInfo &> findStubInfo(StringRef FileName,
-                                            StringRef TargetName);
+                                            StringRef TargetName,
+                                            StringRef KindNameFilter);
   Expected<MemoryRegionInfo &> findGOTEntryInfo(StringRef FileName,
                                                 StringRef TargetName);
 

llvm/tools/llvm-rtdyld/llvm-rtdyld.cpp

@@ -926,7 +926,8 @@ static int linkAndVerify() {
   };
 
   auto GetStubInfo = [&Dyld, &StubMap](StringRef StubContainer,
-                                       StringRef SymbolName)
+                                       StringRef SymbolName,
+                                       StringRef KindNameFilter)
       -> Expected<RuntimeDyldChecker::MemoryRegionInfo> {
     if (!StubMap.count(StubContainer))
       return make_error<StringError>("Stub container not found: " +
@@ -947,6 +948,11 @@ static int linkAndVerify() {
     return StubMemInfo;
   };
 
+  auto GetGOTInfo = [&GetStubInfo](StringRef StubContainer,
+                                   StringRef SymbolName) {
+    return GetStubInfo(StubContainer, SymbolName, "");
+  };
+
   // We will initialize this below once we have the first object file and can
   // know the endianness.
   std::unique_ptr<RuntimeDyldChecker> Checker;
@@ -977,8 +983,7 @@ static int linkAndVerify() {
 
     if (!Checker)
       Checker = std::make_unique<RuntimeDyldChecker>(
-          IsSymbolValid, GetSymbolInfo, GetSectionInfo, GetStubInfo,
-          GetStubInfo,
+          IsSymbolValid, GetSymbolInfo, GetSectionInfo, GetStubInfo, GetGOTInfo,
           Obj.isLittleEndian() ? llvm::endianness::little
                                : llvm::endianness::big,
           TheTriple, MCPU, SubtargetFeatures(), dbgs());