[ORC] Add utilities to RTDyldObjectLinkingLayer2 to simplify symbol flag

management and materialization responsibility registration.

The setOverrideObjectFlagsWithResponsibilityFlags method instructs
RTDyldObjectlinkingLayer2 to override the symbol flags produced by RuntimeDyld with
the flags provided by the MaterializationResponsibility instance. This can be used
to enable symbol visibility (hidden/exported) for COFF object files, which do not
currently support the SF_Exported flag.

The setAutoClaimResponsibilityForObjectSymbols method instructs
RTDyldObjectLinkingLayer2 to claim responsibility for any symbols provided by a
given object file that were not already in the MaterializationResponsibility
instance. Setting this flag allows higher-level program representations (e.g.
LLVM IR) to be added based on only a subset of the symbols they provide, without
having to write intervening layers to scan and add the additional symbols. This
trades diagnostic quality for convenience however: If all symbols are enumerated
up-front then clashes can be detected and reported early. If this option is set,
clashes for the additional symbols may not be detected until late, and detection
may depend on the flow of control through JIT'd code.

llvm-svn: 341154

include/llvm/ExecutionEngine/Orc/Core.h

@@ -139,7 +139,10 @@ public:
   JITDylib &getTargetJITDylib() const { return JD; }
 
   /// Returns the symbol flags map for this responsibility instance.
-  SymbolFlagsMap getSymbols() { return SymbolFlags; }
+  /// Note: The returned flags may have transient flags (Lazy, Materializing)
+  /// set. These should be stripped with JITSymbolFlags::stripTransientFlags
+  /// before using.
+  const SymbolFlagsMap &getSymbols() { return SymbolFlags; }
 
   /// Returns the names of any symbols covered by this
   /// MaterializationResponsibility object that have queries pending. This

include/llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h

@@ -71,8 +71,40 @@ public:
   /// the memory manager, rather than just the sections required for execution.
   ///
   /// This is kludgy, and may be removed in the future.
-  void setProcessAllSections(bool ProcessAllSections) {
+  RTDyldObjectLinkingLayer2 &setProcessAllSections(bool ProcessAllSections) {
     this->ProcessAllSections = ProcessAllSections;
+    return *this;
+  }
+
+  /// Instructs this RTDyldLinkingLayer2 instance to override the symbol flags
+  /// returned by RuntimeDyld for any given object file with the flags supplied
+  /// by the MaterializationResponsibility instance. This is a workaround to
+  /// support symbol visibility in COFF, which does not use the libObject's
+  /// SF_Exported flag. Use only when generating / adding COFF object files.
+  ///
+  /// FIXME: We should be able to remove this if/when COFF properly tracks
+  /// exported symbols.
+  RTDyldObjectLinkingLayer2 &
+  setOverrideObjectFlagsWithResponsibilityFlags(bool OverrideObjectFlags) {
+    this->OverrideObjectFlags = OverrideObjectFlags;
+    return *this;
+  }
+
+  /// If set, this RTDyldObjectLinkingLayer2 instance will claim responsibility
+  /// for any symbols provided by a given object file that were not already in
+  /// the MaterializationResponsibility instance. Setting this flag allows
+  /// higher-level program representations (e.g. LLVM IR) to be added based on
+  /// only a subset of the symbols they provide, without having to write
+  /// intervening layers to scan and add the additional symbols. This trades
+  /// diagnostic quality for convenience however: If all symbols are enumerated
+  /// up-front then clashes can be detected and reported early (and usually
+  /// deterministically). If this option is set, clashes for the additional
+  /// symbols may not be detected until late, and detection may depend on
+  /// the flow of control through JIT'd code. Use with care.
+  RTDyldObjectLinkingLayer2 &
+  setAutoClaimResponsibilityForObjectSymbols(bool AutoClaimObjectSymbols) {
+    this->AutoClaimObjectSymbols = AutoClaimObjectSymbols;
+    return *this;
   }
 
 private:
@@ -80,7 +112,9 @@ private:
   GetMemoryManagerFunction GetMemoryManager;
   NotifyLoadedFunction NotifyLoaded;
   NotifyEmittedFunction NotifyEmitted;
-  bool ProcessAllSections;
+  bool ProcessAllSections = false;
+  bool OverrideObjectFlags = false;
+  bool AutoClaimObjectSymbols = false;
   std::map<VModuleKey, RuntimeDyld *> ActiveRTDylds;
   std::map<VModuleKey, std::shared_ptr<RuntimeDyld::MemoryManager>> MemMgrs;
 };

