From c0c14b45b976ecc18f536db3292a001af4a2d149 Mon Sep 17 00:00:00 2001
From: Ken Jin <28750310+Fidget-Spinner@users.noreply.github.com>
Date: Sat, 18 Oct 2025 01:13:28 +0100
Subject: [PATCH] Restore the optimizer  fully

---
 Python/ceval_macros.h        | 25 +++++++++++++++++--------
 Python/optimizer.c           |  5 +++--
 Python/optimizer_analysis.c  | 14 +++++++-------
 Python/optimizer_bytecodes.c | 22 +++++++++-------------
 Python/optimizer_cases.c.h   | 15 ++++-----------
 5 files changed, 40 insertions(+), 41 deletions(-)

diff --git a/Python/ceval_macros.h b/Python/ceval_macros.h
index ff038026f977..557757f8a259 100644
--- a/Python/ceval_macros.h
+++ b/Python/ceval_macros.h
@@ -143,14 +143,23 @@
 #  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); \
-    }
+    do { \
+        LEAVE_TRACING(); \
+        if (!_PyErr_Occurred(tstate)) { \
+            _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); \
+            } \
+        } \
+        else { \
+            _PyFrame_SetStackPointer(frame, stack_pointer); \
+            _PyJIT_FinalizeTracing(tstate); \
+            stack_pointer = _PyFrame_GetStackPointer(frame); \
+        } \
+    } while (0);
 #  define RECORD_TRACE_NO_DISPATCH() do { \
         if (IS_JIT_TRACING() && add_to_code_trace(tstate, frame, old_code, old_func, this_instr, next_instr, opcode, oparg, _jump_taken)) { \
             BAIL_TRACING_NO_DISPATCH(); \
diff --git a/Python/optimizer.c b/Python/optimizer.c
index 5d75a87ee0de..c9009528e209 100644
--- a/Python/optimizer.c
+++ b/Python/optimizer.c
@@ -794,8 +794,9 @@ _PyJIT_translate_single_bytecode_to_trace(
                 }
                 if (uop == _PUSH_FRAME || uop == _RETURN_VALUE || uop == _RETURN_GENERATOR || uop == _YIELD_VALUE) {
                     PyCodeObject *new_code = (PyCodeObject *)PyStackRef_AsPyObjectBorrow(frame->f_executable);
-                    if (func != NULL) {
-                        operand = (uintptr_t)func;
+                    PyFunctionObject *new_func = (PyCodeObject *)PyStackRef_AsPyObjectBorrow(frame->f_funcobj);
+                    if (new_func != NULL) {
+                        operand = (uintptr_t)new_func;
                     }
                     else if (new_code != NULL) {
                         operand = (uintptr_t)new_code | 1;
diff --git a/Python/optimizer_analysis.c b/Python/optimizer_analysis.c
index 5e18b4e40f3d..7cee20bd8ee2 100644
--- a/Python/optimizer_analysis.c
+++ b/Python/optimizer_analysis.c
@@ -528,13 +528,13 @@ _Py_uop_analyze_and_optimize(
 {
     OPT_STAT_INC(optimizer_attempts);
 
-    // int err = optimize_uops(
-    //     initial_func, buffer,
-    //     length, curr_stacklen, dependencies);
-    //
-    // if (err == 0) {
-    //     return err;
-    // }
+    int err = optimize_uops(
+         initial_func, buffer,
+         length, curr_stacklen, dependencies);
+
+    if (err == 0) {
+        return err;
+    }
 
     assert(length > 0);
 
diff --git a/Python/optimizer_bytecodes.c b/Python/optimizer_bytecodes.c
index f9f8c15ab964..734a9193af7f 100644
--- a/Python/optimizer_bytecodes.c
+++ b/Python/optimizer_bytecodes.c
@@ -794,25 +794,21 @@ dummy_func(void) {
 
     op(_RETURN_GENERATOR, ( -- res)) {
         SYNC_SP();
-        PyCodeObject *co = get_current_code_object(ctx);
-        ctx->frame->stack_pointer = stack_pointer;
-        if (frame_pop(ctx)) {
-            break;
-        }
         stack_pointer = ctx->frame->stack_pointer;
         res = sym_new_unknown(ctx);
-
-        /* Stack space handling */
-        assert(corresponding_check_stack == NULL);
-        assert(co != NULL);
-        int framesize = co->co_framesize;
-        assert(framesize > 0);
-        assert(framesize <= curr_space);
-        curr_space -= framesize;
+        ctx->done = true;
+        ctx->out_of_space = true;
     }
 
     op(_YIELD_VALUE, (unused -- value)) {
+        // TODO (gh-139109): handle this properly in a future optimization.
+        // A possibility to handle underflows is to just restore the current frame information
+        // from whatever is stored in the trace we record at that point of time.
+        // E.g. we record at this YIELD_VALUE, func_obj=x , stack_level=4
+        // We can restore it to there.
         value = sym_new_unknown(ctx);
+        ctx->done = true;
+        ctx->out_of_space = true;
     }
 
     op(_GET_ITER, (iterable -- iter, index_or_null)) {
diff --git a/Python/optimizer_cases.c.h b/Python/optimizer_cases.c.h
index 3d31c7f04dad..f5bb8a5585af 100644
--- a/Python/optimizer_cases.c.h
+++ b/Python/optimizer_cases.c.h
@@ -1171,6 +1171,8 @@
         case _YIELD_VALUE: {
             JitOptRef value;
             value = sym_new_unknown(ctx);
+            ctx->done = true;
+            ctx->out_of_space = true;
             stack_pointer[-1] = value;
             break;
         }
@@ -3011,19 +3013,10 @@
 
         case _RETURN_GENERATOR: {
             JitOptRef res;
-            PyCodeObject *co = get_current_code_object(ctx);
-            ctx->frame->stack_pointer = stack_pointer;
-            if (frame_pop(ctx)) {
-                break;
-            }
             stack_pointer = ctx->frame->stack_pointer;
             res = sym_new_unknown(ctx);
-            assert(corresponding_check_stack == NULL);
-            assert(co != NULL);
-            int framesize = co->co_framesize;
-            assert(framesize > 0);
-            assert(framesize <= curr_space);
-            curr_space -= framesize;
+            ctx->done = true;
+            ctx->out_of_space = true;
             stack_pointer[0] = res;
             stack_pointer += 1;
             assert(WITHIN_STACK_BOUNDS());
-- 
2.47.3