JIT Compilation — Just-In-Time Compilation Explained

class ProfilingInterpreter:
    def __init__(self):
        self.execution_counts = {}
        self.hot_threshold = 100
        self.jit = None

    def interpret(self, bytecode, constants):
        ip = 0
        while ip < len(bytecode):
            op = bytecode[ip]
            self.execution_counts[ip] = self.execution_counts.get(ip, 0) + 1

            if self.execution_counts[ip] > self.hot_threshold and self.jit:
                native_code = self.jit.compile(bytecode, ip)
                if native_code:
                    return self.jit.execute(native_code)

            if op == Opcode.LOAD_CONST:
                ip += 1
                idx = bytecode[ip]
                self.stack.append(constants[idx])
            elif op == Opcode.ADD:
                right = self.stack.pop()
                left = self.stack.pop()
                self.stack.append(left + right)
            ip += 1

class MethodJIT:
    def __init__(self):
        self.compiled_methods = {}

    def should_compile(self, method_name, call_count):
        return call_count > 1000 and method_name not in self.compiled_methods

    def compile(self, method_name, bytecode):
        print(f"JIT compiling: {method_name}")
        # In a real JIT, this would generate native machine code
        native_code = self._emit_native_code(bytecode)
        self.compiled_methods[method_name] = native_code
        return native_code

    def _emit_native_code(self, bytecode):
        # Simplified: returns a callable that executes the bytecode natively
        def native_fn(*args):
            total = 0
            for arg in args:
                total += arg
            return total
        return native_fn

    def execute(self, method_name, *args):
        native_fn = self.compiled_methods.get(method_name)
        if native_fn:
            return native_fn(*args)
        raise RuntimeError(f"Method not compiled: {method_name}")

jit = MethodJIT()
print(jit.should_compile("add", 500))
print(jit.should_compile("add", 1500))
jit.compile("add", [])
result = jit.execute("add", 10, 20)
print(f"Result: {result}")

class TraceJIT:
    def __init__(self):
        self.traces = {}
        self.recording = False
        self.current_trace = []

    def start_recording(self, loop_header):
        self.recording = True
        self.current_trace = [("loop_header", loop_header)]
        print(f"Starting trace at {loop_header}")

    def record_instruction(self, instruction):
        if self.recording:
            self.current_trace.append(instruction)

    def stop_recording(self):
        self.recording = False
        trace = list(self.current_trace)
        self.traces[self.current_trace[0][1]] = self._compile_trace(trace)
        self.current_trace = []
        return trace

    def _compile_trace(self, trace):
        print(f"Compiling trace: {len(trace)} instructions")
        def compiled_trace(state):
            x = state["x"]
            for _ in range(100):
                x = x * 2
            state["x"] = x
            return state
        return compiled_trace

    def execute_trace(self, loop_header, state):
        trace_fn = self.traces.get(loop_header)
        if trace_fn:
            return trace_fn(state)
        return state

trace_jit = TraceJIT()
trace_jit.start_recording("L1")
trace_jit.record_instruction("mul x, 2")
trace_jit.record_instruction("inc i")
trace_jit.record_instruction("cmp i, 100")
trace = trace_jit.stop_recording()
state = trace_jit.execute_trace("L1", {"x": 5})
print(f"Result: {state['x']}")

def check_all_compiled(state):
    """Toggle JIT compilation on and off for testing."""
    if state.get("compiled"):
        return "JIT active"
    return "Interpreter mode"

class Deoptimizer:
    def __init__(self):
        self.assumptions = {}

    def add_assumption(self, code_addr, check_fn):
        self.assumptions[code_addr] = check_fn

    def invalidate(self):
        invalidated = []
        for addr, check_fn in self.assumptions.items():
            if not check_fn():
                invalidated.append(addr)
        for addr in invalidated:
            del self.assumptions[addr]
            print(f"Deoptimized code at {addr}")
        return invalidated