[ORC] Add support for custom generators to the C bindings.

lhames · lhames · commit b6ca0c7dd535 · 2020-10-19T01:59:04.000-07:00
C API clients can now define a custom definition generator by providing a
callback function (to implement DefinitionGenerator::tryToGenerate) and context
object. All arguments for the DefinitionGenerator::tryToGenerate method have
been given C API counterparts, and the API allows for optionally asynchronous
generation.
diff --git a/llvm/include/llvm-c/Orc.h b/llvm/include/llvm-c/Orc.h
@@ -43,6 +43,11 @@ typedef uint64_t LLVMOrcJITTargetAddress;
  */
 typedef struct LLVMOrcOpaqueExecutionSession *LLVMOrcExecutionSessionRef;
 
+/**
+ * Error reporter function.
+ */
+typedef void (*LLVMOrcErrorReporterFunction)(void *Ctx, LLVMErrorRef Err);
+
 /**
  * A reference to an orc::SymbolStringPool.
  */
@@ -55,9 +60,56 @@ typedef struct LLVMOrcOpaqueSymbolStringPoolEntry
     *LLVMOrcSymbolStringPoolEntryRef;
 
 /**
- * Error reporter function.
+ * Lookup kind. This can be used by definition generators when deciding whether
+ * to produce a definition for a requested symbol.
+ *
+ * This enum should be kept in sync with llvm::orc::LookupKind.
  */
-typedef void (*LLVMOrcErrorReporterFunction)(void *Ctx, LLVMErrorRef Err);
+typedef enum {
+  LLVMOrcLookupKindStatic,
+  LLVMOrcLookupKindDLSym
+} LLVMOrcLookupKind;
+
+/**
+ * JITDylib lookup flags. This can be used by definition generators when
+ * deciding whether to produce a definition for a requested symbol.
+ *
+ * This enum should be kept in sync with llvm::orc::JITDylibLookupFlags.
+ */
+typedef enum {
+  LLVMOrcJITDylibLookupFlagsMatchExportedSymbolsOnly,
+  LLVMOrcJITDylibLookupFlagsMatchAllSymbols
+} LLVMOrcJITDylibLookupFlags;
+
+/**
+ * Symbol lookup flags for lookup sets. This should be kept in sync with
+ * llvm::orc::SymbolLookupFlags.
+ */
+typedef enum {
+  LLVMOrcSymbolLookupFlagsRequiredSymbol,
+  LLVMOrcSymbolLookupFlagsWeaklyReferencedSymbol
+} LLVMOrcSymbolLookupFlags;
+
+/**
+ * An element type for a symbol lookup set.
+ */
+typedef struct {
+  LLVMOrcSymbolStringPoolEntryRef Name;
+  LLVMOrcSymbolLookupFlags LookupFlags;
+} LLVMOrcCLookupSetElement;
+
+/**
+ * A set of symbols to look up / generate.
+ *
+ * The list is terminated with an element containing a null pointer for the
+ * Name field.
+ *
+ * If a client creates an instance of this type then they are responsible for
+ * freeing it, and for ensuring that all strings have been retained over the
+ * course of its life. Clients receiving a copy from a callback are not
+ * responsible for managing lifetime or retain counts.
+ */
+typedef LLVMOrcCLookupSetElement *LLVMOrcCLookupSet;
 
 /**
  * A reference to an orc::JITDylib instance.
@@ -75,6 +127,59 @@ typedef struct LLVMOrcOpaqueResourceTracker *LLVMOrcResourceTrackerRef;
 typedef struct LLVMOrcOpaqueDefinitionGenerator
     *LLVMOrcDefinitionGeneratorRef;
 
+/**
+ * An opaque lookup state object. Instances of this type can be captured to
+ * suspend a lookup while a custom generator function attempts to produce a
+ * definition.
+ *
+ * If a client captures a lookup state object then they must eventually call
+ * LLVMOrcLookupStateContinueLookup to restart the lookup. This is required
+ * in order to release memory allocated for the lookup state, even if errors
+ * have occurred while the lookup was suspended (if these errors have made the
+ * lookup impossible to complete then it will issue its own error before
+ * destruction).
+ */
+typedef struct LLVMOrcOpaqueLookupState *LLVMOrcLookupStateRef;
+
+/**
+ * A custom generator function. This can be used to create a custom generator
+ * object using LLVMOrcCreateCustomCAPIDefinitionGenerator. The resulting
+ * object can be attached to a JITDylib, via LLVMOrcJITDylibAddGenerator, to
+ * receive callbacks when lookups fail to match existing definitions.
+ *
+ * GeneratorObj will contain the address of the custom generator object.
+ *
+ * Ctx will contain the context object passed to
+ * LLVMOrcCreateCustomCAPIDefinitionGenerator.
+ *
+ * LookupState will contain a pointer to an LLVMOrcLookupStateRef object. This
+ * can optionally be modified to make the definition generation process
+ * asynchronous: If the LookupStateRef value is copied, and the original
+ * LLVMOrcLookupStateRef set to null, the lookup will be suspended. Once the
+ * asynchronous definition process has been completed clients must call
+ * LLVMOrcLookupStateContinueLookup to continue the lookup (this should be
+ * done unconditionally, even if errors have occurred in the mean time, to
+ * free the lookup state memory and notify the query object of the failures. If
+ * LookupState is captured this function must return LLVMErrorSuccess.
+ *
+ * The Kind argument can be inspected to determine the lookup kind (e.g.
+ * as-if-during-static-link, or as-if-during-dlsym).
+ *
+ * The JD argument specifies which JITDylib the definitions should be generated
+ * into.
+ *
+ * The JDLookupFlags argument can be inspected to determine whether the original
+ * lookup included non-exported symobls.
+ *
+ * Finally, the LookupSet argument contains the set of symbols that could not
+ * be found in JD already (the set of generation candidates).
+ */
+typedef LLVMErrorRef (*LLVMOrcCAPIDefinitionGeneratorTryToGenerateFunction)(
+    LLVMOrcDefinitionGeneratorRef GeneratorObj, void *Ctx,
+    LLVMOrcLookupStateRef *LookupState, LLVMOrcLookupKind Kind,
+    LLVMOrcJITDylibRef JD, LLVMOrcJITDylibLookupFlags JDLookupFlags,
+    LLVMOrcCLookupSet LookupSet);
+
 /**
  * Predicate function for SymbolStringPoolEntries.
  */
