Extract the packer related functions out of the JIT

libserene/include/serene/config.h.in

@@ -12,6 +12,10 @@
 
 #define MAX_PATH_SLOTS 256
 
+#define COMMON_ARGS_COUNT 6
+
+#define PACKED_FUNCTION_NAME_PREFIX "__serene_"
+
 // Should we build the support for MLIR CL OPTIONS?
 #cmakedefine SERENE_WITH_MLIR_CL_OPTION
 

libserene/include/serene/jit/halley.h

@@ -31,10 +31,19 @@
 // to make sure that we did not load it already. If we did just
 // use the existing `JITDylib` for it.
 
-//
 // TODO: [jit] Use Bare JITDylibs for the static and dynamic libs.
 // Hint: Look at `createBareJITDylib` on the `ExecutionSession`
 
+// TODO: [jit] Create a generator class that generates symbols
+// from Serene's code and add them to the JitDylib of a namespace
+// instead of having multiple jitDylibs per NS
+
+// TODO: [jit] Create a another overload of the `lookup` function
+// than returns a pointer to the symbol. This new functions will
+// be used to call those core functions that we already know the
+// signature and we don't want to wrap them. For examples functions
+// in the `serene.core` namespace
+
 #ifndef SERENE_JIT_HALLEY_H
 #define SERENE_JIT_HALLEY_H
 
@@ -87,7 +96,8 @@ class Halley;
 using Engine             = Halley;
 using EnginePtr          = std::unique_ptr<Engine>;
 using MaybeEngine        = llvm::Expected<EnginePtr>;
-using MaybeJitAddress    = llvm::Expected<void *(*)()>;
+using JitWrappedAddress  = void (*)(void **);
+using MaybeJitAddress    = llvm::Expected<JitWrappedAddress>;
 using Dylib              = llvm::orc::JITDylib;
 using DylibPtr           = Dylib *;
 using MaybeDylibPtr      = llvm::Expected<DylibPtr>;
@@ -185,7 +195,7 @@ public:
   /// the namespace from file even if it exists already.
   MaybeDylibPtr loadNamespace(std::string &nsName);
 
-  // TODO: Move all the loader related functions to a Loader class
+  // TODO: [jit] Move all the loader related functions to a Loader class
   /// Load the shared library in the given `path` to the given JITDylib
   /// `jd` via the give ExecutionSession `es`.
   /// This function assumes that the shared lib exists.
@@ -210,46 +220,14 @@ public:
   void setEngine(std::unique_ptr<llvm::orc::LLJIT> e, bool isLazy);
   /// Looks up a packed-argument function with the given sym name and returns a
   /// pointer to it. Propagates errors in case of failure.
-  MaybeJitAddress lookup(const char *nsName, const char *sym);
+  MaybeJitAddress lookup(const char *nsName, const char *sym) const;
   MaybeJitAddress lookup(const types::Symbol &sym) const;
 
   /// Invokes the function with the given name passing it the list of opaque
   /// pointers to the actual arguments.
-  // llvm::Error
-  // invokePacked(llvm::StringRef name,
-  //              llvm::MutableArrayRef<void *> args = llvm::None) const;
-
-  /// Trait that defines how a given type is passed to the JIT code. This
-  /// defaults to passing the address but can be specialized.
-  template <typename T>
-  struct Argument {
-    static void pack(llvm::SmallVectorImpl<void *> &args, T &val) {
-      args.push_back(&val);
-    }
-  };
-
-  /// Tag to wrap an output parameter when invoking a jitted function.
-  template <typename T>
-  struct FnResult {
-    explicit FnResult(T &result) : value(result) {}
-    T &value;
-  };
-
-  /// Helper function to wrap an output operand when using
-  /// ExecutionEngine::invoke.
-  template <typename T>
-  static FnResult<T> result(T &t) {
-    return FnResult<T>(t);
-  }
-
-  // Specialization for output parameter: their address is forwarded directly to
-  // the native code.
-  template <typename T>
-  struct Argument<FnResult<T>> {
-    static void pack(llvm::SmallVectorImpl<void *> &args, FnResult<T> &result) {
-      args.push_back(&result.value);
-    }
-  };
+  llvm::Error
+  invokePacked(const types::Symbol &name,
+               llvm::MutableArrayRef<void *> args = llvm::None) const;
 
   llvm::Error loadModule(const char *nsName, const char *file);
   void dumpToObjectFile(llvm::StringRef filename);

libserene/include/serene/jit/packer.h

