[ORC] Add a symbolAliases function to the Core APIs.

symbolAliases can be used to define symbol aliases within a VSO.

llvm-svn: 335565

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

@@ -245,6 +245,46 @@ absoluteSymbols(SymbolMap Symbols) {
       std::move(Symbols));
 }
 
+struct SymbolAliasMapEntry {
+  SymbolStringPtr Aliasee;
+  JITSymbolFlags AliasFlags;
+};
+
+/// A map of Symbols to (Symbol, Flags) pairs.
+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 {
+public:
+  SymbolAliasesMaterializationUnit(SymbolAliasMap Aliases);
+
+private:
+  void materialize(MaterializationResponsibility R) override;
+  void discard(const VSO &V, SymbolStringPtr Name) override;
+  static SymbolFlagsMap extractFlags(const SymbolAliasMap &Aliases);
+
+  SymbolAliasMap Aliases;
+};
+
+/// Create a SymbolAliasesMaterializationUnit 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:
+/// \code{.cpp}
+///   SymbolStringPtr Baz = ...;
+///   SymbolStringPtr Qux = ...;
+///   if (auto Err = V.define(symbolAliases({
+///       {Baz, { Foo, JITSymbolFlags::Exported }},
+///       {Qux, { Bar, JITSymbolFlags::Weak }}}))
+///     return Err;
+/// \endcode
+inline std::unique_ptr<SymbolAliasesMaterializationUnit>
+symbolAliases(SymbolAliasMap Aliases) {
+  return llvm::make_unique<SymbolAliasesMaterializationUnit>(
+      std::move(Aliases));
+}
+
 /// Base utilities for ExecutionSession.
 class ExecutionSessionBase {
 public:

llvm/lib/ExecutionEngine/Orc/Core.cpp

@@ -344,6 +344,69 @@ AbsoluteSymbolsMaterializationUnit::extractFlags(const SymbolMap &Symbols) {
   return Flags;
 }
 
+SymbolAliasesMaterializationUnit::SymbolAliasesMaterializationUnit(
+    SymbolAliasMap Aliases)
+    : MaterializationUnit(extractFlags(Aliases)), Aliases(std::move(Aliases)) {}
+
+void SymbolAliasesMaterializationUnit::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));
+
+  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);
+      }
+
+      SharedR->resolve(ResolutionMap);
+      SharedR->finalize();
+    } else {
+      auto &ES = SharedR->getTargetVSO().getExecutionSession();
+      ES.reportError(RR.takeError());
+      SharedR->failMaterialization();
+    }
+  };
+
+  auto OnReady = [&ES](Error Err) { ES.reportError(std::move(Err)); };
+
+  SymbolNameSet Aliasees;
+  for (auto &KV : Aliases)
+    Aliasees.insert(KV.second.Aliasee);
+
+  auto Q = std::make_shared<AsynchronousSymbolQuery>(
+      Aliasees, std::move(OnResolve), std::move(OnReady));
+  auto Unresolved = V.lookup(Q, Aliasees);
+
+  if (!Unresolved.empty())
+    ES.failQuery(*Q, make_error<SymbolsNotFound>(std::move(Unresolved)));
+}
+
+void SymbolAliasesMaterializationUnit::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) {
+  SymbolFlagsMap SymbolFlags;
+  for (auto &KV : Aliases)
+    SymbolFlags[KV.first] = KV.second.AliasFlags;
+
+  return SymbolFlags;
+}
+
 Error VSO::defineMaterializing(const SymbolFlagsMap &SymbolFlags) {
   return ES.runSessionLocked([&]() -> Error {
     std::vector<SymbolMap::iterator> AddedSyms;
@@ -858,7 +921,13 @@ void VSO::dump(raw_ostream &OS) {
 Error VSO::defineImpl(MaterializationUnit &MU) {
   SymbolNameSet Duplicates;
   SymbolNameSet MUDefsOverridden;
-  std::vector<SymbolMap::iterator> ExistingDefsOverridden;
+
+  struct ExistingDefOverriddenEntry {
+    SymbolMap::iterator ExistingDefItr;
+    JITSymbolFlags NewFlags;
+  };
+  std::vector<ExistingDefOverriddenEntry> ExistingDefsOverridden;
+
   for (auto &KV : MU.getSymbols()) {
     assert(!KV.second.isLazy() && "Lazy flag should be managed internally.");
     assert(!KV.second.isMaterializing() &&
@@ -879,7 +948,7 @@ Error VSO::defineImpl(MaterializationUnit &MU) {
             (EntryItr->second.getFlags() & JITSymbolFlags::Materializing))
           Duplicates.insert(KV.first);
         else
-          ExistingDefsOverridden.push_back(EntryItr);
+          ExistingDefsOverridden.push_back({EntryItr, NewFlags});
       } else
         MUDefsOverridden.insert(KV.first);
     }
@@ -892,8 +961,8 @@ Error VSO::defineImpl(MaterializationUnit &MU) {
         continue;
 
       bool Found = false;
-      for (const auto &I : ExistingDefsOverridden)
-        if (I->first == KV.first)
+      for (const auto &EDO : ExistingDefsOverridden)
+        if (EDO.ExistingDefItr->first == KV.first)
           Found = true;
 
       if (!Found)
@@ -905,16 +974,18 @@ Error VSO::defineImpl(MaterializationUnit &MU) {
   }
 
   // Update flags on existing defs and call discard on their materializers.
-  for (auto &ExistingDefItr : ExistingDefsOverridden) {
-    assert(ExistingDefItr->second.getFlags().isLazy() &&
-           !ExistingDefItr->second.getFlags().isMaterializing() &&
+  for (auto &EDO : ExistingDefsOverridden) {
+    assert(EDO.ExistingDefItr->second.getFlags().isLazy() &&
+           !EDO.ExistingDefItr->second.getFlags().isMaterializing() &&
            "Overridden existing def should be in the Lazy state");
 
-    auto UMII = UnmaterializedInfos.find(ExistingDefItr->first);
+    EDO.ExistingDefItr->second.setFlags(EDO.NewFlags);
+
+    auto UMII = UnmaterializedInfos.find(EDO.ExistingDefItr->first);
     assert(UMII != UnmaterializedInfos.end() &&
            "Overridden existing def should have an UnmaterializedInfo");
 
-    UMII->second->MU->doDiscard(*this, ExistingDefItr->first);
+    UMII->second->MU->doDiscard(*this, EDO.ExistingDefItr->first);
   }
 
   // Discard overridden symbols povided by MU.

llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp

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