Fix exception tracing

Author: Fidget-Spinner

Include/internal/pycore_optimizer.h

@@ -345,7 +345,8 @@ static inline int is_terminator(const _PyUOpInstruction *uop)
     int opcode = uop->opcode;
     return (
         opcode == _EXIT_TRACE ||
-        opcode == _JUMP_TO_TOP
+        opcode == _JUMP_TO_TOP ||
+        opcode == _DYNAMIC_EXIT
     );
 }
 

Include/internal/pycore_uop_metadata.h

@@ -5446,11 +5446,31 @@ dummy_func(
 
         tier2 op(_GUARD_IP, (ip/4 --)) {
             if (frame->instr_ptr != (_Py_CODEUNIT *)ip) {
+#ifdef Py_DEBUG
+                _Py_CODEUNIT *target = frame->instr_ptr;
+                if (frame->lltrace >= 2) {
+                    printf("GUARD IP EXIT: [UOp ");
+                    _PyUOpPrint(&next_uop[-1]);
+                    printf(", target %d -> %s]\n",
+                        (int)(target - _PyFrame_GetBytecode(frame)),
+                        _PyOpcode_OpName[target->op.code]);
+                }
+#endif
                 GOTO_TIER_ONE(frame->instr_ptr, 1);
             }
         }
 
         tier2 op(_DYNAMIC_EXIT, (ip/4 --)) {
+#ifdef Py_DEBUG
+            _Py_CODEUNIT *target = frame->instr_ptr;
+            if (frame->lltrace >= 2) {
+                printf("GUARD IP EXIT: [UOp ");
+                _PyUOpPrint(&next_uop[-1]);
+                printf(", target %d -> %s]\n",
+                    (int)(target - _PyFrame_GetBytecode(frame)),
+                    _PyOpcode_OpName[target->op.code]);
+            }
+#endif
             GOTO_TIER_ONE(frame->instr_ptr, 1);
         }
 

Python/bytecodes.c

@@ -1198,7 +1198,10 @@ _PyTier2Interpreter(
         if (frame->lltrace >= 4 &&
             next_uop->opcode != _YIELD_VALUE &&
             next_uop->opcode != _FOR_ITER_GEN_FRAME &&
-            next_uop->opcode != _PUSH_FRAME) {
+            next_uop->opcode != _PUSH_FRAME &&
+            next_uop->opcode != _PY_FRAME_KW &&
+            next_uop->opcode != _SAVE_RETURN_OFFSET &&
+            next_uop->opcode != _SAVE_RETURN_OFFSET) {
             dump_stack(frame, stack_pointer);
             if (next_uop->opcode == _START_EXECUTOR) {
                 printf("%4d uop: ", 0);

Python/ceval.c

@@ -587,16 +587,24 @@ _PyJIT_translate_single_bytecode_to_trace(
         goto full;
     }
 
-    // Unsupported opcodes
-    if (opcode == WITH_EXCEPT_START || opcode == RERAISE || opcode == CLEANUP_THROW) {
-        goto full;
-    }
-
     bool needs_guard_ip = _PyOpcode_NeedsGuardIp[opcode] &&
         !(opcode == FOR_ITER_RANGE || opcode == FOR_ITER_LIST || opcode == FOR_ITER_TUPLE) &&
         !(opcode == JUMP_BACKWARD_NO_INTERRUPT || opcode == JUMP_BACKWARD || opcode == JUMP_BACKWARD_JIT) &&
         !(opcode == POP_JUMP_IF_TRUE || opcode == POP_JUMP_IF_FALSE || opcode == POP_JUMP_IF_NONE || opcode == POP_JUMP_IF_NOT_NONE);
 
+
+    // Strange control-flow
+    if (opcode == WITH_EXCEPT_START || opcode == RERAISE || opcode == CLEANUP_THROW || opcode == PUSH_EXC_INFO) {
+        // Rewind to previous instruction and replace with _EXIT_TRACE.
+        _PyUOpInstruction *curr = &trace[trace_length-1];
+        while (curr->opcode != _SET_IP && trace_length > 0) {
+            trace_length--;
+            curr = &trace[trace_length-1];
+        }
+        curr->opcode = _EXIT_TRACE;
+        goto done;
+    }
+
     const struct opcode_macro_expansion *expansion = &_PyOpcode_macro_expansion[opcode];
     RESERVE_RAW(expansion->nuops + needs_guard_ip + 3, "uop and various checks");
 
@@ -618,6 +626,8 @@ _PyJIT_translate_single_bytecode_to_trace(
 
     // Loop back to the start
     if (tstate->interp->jit_tracer_initial_instr == this_instr && tstate->interp->jit_tracer_code_curr_size > 2) {
+        // Undo the last few instructions.
+        trace_length = tstate->interp->jit_tracer_code_curr_size;
         ADD_TO_TRACE(_JUMP_TO_TOP, 0, 0, 0);
         goto done;
     }

Python/executor_cases.c.h

@@ -493,7 +493,7 @@ remove_unneeded_uops(_PyUOpInstruction *buffer, int buffer_size)
                 }
                 /* _PUSH_FRAME doesn't escape or error, but it
                  * does need the IP for the return address */
-                bool needs_ip = (opcode == _PUSH_FRAME || opcode == _YIELD_VALUE);
+                bool needs_ip = (opcode == _PUSH_FRAME || opcode == _YIELD_VALUE || opcode == _DYNAMIC_EXIT);
                 if (_PyUop_Flags[opcode] & HAS_ESCAPES_FLAG) {
                     needs_ip = true;
                     may_have_escaped = true;
@@ -507,6 +507,7 @@ remove_unneeded_uops(_PyUOpInstruction *buffer, int buffer_size)
             }
             case _JUMP_TO_TOP:
             case _EXIT_TRACE:
+            case _DYNAMIC_EXIT:
                 return pc + 1;
         }
     }