[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

Author: lhames

llvm/include/llvm/ExecutionEngine/Orc/Core.h

@@ -134,6 +134,20 @@ class SymbolsNotFound : public ErrorInfo<SymbolsNotFound> {
   SymbolNameSet Symbols;
 };
 
+/// Used to notify clients that a set of symbols could not be removed.
+class SymbolsCouldNotBeRemoved : public ErrorInfo<SymbolsCouldNotBeRemoved> {
+public:
+  static char ID;
+
+  SymbolsCouldNotBeRemoved(SymbolNameSet Symbols);
+  std::error_code convertToErrorCode() const override;
+  void log(raw_ostream &OS) const override;
+  const SymbolNameSet &getSymbols() const { return Symbols; }
+
+private:
+  SymbolNameSet Symbols;
+};
+
 /// Tracks responsibility for materialization, and mediates interactions between
 /// MaterializationUnits and JDs.
 ///
@@ -546,6 +560,18 @@ class JITDylib {
   template <typename MaterializationUnitType>
   Error define(std::unique_ptr<MaterializationUnitType> &MU);
 
+  /// Tries to remove the given symbols.
+  ///
+  /// If any symbols are not defined in this JITDylib this method will return
+  /// a SymbolsNotFound error covering the missing symbols.
+  ///
+  /// If all symbols are found but some symbols are in the process of being
+  /// materialized this method will return a SymbolsCouldNotBeRemoved error.
+  ///
+  /// On success, all symbols are removed. On failure, the JITDylib state is
+  /// left unmodified (no symbols are removed).
+  Error remove(const SymbolNameSet &Names);
+
   /// Search the given JITDylib for the symbols in Symbols. If found, store
   ///        the flags for each symbol in Flags. Returns any unresolved symbols.
   SymbolFlagsMap lookupFlags(const SymbolNameSet &Names);

llvm/lib/ExecutionEngine/Orc/Core.cpp

@@ -136,6 +136,7 @@ namespace orc {
 
 char FailedToMaterialize::ID = 0;
 char SymbolsNotFound::ID = 0;
+char SymbolsCouldNotBeRemoved::ID = 0;
 
 RegisterDependenciesFunction NoDependenciesToRegister =
     RegisterDependenciesFunction();
@@ -242,6 +243,19 @@ void SymbolsNotFound::log(raw_ostream &OS) const {
   OS << "Symbols not found: " << Symbols;
 }
 
+SymbolsCouldNotBeRemoved::SymbolsCouldNotBeRemoved(SymbolNameSet Symbols)
+    : Symbols(std::move(Symbols)) {
+  assert(!this->Symbols.empty() && "Can not fail to resolve an empty set");
+}
+
+std::error_code SymbolsCouldNotBeRemoved::convertToErrorCode() const {
+  return orcError(OrcErrorCode::UnknownORCError);
+}
+
+void SymbolsCouldNotBeRemoved::log(raw_ostream &OS) const {
+  OS << "Symbols could not be removed: " << Symbols;
+}
+
 AsynchronousSymbolQuery::AsynchronousSymbolQuery(
     const SymbolNameSet &Symbols, SymbolsResolvedCallback NotifySymbolsResolved,
     SymbolsReadyCallback NotifySymbolsReady)
@@ -1045,6 +1059,60 @@ void JITDylib::removeFromSearchOrder(JITDylib &JD) {
   });
 }
 
+Error JITDylib::remove(const SymbolNameSet &Names) {
+  return ES.runSessionLocked([&]() -> Error {
+    using SymbolMaterializerItrPair =
+        std::pair<SymbolMap::iterator, UnmaterializedInfosMap::iterator>;
+    std::vector<SymbolMaterializerItrPair> SymbolsToRemove;
+    SymbolNameSet Missing;
+    SymbolNameSet Materializing;
+
+    for (auto &Name : Names) {
+      auto I = Symbols.find(Name);
+
+      // Note symbol missing.
+      if (I == Symbols.end()) {
+        Missing.insert(Name);
+        continue;
+      }
+
+      // Note symbol materializing.
+      if (I->second.getFlags().isMaterializing()) {
+        Materializing.insert(Name);
+        continue;
+      }
+
+      auto UMII = I->second.getFlags().isLazy() ? UnmaterializedInfos.find(Name)
+                                                : UnmaterializedInfos.end();
+      SymbolsToRemove.push_back(std::make_pair(I, UMII));
+    }
+
+    // If any of the symbols are not defined, return an error.
+    if (!Missing.empty())
+      return make_error<SymbolsNotFound>(std::move(Missing));
+
+    // If any of the symbols are currently materializing, return an error.
+    if (!Materializing.empty())
+      return make_error<SymbolsCouldNotBeRemoved>(std::move(Materializing));
+
+    // Remove the symbols.
+    for (auto &SymbolMaterializerItrPair : SymbolsToRemove) {
+      auto UMII = SymbolMaterializerItrPair.second;
+
+      // If there is a materializer attached, call discard.
+      if (UMII != UnmaterializedInfos.end()) {
+        UMII->second->MU->doDiscard(*this, UMII->first);
+        UnmaterializedInfos.erase(UMII);
+      }
+
+      auto SymI = SymbolMaterializerItrPair.first;
+      Symbols.erase(SymI);
+    }
+
+    return Error::success();
+  });
+}
+
 SymbolFlagsMap JITDylib::lookupFlags(const SymbolNameSet &Names) {
   return ES.runSessionLocked([&, this]() {
     SymbolFlagsMap Result;

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}})));