RFC: Native-backed Python Classes for Codon JIT via codon.jitclass

[LeeLee26]

Motivation & Background

Codon already provides codon.convert as a mechanism to expose Python classes to Codon JIT-compiled functions. This approach is suitable for simple interoperability scenarios: a Python object is converted into a Codon-compatible representation, passed into compiled code, and then used within JIT-compiled functions.

To illustrate, consider the following example:

@codon.convert
class A:
    __slots__ = ["x", "y"]

    def __init__(self, x: int, y: int) -> None:
        self.x = x
        self.y = y

    @codon.jit
    def hello(self):
        return self.x + self.y

a = A(1, 2)
a.hello()

When a class is decorated with codon.convert, the Codon JIT side generates a native Codon tuple-like class A[T0, T1]. The @codon.jit decorator on a member function behaves exactly as it does on a standalone function—it treats self as an ordinary parameter.

The expression a = A(1, 2) still executes the original Python __init__ and produces a standard PyObject. When a.hello() is invoked, the Codon JIT first converts the PyObject a into an instance of the generated A[T0, T1] class and then passes that converted object to the JIT-compiled hello function.

However, codon.convert has several inherent limitations when the goal is to model Python classes as native Codon objects:

1. Object identity and lifetime are not naturally preserved
   codon.convert primarily treats Python objects as values that can be transformed into Codon-side data structures. This works well for immutable or value-like objects, but it is not ideal for class instances that are expected to behave as long-lived mutable objects with stable identity.

2. Mutation semantics are limited
   For Python classes with fields, users often expect operations such as

   obj.x = 10
   obj.update()
   print(obj.x)

   to operate on the same underlying object. With conversion-based interop, field mutations may not naturally map to persistent native state unless the object is explicitly reconstructed or round-tripped between Python and Codon representations.

3. Methods are not represented as native class methods
   codon.convert allows converted objects to be consumed by JIT functions, but it does not make the Python class itself a native Codon class with compiled methods, constructors, and field accessors. This prevents object-oriented workloads from being expressed directly in the JIT layer.

4. Performance opportunities are left on the table
   A conversion-based approach introduces boundaries between Python objects and Codon values. For class-heavy workloads, repeatedly converting objects or routing behavior through Python-level wrappers can erode the performance gains expected from JIT compilation.

5. The user model is less direct
   Users who write ordinary Python classes typically expect a decorator-based workflow in which the class itself becomes compiled, rather than manually converting instances or authoring separate JIT functions that operate on them.

Because of these limitations, we propose adding first-class support for JIT-compiled classes through codon.jitclass.

[LeeLee26]

Proposal

I propose adding a new decorator, codon.jitclass, that fully compiles a Python class into a native Codon type, with the Python side functioning solely as a proxy object backed by a Codon-managed native instance.

The earlier background example can be rewritten as follows:

@codon.jitclass
class A:
    x: int
    y: int

    def __init__(self, x: int, y: int) -> None:
        self.x = x
        self.y = y

    def hello(self):
        return self.x + self.y

a = A(1, 2)
print(a.hello())  # 3

a.x = 10
print(a.hello())  # 12

Unlike the codon.convert example, this class A implementation is closer to native Codon code: explicit member-variable declarations replace the __slots__ constraint. When a class is decorated with @codon.jitclass, its source code is handed directly to the Codon JIT for compilation (with the necessary syntactic transformations applied to meet Codon requirements). All member methods are compiled together, and the result is a Python proxy class that wraps every method; calls to those methods invoke the corresponding JIT-compiled Codon class methods directly.

For instance, when executing a = A(1, 2), the constructor of the Codon JIT-generated class A (which we call NativeA) is invoked immediately. This returns a proxy object a that holds a handle to a NativeA instance. From that point forward, every field access such as a.x and every method call such as a.hello() is routed through the proxy to the underlying handle.

Realizing this proposal requires the following components.

1. Class-level JIT compilation

When a class is decorated with @codon.jitclass, its Python AST is transformed into a Codon-compatible class definition. The transformed class is compiled into the existing JIT context.

The decorator supports:

• annotated instance fields,
• __init__ constructors,
• instance methods,
• empty classes,
• method-only classes,
• field getter and setter access from Python.