lib/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.cpp

@@ -69,7 +69,7 @@ RTDyldObjectLinkingLayer2::RTDyldObjectLinkingLayer2(
     NotifyLoadedFunction NotifyLoaded, NotifyEmittedFunction NotifyEmitted)
     : ObjectLayer(ES), GetMemoryManager(GetMemoryManager),
       NotifyLoaded(std::move(NotifyLoaded)),
-      NotifyEmitted(std::move(NotifyEmitted)), ProcessAllSections(false) {}
+      NotifyEmitted(std::move(NotifyEmitted)) {}
 
 void RTDyldObjectLinkingLayer2::emit(MaterializationResponsibility R,
                                      VModuleKey K,
@@ -118,10 +118,38 @@ void RTDyldObjectLinkingLayer2::emit(MaterializationResponsibility R,
       }
     }
 
+    SymbolFlagsMap ExtraSymbolsToClaim;
     SymbolMap Symbols;
-    for (auto &KV : RTDyld->getSymbolTable())
-      if (!InternalSymbols.count(KV.first))
-        Symbols[ES.getSymbolStringPool().intern(KV.first)] = KV.second;
+    for (auto &KV : RTDyld->getSymbolTable()) {
+      // Scan the symbols and add them to the Symbols map for resolution.
+
+      // We never claim internal symbols.
+      if (InternalSymbols.count(KV.first))
+        continue;
+
+      auto InternedName = ES.getSymbolStringPool().intern(KV.first);
+      auto Flags = KV.second.getFlags();
+
+      // Override object flags and claim responsibility for symbols if
+      // requested.
+      if (OverrideObjectFlags || AutoClaimObjectSymbols) {
+        auto I = R.getSymbols().find(InternedName);
+
+        if (OverrideObjectFlags && I != R.getSymbols().end())
+          Flags = JITSymbolFlags::stripTransientFlags(I->second);
+        else if (AutoClaimObjectSymbols && I == R.getSymbols().end())
+          ExtraSymbolsToClaim[InternedName] = Flags;
+      }
+
+      Symbols[InternedName] = JITEvaluatedSymbol(KV.second.getAddress(), Flags);
+    }
+
+    if (!ExtraSymbolsToClaim.empty())
+      if (auto Err = R.defineMaterializing(ExtraSymbolsToClaim)) {
+        ES.reportError(std::move(Err));
+        R.failMaterialization();
+        return;
+      }
 
     R.resolve(Symbols);
   }

unittests/ExecutionEngine/Orc/RTDyldObjectLinkingLayer2Test.cpp

@@ -10,6 +10,7 @@
 #include "OrcTestCommon.h"
 #include "llvm/ExecutionEngine/ExecutionEngine.h"
 #include "llvm/ExecutionEngine/Orc/CompileUtils.h"
+#include "llvm/ExecutionEngine/Orc/IRCompileLayer.h"
 #include "llvm/ExecutionEngine/Orc/LambdaResolver.h"
 #include "llvm/ExecutionEngine/Orc/Legacy.h"
 #include "llvm/ExecutionEngine/Orc/NullResolver.h"
@@ -114,6 +115,128 @@ TEST(RTDyldObjectLinkingLayer2Test, TestSetProcessAllSections) {
       << "Expected to see debug section when ProcessAllSections is true";
 }
 
