Implement the ability to swap functions in the JIT.

This patch allows users of CppInterOp to rewire things like printf to enable
rich displays, for example.

Author: vgvassilev

include/clang/Interpreter/Compatibility.h

@@ -39,11 +39,29 @@ namespace compat {
     cling::utils::Analyze::maybeMangleDeclName(GD, mangledName);
   }
 
+  inline llvm::orc::LLJIT* getExecutionEngine(const cling::Interpreter& I) {
+    // FIXME: This is a horrible hack finding the llvm::orc::LLJIT by computing
+    // the object offsets in Cling. We should add getExecutionEngine interface
+    // to directly.
+
+    // sizeof (m_Opts) + sizeof(m_LLVMContext)
+    const unsigned m_ExecutorOffset = 72;
+    int* IncrementalExecutor = ((int*)(&I)) + m_ExecutorOffset;
+    int* IncrementalJit = *(int**)IncrementalExecutor + 0;
+    int* LLJIT = *(int**)IncrementalJit + 0;
+    return *(llvm::orc::LLJIT**)LLJIT;
+  }
+
   inline llvm::Expected<llvm::JITTargetAddress>
   getSymbolAddress(const cling::Interpreter &I, llvm::StringRef IRName) {
-    return (llvm::JITTargetAddress)I.getAddressOfGlobal(IRName);
-  }
+    if (void* Addr = I.getAddressOfGlobal(IRName))
+      return (llvm::JITTargetAddress)Addr;
 
+    llvm::orc::LLJIT& Jit = *compat::getExecutionEngine(I);
+    llvm::orc::SymbolNameVector Names;
+    Names.push_back(Jit.getExecutionSession().intern(IRName));
+    return llvm::make_error<llvm::orc::SymbolsNotFound>(Names);
+  }
 }
 
 #endif //USE_CLING
