[ORC] Add a 'remove' method to JITDylib to remove symbols.

Symbols can be removed provided that all are present in the JITDylib and none
are currently in the materializing state. On success all requested symbols are
removed. On failure an error is returned and no symbols are removed.

llvm-svn: 343928

1 files changed, 86 insertions, 0 deletions

diff --git a/llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp b/llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp
index d0e4effe75e..47f62a796b8 100644
--- a/llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp
+++ b/llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp
@@ -107,6 +107,92 @@ TEST_F(CoreAPIsStandardTest, EmptyLookup) {
   EXPECT_TRUE(OnReadyRun) << "OnReady was not run for empty query";
 }
 
+TEST_F(CoreAPIsStandardTest, RemoveSymbolsTest) {
+  // Test that:
+  // (1) Missing symbols generate a SymbolsNotFound error.
+  // (2) Materializing symbols generate a SymbolCouldNotBeRemoved error.
+  // (3) Removal of unmaterialized symbols triggers discard on the
+  //     materialization unit.
+  // (4) Removal of symbols destroys empty materialization units.
+  // (5) Removal of materialized symbols works.
+
+  // Foo will be fully materialized.
+  cantFail(JD.define(absoluteSymbols({{Foo, FooSym}})));
+
+  // Bar will be unmaterialized.
+  bool BarDiscarded = false;
+  bool BarMaterializerDestructed = false;
+  cantFail(JD.define(llvm::make_unique<SimpleMaterializationUnit>(
+      SymbolFlagsMap({{Bar, BarSym.getFlags()}}),
+      [this](MaterializationResponsibility R) {
+        ADD_FAILURE() << "Unexpected materialization of \"Bar\"";
+        R.resolve({{Bar, BarSym}});
+        R.emit();
+      },
+      [&](const JITDylib &JD, const SymbolStringPtr &Name) {
+        EXPECT_EQ(Name, Bar) << "Expected \"Bar\" to be discarded";
+        if (Name == Bar)
+          BarDiscarded = true;
+      },
+      [&]() { BarMaterializerDestructed = true; })));
+
+  // Baz will be in the materializing state initially, then
+  // materialized for the final removal attempt.
+  Optional<MaterializationResponsibility> BazR;
+  cantFail(JD.define(llvm::make_unique<SimpleMaterializationUnit>(
+      SymbolFlagsMap({{Baz, BazSym.getFlags()}}),
+      [&](MaterializationResponsibility R) { BazR.emplace(std::move(R)); },
+      [](const JITDylib &JD, const SymbolStringPtr &Name) {
+        ADD_FAILURE() << "\"Baz\" discarded unexpectedly";
+      })));
+
+  bool OnResolvedRun = false;
+  bool OnReadyRun = false;
+  ES.lookup({&JD}, {Foo, Baz},
+            [&](Expected<SymbolMap> Result) {
+              EXPECT_TRUE(!!Result) << "OnResolved failed unexpectedly";
+              consumeError(Result.takeError());
+              OnResolvedRun = true;
+            },
+            [&](Error Err) {
+              EXPECT_FALSE(!!Err) << "OnReady failed unexpectedly";
+              consumeError(std::move(Err));
+              OnReadyRun = true;
+            },
+            NoDependenciesToRegister);
+
+  {
+    // Attempt 1: Search for a missing symbol, Qux.
+    auto Err = JD.remove({Foo, Bar, Baz, Qux});
+    EXPECT_TRUE(!!Err) << "Expected failure";
+    EXPECT_TRUE(Err.isA<SymbolsNotFound>())
+        << "Expected a SymbolsNotFound error";
+  }
+
+  {
+    // Attempt 2: Search for a symbol that is still materializing, Baz.
+    auto Err = JD.remove({Foo, Bar, Baz});
+    EXPECT_TRUE(!!Err) << "Expected failure";
+    EXPECT_TRUE(Err.isA<SymbolsCouldNotBeRemoved>())
+        << "Expected a SymbolsNotFound error";
+  }
+
+  BazR->resolve({{Baz, BazSym}});
+  BazR->emit();
+  {
+    // Attempt 3: Search now that all symbols are fully materialized
+    // (Foo, Baz), or not yet materialized (Bar).
+    auto Err = JD.remove({Foo, Bar, Baz});
+    EXPECT_FALSE(!!Err) << "Expected failure";
+  }
+
+  EXPECT_TRUE(BarDiscarded) << "\"Bar\" should have been discarded";
+  EXPECT_TRUE(BarMaterializerDestructed)
+      << "\"Bar\"'s materializer should have been destructed";
+  EXPECT_TRUE(OnResolvedRun) << "OnResolved should have been run";
+  EXPECT_TRUE(OnReadyRun) << "OnReady should have been run";
+}
+
 TEST_F(CoreAPIsStandardTest, ChainedJITDylibLookup) {
   cantFail(JD.define(absoluteSymbols({{Foo, FooSym}})));