JustJIT

A Just-In-Time compiler for Python that compiles Python bytecode to native machine code using LLVM.

---

Installation

pip install justjit

Quick Start

from justjit import jit

@jit
def fibonacci(n):
    if n <= 1:
        return n
    a, b = 0, 1
    for _ in range(n - 1):
        a, b = b, a + b
    return b

result = fibonacci(40)  # Compiled to native code

Generator Functions

@jit
def count_up(n):
    i = 0
    while i < n:
        yield i
        i += 1

for value in count_up(10):
    print(value)

Async/Await

import asyncio
from justjit import jit

@jit
async def fetch_data():
    await asyncio.sleep(0.1)
    return "data"

result = asyncio.run(fetch_data())

Exception Handling

@jit
def safe_divide(a, b):
    try:
        return a / b
    except ZeroDivisionError:
        return 0

print(safe_divide(10, 2))  # 5.0
print(safe_divide(10, 0))  # 0

Float Mode

@jit(mode='float')
def distance(x1, y1, x2, y2):
    dx = x2 - x1
    dy = y2 - y1
    return (dx * dx + dy * dy) ** 0.5

print(distance(0.0, 0.0, 3.0, 4.0))  # 5.0

List Comprehensions

@jit
def squares(n):
    return [x * x for x in range(n)]

print(squares(5))  # [0, 1, 4, 9, 16]

Features

• Object Mode (default): Full Python semantics via C API calls
• Native Modes: int, float, bool, complex128, vec4f, vec8i for specialized workloads
• Generators: State machine compilation with yield support
• Coroutines: Async/await with awaitable protocol
• Exception Handling: Try/except/finally with stack unwinding
• Inline C/C++: Compile C code at runtime (requires Clang)

Compilation Modes

Mode	Type	Use Case
auto	PyObject*	Default, full Python compatibility
int	i64	Pure integer loops
float	f64	Floating-point math
complex128	{f64, f64}	Complex number operations
vec4f	<4 x f32>	SSE SIMD (4 floats)
vec8i	<8 x i32>	AVX SIMD (8 ints)

@jit(mode='int')
def sum_range(n):
    total = 0
    for i in range(n):
        total += i
    return total

Inline C Compiler

from justjit import C

result = C("""
double square(double x) {
    return x * x;
}
""")

print(result['square'](4.0))  # 16.0

How It Works

Compilation Pipeline

Python Function → Bytecode → CFG Analysis → LLVM IR → Native Code

1. Bytecode Extraction: Uses Python's dis module to extract bytecode instructions
2. CFG Construction: Builds control flow graph, identifies basic blocks and jump targets
3. Stack Depth Analysis: Dataflow analysis to compute stack depth at each instruction
4. IR Generation: Translates each opcode to LLVM IR
5. Optimization: Applies LLVM optimization passes (O0-O3)
6. JIT Compilation: LLVM ORC JIT compiles to native machine code

Object Mode Internals

In object mode, Python operations are compiled to C API calls:

# Python code
result = a + b

; Generated LLVM IR
%result = call ptr @PyNumber_Add(ptr %a, ptr %b)
call void @Py_DECREF(ptr %a)
call void @Py_DECREF(ptr %b)

Every API call is wrapped with NULL-checking for exception handling.

Native Mode Internals

Native modes bypass Python objects entirely:

@jit(mode='int')
def add(a, b):
    return a + b

; Generated LLVM IR - pure LLVM i64
define i64 @add(i64 %a, i64 %b) {
  %result = add i64 %a, %b
  ret i64 %result
}

Generator State Machine

Generators compile to step functions with explicit state:

// Step function signature
PyObject* step(int32_t* state, PyObject** locals, PyObject* sent);

// States: 0=initial, 1..N=resume points, -1=done, -2=error

At each yield:

1. Stack values spill to persistent locals array
2. State advances to next resume point
3. Function returns yielded value

On resume:

1. Switch dispatches to correct state
2. Stack restores from locals array
3. Execution continues after yield

Coroutine Internals

Coroutines extend generators with awaitable protocol:

struct JITCoroutineObject {
    int32_t state;           // Suspension state
    PyObject** locals;       // Persistent variables
    PyObject* awaiting;      // Currently awaited object
    GeneratorStepFunc step;  // Compiled step function
};

await compiles to:

1. GET_AWAITABLE - validate and get iterator
2. SEND - delegate to awaited object
3. END_SEND - extract return value

Inline C Compiler

Uses embedded Clang to compile C/C++ at runtime:

C Code → Clang Frontend → LLVM IR → Same JIT as Python

Python variables are captured as C declarations:

// Python: x = 42, arr = numpy.array([1,2,3])
// Generated C preamble:
long long x = 42LL;
double* arr = (double*)0x7fff12345678ULL;
long long arr_len = 3LL;

Current Status

Working

• Basic control flow (if/else, for, while)
• Arithmetic and comparison operations
• Function calls
• List/dict/set operations
• Try/except handling
• Simple generators

Known Limitations

• Native modes (int, float) only support scalar operations, not array indexing
• Nested list comprehensions in generators may have issues
• Some complex async patterns not fully tested

Building from Source

Requires: Python 3.10+, LLVM 18+, CMake 3.20+, C++17 compiler

# Set LLVM path
export LLVM_DIR=/path/to/llvm/lib/cmake/llvm  # Linux/macOS
$env:LLVM_DIR = "C:\path\to\llvm\lib\cmake\llvm"  # Windows

# Build
pip install -e . --no-build-isolation

License

MIT License - see LICENSE

Links

• Documentation