[JITLink][COFF] Move GetImageBaseSymbol utility into public header. (#195041)

This utility may be useful for people writing
LinkGraphLinkingLayer::Plugins for COFF LinkGraphs, so this commit moves
it a public header where it can easily be reused
(llvm/ExecutionEngine/JITLink/COFF.h).

Also adds unit tests for the utility.

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

@@ -34,6 +34,25 @@ createLinkGraphFromCOFFObject(MemoryBufferRef ObjectBuffer,
 void link_COFF(std::unique_ptr<LinkGraph> G,
                std::unique_ptr<JITLinkContext> Ctx);
 
+/// GetImageBaseSymbol is a function object that finds the __ImageBase symbol
+/// in the given graph if one is present.
+///
+/// The result is cached across calls, and can be reset by calling the reset
+/// method.
+class GetImageBaseSymbol {
+public:
+  GetImageBaseSymbol(StringRef ImageBaseName = "__ImageBase")
+      : ImageBaseName(ImageBaseName) {}
+  Symbol *operator()(LinkGraph &G);
+  void reset(std::optional<Symbol *> CacheValue = std::nullopt) {
+    ImageBase = CacheValue;
+  }
+
+private:
+  StringRef ImageBaseName;
+  std::optional<Symbol *> ImageBase;
+};
+
 } // end namespace jitlink
 } // end namespace llvm
 

llvm/lib/ExecutionEngine/JITLink/COFF.cpp

@@ -131,5 +131,23 @@ void link_COFF(std::unique_ptr<LinkGraph> G,
   }
 }
 
+Symbol *GetImageBaseSymbol::operator()(LinkGraph &G) {
+  if (ImageBase)
+    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;
+
+  return nullptr;
+}
+
 } // end namespace jitlink
 } // end namespace llvm

llvm/lib/ExecutionEngine/JITLink/COFFLinkGraphBuilder.cpp

@@ -630,23 +630,5 @@ COFFLinkGraphBuilder::exportCOMDATSymbol(COFFSymbolIndex SymIndex,
   return GSym;
 }
 
-Symbol *GetImageBaseSymbol::operator()(LinkGraph &G) {
-  if (ImageBase)
-    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;
-
-  return nullptr;
-}
-
 } // namespace jitlink
 } // namespace llvm

llvm/lib/ExecutionEngine/JITLink/COFFLinkGraphBuilder.h

@@ -220,18 +220,6 @@ Error COFFLinkGraphBuilder::forEachRelocation(const object::SectionRef &RelSec,
   return Error::success();
 }
 
-class GetImageBaseSymbol {
-public:
-  GetImageBaseSymbol(StringRef ImageBaseName = "__ImageBase")
-      : ImageBaseName(ImageBaseName) {}
-  Symbol *operator()(LinkGraph &G);
-  void reset() { ImageBase = std::nullopt; }
-
-private:
-  StringRef ImageBaseName;
-  std::optional<Symbol *> ImageBase;
-};
-
 } // end namespace jitlink
 } // end namespace llvm
 

llvm/lib/ExecutionEngine/JITLink/COFF_x86_64.cpp

@@ -15,6 +15,7 @@
 #include "JITLinkGeneric.h"
 #include "SEHFrameSupport.h"
 #include "llvm/BinaryFormat/COFF.h"
+#include "llvm/ExecutionEngine/JITLink/COFF.h"
 #include "llvm/ExecutionEngine/JITLink/x86_64.h"
 #include "llvm/Object/COFF.h"
 #include "llvm/Support/Endian.h"

llvm/unittests/ExecutionEngine/JITLink/CMakeLists.txt

