gh-145064: fix JIT tracing for non-function frames

Author: aviralgarg05

Lib/test/test_capi/test_opt.py

@@ -4065,6 +4065,24 @@ def hot_loop():
         self.assertNotIn('_PyJit_TryInitializeTracing', stderr,
                          f"JIT tracer memory leak detected:\n{stderr}")
 
+    def test_145064(self):
+        # https://github.com/python/cpython/issues/145064
+        result = script_helper.run_python_until_end('-c', textwrap.dedent(f"""
+        SRC = '''
+        class A:
+            def __init__(self, x):
+                self.x = x
+
+        for i in range(8):
+            A(i)
+        '''
+
+        for _ in range({TIER2_THRESHOLD * 6}):
+            ns = {{}}
+            exec(SRC, ns, ns)
+        """), PYTHON_JIT="1", PYTHON_JIT_STRESS="1")
+        self.assertEqual(result[0].rc, 0, result)
+
 def global_identity(x):
     return x
 

Misc/NEWS.d/next/Core_and_Builtins/2026-02-21-18-10-00.gh-issue-145064.p8Nq2v.rst

@@ -0,0 +1 @@
+Fix a JIT crash in ``_Py_uop_frame_new()`` when tracing starts from non-function frames (for example ``_Py_InitCleanup`` shims) by using the traced code object directly and only storing function metadata when available.

Python/optimizer.c

@@ -146,7 +146,7 @@ _PyOptimizer_Optimize(
     assert(!interp->compiling);
     assert(_tstate->jit_tracer_state->initial_state.stack_depth >= 0);
 #ifndef Py_GIL_DISABLED
-    assert(_tstate->jit_tracer_state->initial_state.func != NULL);
+    assert(_tstate->jit_tracer_state->initial_state.code != NULL);
     interp->compiling = true;
     // The first executor in a chain and the MAX_CHAIN_DEPTH'th executor *must*
     // make progress in order to avoid infinite loops or excessively-long
@@ -1020,7 +1020,10 @@ _PyJit_TryInitializeTracing(
     tracer->initial_state.start_instr = start_instr;
     tracer->initial_state.close_loop_instr = close_loop_instr;
     tracer->initial_state.code = (PyCodeObject *)Py_NewRef(code);
-    tracer->initial_state.func = (PyFunctionObject *)Py_NewRef(func);
+    // Tracing can start in shim frames (e.g. _Py_InitCleanup), where
+    // frame->f_funcobj is not a PyFunctionObject.
+    tracer->initial_state.func = PyFunction_Check(func) ?
+        (PyFunctionObject *)Py_NewRef(func) : NULL;
     tracer->initial_state.executor = (_PyExecutorObject *)Py_XNewRef(current_executor);
     tracer->initial_state.exit = exit;
     tracer->initial_state.stack_depth = (int)(stack_pointer - _PyFrame_Stackbase(frame));

Python/optimizer_analysis.c

@@ -458,6 +458,8 @@ optimize_uops(
 {
     assert(!PyErr_Occurred());
     assert(tstate->jit_tracer_state != NULL);
+    PyCodeObject *co = tstate->jit_tracer_state->initial_state.code;
+    assert(co != NULL);
     PyFunctionObject *func = tstate->jit_tracer_state->initial_state.func;
 
     JitOptContext *ctx = &tstate->jit_tracer_state->opt_context;
@@ -473,11 +475,13 @@ optimize_uops(
     }
 
     _Py_uop_abstractcontext_init(ctx, dependencies);
-    _Py_UOpsAbstractFrame *frame = _Py_uop_frame_new(ctx, (PyCodeObject *)func->func_code, NULL, 0);
+    _Py_UOpsAbstractFrame *frame = _Py_uop_frame_new(ctx, co, NULL, 0);
     if (frame == NULL) {
         return 0;
     }
-    frame->func = func;
+    if (func != NULL) {
+        frame->func = func;
+    }
     ctx->curr_frame_depth++;
     ctx->frame = frame;
     _Py_uop_sym_set_stack_depth(ctx, curr_stacklen, frame->stack_pointer);