For example:

@codon.jitclass
class Counter:
    value: int

    def __init__(self, value: int):
        self.value = value

    def inc(self, n: int) -> int:
        self.value += n
        return self.value

is compiled into a native Codon class. Python calls to Counter(...), counter.inc(...), counter.value, and counter.value = ... are dispatched to JIT-generated wrappers.

---

2. Native object construction

Calling a jitclass from Python constructs a real Codon-side object. The constructor arguments are mapped to Codon types using the same runtime type inference machinery employed by codon.jit.

Internally, construction goes through a JIT-generated wrapper similar to:

@export
def __codon_jitclass_new(args: cobj) -> cobj:
    ...
    obj = NativeClass(...)
    return obj.__raw__()

The returned native pointer is stored in a JIT-side handle table. Python receives a lightweight proxy object that refers to this handle.

This means the object is not repeatedly converted back and forth. Instead, the native object remains alive in the JIT runtime and is accessed through stable handles.

---

3. Method dispatch

Each Python method call on a jitclass instance is routed to a compiled Codon method wrapper.

For example, x.hello() is dispatched as:

Python proxy
  -> JIT handle
  -> native Codon object pointer
  -> compiled Codon method wrapper
  -> native method call
  -> Python result conversion

This gives users a natural Python object-oriented interface while executing the actual method body in native Codon code.

---

4. Field access

Field access is supported through generated getter and setter methods.

For a field declaration:

x: int

the implementation internally generates methods equivalent to:

def _codon_jitclass_get_x(self) -> int:
    return self.x

def _codon_jitclass_set_x(self, value: int) -> int:
    self.x = value
    return value

Python descriptors then expose this as normal attribute access:

a.x           # native getter
a.y = 10   # native setter

This allows mutable native state to be observed and updated from Python naturally.

---

5. Lifecycle management

A jitclass instance owns a Python proxy, while the actual object is stored in the Codon JIT runtime.

The proxy supports explicit release and context-manager usage:

a = A(1, 2)
a.close()

with A(1, 2) as a:
    print(a.hello())

Once released, further access raises a JIT error instead of using an invalid native pointer.

The implementation also tracks the active JIT generation. If the JIT context is reset because of a compilation error or runtime failure, old jitclass objects are marked as stale and cannot accidentally call into a new, incompatible JIT context.

---

6. GC root safety

Codon objects are allocated in the Codon runtime and may be managed by the runtime GC. Since __raw__() exposes a borrowed pointer to the native object, the JIT runtime must ensure that the object remains reachable while Python still holds a proxy.

To achieve this, each native jitclass object is registered as a GC root through RAII:

struct JITClassRoot {
    std::unique_ptr<void *> slot;

    explicit JITClassRoot(void *ptr);
    ~JITClassRoot();
};

The root is added when the native object is stored in the JIT handle table and removed automatically when the handle is released. This prevents the Codon GC from collecting native objects that are still reachable from Python.

---

Expected Effect

codon.jitclass provides a more direct and capable object model for Python-to-Codon JIT integration.

Compared with codon.convert, it enables:

1. Native-backed Python objects
   Instances created from a jitclass are backed by real Codon objects rather than temporary converted values.

2. Persistent mutable state
   Field updates and method calls operate on the same native object.

3. Natural Python syntax
   Users can write obj.field, obj.field = value, and obj.method(...) while executing the underlying operations in compiled Codon code.

4. Lower conversion overhead
   The object remains native after construction. Only method arguments and return values cross the Python/Codon boundary.

5. Better support for object-oriented workloads
   Classes with constructors, fields, and methods can be compiled as a unit, making Codon JIT more suitable for simulations, numerical kernels with stateful objects, and class-heavy Python code.

6. Safer runtime behavior
   Handle-based dispatch, explicit release, stale-context detection, and GC rooting provide a controlled ownership model for native objects exposed to Python.

[LeeLee26]

Initial Implementation

1. JIT: add JITCallable to wrap JIT callable functions(1/4) #813
2. JIT: add jitclass for native codon jit class(2/4) #816
3. JIT: add stale jit context detection(3/4) #814
4. JIT: support field access for codon jitclass(4/4) #817

[LeeLee26]

cc @BI71317 @arshajii @inumanag