[JITLink] Add convenience methods to LinkGraph to find symbols by name.

Adds new convenience methods findDefinedSymbolByName, findExternalSymbolByName
and findAbsoluteSymbolByName to the LinkGraph class. These should be used to
find symbols of the given types by name.

COFFLinkGraphBuilder and MachOPlatform are updated to take advantage of the
new methods.

Author: lhames

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

@@ -1387,10 +1387,26 @@ class LinkGraph {
                                  GetExternalSymbolMapEntryValue()));
   }
 
+  /// Returns the external symbol with the given name if one exists, otherwise
+  /// returns nullptr.
+  Symbol *findExternalSymbolByName(const orc::SymbolStringPtrBase &Name) {
+    for (auto *Sym : external_symbols())
+      if (Sym->getName() == Name)
+        return Sym;
+    return nullptr;
+  }
+
   iterator_range<absolute_symbol_iterator> absolute_symbols() {
     return make_range(AbsoluteSymbols.begin(), AbsoluteSymbols.end());
   }
 
+  Symbol *findAbsoluteSymbolByName(const orc::SymbolStringPtrBase &Name) {
+    for (auto *Sym : absolute_symbols())
+      if (Sym->getName() == Name)
+        return Sym;
+    return nullptr;
+  }
+
   iterator_range<defined_symbol_iterator> defined_symbols() {
     auto Secs = sections();
     return make_range(defined_symbol_iterator(Secs.begin(), Secs.end()),
@@ -1403,6 +1419,15 @@ class LinkGraph {
                       const_defined_symbol_iterator(Secs.end(), Secs.end()));
   }
 
+  /// Returns the defined symbol with the given name if one exists, otherwise
+  /// returns nullptr.
+  Symbol *findDefinedSymbolByName(const orc::SymbolStringPtrBase &Name) {
+    for (auto *Sym : defined_symbols())
+      if (Sym->hasName() && Sym->getName() == Name)
+        return Sym;
+    return nullptr;
+  }
+
   /// Make the given symbol external (must not already be external).
   ///
   /// Symbol size, linkage and callability will be left unchanged. Symbol scope

llvm/lib/ExecutionEngine/JITLink/COFFLinkGraphBuilder.cpp

@@ -635,32 +635,16 @@ Symbol *GetImageBaseSymbol::operator()(LinkGraph &G) {
     return *ImageBase;
 
   auto IBN = G.intern(ImageBaseName);
+  ImageBase = G.findExternalSymbolByName(IBN);
+  if (*ImageBase)
+    return *ImageBase;
+  ImageBase = G.findAbsoluteSymbolByName(IBN);
+  if (*ImageBase)
+    return *ImageBase;
+  ImageBase = G.findDefinedSymbolByName(IBN);
+  if (*ImageBase)
+    return *ImageBase;
 
-  // Check external symbols for image base.
-  for (auto *Sym : G.external_symbols()) {
-    if (Sym->getName() == IBN) {
-      ImageBase = Sym;
-      return Sym;
-    }
-  }
-
-  // Check absolute symbols (unexpected, but legal).
-  for (auto *Sym : G.absolute_symbols()) {
-    if (Sym->getName() == IBN) {
-      ImageBase = Sym;
-      return Sym;
-    }
-  }
-
-  // Finally, check defined symbols.
-  for (auto *Sym : G.defined_symbols()) {
-    if (Sym->hasName() && Sym->getName() == IBN) {
-      ImageBase = Sym;
-      return Sym;
-    }
-  }
-
-  ImageBase = nullptr;
   return nullptr;
 }
 

llvm/lib/ExecutionEngine/Orc/MachOPlatform.cpp

@@ -1492,26 +1492,19 @@ Error MachOPlatform::MachOPlatformPlugin::populateObjCRuntimeObject(
     memcpy(SD.Sec.sectname, "__objc_imageinfo", 16);
     strcpy(SD.Sec.segname, "__DATA");
     SD.Sec.size = 8;
-    SD.AddFixups = [&](size_t RecordOffset) {
+    jitlink::Symbol *ObjCImageInfoSym = nullptr;
+    SD.AddFixups = [&, ObjCImageInfoSym](size_t RecordOffset) mutable {
       auto PointerEdge = getPointerEdgeKind(G);
 
       // Look for an existing __objc_imageinfo symbol.
-      jitlink::Symbol *ObjCImageInfoSym = nullptr;
-      for (auto *Sym : G.external_symbols())
-        if (Sym->hasName() && *Sym->getName() == ObjCImageInfoSymbolName) {
-          ObjCImageInfoSym = Sym;
-          break;
-        }
-      if (!ObjCImageInfoSym)
-        for (auto *Sym : G.absolute_symbols())
-          if (Sym->hasName() && *Sym->getName() == ObjCImageInfoSymbolName) {
-            ObjCImageInfoSym = Sym;
-            break;
-          }
-      if (!ObjCImageInfoSym)
-        for (auto *Sym : G.defined_symbols())
-          if (Sym->hasName() && *Sym->getName() == ObjCImageInfoSymbolName) {
-            ObjCImageInfoSym = Sym;
+      if (!ObjCImageInfoSym) {
+        auto Name = G.intern(ObjCImageInfoSymbolName);
+        ObjCImageInfoSym = G.findExternalSymbolByName(Name);
+        if (!ObjCImageInfoSym)
+          ObjCImageInfoSym = G.findAbsoluteSymbolByName(Name);
+        if (!ObjCImageInfoSym) {
+          ObjCImageInfoSym = G.findDefinedSymbolByName(Name);
+          if (ObjCImageInfoSym) {
             std::optional<uint32_t> Flags;
             {
               std::lock_guard<std::mutex> Lock(PluginMutex);
@@ -1525,16 +1518,17 @@ Error MachOPlatform::MachOPlatformPlugin::populateObjCRuntimeObject(
             if (Flags) {
               // We own the definition of __objc_image_info; write the final
               // merged flags value.
-              auto Content = Sym->getBlock().getMutableContent(G);
-              assert(Content.size() == 8 &&
+              auto Content = ObjCImageInfoSym->getBlock().getMutableContent(G);
+              assert(
+                  Content.size() == 8 &&
                   "__objc_image_info size should have been verified already");
               support::endian::write32(&Content[4], *Flags, G.getEndianness());
             }
-            break;
           }
-      if (!ObjCImageInfoSym)
-        ObjCImageInfoSym =
-            &G.addExternalSymbol(ObjCImageInfoSymbolName, 8, false);
+        }
+        if (!ObjCImageInfoSym)
+          ObjCImageInfoSym = &G.addExternalSymbol(std::move(Name), 8, false);
+      }
 
       SecBlock.addEdge(PointerEdge,
                        RecordOffset + ((char *)&SD.Sec.addr - (char *)&SD.Sec),

llvm/unittests/ExecutionEngine/JITLink/LinkGraphTests.cpp

@@ -217,6 +217,51 @@ TEST(LinkGraphTest, ContentAccessAndUpdate) {
                            [](char C) { return C == 0; }));
 }
 
+TEST(LinkGraphTest, FindSymbolsByName) {
+  // Check that we can make defined and absolute symbols external.
+  LinkGraph G("foo", std::make_shared<orc::SymbolStringPool>(),
+              Triple("x86_64-apple-darwin"), SubtargetFeatures(),
+              getGenericEdgeKindName);
+  auto &Sec =
+      G.createSection("__data", orc::MemProt::Read | orc::MemProt::Write);
+
+  auto &B1 =
+      G.createContentBlock(Sec, BlockContent, orc::ExecutorAddr(0x1000), 8, 0);
+
+  // Add an anonymous symbol to make sure that these don't disrupt by-name
+  // lookup of defined symbols.
+  G.addAnonymousSymbol(B1, 0, 0, false, false);
+
+  // Add named defined, external and absolute symbols.
+  auto Foo = G.intern("foo");
+  auto &FooSym = G.addDefinedSymbol(B1, 0, Foo, 4, Linkage::Strong,
+                                    Scope::Default, false, false);
+
+  auto Bar = G.intern("bar");
+  auto &BarSym = G.addExternalSymbol(Bar, 0, false);
+
+  auto Baz = G.intern("baz");
+  auto &BazSym = G.addAbsoluteSymbol(Baz, orc::ExecutorAddr(0x1234), 0,
+                                     Linkage::Strong, Scope::Default, true);
+
+  EXPECT_EQ(G.findDefinedSymbolByName(Foo), &FooSym);
+  EXPECT_EQ(G.findExternalSymbolByName(Foo), nullptr);
+  EXPECT_EQ(G.findAbsoluteSymbolByName(Foo), nullptr);
+
+  EXPECT_EQ(G.findDefinedSymbolByName(Bar), nullptr);
+  EXPECT_EQ(G.findExternalSymbolByName(Bar), &BarSym);
+  EXPECT_EQ(G.findAbsoluteSymbolByName(Bar), nullptr);
+
+  EXPECT_EQ(G.findDefinedSymbolByName(Baz), nullptr);
+  EXPECT_EQ(G.findExternalSymbolByName(Baz), nullptr);
+  EXPECT_EQ(G.findAbsoluteSymbolByName(Baz), &BazSym);
+
+  auto Qux = G.intern("qux");
+  EXPECT_EQ(G.findDefinedSymbolByName(Qux), nullptr);
+  EXPECT_EQ(G.findExternalSymbolByName(Qux), nullptr);
+  EXPECT_EQ(G.findAbsoluteSymbolByName(Qux), nullptr);
+}
+
 TEST(LinkGraphTest, MakeExternal) {
   // Check that we can make defined and absolute symbols external.
   LinkGraph G("foo", std::make_shared<orc::SymbolStringPool>(),