+TEST(RTDyldObjectLinkingLayer2Test, TestOverrideObjectFlags) {
+
+  OrcNativeTarget::initialize();
+
+  std::unique_ptr<TargetMachine> TM(
+      EngineBuilder().selectTarget(Triple("x86_64-unknown-linux-gnu"), "", "",
+                                   SmallVector<std::string, 1>()));
+
+  if (!TM)
+    return;
+
+  // Our compiler is going to modify symbol visibility settings without telling
+  // ORC. This will test our ability to override the flags later.
+  class FunkySimpleCompiler : public SimpleCompiler {
+  public:
+    FunkySimpleCompiler(TargetMachine &TM) : SimpleCompiler(TM) {}
+
+    CompileResult operator()(Module &M) {
+      auto *Foo = M.getFunction("foo");
+      assert(Foo && "Expected function Foo not found");
+      Foo->setVisibility(GlobalValue::HiddenVisibility);
+      return SimpleCompiler::operator()(M);
+    }
+  };
+
+  // Create a module with two void() functions: foo and bar.
+  LLVMContext Context;
+  std::unique_ptr<Module> M;
+  {
+    ModuleBuilder MB(Context, TM->getTargetTriple().str(), "dummy");
+    MB.getModule()->setDataLayout(TM->createDataLayout());
+
+    Function *FooImpl = MB.createFunctionDecl<void()>("foo");
+    BasicBlock *FooEntry = BasicBlock::Create(Context, "entry", FooImpl);
+    IRBuilder<> B1(FooEntry);
+    B1.CreateRetVoid();
+
+    Function *BarImpl = MB.createFunctionDecl<void()>("bar");
+    BasicBlock *BarEntry = BasicBlock::Create(Context, "entry", BarImpl);
+    IRBuilder<> B2(BarEntry);
+    B2.CreateRetVoid();
+
+    M = MB.takeModule();
+  }
+
+  // Create a simple stack and set the override flags option.
+  ExecutionSession ES;
+  auto &JD = ES.createJITDylib("main");
+  auto Foo = ES.getSymbolStringPool().intern("foo");
+  RTDyldObjectLinkingLayer2 ObjLayer(
+      ES, [](VModuleKey) { return std::make_shared<SectionMemoryManager>(); });
+  IRCompileLayer2 CompileLayer(ES, ObjLayer, FunkySimpleCompiler(*TM));
+
+  ObjLayer.setOverrideObjectFlagsWithResponsibilityFlags(true);
+
+  cantFail(CompileLayer.add(JD, ES.allocateVModule(), std::move(M)));
+  ES.lookup({&JD}, {Foo}, [](Expected<SymbolMap> R) { cantFail(std::move(R)); },
+            [](Error Err) { cantFail(std::move(Err)); },
+            NoDependenciesToRegister);
+}
+
+TEST(RTDyldObjectLinkingLayer2Test, TestAutoClaimResponsibilityForSymbols) {
+
+  OrcNativeTarget::initialize();
+
+  std::unique_ptr<TargetMachine> TM(
+      EngineBuilder().selectTarget(Triple("x86_64-unknown-linux-gnu"), "", "",
+                                   SmallVector<std::string, 1>()));
+
+  if (!TM)
+    return;
+
+  // Our compiler is going to add a new symbol without telling ORC.
+  // This will test our ability to auto-claim responsibility later.
+  class FunkySimpleCompiler : public SimpleCompiler {
+  public:
+    FunkySimpleCompiler(TargetMachine &TM) : SimpleCompiler(TM) {}
+
+    CompileResult operator()(Module &M) {
+      Function *BarImpl =
+          Function::Create(TypeBuilder<void(), false>::get(M.getContext()),
+                           GlobalValue::ExternalLinkage, "bar", &M);
+      BasicBlock *BarEntry =
+          BasicBlock::Create(M.getContext(), "entry", BarImpl);
+      IRBuilder<> B(BarEntry);
+      B.CreateRetVoid();
+
+      return SimpleCompiler::operator()(M);
+    }
+  };
+
+  // Create a module with two void() functions: foo and bar.
+  LLVMContext Context;
+  std::unique_ptr<Module> M;
+  {
+    ModuleBuilder MB(Context, TM->getTargetTriple().str(), "dummy");
+    MB.getModule()->setDataLayout(TM->createDataLayout());
+
+    Function *FooImpl = MB.createFunctionDecl<void()>("foo");
+    BasicBlock *FooEntry = BasicBlock::Create(Context, "entry", FooImpl);
+    IRBuilder<> B(FooEntry);
+    B.CreateRetVoid();
+
+    M = MB.takeModule();
+  }
+
+  // Create a simple stack and set the override flags option.
+  ExecutionSession ES;
+  auto &JD = ES.createJITDylib("main");
+  auto Foo = ES.getSymbolStringPool().intern("foo");
+  RTDyldObjectLinkingLayer2 ObjLayer(
+      ES, [](VModuleKey) { return std::make_shared<SectionMemoryManager>(); });
+  IRCompileLayer2 CompileLayer(ES, ObjLayer, FunkySimpleCompiler(*TM));
+
+  ObjLayer.setAutoClaimResponsibilityForObjectSymbols(true);
+
+  cantFail(CompileLayer.add(JD, ES.allocateVModule(), std::move(M)));
+  ES.lookup({&JD}, {Foo}, [](Expected<SymbolMap> R) { cantFail(std::move(R)); },
+            [](Error Err) { cantFail(std::move(Err)); },
+            NoDependenciesToRegister);
+}
+
 TEST(RTDyldObjectLinkingLayer2Test, NoDuplicateFinalization) {
   // Create a pair of modules that will trigger recursive finalization:
   // Module 1: