From be1066de56118a93f11ea53ea1ac6eb95b8a90f7 Mon Sep 17 00:00:00 2001
From: Lang Hames <lhames@gmail.com>
Date: Thu, 2 Aug 2018 20:13:58 +0000
Subject: [PATCH] [ORC] Add a re-exports fallback definition generator.

An instance of ReexportsFallbackDefinitionGenerator can be attached to a VSO
(via setFallbackDefinitionGenerator) to re-export symbols on demandy from a
backing VSO.

llvm-svn: 338764
---
 llvm/include/llvm/ExecutionEngine/Orc/Core.h  | 11 +++++++++
 llvm/lib/ExecutionEngine/Orc/Core.cpp         | 24 +++++++++++++++++++
 .../ExecutionEngine/Orc/CoreAPIsTest.cpp      | 22 +++++++++++++++++
 3 files changed, 57 insertions(+)

diff --git a/llvm/include/llvm/ExecutionEngine/Orc/Core.h b/llvm/include/llvm/ExecutionEngine/Orc/Core.h
index fd03687cfc21..8456dff7055c 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/Core.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/Core.h
@@ -344,6 +344,17 @@ reexports(VSO &SourceV, SymbolAliasMap Aliases) {
 Expected<SymbolAliasMap>
 buildSimpleReexportsAliasMap(VSO &SourceV, const SymbolNameSet &Symbols);
 
+class ReexportsFallbackDefinitionGenerator {
+public:
+  using SymbolPredicate = std::function<bool(SymbolStringPtr)>;
+  ReexportsFallbackDefinitionGenerator(VSO &BackingVSO, SymbolPredicate Allow);
+  SymbolNameSet operator()(VSO &V, const SymbolNameSet &Names);
+
+private:
+  VSO &BackingVSO;
+  SymbolPredicate Allow;
+};
+
 /// Base utilities for ExecutionSession.
 class ExecutionSessionBase {
   // FIXME: Remove this when we remove the old ORC layers.
diff --git a/llvm/lib/ExecutionEngine/Orc/Core.cpp b/llvm/lib/ExecutionEngine/Orc/Core.cpp
index 4325d57f73d0..3e5ba0399db0 100644
--- a/llvm/lib/ExecutionEngine/Orc/Core.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/Core.cpp
@@ -886,6 +886,30 @@ buildSimpleReexportsAliasMap(VSO &SourceV, const SymbolNameSet &Symbols) {
   return Result;
 }
 
+ReexportsFallbackDefinitionGenerator::ReexportsFallbackDefinitionGenerator(
+    VSO &BackingVSO, SymbolPredicate Allow)
+    : BackingVSO(BackingVSO), Allow(std::move(Allow)) {}
+
+SymbolNameSet ReexportsFallbackDefinitionGenerator::
+operator()(VSO &V, const SymbolNameSet &Names) {
+  orc::SymbolNameSet Added;
+  orc::SymbolAliasMap AliasMap;
+
+  auto Flags = BackingVSO.lookupFlags(Names);
+
+  for (auto &KV : Flags) {
+    if (!Allow(KV.first))
+      continue;
+    AliasMap[KV.first] = SymbolAliasMapEntry(KV.first, KV.second);
+    Added.insert(KV.first);
+  }
+
+  if (!Added.empty())
+    cantFail(V.define(reexports(BackingVSO, AliasMap)));
+
+  return Added;
+}
+
 Error VSO::defineMaterializing(const SymbolFlagsMap &SymbolFlags) {
   return ES.runSessionLocked([&]() -> Error {
     std::vector<SymbolMap::iterator> AddedSyms;
diff --git a/llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp b/llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp
index baa1e3b5e8cb..68c081dd5208 100644
--- a/llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp
+++ b/llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp
@@ -273,6 +273,28 @@ TEST_F(CoreAPIsStandardTest, TestThatReExportsDontUnnecessarilyMaterialize) {
   EXPECT_FALSE(BarMaterialized) << "Bar should not have been materialized";
 }
 
+TEST_F(CoreAPIsStandardTest, TestReexportsFallbackGenerator) {
+  // Test that a re-exports fallback generator can dynamically generate
+  // reexports.
+
+  auto &V2 = ES.createVSO("V2");
+  cantFail(V2.define(absoluteSymbols({{Foo, FooSym}, {Bar, BarSym}})));
+
+  auto Filter = [this](SymbolStringPtr Name) { return Name != Bar; };
+
+  V.setFallbackDefinitionGenerator(
+      ReexportsFallbackDefinitionGenerator(V2, Filter));
+
+  auto Flags = V.lookupFlags({Foo, Bar, Baz});
+  EXPECT_EQ(Flags.size(), 1U) << "Unexpected number of results";
+  EXPECT_EQ(Flags[Foo], FooSym.getFlags()) << "Unexpected flags for Foo";
+
+  auto Result = cantFail(lookup({&V}, Foo));
+
+  EXPECT_EQ(Result.getAddress(), FooSym.getAddress())
+      << "Incorrect reexported symbol address";
+}
+
 TEST_F(CoreAPIsStandardTest, TestTrivialCircularDependency) {
   Optional<MaterializationResponsibility> FooR;
   auto FooMU = llvm::make_unique<SimpleMaterializationUnit>(