[ORC] Re-apply r336760 with fixes.

llvm-svn: 337637

include/llvm/ExecutionEngine/Orc/Core.h

@@ -179,8 +179,11 @@ public:
   /// threads, or different kinds of materialization processes.
   MaterializationResponsibility delegate(const SymbolNameSet &Symbols);
 
-  /// Add dependencies for the symbols in this dylib.
-  void addDependencies(const SymbolDependenceMap &Dependencies);
+  void addDependencies(const SymbolStringPtr &Name,
+                       const SymbolDependenceMap &Dependencies);
+
+  /// Add dependencies that apply to all symbols covered by this instance.
+  void addDependenciesForAll(const SymbolDependenceMap &Dependencies);
 
 private:
   /// Create a MaterializationResponsibility for the given VSO and

lib/ExecutionEngine/Orc/Core.cpp

@@ -11,6 +11,7 @@
 #include "llvm/Config/llvm-config.h"
 #include "llvm/ExecutionEngine/Orc/OrcError.h"
 #include "llvm/IR/Mangler.h"
+#include "llvm/Support/Debug.h"
 #include "llvm/Support/Format.h"
 
 #if LLVM_ENABLE_THREADS
@@ -685,6 +686,13 @@ MaterializationResponsibility::delegate(const SymbolNameSet &Symbols) {
 }
 
 void MaterializationResponsibility::addDependencies(
+    const SymbolStringPtr &Name, const SymbolDependenceMap &Dependencies) {
+  assert(SymbolFlags.count(Name) &&
+         "Symbol not covered by this MaterializationResponsibility instance");
+  V.addDependencies(Name, Dependencies);
+}
+
+void MaterializationResponsibility::addDependenciesForAll(
     const SymbolDependenceMap &Dependencies) {
   for (auto &KV : SymbolFlags)
     V.addDependencies(KV.first, Dependencies);
@@ -797,8 +805,25 @@ void ReExportsMaterializationUnit::materialize(
 
     QueryInfos.pop_back();
 
-    auto RegisterDependencies = [&](const SymbolDependenceMap &Deps) {
-      R.addDependencies(Deps);
+    auto RegisterDependencies = [QueryInfo,
+                                 &SrcV](const SymbolDependenceMap &Deps) {
+      // If there were no materializing symbols, just bail out.
+      if (Deps.empty())
+        return;
+
+      // Otherwise the only deps should be on SrcV.
+      assert(Deps.size() == 1 && Deps.count(&SrcV) &&
+             "Unexpected dependencies for reexports");
+
+      auto &SrcVDeps = Deps.find(&SrcV)->second;
+      SymbolDependenceMap PerAliasDepsMap;
+      auto &PerAliasDeps = PerAliasDepsMap[&SrcV];
+
+      for (auto &KV : QueryInfo->Aliases)
+        if (SrcVDeps.count(KV.second.Aliasee)) {
+          PerAliasDeps = {KV.second.Aliasee};
+          QueryInfo->R.addDependencies(KV.first, PerAliasDepsMap);
+        }
     };
 
     auto OnResolve = [QueryInfo](Expected<SymbolMap> Result) {
@@ -979,6 +1004,15 @@ void VSO::addDependencies(const SymbolStringPtr &Name,
     auto &DepsOnOtherVSO = MI.UnfinalizedDependencies[&OtherVSO];
 
     for (auto &OtherSymbol : KV.second) {
+#ifndef NDEBUG
+      // Assert that this symbol exists and has not been finalized already.
+      auto SymI = OtherVSO.Symbols.find(OtherSymbol);
+      assert(SymI != OtherVSO.Symbols.end() &&
+             (SymI->second.getFlags().isLazy() ||
+              SymI->second.getFlags().isMaterializing()) &&
+             "Dependency on finalized symbol");
+#endif
+
       auto &OtherMI = OtherVSO.MaterializingInfos[OtherSymbol];
 
       if (OtherMI.IsFinalized)

lib/ExecutionEngine/Orc/Legacy.cpp

@@ -31,7 +31,7 @@ JITSymbolResolverAdapter::lookup(const LookupSet &Symbols) {
 
   auto RegisterDependencies = [&](const SymbolDependenceMap &Deps) {
     if (MR)
-      MR->addDependencies(Deps);
+      MR->addDependenciesForAll(Deps);
   };
 
   auto InternedResult =

lib/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.cpp

@@ -26,7 +26,7 @@ public:
       InternedSymbols.insert(ES.getSymbolStringPool().intern(S));
 
     auto RegisterDependencies = [&](const SymbolDependenceMap &Deps) {
-      MR.addDependencies(Deps);
+      MR.addDependenciesForAll(Deps);
     };
 
     auto InternedResult =

unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp

@@ -231,6 +231,48 @@ TEST_F(CoreAPIsStandardTest, TestChainedAliases) {
       << "\"Baz\"'s address should match \"Foo\"'s";
 }
 
+TEST_F(CoreAPIsStandardTest, TestBasicReExports) {
+  // Test that the basic use case of re-exporting a single symbol from another
+  // VSO works.
+  cantFail(V.define(absoluteSymbols({{Foo, FooSym}})));
+
+  auto &V2 = ES.createVSO("V2");
+
+  cantFail(V2.define(reexports(V, {{Bar, {Foo, BarSym.getFlags()}}})));
+
+  auto Result = cantFail(lookup({&V2}, Bar));
+  EXPECT_EQ(Result.getAddress(), FooSym.getAddress())
+      << "Re-export Bar for symbol Foo should match FooSym's address";
+}
+
+TEST_F(CoreAPIsStandardTest, TestThatReExportsDontUnnecessarilyMaterialize) {
+  // Test that re-exports do not materialize symbols that have not been queried
+  // for.
+  cantFail(V.define(absoluteSymbols({{Foo, FooSym}})));
+
+  bool BarMaterialized = false;
+  auto BarMU = llvm::make_unique<SimpleMaterializationUnit>(
+      SymbolFlagsMap({{Bar, BarSym.getFlags()}}),
+      [&](MaterializationResponsibility R) {
+        BarMaterialized = true;
+        R.resolve({{Bar, BarSym}});
+        R.finalize();
+      });
+
+  cantFail(V.define(BarMU));
+
+  auto &V2 = ES.createVSO("V2");
+
+  cantFail(V2.define(reexports(
+      V, {{Baz, {Foo, BazSym.getFlags()}}, {Qux, {Bar, QuxSym.getFlags()}}})));
+
+  auto Result = cantFail(lookup({&V2}, Baz));
+  EXPECT_EQ(Result.getAddress(), FooSym.getAddress())
+      << "Re-export Baz for symbol Foo should match FooSym's address";
+
+  EXPECT_FALSE(BarMaterialized) << "Bar should not have been materialized";
+}
+
 TEST_F(CoreAPIsStandardTest, TestTrivialCircularDependency) {
   Optional<MaterializationResponsibility> FooR;
   auto FooMU = llvm::make_unique<SimpleMaterializationUnit>(
@@ -338,15 +380,15 @@ TEST_F(CoreAPIsStandardTest, TestCircularDependenceInOneVSO) {
             NoDependenciesToRegister);
 
   // Add a circular dependency: Foo -> Bar, Bar -> Baz, Baz -> Foo.
-  FooR->addDependencies({{&V, SymbolNameSet({Bar})}});
-  BarR->addDependencies({{&V, SymbolNameSet({Baz})}});
-  BazR->addDependencies({{&V, SymbolNameSet({Foo})}});
+  FooR->addDependenciesForAll({{&V, SymbolNameSet({Bar})}});
+  BarR->addDependenciesForAll({{&V, SymbolNameSet({Baz})}});
+  BazR->addDependenciesForAll({{&V, SymbolNameSet({Foo})}});
 
   // Add self-dependencies for good measure. This tests that the implementation
   // of addDependencies filters these out.
-  FooR->addDependencies({{&V, SymbolNameSet({Foo})}});
-  BarR->addDependencies({{&V, SymbolNameSet({Bar})}});
-  BazR->addDependencies({{&V, SymbolNameSet({Baz})}});
+  FooR->addDependenciesForAll({{&V, SymbolNameSet({Foo})}});
+  BarR->addDependenciesForAll({{&V, SymbolNameSet({Bar})}});
+  BazR->addDependenciesForAll({{&V, SymbolNameSet({Baz})}});
 
   // Check that nothing has been resolved yet.
   EXPECT_FALSE(FooResolved) << "\"Foo\" should not be resolved yet";