From 5ff5a30beecc65cb72b50047990dd684977f3001 Mon Sep 17 00:00:00 2001
From: Lang Hames <lhames@gmail.com>
Date: Sun, 21 Jan 2018 03:20:39 +0000
Subject: [PATCH] [ORC] Add a lookupFlags method to VSO.

lookupFlags returns a SymbolFlagsMap for the requested symbols, along with a
set containing the SymbolStringPtr for any symbol not found in the VSO.

The JITSymbolFlags for each symbol will have been stripped of its transient
JIT-state flags (i.e. NotMaterialized, Materializing).

Calling lookupFlags does not trigger symbol materialization.

llvm-svn: 323060
---
 llvm/include/llvm/ExecutionEngine/JITSymbol.h |  5 ++
 llvm/include/llvm/ExecutionEngine/Orc/Core.h  | 11 ++++
 llvm/lib/ExecutionEngine/Orc/Core.cpp         | 20 ++++++++
 .../ExecutionEngine/Orc/CoreAPIsTest.cpp      | 50 +++++++++++++++++++
 4 files changed, 86 insertions(+)

diff --git a/llvm/include/llvm/ExecutionEngine/JITSymbol.h b/llvm/include/llvm/ExecutionEngine/JITSymbol.h
index 6aba236b6196..60b4f8fbc74b 100644
--- a/llvm/include/llvm/ExecutionEngine/JITSymbol.h
+++ b/llvm/include/llvm/ExecutionEngine/JITSymbol.h
@@ -56,6 +56,11 @@ public:
     Materializing = 1U << 6
   };
 
+  static JITSymbolFlags stripTransientFlags(JITSymbolFlags Orig) {
+    return static_cast<FlagNames>(Orig.Flags &
+                                  ~(NotMaterialized | Materializing));
+  }
+
   /// @brief Default-construct a JITSymbolFlags instance.
   JITSymbolFlags() = default;
 
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/Core.h b/llvm/include/llvm/ExecutionEngine/Orc/Core.h
index 56615a33475a..443125b79519 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/Core.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/Core.h
@@ -138,6 +138,11 @@ public:
       std::map<SymbolStringPtr, RelativeLinkageStrength>;
   using SourceWorkMap = std::map<SymbolSource *, SymbolNameSet>;
 
+  struct LookupFlagsResult {
+    SymbolFlagsMap SymbolFlags;
+    SymbolNameSet SymbolsNotFound;
+  };
+
   struct LookupResult {
     SourceWorkMap MaterializationWork;
     SymbolNameSet UnresolvedSymbols;
@@ -175,6 +180,12 @@ public:
   /// @brief Finalize the given symbols.
   void finalize(SymbolNameSet SymbolsToFinalize);
 
+  /// @brief Look up the flags for the given symbols.
+  ///
+  /// Returns the flags for the give symbols, together with the set of symbols
+  /// not found.
+  LookupFlagsResult lookupFlags(SymbolNameSet Symbols);
+
   /// @brief Apply the given query to the given symbols in this VSO.
   ///
   /// For symbols in this VSO that have already been materialized, their address
diff --git a/llvm/lib/ExecutionEngine/Orc/Core.cpp b/llvm/lib/ExecutionEngine/Orc/Core.cpp
index c99c8ec120fd..232c880b9dd6 100644
--- a/llvm/lib/ExecutionEngine/Orc/Core.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/Core.cpp
@@ -287,6 +287,26 @@ void VSO::finalize(SymbolNameSet SymbolsToFinalize) {
   }
 }
 
+VSO::LookupFlagsResult VSO::lookupFlags(SymbolNameSet Names) {
+  SymbolFlagsMap FlagsFound;
+
+  for (SymbolNameSet::iterator I = Names.begin(), E = Names.end(); I != E;) {
+    auto Tmp = I++;
+    auto SymI = Symbols.find(*Tmp);
+
+    // If the symbol isn't in this dylib then just continue.
+    if (SymI == Symbols.end())
+      continue;
+
+    Names.erase(Tmp);
+
+    FlagsFound[SymI->first] =
+        JITSymbolFlags::stripTransientFlags(SymI->second.getFlags());
+  }
+
+  return {std::move(FlagsFound), std::move(Names)};
+}
+
 VSO::LookupResult VSO::lookup(AsynchronousSymbolQuery &Query,
                               SymbolNameSet Names) {
   SourceWorkMap MaterializationWork;
diff --git a/llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp b/llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp
index a198ed0b96e5..2b63dbf7c1fb 100644
--- a/llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp
+++ b/llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp
@@ -130,6 +130,56 @@ TEST(CoreAPIsTest, SimpleAsynchronousSymbolQueryAgainstVSO) {
   EXPECT_TRUE(OnReadyRun) << "OnReady was not run";
 }
 
+TEST(CoreAPIsTest, LookupFlagsTest) {
+
+  // Test that lookupFlags works on a predefined symbol, and does not trigger
+  // materialization of a lazy symbol.
+
+  SymbolStringPool SP;
+  auto Foo = SP.intern("foo");
+  auto Bar = SP.intern("bar");
+  auto Baz = SP.intern("baz");
+
+  VSO V;
+
+  auto Source = std::make_shared<SimpleSource>(
+      [](VSO &V, SymbolNameSet Symbols) -> Error {
+        llvm_unreachable("Symbol materialized on flags lookup");
+      },
+      [](VSO &V, SymbolStringPtr Name) -> Error {
+        llvm_unreachable("Symbol finalized on flags lookup");
+      });
+
+  JITSymbolFlags FooFlags = JITSymbolFlags::Exported;
+  JITSymbolFlags BarFlags = static_cast<JITSymbolFlags::FlagNames>(
+      JITSymbolFlags::Exported | JITSymbolFlags::Weak);
+
+  SymbolMap InitialDefs;
+  InitialDefs[Foo] = JITEvaluatedSymbol(0xdeadbeef, FooFlags);
+  cantFail(V.define(std::move(InitialDefs)));
+
+  SymbolFlagsMap InitialLazyDefs({{Bar, BarFlags}});
+  cantFail(V.defineLazy(InitialLazyDefs, *Source));
+
+  SymbolNameSet Names({Foo, Bar, Baz});
+
+  auto LFR = V.lookupFlags(Names);
+
+  EXPECT_EQ(LFR.SymbolsNotFound.size(), 1U) << "Expected one not-found symbol";
+  EXPECT_EQ(*LFR.SymbolsNotFound.begin(), Baz)
+      << "Expected Baz to be not-found";
+  EXPECT_EQ(LFR.SymbolFlags.size(), 2U)
+      << "Returned symbol flags contains unexpected results";
+  EXPECT_EQ(LFR.SymbolFlags.count(Foo), 1U)
+      << "Missing lookupFlags result for Foo";
+  EXPECT_EQ(LFR.SymbolFlags[Foo], FooFlags)
+      << "Incorrect flags returned for Foo";
+  EXPECT_EQ(LFR.SymbolFlags.count(Bar), 1U)
+      << "Missing  lookupFlags result for Bar";
+  EXPECT_EQ(LFR.SymbolFlags[Bar], BarFlags)
+      << "Incorrect flags returned for Bar";
+}
+
 TEST(CoreAPIsTest, AddAndMaterializeLazySymbol) {
 
   constexpr JITTargetAddress FakeFooAddr = 0xdeadbeef;