Handle recursive tracing and CALL_ALLOC_AND_ENTER_INIT

Author: Fidget-Spinner

Python/bytecodes.c

@@ -5268,6 +5268,7 @@ dummy_func(
 
         tier2 op(_EXIT_TRACE, (exit_p/4 --)) {
             _PyExitData *exit = (_PyExitData *)exit_p;
+            assert(!exit->is_dynamic);
         #if defined(Py_DEBUG) && !defined(_Py_JIT)
             _Py_CODEUNIT *target = _PyFrame_GetBytecode(frame) + exit->target;
             OPT_HIST(trace_uop_execution_counter, trace_run_length_hist);
@@ -5281,7 +5282,6 @@ dummy_func(
             }
         #endif
             tstate->jit_exit = exit;
-            assert(!exit->is_dynamic);
             TIER2_TO_TIER2(exit->executor);
         }
 
@@ -5441,6 +5441,9 @@ dummy_func(
                 TIER2_TO_TIER2(exit->executor);
             }
             else {
+                if (frame->owner >= FRAME_OWNED_BY_INTERPRETER) {
+                    GOTO_TIER_ONE(target, 0);
+                }
                 if (!backoff_counter_triggers(temperature)) {
                     exit->temperature = advance_backoff_counter(temperature);
                     GOTO_TIER_ONE(target, 0);

Python/ceval_macros.h

@@ -138,13 +138,15 @@
 #if _Py_TAIL_CALL_INTERP || USE_COMPUTED_GOTOS
 #  define IS_JIT_TRACING() (tstate->interp->jit_is_tracing)
 #  define ENTER_TRACING() \
-    assert(!IS_JIT_TRACING()); \
     DISPATCH_TABLE_VAR = TRACING_DISPATCH_TABLE; \
     tstate->interp->jit_is_tracing = true;
 #  define LEAVE_TRACING() \
     assert(IS_JIT_TRACING()); \
     DISPATCH_TABLE_VAR = DISPATCH_TABLE; \
     tstate->interp->jit_is_tracing = false;
+// This handles recursive tracing over C calls.
+#  define RELOAD_TRACING() \
+    DISPATCH_TABLE_VAR = tstate->interp->jit_is_tracing ? TRACING_DISPATCH_TABLE : DISPATCH_TABLE;
 #  define BAIL_TRACING_NO_DISPATCH() \
     LEAVE_TRACING(); \
     _PyFrame_SetStackPointer(frame, stack_pointer); \
@@ -155,7 +157,7 @@
         JUMP_TO_LABEL(error); \
     }
 #  define RECORD_TRACE_NO_DISPATCH() do { \
-        if (add_to_code_trace(tstate, frame, old_code, old_func, this_instr, next_instr, opcode, oparg, _jump_taken)) { \
+        if (tstate->interp->jit_is_tracing && add_to_code_trace(tstate, frame, old_code, old_func, this_instr, next_instr, opcode, oparg, _jump_taken)) { \
             BAIL_TRACING_NO_DISPATCH(); \
         } \
     } while (0);
@@ -222,6 +224,7 @@ do { \
 #define TRACING_DISPATCH() \
     { \
         assert(frame->stackpointer == NULL); \
+        RELOAD_TRACING(); \
         RECORD_TRACE_NO_DISPATCH(); \
         NEXTOPARG(); \
         PRE_DISPATCH_GOTO(); \

Python/executor_cases.c.h

@@ -605,6 +605,10 @@ _PyJIT_translate_single_bytecode_to_trace(
         return 1;
     }
 
+    if (opcode == JUMP_FORWARD) {
+        return 1;
+    }
+
     if (opcode == ENTER_EXECUTOR) {
         ADD_TO_TRACE(_CHECK_VALIDITY, 0, 0, target);
         ADD_TO_TRACE(_SET_IP, 0, (uintptr_t)target_instr, target);
@@ -623,7 +627,10 @@ _PyJIT_translate_single_bytecode_to_trace(
 
     // Strange control-flow, unsupported opcode, etc.
     if (jump_taken ||
-        opcode == WITH_EXCEPT_START || opcode == RERAISE || opcode == CLEANUP_THROW || opcode == PUSH_EXC_INFO) {
+        opcode == WITH_EXCEPT_START || opcode == RERAISE || opcode == CLEANUP_THROW || opcode == PUSH_EXC_INFO ||
+        frame->owner >= FRAME_OWNED_BY_INTERPRETER ||
+        // This can be supported, but requires a tracing shim frame.
+        opcode == CALL_ALLOC_AND_ENTER_INIT) {
     unsupported:
             {
                 // Rewind to previous instruction and replace with _EXIT_TRACE.
@@ -814,6 +821,7 @@ _PyJIT_translate_single_bytecode_to_trace(
     if (!is_terminator(&tstate->interp->jit_tracer_code_buffer[trace_length-1])) {
         // Undo the last few instructions.
         trace_length = tstate->interp->jit_tracer_code_curr_size;
+        max_length = tstate->interp->jit_tracer_code_max_size;
         // We previously reversed one.
         max_length += 1;
         ADD_TO_TRACE(_EXIT_TRACE, 0, 0, target);