Merge commit '9c5001e45491' from llvm.org/main into next

Author: git apple-llvm automerger

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