codon build command: add --cir output type option (#649)

* Fix argument ordering in JIT

* Format

* Update JIT tests

* Fix JIT test

.github/workflows/ci.yml

@@ -187,7 +187,7 @@ jobs:
       - name: Prepare Artifacts
         run: |
           cp -rf codon-deploy/python/dist .
-          rm -rf codon-deploy/lib/libfmt.a codon-deploy/lib/pkgconfig codon-deploy/lib/cmake codon-deploy/python/codon.egg-info codon-deploy/python/dist codon-deploy/python/build
+          rm -rf codon-deploy/lib/libfmt.a codon-deploy/lib/pkgconfig codon-deploy/lib/cmake codon-deploy/python/codon_jit.egg-info codon-deploy/python/build
           tar -czf ${CODON_BUILD_ARCHIVE} codon-deploy
           du -sh codon-deploy
 

CMakeLists.txt

@@ -1,10 +1,10 @@
 cmake_minimum_required(VERSION 3.14)
 project(
   Codon
-  VERSION "0.18.1"
+  VERSION "0.18.2"
   HOMEPAGE_URL "https://github.com/exaloop/codon"
   DESCRIPTION "high-performance, extensible Python compiler")
-set(CODON_JIT_PYTHON_VERSION "0.3.1")
+set(CODON_JIT_PYTHON_VERSION "0.3.2")
 configure_file("${PROJECT_SOURCE_DIR}/cmake/config.h.in"
                "${PROJECT_SOURCE_DIR}/codon/config/config.h")
 configure_file("${PROJECT_SOURCE_DIR}/cmake/config.py.in"
@@ -48,10 +48,8 @@ include(${CMAKE_SOURCE_DIR}/cmake/deps.cmake)
 set(CMAKE_BUILD_WITH_INSTALL_RPATH ON)
 if(APPLE)
   set(CMAKE_INSTALL_RPATH "@loader_path;@loader_path/../lib/codon")
-  set(STATIC_LIBCPP "")
 else()
   set(CMAKE_INSTALL_RPATH "$ORIGIN:$ORIGIN/../lib/codon")
-  set(STATIC_LIBCPP "-static-libstdc++")
 endif()
 
 add_executable(peg2cpp codon/util/peg2cpp.cpp)
@@ -138,7 +136,7 @@ target_include_directories(codonrt PRIVATE ${backtrace_SOURCE_DIR}
                                            ${highway_SOURCE_DIR}
                                            "${gc_SOURCE_DIR}/include"
                                            "${fast_float_SOURCE_DIR}/include" runtime)
-target_link_libraries(codonrt PRIVATE fmt omp backtrace ${STATIC_LIBCPP} LLVMSupport)
+target_link_libraries(codonrt PRIVATE fmt omp backtrace LLVMSupport)
 if(APPLE)
   target_link_libraries(
     codonrt
@@ -434,11 +432,7 @@ llvm_map_components_to_libnames(
   TransformUtils
   Vectorize
   Passes)
-if(APPLE)
-  target_link_libraries(codonc PRIVATE ${LLVM_LIBS} fmt dl codonrt)
-else()
-  target_link_libraries(codonc PRIVATE ${STATIC_LIBCPP} ${LLVM_LIBS} fmt dl codonrt)
-endif()
+target_link_libraries(codonc PRIVATE ${LLVM_LIBS} fmt dl codonrt)
 
 # Gather headers
 add_custom_target(
@@ -482,13 +476,13 @@ add_dependencies(libs codonrt codonc)
 
 # Codon command-line tool
 add_executable(codon codon/app/main.cpp)
-target_link_libraries(codon PUBLIC ${STATIC_LIBCPP} fmt codonc codon_jupyter Threads::Threads)
+target_link_libraries(codon PUBLIC fmt codonc codon_jupyter Threads::Threads)
 
 # Codon test Download and unpack googletest at configure time
 include(FetchContent)
 FetchContent_Declare(
   googletest
-  URL https://github.com/google/googletest/archive/609281088cfefc76f9d0ce82e1ff6c30cc3591e5.zip
+  URL https://github.com/google/googletest/archive/03597a01ee50ed33e9dfd640b249b4be3799d395.zip
 )
 # For Windows: Prevent overriding the parent project's compiler/linker settings
 set(gtest_force_shared_crt ON CACHE BOOL "" FORCE)

README.md

@@ -149,7 +149,7 @@ print(total)
 ```
 
 Note that Codon automatically turns the `total += 1` statement in the loop body into an atomic
-reduction to avoid race conditions. Learn more in the [multitheading docs](advanced/parallel.md).
+reduction to avoid race conditions. Learn more in the [multithreading docs](https://docs.exaloop.io/codon/advanced/parallel).
 
 Codon also supports writing and executing GPU kernels. Here's an example that computes the
 [Mandelbrot set](https://en.wikipedia.org/wiki/Mandelbrot_set):

cmake/deps.cmake

@@ -1,8 +1,8 @@
-set(CPM_DOWNLOAD_VERSION 0.32.3)
+set(CPM_DOWNLOAD_VERSION 0.40.8)
 set(CPM_DOWNLOAD_LOCATION "${CMAKE_BINARY_DIR}/cmake/CPM_${CPM_DOWNLOAD_VERSION}.cmake")
 if(NOT (EXISTS ${CPM_DOWNLOAD_LOCATION}))
     message(STATUS "Downloading CPM.cmake...")
-    file(DOWNLOAD https://github.com/TheLartians/CPM.cmake/releases/download/v${CPM_DOWNLOAD_VERSION}/CPM.cmake ${CPM_DOWNLOAD_LOCATION})
+    file(DOWNLOAD https://github.com/cpm-cmake/CPM.cmake/releases/download/v${CPM_DOWNLOAD_VERSION}/CPM.cmake ${CPM_DOWNLOAD_LOCATION})
 endif()
 include(${CPM_DOWNLOAD_LOCATION})
 
@@ -77,9 +77,9 @@ endif()
 
 CPMAddPackage(
     NAME bdwgc
-    GITHUB_REPOSITORY "ivmai/bdwgc"
+    GITHUB_REPOSITORY "exaloop/bdwgc"
     VERSION 8.0.5
-    GIT_TAG d0ba209660ea8c663e06d9a68332ba5f42da54ba
+    GIT_TAG e16c67244aff26802203060422545d38305e0160
     EXCLUDE_FROM_ALL YES
     OPTIONS "CMAKE_POSITION_INDEPENDENT_CODE ON"
             "BUILD_SHARED_LIBS OFF"
@@ -169,7 +169,7 @@ if(NOT APPLE)
     CPMAddPackage(
         NAME openblas
         GITHUB_REPOSITORY "OpenMathLib/OpenBLAS"
-        GIT_TAG v0.3.28
+        GIT_TAG v0.3.29
         EXCLUDE_FROM_ALL YES
         OPTIONS "DYNAMIC_ARCH ON"
                 "BUILD_TESTING OFF"

codon/app/main.cpp

@@ -11,6 +11,7 @@
 #include <unordered_map>
 #include <vector>
 
+#include "codon/cir/util/format.h"
 #include "codon/compiler/compiler.h"
 #include "codon/compiler/error.h"
 #include "codon/compiler/jit.h"
@@ -87,7 +88,7 @@ void initLogFlags(const llvm::cl::opt<std::string> &log) {
     codon::getLogger().parse(std::string(d));
 }
 
-enum BuildKind { LLVM, Bitcode, Object, Executable, Library, PyExtension, Detect };
+enum BuildKind { LLVM, Bitcode, Object, Executable, Library, PyExtension, Detect, CIR };
 enum OptMode { Debug, Release };
 enum Numerics { C, Python };
 } // namespace
@@ -333,6 +334,7 @@ int buildMode(const std::vector<const char *> &args, const std::string &argv0) {
           clEnumValN(Executable, "exe", "Generate executable"),
           clEnumValN(Library, "lib", "Generate shared library"),
           clEnumValN(PyExtension, "pyext", "Generate Python extension module"),
+          clEnumValN(CIR, "cir", "Generate Codon Intermediate Representation"),
           clEnumValN(Detect, "detect",
                      "Detect output type based on output file extension")),
       llvm::cl::init(Detect));
@@ -372,6 +374,9 @@ int buildMode(const std::vector<const char *> &args, const std::string &argv0) {
   case BuildKind::Detect:
     extension = "";
     break;
+  case BuildKind::CIR:
+    extension = ".cir";
+    break;
   default:
     seqassertn(0, "unknown build kind");
   }
@@ -401,6 +406,11 @@ int buildMode(const std::vector<const char *> &args, const std::string &argv0) {
     compiler->getLLVMVisitor()->writeToPythonExtension(*compiler->getCache()->pyModule,
                                                        filename);
     break;
+  case BuildKind::CIR: {
+    std::ofstream out(filename);
+    codon::ir::util::format(out, compiler->getModule());
+    break;
+  }
   case BuildKind::Detect:
     compiler->getLLVMVisitor()->compile(filename, argv0, libsVec, lflags);
     break;

codon/compiler/jit.cpp

@@ -263,21 +263,21 @@ ir::types::Type *JIT::PythonData::getCObjType(ir::Module *M) {
   return cobj;
 }
 
-JITResult JIT::executeSafe(const std::string &code, const std::string &file, int line,
-                           bool debug) {
+JIT::JITResult JIT::executeSafe(const std::string &code, const std::string &file,
+                                int line, bool debug) {
   auto result = execute(code, file, line, debug);
   if (auto err = result.takeError()) {
     auto errorInfo = llvm::toString(std::move(err));
     return JITResult::error(errorInfo);
   }
-  return JITResult::success(nullptr);
+  return JITResult::success();
 }
 
-JITResult JIT::executePython(const std::string &name,
-                             const std::vector<std::string> &types,
-                             const std::string &pyModule,
-                             const std::vector<std::string> &pyVars, void *arg,
-                             bool debug) {
+JIT::JITResult JIT::executePython(const std::string &name,
+                                  const std::vector<std::string> &types,
+                                  const std::string &pyModule,
+                                  const std::vector<std::string> &pyVars, void *arg,
+                                  bool debug) {
   auto key = buildKey(name, types);
   auto &cache = pydata->cache;
   auto it = cache.find(key);
@@ -322,26 +322,48 @@ JITResult JIT::executePython(const std::string &name,
   }
 }
 
-JIT *jitInit(const std::string &name) {
-  auto jit = new JIT(name);
+} // namespace jit
+} // namespace codon
+
+void *jit_init(char *name) {
+  auto jit = new codon::jit::JIT(std::string(name));
   llvm::cantFail(jit->init());
   return jit;
 }
 
-JITResult jitExecutePython(JIT *jit, const std::string &name,
-                           const std::vector<std::string> &types,
-                           const std::string &pyModule,
-                           const std::vector<std::string> &pyVars, void *arg,
-                           bool debug) {
-  return jit->executePython(name, types, pyModule, pyVars, arg, debug);
+void jit_exit(void *jit) { delete ((codon::jit::JIT *)jit); }
+
+CJITResult jit_execute_python(void *jit, char *name, char **types, size_t types_size,
+                              char *pyModule, char **py_vars, size_t py_vars_size,
+                              void *arg, uint8_t debug) {
+  std::vector<std::string> cppTypes;
+  cppTypes.reserve(types_size);
+  for (size_t i = 0; i < types_size; i++)
+    cppTypes.emplace_back(types[i]);
+  std::vector<std::string> cppPyVars;
+  cppPyVars.reserve(py_vars_size);
+  for (size_t i = 0; i < py_vars_size; i++)
+    cppPyVars.emplace_back(py_vars[i]);
+  auto t = ((codon::jit::JIT *)jit)
+               ->executePython(std::string(name), cppTypes, std::string(pyModule),
+                               cppPyVars, arg, bool(debug));
+  void *result = t.result;
+  char *message =
+      t.message.empty() ? nullptr : strndup(t.message.c_str(), t.message.size());
+  return {result, message};
 }
 
-JITResult jitExecuteSafe(JIT *jit, const std::string &code, const std::string &file,
-                         int line, bool debug) {
-  return jit->executeSafe(code, file, line, debug);
+CJITResult jit_execute_safe(void *jit, char *code, char *file, int32_t line,
+                            uint8_t debug) {
+  auto t = ((codon::jit::JIT *)jit)
+               ->executeSafe(std::string(code), std::string(file), line, bool(debug));
+  void *result = t.result;
+  char *message =
+      t.message.empty() ? nullptr : strndup(t.message.c_str(), t.message.size());
+  return {result, message};
 }
 
-std::string getJITLibrary() { return ast::library_path(); }
-
-} // namespace jit
-} // namespace codon
+char *get_jit_library() {
+  auto t = codon::ast::library_path();
+  return strndup(t.c_str(), t.size());
+}

codon/compiler/jit.h

@@ -31,6 +31,15 @@ public:
     ir::types::Type *getCObjType(ir::Module *M);
   };
 
+  struct JITResult {
+    void *result;
+    std::string message;
+
+    operator bool() const { return message.empty(); }
+    static JITResult success(void *result = nullptr) { return {result, ""}; }
+    static JITResult error(const std::string &message) { return {nullptr, message}; }
+  };
+
 private:
   std::unique_ptr<Compiler> compiler;
   std::unique_ptr<Engine> engine;

codon/compiler/jit_extern.h

@@ -2,35 +2,30 @@
 
 #pragma once
 
-#include <string>
-#include <vector>
+#include <stddef.h>
+#include <stdint.h>
 
-namespace codon {
-namespace jit {
+#ifdef __cplusplus
+extern "C" {
+#endif
 
-class JIT;
-
-struct JITResult {
+struct CJITResult {
   void *result;
-  std::string message;
-
-  operator bool() const { return message.empty(); }
-  static JITResult success(void *result) { return {result, ""}; }
-  static JITResult error(const std::string &message) { return {nullptr, message}; }
+  char *error;
 };
 
-JIT *jitInit(const std::string &name);
+void *jit_init(char *name);
+void jit_exit(void *jit);
 
-JITResult jitExecutePython(JIT *jit, const std::string &name,
-                           const std::vector<std::string> &types,
-                           const std::string &pyModule,
-                           const std::vector<std::string> &pyVars, void *arg,
-                           bool debug);
+struct CJITResult jit_execute_python(void *jit, char *name, char **types,
+                                     size_t types_size, char *pyModule, char **py_vars,
+                                     size_t py_vars_size, void *arg, uint8_t debug);
 
-JITResult jitExecuteSafe(JIT *jit, const std::string &code, const std::string &file,
-                         int line, bool debug);
+struct CJITResult jit_execute_safe(void *jit, char *code, char *file, int32_t line,
+                                   uint8_t debug);
 
-std::string getJITLibrary();
+char *get_jit_library();
 
-} // namespace jit
-} // namespace codon
+#ifdef __cplusplus
+}
+#endif

docs/README.md

@@ -124,7 +124,7 @@ print(total)
 ```
 
 Note that Codon automatically turns the `total += 1` statement in the loop body into an atomic
-reduction to avoid race conditions. Learn more in the [multitheading docs](advanced/parallel.md).
+reduction to avoid race conditions. Learn more in the [multithreading docs](advanced/parallel.md).
 
 Codon also supports writing and executing GPU kernels. Here's an example that computes the
 [Mandelbrot set](https://en.wikipedia.org/wiki/Mandelbrot_set):

jit/codon/__init__.py

@@ -1,5 +1,7 @@
 # Copyright (C) 2022-2025 Exaloop Inc. <https://exaloop.io>
 
-__all__ = ["jit", "convert", "JITError"]
+__all__ = [
+    "jit", "convert", "JITError", "JITWrapper", "_jit_register_fn", "_jit"
+]
 
-from .decorator import jit, convert, execute, JITError
+from .decorator import jit, convert, execute, JITError, JITWrapper, _jit_register_fn, _jit_callback_fn, _jit

jit/codon/decorator.py

@@ -23,16 +23,14 @@ if "CODON_PATH" not in os.environ:
     if codon_lib_path:
         codon_path.append(Path(codon_lib_path).parent / "stdlib")
     codon_path.append(
-        Path(os.path.expanduser("~")) / ".codon" / "lib" / "codon" / "stdlib"
-    )
+        Path(os.path.expanduser("~")) / ".codon" / "lib" / "codon" / "stdlib")
     for path in codon_path:
         if path.exists():
             os.environ["CODON_PATH"] = str(path.resolve())
             break
     else:
         raise RuntimeError(
-            "Cannot locate Codon. Please install Codon or set CODON_PATH."
-        )
+            "Cannot locate Codon. Please install Codon or set CODON_PATH.")
 
 pod_conversions = {
     type(None): "pyobj",
@@ -61,7 +59,6 @@ pod_conversions = {
 custom_conversions = {}
 _error_msgs = set()
 
-
 def _common_type(t, debug, sample_size):
     sub, is_optional = None, False
     for i in itertools.islice(t, sample_size):
@@ -76,7 +73,6 @@ def _common_type(t, debug, sample_size):
         sub = "Optional[{}]".format(sub)
     return sub if sub else "pyobj"
 
-
 def _codon_type(arg, **kwargs):
     t = type(arg)
 
@@ -88,11 +84,11 @@ def _codon_type(arg, **kwargs):
     if issubclass(t, set):
         return "Set[{}]".format(_common_type(arg, **kwargs))
     if issubclass(t, dict):
-        return "Dict[{},{}]".format(
-            _common_type(arg.keys(), **kwargs), _common_type(arg.values(), **kwargs)
-        )
+        return "Dict[{},{}]".format(_common_type(arg.keys(), **kwargs),
+                                    _common_type(arg.values(), **kwargs))
     if issubclass(t, tuple):
-        return "Tuple[{}]".format(",".join(_codon_type(a, **kwargs) for a in arg))
+        return "Tuple[{}]".format(",".join(
+            _codon_type(a, **kwargs) for a in arg))
     if issubclass(t, np.ndarray):
         if arg.dtype == np.bool_:
             dtype = "bool"
@@ -134,7 +130,8 @@ def _codon_type(arg, **kwargs):
 
     s = custom_conversions.get(t, "")
     if s:
-        j = ",".join(_codon_type(getattr(arg, slot), **kwargs) for slot in t.__slots__)
+        j = ",".join(
+            _codon_type(getattr(arg, slot), **kwargs) for slot in t.__slots__)
         return "{}[{}]".format(s, j)
 
     debug = kwargs.get("debug", None)
@@ -145,28 +142,22 @@ def _codon_type(arg, **kwargs):
             _error_msgs.add(msg)
     return "pyobj"
 
-
 def _codon_types(args, **kwargs):
     return tuple(_codon_type(arg, **kwargs) for arg in args)
 
-
 def _reset_jit():
     global _jit
     _jit = JITWrapper()
-    init_code = (
-        "from internal.python import "
-        "setup_decorator, PyTuple_GetItem, PyObject_GetAttrString\n"
-        "setup_decorator()\n"
-        "import numpy as np\n"
-        "import numpy.pybridge\n"
-    )
+    init_code = ("from internal.python import "
+                 "setup_decorator, PyTuple_GetItem, PyObject_GetAttrString\n"
+                 "setup_decorator()\n"
+                 "import numpy as np\n"
+                 "import numpy.pybridge\n")
     _jit.execute(init_code, "", 0, False)
     return _jit
 
-
 _jit = _reset_jit()
 
-
 class RewriteFunctionArgs(ast.NodeTransformer):
     def __init__(self, args):
         self.args = args
@@ -176,7 +167,6 @@ class RewriteFunctionArgs(ast.NodeTransformer):
             node.args.args.append(ast.arg(arg=a, annotation=None))
         return node
 
-
 def _obj_to_str(obj, **kwargs) -> str:
     if inspect.isclass(obj):
         lines = inspect.getsourcelines(obj)[0]
@@ -185,8 +175,10 @@ def _obj_to_str(obj, **kwargs) -> str:
         obj_name = obj.__name__
     elif callable(obj) or isinstance(obj, str):
         is_str = isinstance(obj, str)
-        lines = [i + '\n' for i in obj.split('\n')] if is_str else inspect.getsourcelines(obj)[0]
-        if not is_str: lines = lines[1:]
+        lines = [i + '\n' for i in obj.split('\n')
+                 ] if is_str else inspect.getsourcelines(obj)[0]
+        if not is_str:
+            lines = lines[1:]
         obj_str = textwrap.dedent(''.join(lines))
 
         pyvars = kwargs.get("pyvars", None)
@@ -195,8 +187,7 @@ def _obj_to_str(obj, **kwargs) -> str:
                 if not isinstance(i, str):
                     raise ValueError("pyvars only takes string literals")
             node = ast.fix_missing_locations(
-                RewriteFunctionArgs(pyvars).visit(ast.parse(obj_str))
-            )
+                RewriteFunctionArgs(pyvars).visit(ast.parse(obj_str)))
             obj_str = astunparse.unparse(node)
         if is_str:
             try:
@@ -206,28 +197,23 @@ def _obj_to_str(obj, **kwargs) -> str:
         else:
             obj_name = obj.__name__
     else:
-        raise TypeError("Function or class expected, got " + type(obj).__name__)
+        raise TypeError("Function or class expected, got " +
+                        type(obj).__name__)
     return obj_name, obj_str.replace("_@par", "@par")
 
-
 def _parse_decorated(obj, **kwargs):
-    return  _obj_to_str(obj, **kwargs)
-
+    return _obj_to_str(obj, **kwargs)
 
 def convert(t):
     if not hasattr(t, "__slots__"):
-        raise JITError("class '{}' does not have '__slots__' attribute".format(str(t)))
+        raise JITError("class '{}' does not have '__slots__' attribute".format(
+            str(t)))
 
     name = t.__name__
     slots = t.__slots__
-    code = (
-        "@tuple\n"
-        "class "
-        + name
-        + "["
-        + ",".join("T{}".format(i) for i in range(len(slots)))
-        + "]:\n"
-    )
+    code = ("@tuple\n"
+            "class " + name + "[" +
+            ",".join("T{}".format(i) for i in range(len(slots))) + "]:\n")
     for i, slot in enumerate(slots):
         code += "    {}: T{}\n".format(slot, i)
 
@@ -235,17 +221,14 @@ def convert(t):
     code += "    def __from_py__(p: cobj):\n"
     for i, slot in enumerate(slots):
         code += "        a{} = T{}.__from_py__(PyObject_GetAttrString(p, '{}'.ptr))\n".format(
-            i, i, slot
-        )
+            i, i, slot)
     code += "        return {}({})\n".format(
-        name, ", ".join("a{}".format(i) for i in range(len(slots)))
-    )
+        name, ", ".join("a{}".format(i) for i in range(len(slots))))
 
     _jit.execute(code, "", 0, False)
     custom_conversions[t] = name
     return t
 
-
 def _jit_register_fn(f, pyvars, debug):
     try:
         obj_name, obj_str = _parse_decorated(f, pyvars=pyvars)
@@ -258,29 +241,46 @@ def _jit_register_fn(f, pyvars, debug):
         _reset_jit()
         raise
 
-def _jit_callback_fn(obj_name, module, debug=None, sample_size=5, pyvars=None, *args, **kwargs):
-    try:
+def _jit_callback_fn(fn,
+                     obj_name,
+                     module,
+                     debug=None,
+                     sample_size=5,
+                     pyvars=None,
+                     *args,
+                     **kwargs):
+    if fn is not None:
+        sig = inspect.signature(fn)
+        bound_args = sig.bind(*args, **kwargs)
+        bound_args.apply_defaults()
+        args = tuple(bound_args.arguments[param] for param in sig.parameters)
+    else:
         args = (*args, *kwargs.values())
+
+    try:
         types = _codon_types(args, debug=debug, sample_size=sample_size)
         if debug:
-            print("[python] {}({})".format(obj_name, list(types)), file=sys.stderr)
-        return _jit.run_wrapper(
-            obj_name, list(types), module, list(pyvars), args, 1 if debug else 0
-        )
+            print("[python] {}({})".format(obj_name, list(types)),
+                  file=sys.stderr)
+        return _jit.run_wrapper(obj_name, list(types), module, list(pyvars),
+                                args, 1 if debug else 0)
     except JITError:
         _reset_jit()
         raise
 
 def _jit_str_fn(fstr, debug=None, sample_size=5, pyvars=None):
     obj_name = _jit_register_fn(fstr, pyvars, debug)
-    def wrapped(*args, **kwargs):
-        return _jit_callback_fn(obj_name, "__main__", debug, sample_size, pyvars, *args, **kwargs)
-    return wrapped
 
+    def wrapped(*args, **kwargs):
+        return _jit_callback_fn(None, obj_name, "__main__", debug, sample_size,
+                                pyvars, *args, **kwargs)
+
+    return wrapped
 
 def jit(fn=None, debug=None, sample_size=5, pyvars=None):
     if not pyvars:
         pyvars = []
+
     if not isinstance(pyvars, list):
         raise ArgumentError("pyvars must be a list")
 
@@ -289,12 +289,15 @@ def jit(fn=None, debug=None, sample_size=5, pyvars=None):
 
     def _decorate(f):
         obj_name = _jit_register_fn(f, pyvars, debug)
+
         @functools.wraps(f)
         def wrapped(*args, **kwargs):
-            return _jit_callback_fn(obj_name, f.__module__, debug, sample_size, pyvars, *args, **kwargs)
-        return wrapped
-    return _decorate(fn) if fn else _decorate
+            return _jit_callback_fn(f, obj_name, f.__module__, debug,
+                                    sample_size, pyvars, *args, **kwargs)
 
+        return wrapped
+
+    return _decorate(fn) if fn else _decorate
 
 def execute(code, debug=False):
     try:

jit/codon/jit.pxd

@@ -1,16 +1,22 @@
 # Copyright (C) 2022-2025 Exaloop Inc. <https://exaloop.io>
 
-from libcpp.string cimport string
-from libcpp.vector cimport vector
+from libc.stdint cimport int32_t, uint8_t
 
-cdef extern from "codon/compiler/jit_extern.h" namespace "codon::jit":
-    cdef cppclass JIT
-    cdef cppclass JITResult:
+cdef extern from "codon/compiler/jit_extern.h":
+    cdef struct CJITResult:
         void *result
-        string message
-        bint operator bool()
+        char *error
 
-    JIT *jitInit(string)
-    JITResult jitExecuteSafe(JIT*, string, string, int, char)
-    JITResult jitExecutePython(JIT*, string, vector[string], string, vector[string], object, char)
-    string getJITLibrary()
+    void *jit_init(char *name)
+    void jit_exit(void *jit)
+
+    cdef char *get_jit_library()
+
+    cdef CJITResult jit_execute_safe(
+        void *jit, char *code, char *file, int32_t line, uint8_t debug
+    )
+    cdef CJITResult jit_execute_python(
+        void *jit, char *name, char **types, size_t types_size,
+        char *pyModule, char **py_vars, size_t py_vars_size,
+        void *arg, uint8_t debug
+    )

jit/codon/jit.pyx

@@ -1,45 +1,82 @@
 # Copyright (C) 2022-2025 Exaloop Inc. <https://exaloop.io>
 
-# distutils: language=c++
+# distutils: language=c
 # cython: language_level=3
 # cython: c_string_type=unicode
 # cython: c_string_encoding=utf8
 
-from libcpp.string cimport string
-from libcpp.vector cimport vector
 cimport codon.jit
+from libc.stdlib cimport malloc, calloc, free
+from libc.string cimport strcpy
+from libc.stdint cimport int32_t, uint8_t
 
 
 class JITError(Exception):
     pass
 
 
+cdef str get_free_str(char *s):
+    cdef bytes py_s
+    try:
+        py_s = s
+        return py_s.decode('utf-8')
+    finally:
+        free(s)
+
+
 cdef class JITWrapper:
-    cdef codon.jit.JIT* jit
+    cdef void* jit
 
     def __cinit__(self):
-        self.jit = codon.jit.jitInit(b"codon jit")
+        self.jit = codon.jit.jit_init(b"codon jit")
 
     def __dealloc__(self):
-        del self.jit
+        codon.jit.jit_exit(self.jit)
 
-    def execute(self, code: str, filename: str, fileno: int, debug: char) -> str:
-        result = codon.jit.jitExecuteSafe(self.jit, code, filename, fileno, <char>debug)
-        if <bint>result:
+    def execute(self, code: str, filename: str, fileno: int, debug) -> str:
+        result = codon.jit.jit_execute_safe(
+            self.jit, code.encode('utf-8'), filename.encode('utf-8'), fileno, <uint8_t>debug
+        )
+        if result.error is NULL:
             return None
         else:
-            raise JITError(result.message)
+            msg = get_free_str(result.error)
+            raise JITError(msg)
+
+    def run_wrapper(self, name: str, types: list[str], module: str, pyvars: list[str], args, debug) -> object:
+        cdef char** c_types = <char**>calloc(len(types), sizeof(char*))
+        cdef char** c_pyvars = <char**>calloc(len(pyvars), sizeof(char*))
+        if not c_types or not c_pyvars:
+            raise JITError("Cython allocation failed")
+        try:
+            for i, s in enumerate(types):
+                bytes = s.encode('utf-8')
+                c_types[i] = <char*>malloc(len(bytes) + 1)
+                strcpy(c_types[i], bytes)
+            for i, s in enumerate(pyvars):
+                bytes = s.encode('utf-8')
+                c_pyvars[i] = <char*>malloc(len(bytes) + 1)
+                strcpy(c_pyvars[i], bytes)
+
+            result = codon.jit.jit_execute_python(
+                self.jit, name.encode('utf-8'), c_types, len(types),
+                module.encode('utf-8'), c_pyvars, len(pyvars),
+                <void *>args, <uint8_t>debug
+            )
+            if result.error is NULL:
+                return <object>result.result
+            else:
+                msg = get_free_str(result.error)
+                raise JITError(msg)
+        finally:
+            for i in range(len(types)):
+                free(c_types[i])
+            free(c_types)
+            for i in range(len(pyvars)):
+                free(c_pyvars[i])
+            free(c_pyvars)
 
-    def run_wrapper(self, name: str, types: list[str], module: str, pyvars: list[str], args, debug: char) -> object:
-        cdef vector[string] types_vec = types
-        cdef vector[string] pyvars_vec = pyvars
-        result = codon.jit.jitExecutePython(
-            self.jit, name, types_vec, module, pyvars_vec, <object>args, <char>debug
-        )
-        if <bint>result:
-            return <object>result.result
-        else:
-            raise JITError(result.message)
 
 def codon_library():
-    return codon.jit.getJITLibrary()
+    cdef char* c = codon.jit.get_jit_library()
+    return get_free_str(c)

jit/setup.py

@@ -67,7 +67,7 @@ jit_extension = Extension(
     "codon.codon_jit",
     sources=["codon/jit.pyx"],
     libraries=libraries,
-    language="c++",
+    language="c",
     extra_compile_args=["-w"],
     extra_link_args=linker_args,
     include_dirs=[str(codon_path / "include")],

stdlib/internal/python.codon

@@ -24,6 +24,7 @@ PyFloat_AsDouble = Function[[cobj], float](cobj())
 PyFloat_FromDouble = Function[[float], cobj](cobj())
 PyBool_FromLong = Function[[int], cobj](cobj())
 PyBytes_AsString = Function[[cobj], cobj](cobj())
+PyBytes_Size = Function[[cobj], int](cobj())
 PyList_New = Function[[int], cobj](cobj())
 PyList_Size = Function[[cobj], int](cobj())
 PyList_GetItem = Function[[cobj, int], cobj](cobj())
@@ -130,6 +131,7 @@ PyLong_Type = cobj()
 PyFloat_Type = cobj()
 PyBool_Type = cobj()
 PyUnicode_Type = cobj()
+PyBytes_Type = cobj()
 PyComplex_Type = cobj()
 PyList_Type = cobj()
 PyDict_Type = cobj()
@@ -213,6 +215,7 @@ def init_handles_dlopen(py_handle: cobj):
     global PyFloat_FromDouble
     global PyBool_FromLong
     global PyBytes_AsString
+    global PyBytes_Size
     global PyList_New
     global PyList_Size
     global PyList_GetItem
@@ -303,6 +306,7 @@ def init_handles_dlopen(py_handle: cobj):
     global PyFloat_Type
     global PyBool_Type
     global PyUnicode_Type
+    global PyBytes_Type
     global PyComplex_Type
     global PyList_Type
     global PyDict_Type
@@ -347,6 +351,7 @@ def init_handles_dlopen(py_handle: cobj):
     PyFloat_FromDouble = dlsym(py_handle, "PyFloat_FromDouble")
     PyBool_FromLong = dlsym(py_handle, "PyBool_FromLong")
     PyBytes_AsString = dlsym(py_handle, "PyBytes_AsString")
+    PyBytes_Size = dlsym(py_handle, "PyBytes_Size")
     PyList_New = dlsym(py_handle, "PyList_New")
     PyList_Size = dlsym(py_handle, "PyList_Size")
     PyList_GetItem = dlsym(py_handle, "PyList_GetItem")
@@ -437,6 +442,7 @@ def init_handles_dlopen(py_handle: cobj):
     PyFloat_Type = dlsym(py_handle, "PyFloat_Type")
     PyBool_Type = dlsym(py_handle, "PyBool_Type")
     PyUnicode_Type = dlsym(py_handle, "PyUnicode_Type")
+    PyBytes_Type = dlsym(py_handle, "PyBytes_Type")
     PyComplex_Type = dlsym(py_handle, "PyComplex_Type")
     PyList_Type = dlsym(py_handle, "PyList_Type")
     PyDict_Type = dlsym(py_handle, "PyDict_Type")
@@ -482,6 +488,7 @@ def init_handles_static():
     from C import PyFloat_FromDouble(float) -> cobj as _PyFloat_FromDouble
     from C import PyBool_FromLong(int) -> cobj as _PyBool_FromLong
     from C import PyBytes_AsString(cobj) -> cobj as _PyBytes_AsString
+    from C import PyBytes_Size(cobj) -> int as _PyBytes_Size
     from C import PyList_New(int) -> cobj as _PyList_New
     from C import PyList_Size(cobj) -> int as _PyList_Size
     from C import PyList_GetItem(cobj, int) -> cobj as _PyList_GetItem
@@ -572,6 +579,7 @@ def init_handles_static():
     from C import PyFloat_Type: cobj as _PyFloat_Type
     from C import PyBool_Type: cobj as _PyBool_Type
     from C import PyUnicode_Type: cobj as _PyUnicode_Type
+    from C import PyBytes_Type: cobj as _PyBytes_Type
     from C import PyComplex_Type: cobj as _PyComplex_Type
     from C import PyList_Type: cobj as _PyList_Type
     from C import PyDict_Type: cobj as _PyDict_Type
@@ -616,6 +624,7 @@ def init_handles_static():
     global PyFloat_FromDouble
     global PyBool_FromLong
     global PyBytes_AsString
+    global PyBytes_Size
     global PyList_New
     global PyList_Size
     global PyList_GetItem
@@ -706,6 +715,7 @@ def init_handles_static():
     global PyFloat_Type
     global PyBool_Type
     global PyUnicode_Type
+    global PyBytes_Type
     global PyComplex_Type
     global PyList_Type
     global PyDict_Type
@@ -750,6 +760,7 @@ def init_handles_static():
     PyFloat_FromDouble = _PyFloat_FromDouble
     PyBool_FromLong = _PyBool_FromLong
     PyBytes_AsString = _PyBytes_AsString
+    PyBytes_Size = _PyBytes_Size
     PyList_New = _PyList_New
     PyList_Size = _PyList_Size
     PyList_GetItem = _PyList_GetItem
@@ -840,6 +851,7 @@ def init_handles_static():
     PyFloat_Type = __ptr__(_PyFloat_Type).as_byte()
     PyBool_Type = __ptr__(_PyBool_Type).as_byte()
     PyUnicode_Type = __ptr__(_PyUnicode_Type).as_byte()
+    PyBytes_Type = __ptr__(_PyBytes_Type).as_byte()
     PyComplex_Type = __ptr__(_PyComplex_Type).as_byte()
     PyList_Type = __ptr__(_PyList_Type).as_byte()
     PyDict_Type = __ptr__(_PyDict_Type).as_byte()
@@ -1174,7 +1186,7 @@ class pyobj:
         return pyobj.to_str(self.p, errors, empty)
 
     def to_str(p: cobj, errors: str, empty: str = "") -> str:
-        obj = PyUnicode_AsEncodedString(p, "utf-8".c_str(), errors.c_str())
+        obj = PyUnicode_AsEncodedString(p, "utf-8".ptr, errors.c_str() if errors else "".ptr)
         if obj == cobj():
             return empty
         bts = PyBytes_AsString(obj)
@@ -1292,8 +1304,11 @@ class _PyObject_Struct:
 def _conversion_error(name: Static[str]):
     raise PyError("conversion error: Python object did not have type '" + name + "'")
 
+def _get_type(o: cobj):
+    return Ptr[_PyObject_Struct](o)[0].pytype
+
 def _ensure_type(o: cobj, t: cobj, name: Static[str]):
-    if Ptr[_PyObject_Struct](o)[0].pytype != t:
+    if _get_type(o) != t:
         _conversion_error(name)
 
 
@@ -1350,7 +1365,14 @@ class str:
         return pyobj.exc_wrap(PyUnicode_DecodeFSDefaultAndSize(self.ptr, self.len))
 
     def __from_py__(s: cobj) -> str:
-        return pyobj.exc_wrap(pyobj.to_str(s, "strict"))
+        if _get_type(s) == PyBytes_Type:
+            n = PyBytes_Size(s)
+            p0 = PyBytes_AsString(s)
+            p1 = cobj(n)
+            str.memcpy(p1, p0, n)
+            return str(p1, n)
+        else:
+            return pyobj.exc_wrap(pyobj.to_str(s, "strict"))
 
 @extend
 class complex:

stdlib/numpy/pybridge.codon

@@ -245,7 +245,7 @@ class ndarray:
             k = 0
             for idx in util.multirange(shape):
                 off = 0
-                for i in range(ndim):
+                for i in staticrange(ndim):
                     off += idx[i] * strides[i]
                 e = Ptr[cobj](arr_data + off)[0]
                 if hasattr(dtype, "__from_py__"):
@@ -263,7 +263,7 @@ class ndarray:
                 k = 0
                 for idx in util.multirange(shape):
                     off = 0
-                    for i in range(ndim):
+                    for i in staticrange(ndim):
                         off += idx[i] * strides[i]
                     e = Ptr[dtype](arr_data + off)[0]
                     data[k] = e

test/python/cython_jit.py

@@ -181,3 +181,16 @@ def test_ndarray():
     assert np.datetime_data(y.dtype) == ('s', 2)
 
 test_ndarray()
+
+@codon.jit
+def e(x=2, y=99):
+    return 2*x + y
+
+def test_arg_order():
+    assert e(1, 2) == 4
+    assert e(1) == 101
+    assert e(y=10, x=1) == 12
+    assert e(x=1) == 101
+    assert e() == 103
+
+test_arg_order()

test/python/pyext.py

@@ -88,6 +88,7 @@ def test_codon_extensions(m):
         assert m.f4(a=2.2) == (2.2, 2.22)
         assert m.f4(b=3.3) == (1.11, 3.3)
         assert m.f4('foo') == ('foo', 'foo')
+        assert m.f4(b'foo') == ('foo', 'foo')
         assert m.f4({1}) == {1}
         assert m.f5() is None
         assert equal(m.f6(1.9, 't'), 1.9, 1.9, 't')