@@ -11,6 +11,7 @@ add_llvm_unittest(JITLinkTests
     AArch32Tests.cpp
     AArch32ErrorTests.cpp
     AArch64Tests.cpp
+    COFFLinkGraphTests.cpp
     EHFrameSupportTests.cpp
     JITLinkTestUtils.cpp
     LinkGraphTests.cpp

llvm/unittests/ExecutionEngine/JITLink/COFFLinkGraphTests.cpp

@@ -0,0 +1,109 @@
+//===-- COFFLinkGraphTests.cpp - Unit tests for COFF LinkGraph utils ------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "JITLinkTestUtils.h"
+
+#include "llvm/ExecutionEngine/JITLink/COFF.h"
+#include "llvm/ExecutionEngine/JITLink/JITLink.h"
+
+#include "gtest/gtest.h"
+
+using namespace llvm;
+using namespace llvm::jitlink;
+
+TEST(COFFLinkGraphTest, GetImageBaseSymbolReturnsNullWhenMissing) {
+  LinkGraph G("foo", std::make_shared<orc::SymbolStringPool>(),
+              Triple("x86_64-pc-windows-msvc"), SubtargetFeatures(),
+              getGenericEdgeKindName);
+
+  GetImageBaseSymbol GetIB;
+  EXPECT_EQ(GetIB(G), nullptr);
+}
+
+TEST(COFFLinkGraphTest, GetImageBaseSymbolFindsExternal) {
+  LinkGraph G("foo", std::make_shared<orc::SymbolStringPool>(),
+              Triple("x86_64-pc-windows-msvc"), SubtargetFeatures(),
+              getGenericEdgeKindName);
+
+  auto &ExtSym = G.addExternalSymbol(G.intern("__ImageBase"), 0, false);
+  GetImageBaseSymbol GetIB;
+  EXPECT_EQ(GetIB(G), &ExtSym);
+}
+
+TEST(COFFLinkGraphTest, GetImageBaseSymbolFindsAbsolute) {
+  LinkGraph G("foo", std::make_shared<orc::SymbolStringPool>(),
+              Triple("x86_64-pc-windows-msvc"), SubtargetFeatures(),
+              getGenericEdgeKindName);
+
+  auto &AbsSym =
+      G.addAbsoluteSymbol(G.intern("__ImageBase"), orc::ExecutorAddr(0x1000), 0,
+                          Linkage::Strong, Scope::Default, true);
+  GetImageBaseSymbol GetIB;
+  EXPECT_EQ(GetIB(G), &AbsSym);
+}
+
+TEST(COFFLinkGraphTest, GetImageBaseSymbolFindsDefined) {
+  LinkGraph G("foo", std::make_shared<orc::SymbolStringPool>(),
+              Triple("x86_64-pc-windows-msvc"), SubtargetFeatures(),
+              getGenericEdgeKindName);
+
+  auto &Sec =
+      G.createSection("__data", orc::MemProt::Read | orc::MemProt::Write);
+  auto &B =
+      G.createContentBlock(Sec, BlockContent, orc::ExecutorAddr(0x2000), 8, 0);
+  auto &DefSym =
+      G.addDefinedSymbol(B, 0, G.intern("__ImageBase"), 4, Linkage::Strong,
+                         Scope::Default, false, false);
+  GetImageBaseSymbol GetIB;
+  EXPECT_EQ(GetIB(G), &DefSym);
+}
+
+TEST(COFFLinkGraphTest, GetImageBaseSymbolCachesResult) {
+  LinkGraph G("foo", std::make_shared<orc::SymbolStringPool>(),
+              Triple("x86_64-pc-windows-msvc"), SubtargetFeatures(),
+              getGenericEdgeKindName);
+
+  auto &ExtSym = G.addExternalSymbol(G.intern("__ImageBase"), 0, false);
+  GetImageBaseSymbol GetIB;
+  EXPECT_EQ(GetIB(G), &ExtSym);
+
+  // Remove the symbol -- cached result should still be returned.
+  G.removeExternalSymbol(ExtSym);
+  EXPECT_EQ(GetIB(G), &ExtSym);
+}
+
+TEST(COFFLinkGraphTest, GetImageBaseSymbolReset) {
+  LinkGraph G("foo", std::make_shared<orc::SymbolStringPool>(),
+              Triple("x86_64-pc-windows-msvc"), SubtargetFeatures(),
+              getGenericEdgeKindName);
+
+  GetImageBaseSymbol GetIB;
+  EXPECT_EQ(GetIB(G), nullptr);
+
+  // Add a symbol and reset -- should now find it.
+  auto &ExtSym = G.addExternalSymbol(G.intern("__ImageBase"), 0, false);
+  GetIB.reset();
+  EXPECT_EQ(GetIB(G), &ExtSym);
+
+  // Reset with an explicit cache value of nullptr -- should return nullptr
+  // without searching.
+  GetIB.reset(nullptr);
+  EXPECT_EQ(GetIB(G), nullptr);
+}
+
+TEST(COFFLinkGraphTest, GetImageBaseSymbolCustomName) {
+  LinkGraph G("foo", std::make_shared<orc::SymbolStringPool>(),
+              Triple("x86_64-pc-windows-msvc"), SubtargetFeatures(),
+              getGenericEdgeKindName);
+
+  G.addExternalSymbol(G.intern("__ImageBase"), 0, false);
+  auto &CustomSym = G.addExternalSymbol(G.intern("__CustomBase"), 0, false);
+
+  GetImageBaseSymbol GetIB("__CustomBase");
+  EXPECT_EQ(GetIB(G), &CustomSym);
+}