[ORC] Generalize alias materialization to support re-exports (i.e. aliasing of

symbols in another VSO).

Also fixes a bug where chained aliases within a single VSO would deadlock on
materialization.

llvm-svn: 336741

include/llvm/ExecutionEngine/Orc/Core.h

@@ -251,6 +251,10 @@ absoluteSymbols(SymbolMap Symbols) {
 }
 
 struct SymbolAliasMapEntry {
+  SymbolAliasMapEntry() = default;
+  SymbolAliasMapEntry(SymbolStringPtr Aliasee, JITSymbolFlags AliasFlags)
+      : Aliasee(std::move(Aliasee)), AliasFlags(AliasFlags) {}
+
   SymbolStringPtr Aliasee;
   JITSymbolFlags AliasFlags;
 };
@@ -260,19 +264,27 @@ using SymbolAliasMap = std::map<SymbolStringPtr, SymbolAliasMapEntry>;
 
 /// A materialization unit for symbol aliases. Allows existing symbols to be
 /// aliased with alternate flags.
-class SymbolAliasesMaterializationUnit : public MaterializationUnit {
+class ReExportsMaterializationUnit : public MaterializationUnit {
 public:
-  SymbolAliasesMaterializationUnit(SymbolAliasMap Aliases);
+  /// SourceVSO is allowed to be nullptr, in which case the source VSO is
+  /// taken to be whatever VSO these definitions are materialized in. This
+  /// is useful for defining aliases within a VSO.
+  ///
+  /// Note: Care must be taken that no sets of aliases form a cycle, as such
+  ///       a cycle will result in a deadlock when any symbol in the cycle is
+  ///       resolved.
+  ReExportsMaterializationUnit(VSO *SourceVSO, SymbolAliasMap Aliases);
 
 private:
   void materialize(MaterializationResponsibility R) override;
   void discard(const VSO &V, SymbolStringPtr Name) override;
   static SymbolFlagsMap extractFlags(const SymbolAliasMap &Aliases);
 
+  VSO *SourceVSO = nullptr;
   SymbolAliasMap Aliases;
 };
 
-/// Create a SymbolAliasesMaterializationUnit with the given aliases.
+/// Create a ReExportsMaterializationUnit with the given aliases.
 /// Useful for defining symbol aliases.: E.g., given a VSO V containing symbols
 /// "foo" and "bar", we can define aliases "baz" (for "foo") and "qux" (for
 /// "bar") with:
@@ -284,12 +296,25 @@ private:
 ///       {Qux, { Bar, JITSymbolFlags::Weak }}}))
 ///     return Err;
 /// \endcode
-inline std::unique_ptr<SymbolAliasesMaterializationUnit>
+inline std::unique_ptr<ReExportsMaterializationUnit>
 symbolAliases(SymbolAliasMap Aliases) {
-  return llvm::make_unique<SymbolAliasesMaterializationUnit>(
-      std::move(Aliases));
+  return llvm::make_unique<ReExportsMaterializationUnit>(nullptr,
+                                                         std::move(Aliases));
 }
 
+/// Create a materialization unit for re-exporting symbols from another VSO
+/// with alternative names/flags.
+inline std::unique_ptr<ReExportsMaterializationUnit>
+reexports(VSO &SourceV, SymbolAliasMap Aliases) {
+  return llvm::make_unique<ReExportsMaterializationUnit>(&SourceV,
+                                                         std::move(Aliases));
+}
+
+/// Build a SymbolAliasMap for the common case where you want to re-export
+/// symbols from another VSO with the same linkage/flags.
+Expected<SymbolAliasMap>
+buildSimpleReexportsAliasMap(VSO &SourceV, const SymbolNameSet &Symbols);
+
 /// Base utilities for ExecutionSession.
 class ExecutionSessionBase {
 public:

lib/ExecutionEngine/Orc/Core.cpp

@@ -363,62 +363,143 @@ AbsoluteSymbolsMaterializationUnit::extractFlags(const SymbolMap &Symbols) {
   return Flags;
 }
 
-SymbolAliasesMaterializationUnit::SymbolAliasesMaterializationUnit(
-    SymbolAliasMap Aliases)
-    : MaterializationUnit(extractFlags(Aliases)), Aliases(std::move(Aliases)) {}
+ReExportsMaterializationUnit::ReExportsMaterializationUnit(
+    VSO *SourceVSO, SymbolAliasMap Aliases)
+    : MaterializationUnit(extractFlags(Aliases)), SourceVSO(SourceVSO),
+      Aliases(std::move(Aliases)) {}
 
-void SymbolAliasesMaterializationUnit::materialize(
+void ReExportsMaterializationUnit::materialize(
     MaterializationResponsibility R) {
-  auto &V = R.getTargetVSO();
-  auto &ES = V.getExecutionSession();
 
-  // FIXME: Use a unique_ptr when we move to C++14 and have generalized lambda
-  // capture.
-  auto SharedR = std::make_shared<MaterializationResponsibility>(std::move(R));
+  VSO &SrcV = SourceVSO ? *SourceVSO : R.getTargetVSO();
+  auto &ES = SrcV.getExecutionSession();
 
-  auto OnResolve = [this, SharedR](
-                       Expected<AsynchronousSymbolQuery::ResolutionResult> RR) {
-    if (RR) {
-      SymbolMap ResolutionMap;
-      for (auto &KV : Aliases) {
-        assert(RR->Symbols.count(KV.second.Aliasee) &&
-               "Result map missing entry?");
-        ResolutionMap[KV.first] = JITEvaluatedSymbol(
-            RR->Symbols[KV.second.Aliasee].getAddress(), KV.second.AliasFlags);
-      }
+  // Find the set of requested aliases and aliasees. Return any unrequested
+  // aliases back to the VSO so as to not prematurely materialize any aliasees.
+  auto RequestedSymbols = R.getRequestedSymbols();
+  SymbolAliasMap RequestedAliases;
 
-      SharedR->resolve(ResolutionMap);
-      SharedR->finalize();
-    } else {
-      auto &ES = SharedR->getTargetVSO().getExecutionSession();
-      ES.reportError(RR.takeError());
-      SharedR->failMaterialization();
-    }
+  for (auto &Name : RequestedSymbols) {
+    auto I = Aliases.find(Name);
+    assert(I != Aliases.end() && "Symbol not found in aliases map?");
+    RequestedAliases[Name] = std::move(I->second);
+    Aliases.erase(I);
+  }
+
+  if (!Aliases.empty()) {
+    if (SourceVSO)
+      R.replace(reexports(*SourceVSO, std::move(Aliases)));
+    else
+      R.replace(symbolAliases(std::move(Aliases)));
+  }
+
+  // The OnResolveInfo struct will hold the aliases and responsibilty for each
+  // query in the list.
+  struct OnResolveInfo {
+    OnResolveInfo(MaterializationResponsibility R, SymbolAliasMap Aliases)
+        : R(std::move(R)), Aliases(std::move(Aliases)) {}
+
+    MaterializationResponsibility R;
+    SymbolAliasMap Aliases;
   };
 
-  auto OnReady = [&ES](Error Err) { ES.reportError(std::move(Err)); };
+  // Build a list of queries to issue. In each round we build the largest set of
+  // aliases that we can resolve without encountering a chain definition of the
+  // form Foo -> Bar, Bar -> Baz. Such a form would deadlock as the query would
+  // be waitin on a symbol that it itself had to resolve. Usually this will just
+  // involve one round and a single query.
 
-  SymbolNameSet Aliasees;
-  for (auto &KV : Aliases)
-    Aliasees.insert(KV.second.Aliasee);
+  std::vector<std::pair<SymbolNameSet, std::shared_ptr<OnResolveInfo>>>
+      QueryInfos;
+  while (!RequestedAliases.empty()) {
+    SymbolNameSet ResponsibilitySymbols;
+    SymbolNameSet QuerySymbols;
+    SymbolAliasMap QueryAliases;
 
-  auto Q = std::make_shared<AsynchronousSymbolQuery>(
-      Aliasees, std::move(OnResolve), std::move(OnReady));
-  auto Unresolved = V.lookup(Q, Aliasees);
+    for (auto I = RequestedAliases.begin(), E = RequestedAliases.end();
+         I != E;) {
+      auto Tmp = I++;
 
-  if (!Unresolved.empty())
-    ES.failQuery(*Q, make_error<SymbolsNotFound>(std::move(Unresolved)));
+      // Chain detected. Skip this symbol for this round.
+      if (&SrcV == &R.getTargetVSO() &&
+          (QueryAliases.count(Tmp->second.Aliasee) ||
+           RequestedAliases.count(Tmp->second.Aliasee)))
+        continue;
+
+      ResponsibilitySymbols.insert(Tmp->first);
+      QuerySymbols.insert(Tmp->second.Aliasee);
+      QueryAliases[Tmp->first] = std::move(Tmp->second);
+      RequestedAliases.erase(Tmp);
+    }
+    assert(!QuerySymbols.empty() && "Alias cycle detected!");
+
+    auto QueryInfo = std::make_shared<OnResolveInfo>(
+        R.delegate(ResponsibilitySymbols), std::move(QueryAliases));
+    QueryInfos.push_back(
+        make_pair(std::move(QuerySymbols), std::move(QueryInfo)));
+  }
+
+  // Issue the queries.
+  while (!QueryInfos.empty()) {
+    auto QuerySymbols = std::move(QueryInfos.back().first);
+    auto QueryInfo = std::move(QueryInfos.back().second);
+
+    QueryInfos.pop_back();
+
+    auto OnResolve =
+        [QueryInfo,
+         &SrcV](Expected<AsynchronousSymbolQuery::ResolutionResult> RR) {
+          if (RR) {
+            SymbolMap ResolutionMap;
+            SymbolNameSet Resolved;
+            for (auto &KV : QueryInfo->Aliases) {
+              assert(RR->Symbols.count(KV.second.Aliasee) &&
+                     "Result map missing entry?");
+              ResolutionMap[KV.first] = JITEvaluatedSymbol(
+                  RR->Symbols[KV.second.Aliasee].getAddress(),
+                  KV.second.AliasFlags);
+
+              // FIXME: We're creating a SymbolFlagsMap and a std::map of
+              // std::sets just to add one dependency here. This needs a
+              // re-think.
+              Resolved.insert(KV.first);
+            }
+            QueryInfo->R.resolve(ResolutionMap);
+
+            SymbolDependenceMap Deps;
+            Deps[&SrcV] = std::move(Resolved);
+            QueryInfo->R.addDependencies(Deps);
+
+            QueryInfo->R.finalize();
+          } else {
+            auto &ES = QueryInfo->R.getTargetVSO().getExecutionSession();
+            ES.reportError(RR.takeError());
+            QueryInfo->R.failMaterialization();
+          }
+        };
+
+    auto OnReady = [&ES](Error Err) { ES.reportError(std::move(Err)); };
+
+    auto Q = std::make_shared<AsynchronousSymbolQuery>(
+        QuerySymbols, std::move(OnResolve), std::move(OnReady));
+
+    auto Unresolved = SrcV.lookup(Q, std::move(QuerySymbols));
+
+    if (!Unresolved.empty()) {
+      ES.failQuery(*Q, make_error<SymbolsNotFound>(std::move(Unresolved)));
+      return;
+    }
+  }
 }
 
-void SymbolAliasesMaterializationUnit::discard(const VSO &V,
-                                               SymbolStringPtr Name) {
+void ReExportsMaterializationUnit::discard(const VSO &V, SymbolStringPtr Name) {
   assert(Aliases.count(Name) &&
          "Symbol not covered by this MaterializationUnit");
   Aliases.erase(Name);
 }
 
 SymbolFlagsMap
-SymbolAliasesMaterializationUnit::extractFlags(const SymbolAliasMap &Aliases) {
+ReExportsMaterializationUnit::extractFlags(const SymbolAliasMap &Aliases) {
   SymbolFlagsMap SymbolFlags;
   for (auto &KV : Aliases)
     SymbolFlags[KV.first] = KV.second.AliasFlags;
@@ -426,6 +507,23 @@ SymbolAliasesMaterializationUnit::extractFlags(const SymbolAliasMap &Aliases) {
   return SymbolFlags;
 }
 
+Expected<SymbolAliasMap>
+buildSimpleReexportsAliasMap(VSO &SourceV, const SymbolNameSet &Symbols) {
+  SymbolFlagsMap Flags;
+  auto Unresolved = SourceV.lookupFlags(Flags, Symbols);
+
+  if (!Unresolved.empty())
+    return make_error<SymbolsNotFound>(std::move(Unresolved));
+
+  SymbolAliasMap Result;
+  for (auto &Name : Symbols) {
+    assert(Flags.count(Name) && "Missing entry in flags map");
+    Result[Name] = SymbolAliasMapEntry(Name, Flags[Name]);
+  }
+
+  return Result;
+}
+
 Error VSO::defineMaterializing(const SymbolFlagsMap &SymbolFlags) {
   return ES.runSessionLocked([&]() -> Error {
     std::vector<SymbolMap::iterator> AddedSyms;

unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp

@@ -259,7 +259,7 @@ TEST(CoreAPIsTest, LookupFlagsTest) {
   EXPECT_EQ(SymbolFlags[Bar], BarFlags) << "Incorrect flags returned for Bar";
 }
 
-TEST(CoreAPIsTest, TestAliases) {
+TEST(CoreAPIsTest, TestBasicAliases) {
   ExecutionSession ES;
   auto &V = ES.createVSO("V");
 
@@ -288,6 +288,31 @@ TEST(CoreAPIsTest, TestAliases) {
       << "The \"Qux\" alias should have been overriden";
 }
 
+TEST(CoreAPIsTest, TestChainedAliases) {
+  ExecutionSession ES;
+  auto &V = ES.createVSO("V");
+
+  auto Foo = ES.getSymbolStringPool().intern("foo");
+  auto FooSym = JITEvaluatedSymbol(1U, JITSymbolFlags::Exported);
+
+  auto Bar = ES.getSymbolStringPool().intern("bar");
+
+  auto Baz = ES.getSymbolStringPool().intern("baz");
+
+  cantFail(V.define(absoluteSymbols({{Foo, FooSym}})));
+  cantFail(V.define(symbolAliases({{Baz, {Bar, JITSymbolFlags::Exported}},
+                                   {Bar, {Foo, JITSymbolFlags::Exported}}})));
+
+  auto Result = lookup({&V}, {Bar, Baz});
+  EXPECT_TRUE(!!Result) << "Unexpected lookup failure";
+  EXPECT_EQ(Result->count(Bar), 1U) << "No result for \"bar\"";
+  EXPECT_EQ(Result->count(Baz), 1U) << "No result for \"baz\"";
+  EXPECT_EQ((*Result)[Bar].getAddress(), FooSym.getAddress())
+      << "\"Bar\"'s address should match \"Foo\"'s";
+  EXPECT_EQ((*Result)[Baz].getAddress(), FooSym.getAddress())
+      << "\"Baz\"'s address should match \"Foo\"'s";
+}
+
 TEST(CoreAPIsTest, TestTrivialCircularDependency) {
   ExecutionSession ES;