@@ -0,0 +1,66 @@
+/* -*- C++ -*-
+ * Serene Programming Language
+ *
+ * Copyright (c) 2019-2022 Sameer Rahmani <lxsameer@gnu.org>
+ *
+ * This program is free software; you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, version 2.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+#ifndef SERENE_JIT_PACKER_H
+#define SERENE_JIT_PACKER_H
+
+#include <llvm/ADT/SmallVector.h>
+#include <llvm/ADT/StringRef.h>
+#include <llvm/IR/Module.h>
+
+namespace serene::jit {
+
+struct Packer {
+  /// Trait that defines how a given type is passed to the JIT code. This
+  /// defaults to passing the address but can be specialized.
+  template <typename T>
+  struct Argument {
+    static void pack(llvm::SmallVectorImpl<void *> &args, T &val) {
+      args.push_back(&val);
+    }
+  };
+
+  /// Tag to wrap an output parameter when invoking a jitted function.
+  template <typename T>
+  struct FnResult {
+    explicit FnResult(T &result) : value(result) {}
+    T &value;
+  };
+
+  /// Helper function to wrap an output operand when using
+  /// ExecutionEngine::invoke.
+  template <typename T>
+  static FnResult<T> result(T &t) {
+    return FnResult<T>(t);
+  }
+
+  // Specialization for output parameter: their address is forwarded directly to
+  // the native code.
+  template <typename T>
+  struct Argument<FnResult<T>> {
+    static void pack(llvm::SmallVectorImpl<void *> &args, FnResult<T> &result) {
+      args.push_back(&result.value);
+    }
+  };
+};
+
+std::string makePackedFunctionName(llvm::StringRef name);
+void packFunctionArguments(llvm::Module *module);
+
+} // namespace serene::jit
+#endif

libserene/lib/CMakeLists.txt

@@ -27,7 +27,8 @@ add_library(serene
   context.cpp
   fs.cpp
 
-  jit/halley.cpp)
+  jit/halley.cpp
+  jit/packer.cpp)
 
 # Create an ALIAS target. This way if we mess up the name
 # there will be an cmake error inseat of a linker error which is harder

libserene/lib/jit/halley.cpp

@@ -98,6 +98,7 @@ ObjectCache::getObject(const llvm::Module *m) {
                " in cache. Compiling.");
     return nullptr;
   }
+
   HALLEY_LOG("Object for " + m->getModuleIdentifier() + " loaded from cache.");
   return llvm::MemoryBuffer::getMemBuffer(i->second->getMemBufferRef());
 }
@@ -405,7 +406,11 @@ llvm::Error Halley::createEmptyNS(const char *name) {
   return llvm::Error::success();
 };
 
-MaybeJitAddress Halley::lookup(const char *nsName, const char *sym) {
+MaybeJitAddress Halley::lookup(const types::Symbol &sym) const {
+  return lookup(sym.ns->data, sym.name->data);
+}
+
+MaybeJitAddress Halley::lookup(const char *nsName, const char *sym) const {
   assert(sym != nullptr && "'sym' is null: lookup");
   assert(nsName != nullptr && "'nsName' is null: lookup");
 
@@ -415,7 +420,7 @@ MaybeJitAddress Halley::lookup(const char *nsName, const char *sym) {
   std::string fqsym = (ns + "/" + s).str();
 
   HALLEY_LOG("Looking up symbol: " << fqsym);
-  auto *dylib = jitDylibs[nsName].back();
+  auto *dylib = const_cast<Halley *>(this)->jitDylibs[nsName].back();
 
   if (dylib == nullptr) {
     return tempError(*ctx, "No dylib " + s);
@@ -437,7 +442,7 @@ MaybeJitAddress Halley::lookup(const char *nsName, const char *sym) {
   auto rawFPtr = expectedSymbol->getAddress();
 
   // NOLINTNEXTLINE(performance-no-int-to-ptr)
-  auto fptr = reinterpret_cast<void *(*)()>(rawFPtr);
+  auto fptr = reinterpret_cast<JitWrappedAddress>(rawFPtr);
 
   if (fptr == nullptr) {
     return tempError(*ctx, "Lookup function is null!");
@@ -532,7 +537,6 @@ Halley::loadNamespaceFrom<fs::NSFileType::ObjectFile>(NSLoadRequest &req) {
     return err;
   }
 
-  llvm::outs() << "ok\n";
   return jd;
 };
 
@@ -828,6 +832,19 @@ llvm::Error Halley::createCurrentProcessJD() {
   return llvm::Error::success();
 };
 
+llvm::Error Halley::invokePacked(const types::Symbol &name,
+                                 llvm::MutableArrayRef<void *> args) const {
+  auto expectedFPtr = lookup(name);
+  if (!expectedFPtr) {
+    return expectedFPtr.takeError();
+  }
+
+  auto *fptr = *expectedFPtr;
+  (*fptr)(args.data());
+
+  return llvm::Error::success();
+}
+
 MaybeEngine makeHalleyJIT(std::unique_ptr<SereneContext> ctx) {
   llvm::orc::JITTargetMachineBuilder jtmb(ctx->triple);
   auto maybeJIT = Halley::make(std::move(ctx), std::move(jtmb));

libserene/lib/jit/packer.cpp

@@ -0,0 +1,96 @@
+/* -*- C++ -*-
+ * Serene Programming Language
+ *
+ * Copyright (c) 2019-2022 Sameer Rahmani <lxsameer@gnu.org>
+ *
+ * This program is free software; you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, version 2.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+#include "serene/jit/packer.h"
+
+#include "serene/config.h"
+
+#include <llvm/IR/IRBuilder.h>
+
+namespace serene::jit {
+
+std::string makePackedFunctionName(llvm::StringRef name) {
+  // TODO: move the "_serene_" constant to a macro or something
+  return PACKED_FUNCTION_NAME_PREFIX + name.str();
+}
+
+void packFunctionArguments(llvm::Module *module) {
+  auto &ctx = module->getContext();
+  llvm::IRBuilder<> builder(ctx);
+  llvm::DenseSet<llvm::Function *> interfaceFunctions;
+  for (auto &func : module->getFunctionList()) {
+    if (func.isDeclaration()) {
+      continue;
+    }
+    if (interfaceFunctions.count(&func) != 0) {
+      continue;
+    }
+
+    // Given a function `foo(<...>)`, define the interface function
+    // `serene_foo(i8**)`.
+    auto *newType = llvm::FunctionType::get(
+        builder.getVoidTy(), builder.getInt8PtrTy()->getPointerTo(),
+        /*isVarArg=*/false);
+    auto newName = makePackedFunctionName(func.getName());
+    auto funcCst = module->getOrInsertFunction(newName, newType);
+    llvm::Function *interfaceFunc =
+        llvm::cast<llvm::Function>(funcCst.getCallee());
+    interfaceFunctions.insert(interfaceFunc);
+
+    // Extract the arguments from the type-erased argument list and cast them to
+    // the proper types.
+    auto *bb = llvm::BasicBlock::Create(ctx);
+    bb->insertInto(interfaceFunc);
+    builder.SetInsertPoint(bb);
+    llvm::Value *argList = interfaceFunc->arg_begin();
+    llvm::SmallVector<llvm::Value *, COMMON_ARGS_COUNT> args;
+    args.reserve(llvm::size(func.args()));
+    for (const auto &indexedArg : llvm::enumerate(func.args())) {
+      llvm::Value *argIndex = llvm::Constant::getIntegerValue(
+          builder.getInt64Ty(), llvm::APInt(I64_SIZE, indexedArg.index()));
+      llvm::Value *argPtrPtr =
+          builder.CreateGEP(builder.getInt8PtrTy(), argList, argIndex);
+      llvm::Value *argPtr =
+          builder.CreateLoad(builder.getInt8PtrTy(), argPtrPtr);
+      llvm::Type *argTy = indexedArg.value().getType();
+      argPtr            = builder.CreateBitCast(argPtr, argTy->getPointerTo());
+      llvm::Value *arg  = builder.CreateLoad(argTy, argPtr);
+      args.push_back(arg);
+    }
+
+    // Call the implementation function with the extracted arguments.
+    llvm::Value *result = builder.CreateCall(&func, args);
+
+    // Assuming the result is one value, potentially of type `void`.
+    if (!result->getType()->isVoidTy()) {
+      llvm::Value *retIndex = llvm::Constant::getIntegerValue(
+          builder.getInt64Ty(), llvm::APInt(I64_SIZE, llvm::size(func.args())));
+      llvm::Value *retPtrPtr =
+          builder.CreateGEP(builder.getInt8PtrTy(), argList, retIndex);
+      llvm::Value *retPtr =
+          builder.CreateLoad(builder.getInt8PtrTy(), retPtrPtr);
+      retPtr = builder.CreateBitCast(retPtr, result->getType()->getPointerTo());
+      builder.CreateStore(result, retPtr);
+    }
+
+    // The interface function returns void.
+    builder.CreateRetVoid();
+  }
+};
+
+} // namespace serene::jit

serenec/serenec.cpp

@@ -318,13 +318,14 @@ int main(int argc, char *argv[]) {
   }
   auto *c = *bt;
 
-  void *res = c();
-  // for (int i = 0; i <= 10; i++) {
-  //   printf(">> %02x", *(c + i));
-  // }
-  printf("Res >> %p\n", res);
-  llvm::outs() << "Res: " << *((int *)res) << "\n";
-  (void)res;
+  (void)c;
+  // void *res = c();
+  // // for (int i = 0; i <= 10; i++) {
+  // //   printf(">> %02x", *(c + i));
+  // // }
+  // printf("Res >> %p\n", res);
+  // llvm::outs() << "Res: " << *((int *)res) << "\n";
+  // (void)res;
 
   // // TODO: handle the outputDir by not forcing it. it should be
   // //       default to the current working dir