@@ -92,9 +110,10 @@ namespace compat {
   //            getSymbolAddress, getSymbolAddressFromLinkerName
   // Clang 15 - Add new Interpreter methods: Undo
 
-  inline const llvm::orc::LLJIT *getExecutionEngine(clang::Interpreter &I)  {
+  inline llvm::orc::LLJIT* getExecutionEngine(clang::Interpreter& I) {
 #if CLANG_VERSION_MAJOR >= 14
-    return &llvm::cantFail(I.getExecutionEngine());
+    auto* engine = &llvm::cantFail(I.getExecutionEngine());
+    return const_cast<llvm::orc::LLJIT*>(engine);
 #else
     assert(0 && "Not implemented in Clang <14!");
     return nullptr;

include/clang/Interpreter/CppInterOp.h

@@ -455,6 +455,17 @@ namespace Cpp {
   /// function.
   bool LoadLibrary(const char *lib_path, bool lookup = true);
 
+  /// Inserts or replaces a symbol in the JIT with the one provided. This is
+  /// useful for providing our own implementations of facilities such as printf.
+  ///
+  ///\param[in] linker_mangled_name - the name of the symbol to be inserted or
+  ///           replaced.
+  ///\param[in] address - the new address of the symbol.
+  ///
+  ///\returns true on failure.
+  bool InsertOrReplaceJitSymbol(const char* linker_mangled_name,
+                                uint64_t address);
+
   /// Tries to load provided objects in a string format (prettyprint).
   std::string ObjToString(const char *type, void *obj);
 

include/clang/Interpreter/CppInterOpInterpreter.h

@@ -154,6 +154,7 @@ class Interpreter {
   ~Interpreter() {}
 
   operator const clang::Interpreter&() const { return *inner; }
+  operator clang::Interpreter&() { return *inner; }
 
   ///\brief Describes the return result of the different routines that do the
   /// incremental compilation.

lib/Interpreter/CppInterOp.cpp

@@ -2563,6 +2563,57 @@ namespace Cpp {
     return res == compat::Interpreter::kSuccess;
   }
 
+  bool InsertOrReplaceJitSymbol(const char* linker_mangled_name,
+                                uint64_t address) {
+    using namespace llvm;
+    using namespace llvm::orc;
+
+    auto& I = getInterp();
+    auto Symbol = compat::getSymbolAddress(I, linker_mangled_name);
+    if (Error Err = Symbol.takeError()) {
+      logAllUnhandledErrors(std::move(Err), errs(),
+                            "[InsertOrReplaceJitSymbol] error: ");
+      return true;
+    }
+
+    // Nothing to define, we are redefining the same function.
+    if (*Symbol && *Symbol == address) {
+      errs() << "[InsertOrReplaceJitSymbol] warning: redefining '"
+             << linker_mangled_name << "' with the same address\n";
+      return true;
+    }
+
+    // Let's inject it.
+    bool Inserted;
+    SymbolMap::iterator It;
+    llvm::orc::SymbolMap InjectedSymbols;
+
+    llvm::orc::LLJIT& Jit = *compat::getExecutionEngine(I);
+    JITDylib& DyLib = Jit.getMainJITDylib();
+
+    std::tie(It, Inserted) = InjectedSymbols.try_emplace(
+        Jit.getExecutionSession().intern(linker_mangled_name),
+        JITEvaluatedSymbol(address, JITSymbolFlags::Exported));
+    assert(Inserted && "Why wasn't this found in the initial Jit lookup?");
+
+    // We want to replace a symbol with a custom provided one.
+    if (Symbol && address)
+      // The symbol be in the DyLib or in-process.
+      if (auto Err = DyLib.remove({It->first})) {
+        logAllUnhandledErrors(std::move(Err), errs(),
+                              "[InsertOrReplaceJitSymbol] error: ");
+        return true;
+      }
+
+    if (Error Err = DyLib.define(absoluteSymbols({*It}))) {
+      logAllUnhandledErrors(std::move(Err), errs(),
+                            "[InsertOrReplaceJitSymbol] error: ");
+      return true;
+    }
+
+    return false;
+  }
+
   std::string ObjToString(const char *type, void *obj) {
     return getInterp().toString(type, obj);
   }

unittests/CppInterOp/CMakeLists.txt

@@ -6,6 +6,7 @@ add_cppinterop_unittest(CppInterOpTests
   EnumReflectionTest.cpp
   FunctionReflectionTest.cpp
   InterpreterTest.cpp
+  JitTest.cpp
   ScopeReflectionTest.cpp
   TypeReflectionTest.cpp
   Utils.cpp

unittests/CppInterOp/JitTest.cpp

@@ -0,0 +1,32 @@
+#include "Utils.h"
+
+#include "clang/Interpreter/CppInterOp.h"
+
+#include "gtest/gtest.h"
+
+using namespace TestUtils;
+
+static int printf_jit(const char* format, ...) {
+  llvm::errs() << "printf_jit called!\n";
+  return 0;
+}
+
+TEST(JitTest, InsertOrReplaceJitSymbol) {
+  std::vector<Decl*> Decls;
+  std::string code = R"(
+    extern "C" int printf(const char*,...);
+    )";
+
+  GetAllTopLevelDecls(code, Decls);
+
+  EXPECT_FALSE(Cpp::InsertOrReplaceJitSymbol("printf", (uint64_t)&printf_jit));
+
+  testing::internal::CaptureStderr();
+  Cpp::Process("printf(\"Blah\");");
+  std::string cerrs = testing::internal::GetCapturedStderr();
+  EXPECT_STREQ(cerrs.c_str(), "printf_jit called!\n");
+
+  EXPECT_TRUE(
+      Cpp::InsertOrReplaceJitSymbol("non_existent", (uint64_t)&printf_jit));
+  EXPECT_TRUE(Cpp::InsertOrReplaceJitSymbol("non_existent", 0));
+}