@@ -156,6 +261,11 @@ void LLVMOrcSymbolStringPoolClearDeadEntries(LLVMOrcSymbolStringPoolRef SSP);
 LLVMOrcSymbolStringPoolEntryRef
 LLVMOrcExecutionSessionIntern(LLVMOrcExecutionSessionRef ES, const char *Name);
 
+/**
+ * Increments the ref-count for a SymbolStringPool entry.
+ */
+void LLVMOrcRetainSymbolStringPoolEntry(LLVMOrcSymbolStringPoolEntryRef S);
+
 /**
  * Reduces the ref-count for of a SymbolStringPool entry.
  */
@@ -254,6 +364,12 @@ LLVMErrorRef LLVMOrcJITDylibClear(LLVMOrcJITDylibRef JD);
 void LLVMOrcJITDylibAddGenerator(LLVMOrcJITDylibRef JD,
                                  LLVMOrcDefinitionGeneratorRef DG);
 
+/**
+ * Create a custom generator.
+ */
+LLVMOrcDefinitionGeneratorRef LLVMOrcCreateCustomCAPIDefinitionGenerator(
+    void *Ctx, LLVMOrcCAPIDefinitionGeneratorTryToGenerateFunction F);
+
 /**
  * Get a DynamicLibrarySearchGenerator that will reflect process symbols into
  * the JITDylib. On success the resulting generator is owned by the client.
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/Core.h b/llvm/include/llvm/ExecutionEngine/Orc/Core.h
@@ -845,15 +845,22 @@ class AsynchronousSymbolQuery {
 /// DefinitionGenerators can optionally take ownership of a LookupState object
 /// to suspend a lookup-in-progress while they search for definitions.
 class LookupState {
+  friend class OrcV2CAPIHelper;
   friend class ExecutionSession;
 
 public:
+  ~LookupState();
+
   /// Continue the lookup. This can be called by DefinitionGenerators
   /// to re-start a captured query-application operation.
   void continueLookup(Error Err);
 
 private:
   LookupState(std::unique_ptr<InProgressLookupState> IPLS);
+
+  // For C API.
+  void reset(InProgressLookupState *IPLS);
+
   std::unique_ptr<InProgressLookupState> IPLS;
 };
 
diff --git a/llvm/lib/ExecutionEngine/Orc/Core.cpp b/llvm/lib/ExecutionEngine/Orc/Core.cpp
@@ -577,6 +577,10 @@ Error ReexportsGenerator::tryToGenerate(LookupState &LS, LookupKind K,
 LookupState::LookupState(std::unique_ptr<InProgressLookupState> IPLS)
     : IPLS(std::move(IPLS)) {}
 
+void LookupState::reset(InProgressLookupState *IPLS) { this->IPLS.reset(IPLS); }
+
+LookupState::~LookupState() {}
+
 void LookupState::continueLookup(Error Err) {
   assert(IPLS && "Cannot call continueLookup on empty LookupState");
   auto &ES = IPLS->SearchOrder.begin()->first->getExecutionSession();
diff --git a/llvm/lib/ExecutionEngine/Orc/OrcV2CBindings.cpp b/llvm/lib/ExecutionEngine/Orc/OrcV2CBindings.cpp
@@ -18,6 +18,8 @@ using namespace llvm::orc;
 namespace llvm {
 namespace orc {
 
+class InProgressLookupState;
+
 class OrcV2CAPIHelper {
 public:
   using PoolEntry = SymbolStringPtr::PoolEntry;
@@ -33,14 +35,27 @@ class OrcV2CAPIHelper {
     return S.S;
   }
 
+  static void retainPoolEntry(PoolEntryPtr P) {
+    SymbolStringPtr S(P);
+    S.S = nullptr;
+  }
+
   static void releasePoolEntry(PoolEntryPtr P) {
     SymbolStringPtr S;
     S.S = P;
   }
+
+  static InProgressLookupState *extractLookupState(LookupState &LS) {
+    return LS.IPLS.release();
+  }
+
+  static void resetLookupState(LookupState &LS, InProgressLookupState *IPLS) {
+    return LS.reset(IPLS);
+  }
 };
 
-} // end namespace orc
-} // end namespace llvm
+} // namespace orc
+} // namespace llvm
 
 DEFINE_SIMPLE_CONVERSION_FUNCTIONS(ExecutionSession, LLVMOrcExecutionSessionRef)
 DEFINE_SIMPLE_CONVERSION_FUNCTIONS(SymbolStringPool, LLVMOrcSymbolStringPoolRef)
@@ -50,6 +65,7 @@ DEFINE_SIMPLE_CONVERSION_FUNCTIONS(JITDylib, LLVMOrcJITDylibRef)
 DEFINE_SIMPLE_CONVERSION_FUNCTIONS(ResourceTracker, LLVMOrcResourceTrackerRef)
 DEFINE_SIMPLE_CONVERSION_FUNCTIONS(DefinitionGenerator,
                                    LLVMOrcDefinitionGeneratorRef)
+DEFINE_SIMPLE_CONVERSION_FUNCTIONS(InProgressLookupState, LLVMOrcLookupStateRef)
 DEFINE_SIMPLE_CONVERSION_FUNCTIONS(ThreadSafeContext,
                                    LLVMOrcThreadSafeContextRef)
 DEFINE_SIMPLE_CONVERSION_FUNCTIONS(ThreadSafeModule, LLVMOrcThreadSafeModuleRef)
@@ -60,6 +76,83 @@ DEFINE_SIMPLE_CONVERSION_FUNCTIONS(LLJIT, LLVMOrcLLJITRef)
 
 DEFINE_SIMPLE_CONVERSION_FUNCTIONS(TargetMachine, LLVMTargetMachineRef)
 
+namespace llvm {
+namespace orc {
+
+class CAPIDefinitionGenerator final : public DefinitionGenerator {
+public:
+  CAPIDefinitionGenerator(
+      void *Ctx,
+      LLVMOrcCAPIDefinitionGeneratorTryToGenerateFunction TryToGenerate)
+      : Ctx(Ctx), TryToGenerate(TryToGenerate) {}
+
+  Error tryToGenerate(LookupState &LS, LookupKind K, JITDylib &JD,
+                      JITDylibLookupFlags JDLookupFlags,
+                      const SymbolLookupSet &LookupSet) override {
+
+    // Take the lookup state.
+    LLVMOrcLookupStateRef LSR = ::wrap(OrcV2CAPIHelper::extractLookupState(LS));
+
+    // Translate the lookup kind.
+    LLVMOrcLookupKind CLookupKind;
+    switch (K) {
+    case LookupKind::Static:
+      CLookupKind = LLVMOrcLookupKindStatic;
+      break;
+    case LookupKind::DLSym:
+      CLookupKind = LLVMOrcLookupKindDLSym;
+      break;
+    }
+
+    // Translate the JITDylibSearchFlags.
+    LLVMOrcJITDylibLookupFlags CJDLookupFlags;
+    switch (JDLookupFlags) {
+    case JITDylibLookupFlags::MatchExportedSymbolsOnly:
+      CJDLookupFlags = LLVMOrcJITDylibLookupFlagsMatchExportedSymbolsOnly;
+      break;
+    case JITDylibLookupFlags::MatchAllSymbols:
+      CJDLookupFlags = LLVMOrcJITDylibLookupFlagsMatchAllSymbols;
+      break;
+    }
+
+    // Translate the lookup set.
+    std::vector<LLVMOrcCLookupSetElement> CLookupSet;
+    CLookupSet.reserve(LookupSet.size());
+    for (auto &KV : LookupSet) {
+      LLVMOrcSymbolLookupFlags SLF;
+      switch (KV.second) {
+      case SymbolLookupFlags::RequiredSymbol:
+        SLF = LLVMOrcSymbolLookupFlagsRequiredSymbol;
+        break;
+      case SymbolLookupFlags::WeaklyReferencedSymbol:
+        SLF = LLVMOrcSymbolLookupFlagsWeaklyReferencedSymbol;
+        break;
+      }
+
+      CLookupSet.push_back(
+          {::wrap(OrcV2CAPIHelper::getRawPoolEntryPtr(KV.first)), SLF});
+    }
+    CLookupSet.push_back({nullptr, LLVMOrcSymbolLookupFlagsRequiredSymbol});
+
+    // Run the C TryToGenerate function.
+    auto Err =
+        unwrap(TryToGenerate(::wrap(this), Ctx, &LSR, CLookupKind, ::wrap(&JD),
+                             CJDLookupFlags, CLookupSet.data()));
+
+    // Restore the lookup state.
+    OrcV2CAPIHelper::resetLookupState(LS, ::unwrap(LSR));
+
+    return Err;
+  }
+
+private:
+  void *Ctx;
+  LLVMOrcCAPIDefinitionGeneratorTryToGenerateFunction TryToGenerate;
+};
+
+} // end namespace orc
+} // end namespace llvm
+
 void LLVMOrcExecutionSessionSetErrorReporter(
     LLVMOrcExecutionSessionRef ES, LLVMOrcErrorReporterFunction ReportError,
     void *Ctx) {
@@ -82,6 +175,10 @@ LLVMOrcExecutionSessionIntern(LLVMOrcExecutionSessionRef ES, const char *Name) {
       OrcV2CAPIHelper::releaseSymbolStringPtr(unwrap(ES)->intern(Name)));
 }
 
+void LLVMOrcRetainSymbolStringPoolEntry(LLVMOrcSymbolStringPoolEntryRef S) {
+  OrcV2CAPIHelper::retainPoolEntry(unwrap(S));
+}
+
 void LLVMOrcReleaseSymbolStringPoolEntry(LLVMOrcSymbolStringPoolEntryRef S) {
   OrcV2CAPIHelper::releasePoolEntry(unwrap(S));
 }
@@ -117,6 +214,11 @@ LLVMErrorRef LLVMOrcResourceTrackerRemove(LLVMOrcResourceTrackerRef RT) {
   return wrap(TmpRT->remove());
 }
 
+void LLVMOrcDisposeDefinitionGenerator(
+    LLVMOrcDefinitionGeneratorRef DG) {
+  delete unwrap(DG);
+}
+
 LLVMOrcJITDylibRef
 LLVMOrcExecutionSessionCreateBareJITDylib(LLVMOrcExecutionSessionRef ES,
                                           const char *Name) {
@@ -140,11 +242,6 @@ LLVMOrcExecutionSessionGetJITDylibByName(LLVMOrcExecutionSessionRef ES,
   return wrap(unwrap(ES)->getJITDylibByName(Name));
 }
 
-void LLVMOrcDisposeDefinitionGenerator(
-    LLVMOrcDefinitionGeneratorRef DG) {
-  delete unwrap(DG);
-}
-
 LLVMErrorRef LLVMOrcJITDylibClear(LLVMOrcJITDylibRef JD) {
   return wrap(unwrap(JD)->clear());
 }
@@ -154,6 +251,16 @@ void LLVMOrcJITDylibAddGenerator(LLVMOrcJITDylibRef JD,
   unwrap(JD)->addGenerator(std::unique_ptr<DefinitionGenerator>(unwrap(DG)));
 }
 
+void LLVMOrcDisposeDefinitionGenerator(LLVMOrcDefinitionGeneratorRef DG) {
+  std::unique_ptr<DefinitionGenerator> TmpDG(unwrap(DG));
+}
+
+LLVMOrcDefinitionGeneratorRef LLVMOrcCreateCustomCAPIDefinitionGenerator(
+    void *Ctx, LLVMOrcCAPIDefinitionGeneratorTryToGenerateFunction F) {
+  auto DG = std::make_unique<CAPIDefinitionGenerator>(Ctx, F);
+  return wrap(DG.release());
+}
+
 LLVMErrorRef LLVMOrcCreateDynamicLibrarySearchGeneratorForProcess(
     LLVMOrcDefinitionGeneratorRef *Result, char GlobalPrefix,
     LLVMOrcSymbolPredicate Filter, void *FilterCtx) {
