[ORC] Remove symbols from dependency lists when failing materialization.

When failing materialization of a symbol X, remove X from the dependants list
of any of X's dependencies. This ensures that when X's dependencies are
emitted (or fail themselves) they do not try to access the no-longer-existing
MaterializationInfo for X.

llvm-svn: 359252

3 files changed, 68 insertions, 3 deletions

diff --git a/llvm/lib/ExecutionEngine/Orc/Core.cpp b/llvm/lib/ExecutionEngine/Orc/Core.cpp
index cc52b564bec..4aa763feb79 100644
--- a/llvm/lib/ExecutionEngine/Orc/Core.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/Core.cpp
@@ -390,8 +390,8 @@ SymbolNameSet MaterializationResponsibility::getRequestedSymbols() const {
 }
 
 void MaterializationResponsibility::resolve(const SymbolMap &Symbols) {
-  LLVM_DEBUG(dbgs() << "In " << JD.getName() << " resolving " << Symbols
-                    << "\n");
+  LLVM_DEBUG(
+      { dbgs() << "In " << JD.getName() << " resolving " << Symbols << "\n"; });
 #ifndef NDEBUG
   for (auto &KV : Symbols) {
     auto I = SymbolFlags.find(KV.first);
@@ -412,6 +412,11 @@ void MaterializationResponsibility::resolve(const SymbolMap &Symbols) {
 }
 
 void MaterializationResponsibility::emit() {
+
+  LLVM_DEBUG({
+    dbgs() << "In " << JD.getName() << " emitting " << SymbolFlags << "\n";
+  });
+
 #ifndef NDEBUG
   for (auto &KV : SymbolFlags)
     assert(!KV.second.isMaterializing() &&
@@ -441,6 +446,11 @@ Error MaterializationResponsibility::defineMaterializing(
 
 void MaterializationResponsibility::failMaterialization() {
 
+  LLVM_DEBUG({
+    dbgs() << "In " << JD.getName() << " failing materialization for "
+           << SymbolFlags << "\n";
+  });
+
   SymbolNameSet FailedSymbols;
   for (auto &KV : SymbolFlags)
     FailedSymbols.insert(KV.first);
@@ -1025,6 +1035,23 @@ void JITDylib::notifyFailed(const SymbolNameSet &FailedSymbols) {
       if (MII == MaterializingInfos.end())
         continue;
 
+      // Remove this symbol from the dependants list of any dependencies.
+      for (auto &KV : MII->second.UnemittedDependencies) {
+        auto *DependencyJD = KV.first;
+        auto &Dependencies = KV.second;
+        for (auto &DependencyName : Dependencies) {
+          auto DependencyMII =
+              DependencyJD->MaterializingInfos.find(DependencyName);
+          assert(DependencyMII != DependencyJD->MaterializingInfos.end() &&
+                 "Unemitted dependency must have a MaterializingInfo entry");
+          assert(DependencyMII->second.Dependants.count(this) &&
+                 "Dependency's dependants list does not contain this JITDylib");
+          assert(DependencyMII->second.Dependants[this].count(Name) &&
+                 "Dependency's dependants list does not contain dependant");
+          DependencyMII->second.Dependants[this].erase(Name);
+        }
+      }
+
       // Copy all the queries to the FailedQueries list, then abandon them.
       // This has to be a copy, and the copy has to come before the abandon
       // operation: Each Q.detach() call will reach back into this
diff --git a/llvm/unittests/ExecutionEngine/Orc/CMakeLists.txt b/llvm/unittests/ExecutionEngine/Orc/CMakeLists.txt
index 019437d4ad5..76b18fd1ba9 100644
--- a/llvm/unittests/ExecutionEngine/Orc/CMakeLists.txt
+++ b/llvm/unittests/ExecutionEngine/Orc/CMakeLists.txt
@@ -30,4 +30,6 @@ add_llvm_unittest(OrcJITTests
   ThreadSafeModuleTest.cpp
   )
 
-target_link_libraries(OrcJITTests PRIVATE ${ORC_JIT_TEST_LIBS})
+target_link_libraries(OrcJITTests PRIVATE
+                        LLVMTestingSupport
+                        ${ORC_JIT_TEST_LIBS})
diff --git a/llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp b/llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp
index 84b19ec2f90..7b29f4c90da 100644
--- a/llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp
+++ b/llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp
@@ -10,6 +10,7 @@
 #include "llvm/Config/llvm-config.h"
 #include "llvm/ExecutionEngine/Orc/Core.h"
 #include "llvm/ExecutionEngine/Orc/OrcError.h"
+#include "llvm/Testing/Support/Error.h"
 
 #include <set>
 #include <thread>
@@ -730,6 +731,41 @@ TEST_F(CoreAPIsStandardTest, FailResolution) {
   }
 }
 
+TEST_F(CoreAPIsStandardTest, FailEmissionEarly) {
+
+  cantFail(JD.define(absoluteSymbols({{Baz, BazSym}})));
+
+  auto MU = llvm::make_unique<SimpleMaterializationUnit>(
+      SymbolFlagsMap({{Foo, FooSym.getFlags()}, {Bar, BarSym.getFlags()}}),
+      [&](MaterializationResponsibility R) {
+        R.resolve(SymbolMap({{Foo, FooSym}, {Bar, BarSym}}));
+
+        ES.lookup(
+            JITDylibSearchList({{&JD, false}}), SymbolNameSet({Baz}),
+            [&R](Expected<SymbolMap> Result) {
+              // Called when "baz" is resolved. We don't actually depend
+              // on or care about baz, but use it to trigger failure of
+              // this materialization before Baz has been finalized in
+              // order to test that error propagation is correct in this
+              // scenario.
+              cantFail(std::move(Result));
+              R.failMaterialization();
+            },
+            [](Error Err) { cantFail(std::move(Err)); },
+            [&](const SymbolDependenceMap &Deps) {
+              R.addDependenciesForAll(Deps);
+            });
+      });
+
+  cantFail(JD.define(MU));
+
+  SymbolNameSet Names({Foo, Bar});
+  auto Result = ES.lookup(JITDylibSearchList({{&JD, false}}), Names);
+
+  EXPECT_THAT_EXPECTED(std::move(Result), Failed())
+      << "Unexpected success while trying to test error propagation";
+}
+
 TEST_F(CoreAPIsStandardTest, TestLookupWithUnthreadedMaterialization) {
   auto MU = llvm::make_unique<SimpleMaterializationUnit>(
       SymbolFlagsMap({{Foo, JITSymbolFlags::Exported}}),