From 321d35cff108daf23f8c204eae9cacd95f916e1f Mon Sep 17 00:00:00 2001
From: Ken Jin <kenjin4096@gmail.com>
Date: Mon, 23 Mar 2026 00:10:34 +0800
Subject: [PATCH] Rewind optimizer when out of space

---
 Include/internal/pycore_optimizer.h | 10 +++++++
 Python/optimizer_analysis.c         | 45 ++++++++++++++++++++---------
 2 files changed, 41 insertions(+), 14 deletions(-)

diff --git a/Include/internal/pycore_optimizer.h b/Include/internal/pycore_optimizer.h
index 94c79edc855bbf..a79d502a7a7395 100644
--- a/Include/internal/pycore_optimizer.h
+++ b/Include/internal/pycore_optimizer.h
@@ -67,6 +67,16 @@ uop_buffer_last(_PyJitUopBuffer *trace)
     return trace->next-1;
 }
 
+static inline bool
+uop_buffer_rewind(_PyJitUopBuffer *trace)
+{
+   if (trace->next <= trace->start) {
+       return false;
+   }
+    trace->next--;
+    return true;
+}
+
 static inline int
 uop_buffer_length(_PyJitUopBuffer *trace)
 {
diff --git a/Python/optimizer_analysis.c b/Python/optimizer_analysis.c
index 520420e9878575..c04b2222b68319 100644
--- a/Python/optimizer_analysis.c
+++ b/Python/optimizer_analysis.c
@@ -488,17 +488,9 @@ optimize_uops(
     _PyUOpInstruction *this_instr = NULL;
     JitOptRef *stack_pointer = ctx->frame->stack_pointer;
 
-    for (int i = 0; i < trace_len; i++) {
+    int i = 0;
+    for (; i < trace_len && !ctx->done; i++) {
         this_instr = &trace[i];
-        if (ctx->done) {
-            // Don't do any more optimization, but
-            // we still need to reach a terminator for corrctness.
-            *(ctx->out_buffer.next++) = *this_instr;
-            if (is_terminator_uop(this_instr)) {
-                break;
-            }
-            continue;
-        }
 
         int oparg = this_instr->oparg;
         opcode = this_instr->opcode;
@@ -531,10 +523,6 @@ optimize_uops(
             assert(STACK_LEVEL() >= 0);
         }
     }
-    if (ctx->out_of_space) {
-        DPRINTF(3, "\n");
-        DPRINTF(1, "Out of space in abstract interpreter\n");
-    }
     if (ctx->contradiction) {
         // Attempted to push a "bottom" (contradiction) symbol onto the stack.
         // This means that the abstract interpreter has optimized to trace
@@ -548,6 +536,35 @@ optimize_uops(
         OPT_STAT_INC(optimizer_contradiction);
         return 0;
     }
+    if (ctx->out_of_space) {
+        DPRINTF(3, "\n");
+        DPRINTF(1, "Out of space in abstract interpreter at length %d\n",
+            uop_buffer_length(&ctx->out_buffer));
+        // Rewind to previous instruction and replace with _EXIT_TRACE.
+        _PyUOpInstruction *curr = uop_buffer_last(&ctx->out_buffer);
+        while (curr->opcode != _SET_IP) {
+            if (!uop_buffer_rewind(&ctx->out_buffer)) {
+                // Reached the start.
+                return 0;
+            }
+            curr = uop_buffer_last(&ctx->out_buffer);
+        }
+        assert(curr->opcode == _SET_IP);
+        int32_t old_target = (int32_t)uop_get_target(curr);
+        curr->opcode = _EXIT_TRACE;
+        curr->format = UOP_FORMAT_TARGET;
+        curr->target = old_target;
+        DPRINTF(1, "Rewound to length %d\n", uop_buffer_length(&ctx->out_buffer));
+    }
+    else {
+        // Don't do any more optimization, but
+        // we still need to reach a terminator for correctness.
+        while (!is_terminator_uop(uop_buffer_last(&ctx->out_buffer))) {
+            this_instr = &trace[i];
+            *(ctx->out_buffer.next++) = *this_instr;
+            i++;
+        }
+    }
 
     /* Either reached the end or cannot optimize further, but there
      * would be no benefit in retrying later */