[llvm] 56af226 - [ORC] Reset MaterializationUnit::InitSymbol if the init symbol is discarded.

[Lang Hames via llvm-commits]

Author: Lang Hames
Date: 2022-10-23T13:55:55-07:00
New Revision: 56af226b63f86ed1317020ebfe7f8bdd6704998f

URL: https://github.com/llvm/llvm-project/commit/56af226b63f86ed1317020ebfe7f8bdd6704998f
DIFF: https://github.com/llvm/llvm-project/commit/56af226b63f86ed1317020ebfe7f8bdd6704998f.diff

LOG: [ORC] Reset MaterializationUnit::InitSymbol if the init symbol is discarded.

Discarding the init symbol is expected to be uncommon (it represents metadata
in the MaterializationUnit that is relevant to dlopen, and this will not
usually be fully duplicated in some other location), however if a client has
marked an InitSymbol as weak and it is selected to be discarded then we should
keep the data structure consistent.

Added: 
    

Modified: 
    llvm/include/llvm/ExecutionEngine/Orc/Core.h
    llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/ExecutionEngine/Orc/Core.h b/llvm/include/llvm/ExecutionEngine/Orc/Core.h
index f86ef14ffb75b..d5e5222d84762 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/Core.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/Core.h
@@ -707,6 +707,13 @@ class MaterializationUnit {
   /// has been overridden.
   void doDiscard(const JITDylib &JD, const SymbolStringPtr &Name) {
     SymbolFlags.erase(Name);
+    if (InitSymbol == Name) {
+      DEBUG_WITH_TYPE("orc", {
+        dbgs() << "In " << JD.getName() << ", " << getName()
+               << ": discarding init symbol \"" << *Name << "\"\n";
+      });
+      InitSymbol = nullptr;
+    }
     discard(JD, std::move(Name));
   }
 

diff  --git a/llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp b/llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp
index 7ef20cdc49d5b..8022f03172326 100644
--- a/llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp
+++ b/llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp
@@ -237,6 +237,27 @@ TEST_F(CoreAPIsStandardTest, RemoveSymbolsTest) {
   EXPECT_TRUE(OnCompletionRun) << "OnCompletion should have been run";
 }
 
+TEST_F(CoreAPIsStandardTest, DiscardInitSymbol) {
+  SymbolStringPtr ForwardedDiscardSym = nullptr;
+
+  auto MU = std::make_unique<SimpleMaterializationUnit>(
+      SymbolFlagsMap({{Foo, FooSym.getFlags()}, {Bar, BarSym.getFlags()}}),
+      [](std::unique_ptr<MaterializationResponsibility> R) {
+        llvm_unreachable("Materialize called unexpectedly?");
+      },
+      Foo,
+      [&](const JITDylib &, SymbolStringPtr Sym) {
+        ForwardedDiscardSym = std::move(Sym);
+      });
+
+  MU->doDiscard(JD, Foo);
+
+  EXPECT_EQ(ForwardedDiscardSym, Foo);
+  EXPECT_EQ(MU->getSymbols().size(), 1U);
+  EXPECT_TRUE(MU->getSymbols().count(Bar));
+  EXPECT_EQ(MU->getInitializerSymbol(), nullptr);
+}
+
 TEST_F(CoreAPIsStandardTest, LookupWithHiddenSymbols) {
   auto BarHiddenFlags = BarSym.getFlags() & ~JITSymbolFlags::Exported;
   auto BarHiddenSym = JITEvaluatedSymbol(BarSym.getAddress(), BarHiddenFlags);