Fix bug where code/func get freed halfway

Author: Fidget-Spinner

Include/internal/pycore_interp_structs.h

@@ -951,8 +951,8 @@ struct _is {
     _Py_CODEUNIT *jit_tracer_initial_instr;
     int jit_tracer_initial_stack_depth;
     int jit_tracer_initial_chain_depth;
-    PyCodeObject *jit_tracer_initial_code; // Borrowed
-    PyFunctionObject *jit_tracer_initial_func; // Borrowed
+    PyCodeObject *jit_tracer_initial_code; // Strong
+    PyFunctionObject *jit_tracer_initial_func; // Strong
     struct _PyExitData *jit_tracer_previous_exit;
     bool jit_completed_loop;
     bool jit;

Include/internal/pycore_optimizer.h

@@ -372,6 +372,8 @@ _PyJIT_translate_single_bytecode_to_trace(
 
 void
 _PyJIT_InitializeTracing(PyThreadState *tstate, _PyInterpreterFrame *frame, _Py_CODEUNIT *next_instr, int curr_stackdepth, int chain_depth, _PyExitData *exit);
+
+void _PyJIT_FinalizeTracing(PyThreadState *tstate);
 #ifdef __cplusplus
 }
 #endif

Python/ceval_macros.h

@@ -141,21 +141,22 @@
 #  define LEAVE_TRACING() DISPATCH_TABLE_VAR = DISPATCH_TABLE;
 #  define BAIL_TRACING_NO_DISPATCH() \
     LEAVE_TRACING(); \
+    _PyFrame_SetStackPointer(frame, stack_pointer); \
     int _err = _PyOptimizer_Optimize(frame, tstate); \
+    _PyJIT_FinalizeTracing(tstate); \
+    stack_pointer = _PyFrame_GetStackPointer(frame); \
     if (_err < 0) { \
         JUMP_TO_LABEL(error); \
     }
 #  define BAIL_TRACING() \
     BAIL_TRACING_NO_DISPATCH() \
     DISPATCH();
 #  define RECORD_TRACE() do { \
-        frame->instr_ptr = next_instr; \
         if (add_to_code_trace(tstate, frame, old_code, old_func, this_instr, next_instr, opcode, oparg, _jump_taken)) { \
             BAIL_TRACING(); \
         } \
     } while (0);
 #  define RECORD_TRACE_NO_DISPATCH() do { \
-        frame->instr_ptr = next_instr; \
         if (add_to_code_trace(tstate, frame, old_code, old_func, this_instr, next_instr, opcode, oparg, _jump_taken)) { \
             BAIL_TRACING_NO_DISPATCH(); \
         } \
@@ -413,6 +414,7 @@ _PyFrame_SetStackPointer(frame, stack_pointer)
 
 #define TIER1_TO_TIER2(EXECUTOR)                        \
 do {                                                   \
+    LEAVE_TRACING(); \
     OPT_STAT_INC(traces_executed);                     \
     next_instr = _Py_jit_entry((EXECUTOR), frame, stack_pointer, tstate); \
     frame = tstate->current_frame;                     \
@@ -428,9 +430,6 @@ do {                                                   \
         assert(tstate->interp->jit_tracer_code_curr_size == 2); \
         ENTER_TRACING(); \
     } \
-    else { \
-        LEAVE_TRACING(); \
-    } \
     DISPATCH();                                        \
 } while (0)
 

Python/generated_cases.c.h

@@ -129,13 +129,18 @@ _PyOptimizer_Optimize(
     // side-exit chains. We can only insert the executor into the bytecode if
     // this is true, since a deopt won't infinitely re-enter the executor:
     bool progress_needed = (chain_depth % MAX_CHAIN_DEPTH) == 0;
-    PyCodeObject *code = (PyCodeObject *)tstate->interp->jit_tracer_initial_func->func_code;
-    assert(PyCode_Check(code));
+    PyCodeObject *code = (PyCodeObject *)tstate->interp->jit_tracer_initial_code;
     _Py_CODEUNIT *start = tstate->interp->jit_tracer_initial_instr;
     if (progress_needed && !has_space_for_executor(code, start)) {
         interp->compiling = false;
         return 0;
     }
+    // We are the only one still holding a reference to this code object that
+    // is practically dead.
+    if (_PyObject_IsUniquelyReferenced(code) || _PyObject_IsUniquelyReferenced(tstate->interp->jit_tracer_initial_func)) {
+        interp->compiling = false;
+        return 0;
+    }
     _PyExecutorObject *executor;
     int err = uop_optimize(frame, tstate, &executor, progress_needed);
     if (err <= 0) {
@@ -840,15 +845,24 @@ _PyJIT_InitializeTracing(PyThreadState *tstate, _PyInterpreterFrame *frame, _Py_
     tstate->interp->jit_tracer_code_curr_size = 2;
     tstate->interp->jit_tracer_code_max_size = UOP_MAX_TRACE_LENGTH;
     tstate->interp->jit_tracer_initial_instr = next_instr;
-    tstate->interp->jit_tracer_initial_code = code;
-    tstate->interp->jit_tracer_initial_func = _PyFrame_GetFunction(frame);
+    tstate->interp->jit_tracer_initial_code = (PyCodeObject *)Py_NewRef(code);
+    tstate->interp->jit_tracer_initial_func = (PyFunctionObject *)Py_NewRef(_PyFrame_GetFunction(frame));
     tstate->interp->jit_tracer_previous_exit = exit;
     memset(&tstate->interp->jit_tracer_dependencies.bits, 0, sizeof(tstate->interp->jit_tracer_dependencies.bits));
     tstate->interp->jit_completed_loop = false;
     tstate->interp->jit_tracer_initial_stack_depth = curr_stackdepth;
     tstate->interp->jit_tracer_initial_chain_depth = chain_depth;
 }
 
+void
+_PyJIT_FinalizeTracing(PyThreadState *tstate)
+{
+    Py_CLEAR(tstate->interp->jit_tracer_initial_code);
+    Py_CLEAR(tstate->interp->jit_tracer_initial_func);
+    tstate->interp->jit_tracer_code_curr_size = 2;
+}
+
+
 #undef RESERVE
 #undef RESERVE_RAW
 #undef INSTR_IP