GH-104584: Miscellaneous fixes for -Xuops (GH-106908)

10 files changed, 67 insertions, 18 deletions

diff --git a/Lib/dis.py b/Lib/dis.py
index f7a31f2..184eb14 100644
--- a/Lib/dis.py
+++ b/Lib/dis.py
@@ -430,7 +430,8 @@ def get_instructions(x, *, first_line=None, show_caches=False, adaptive=False):
                                    co.co_names, co.co_consts,
                                    linestarts, line_offset,
                                    co_positions=co.co_positions(),
-                                   show_caches=show_caches)
+                                   show_caches=show_caches,
+                                   original_code=co.co_code)
 
 def _get_const_value(op, arg, co_consts):
     """Helper to get the value of the const in a hasconst op.
@@ -504,7 +505,7 @@ def _get_instructions_bytes(code, varname_from_oparg=None,
                             names=None, co_consts=None,
                             linestarts=None, line_offset=0,
                             exception_entries=(), co_positions=None,
-                            show_caches=False):
+                            show_caches=False, original_code=None):
     """Iterate over the instructions in a bytecode string.
 
     Generates a sequence of Instruction namedtuples giving the details of each
@@ -513,14 +514,18 @@ def _get_instructions_bytes(code, varname_from_oparg=None,
     arguments.
 
     """
+    # Use the basic, unadaptive code for finding labels and actually walking the
+    # bytecode, since replacements like ENTER_EXECUTOR and INSTRUMENTED_* can
+    # mess that logic up pretty badly:
+    original_code = original_code or code
     co_positions = co_positions or iter(())
     get_name = None if names is None else names.__getitem__
-    labels = set(findlabels(code))
+    labels = set(findlabels(original_code))
     for start, end, target, _, _ in exception_entries:
         for i in range(start, end):
             labels.add(target)
     starts_line = None
-    for offset, start_offset, op, arg in _unpack_opargs(code):
+    for offset, start_offset, op, arg in _unpack_opargs(original_code):
         if linestarts is not None:
             starts_line = linestarts.get(offset, None)
             if starts_line is not None:
@@ -531,6 +536,7 @@ def _get_instructions_bytes(code, varname_from_oparg=None,
         positions = Positions(*next(co_positions, ()))
         deop = _deoptop(op)
         caches = _inline_cache_entries[deop]
+        op = code[offset]
         if arg is not None:
             #  Set argval to the dereferenced value of the argument when
             #  available, and argrepr to the string representation of argval.
@@ -591,7 +597,6 @@ def _get_instructions_bytes(code, varname_from_oparg=None,
         yield Instruction(_all_opname[op], op,
                           arg, argval, argrepr,
                           offset, start_offset, starts_line, is_jump_target, positions)
-        caches = _inline_cache_entries[deop]
         if not caches:
             continue
         if not show_caches:
@@ -622,7 +627,8 @@ def disassemble(co, lasti=-1, *, file=None, show_caches=False, adaptive=False):
                        lasti, co._varname_from_oparg,
                        co.co_names, co.co_consts, linestarts, file=file,
                        exception_entries=exception_entries,
-                       co_positions=co.co_positions(), show_caches=show_caches)
+                       co_positions=co.co_positions(), show_caches=show_caches,
+                       original_code=co.co_code)
 
 def _disassemble_recursive(co, *, file=None, depth=None, show_caches=False, adaptive=False):
     disassemble(co, file=file, show_caches=show_caches, adaptive=adaptive)
@@ -640,7 +646,7 @@ def _disassemble_recursive(co, *, file=None, depth=None, show_caches=False, adap
 def _disassemble_bytes(code, lasti=-1, varname_from_oparg=None,
                        names=None, co_consts=None, linestarts=None,
                        *, file=None, line_offset=0, exception_entries=(),
-                       co_positions=None, show_caches=False):
+                       co_positions=None, show_caches=False, original_code=None):
     # Omit the line number column entirely if we have no line number info
     show_lineno = bool(linestarts)
     if show_lineno:
@@ -661,7 +667,8 @@ def _disassemble_bytes(code, lasti=-1, varname_from_oparg=None,
                                          line_offset=line_offset,
                                          exception_entries=exception_entries,
                                          co_positions=co_positions,
-                                         show_caches=show_caches):
+                                         show_caches=show_caches,
+                                         original_code=original_code):
         new_source_line = (show_lineno and
                            instr.starts_line is not None and
                            instr.offset > 0)
@@ -823,7 +830,8 @@ class Bytecode:
                                        line_offset=self._line_offset,
                                        exception_entries=self.exception_entries,
                                        co_positions=co.co_positions(),
-                                       show_caches=self.show_caches)
+                                       show_caches=self.show_caches,
+                                       original_code=co.co_code)
 
     def __repr__(self):
         return "{}({!r})".format(self.__class__.__name__,
@@ -859,7 +867,8 @@ class Bytecode:
                                lasti=offset,
                                exception_entries=self.exception_entries,
                                co_positions=co.co_positions(),
-                               show_caches=self.show_caches)
+                               show_caches=self.show_caches,
+                               original_code=co.co_code)
             return output.getvalue()
 
 
diff --git a/Lib/test/test_capi/test_misc.py b/Lib/test/test_capi/test_misc.py
index 4e519fa..e40ffdf 100644
--- a/Lib/test/test_capi/test_misc.py
+++ b/Lib/test/test_capi/test_misc.py
@@ -2368,12 +2368,16 @@ def clear_executors(func):
 class TestOptimizerAPI(unittest.TestCase):
 
     def test_get_set_optimizer(self):
-        self.assertEqual(_testinternalcapi.get_optimizer(), None)
+        old = _testinternalcapi.get_optimizer()
         opt = _testinternalcapi.get_counter_optimizer()
-        _testinternalcapi.set_optimizer(opt)
-        self.assertEqual(_testinternalcapi.get_optimizer(), opt)
-        _testinternalcapi.set_optimizer(None)
-        self.assertEqual(_testinternalcapi.get_optimizer(), None)
+        try:
+            _testinternalcapi.set_optimizer(opt)
+            self.assertEqual(_testinternalcapi.get_optimizer(), opt)
+            _testinternalcapi.set_optimizer(None)
+            self.assertEqual(_testinternalcapi.get_optimizer(), None)
+        finally:
+            _testinternalcapi.set_optimizer(old)
+
 
     def test_counter_optimizer(self):
         # Generate a new function at each call
diff --git a/Lib/test/test_dis.py b/Lib/test/test_dis.py
index 8597b8f..642b261 100644
--- a/Lib/test/test_dis.py
+++ b/Lib/test/test_dis.py
@@ -1244,10 +1244,14 @@ class DisTests(DisTestBase):
     @cpython_only
     @requires_specialization
     def test_loop_quicken(self):
+        import _testinternalcapi
         # Loop can trigger a quicken where the loop is located
         self.code_quicken(loop_test, 1)
         got = self.get_disassembly(loop_test, adaptive=True)
-        self.do_disassembly_compare(got, dis_loop_test_quickened_code)
+        expected = dis_loop_test_quickened_code
+        if _testinternalcapi.get_optimizer():
+            expected = expected.replace("JUMP_BACKWARD ", "ENTER_EXECUTOR")
+        self.do_disassembly_compare(got, expected)
 
     @cpython_only
     def test_extended_arg_quick(self):
diff --git a/Misc/NEWS.d/next/Core and Builtins/2023-07-20-01-15-58.gh-issue-106908.cDmcVI.rst b/Misc/NEWS.d/next/Core and Builtins/2023-07-20-01-15-58.gh-issue-106908.cDmcVI.rst
new file mode 100644
index 0000000..9c9b845
--- /dev/null
+++ b/Misc/NEWS.d/next/Core and Builtins/2023-07-20-01-15-58.gh-issue-106908.cDmcVI.rst
@@ -0,0 +1,3 @@
+Fix various hangs, reference leaks, test failures, and tracing/introspection
+bugs when running with :envvar:`PYTHONUOPS` or :option:`-X uops <-X>`
+enabled.
diff --git a/Python/bytecodes.c b/Python/bytecodes.c
index ea136a3..81d6f80 100644
--- a/Python/bytecodes.c
+++ b/Python/bytecodes.c
@@ -2182,7 +2182,14 @@ dummy_func(
             JUMPBY(1-oparg);
             #if ENABLE_SPECIALIZATION
             here[1].cache += (1 << OPTIMIZER_BITS_IN_COUNTER);
-            if (here[1].cache > tstate->interp->optimizer_backedge_threshold) {
+            if (here[1].cache > tstate->interp->optimizer_backedge_threshold &&
+                // Double-check that the opcode isn't instrumented or something:
+                here->op.code == JUMP_BACKWARD &&
+                // _PyOptimizer_BackEdge is going to change frame->prev_instr,
+                // which breaks line event calculations:
+                next_instr->op.code != INSTRUMENTED_LINE
+                )
+            {
                 OBJECT_STAT_INC(optimization_attempts);
                 frame = _PyOptimizer_BackEdge(frame, here, next_instr, stack_pointer);
                 if (frame == NULL) {
diff --git a/Python/executor_cases.c.h b/Python/executor_cases.c.h
index e1f8b9f..90607e8 100644
--- a/Python/executor_cases.c.h
+++ b/Python/executor_cases.c.h
@@ -2017,6 +2017,7 @@
             STACK_SHRINK(oparg);
             STACK_SHRINK(1);
             stack_pointer[-1] = res;
+            CHECK_EVAL_BREAKER();
             break;
         }
 
@@ -2038,6 +2039,7 @@
             STACK_SHRINK(oparg);
             STACK_SHRINK(1);
             stack_pointer[-1] = res;
+            CHECK_EVAL_BREAKER();
             break;
         }
 
@@ -2089,6 +2091,7 @@
             STACK_SHRINK(oparg);
             STACK_SHRINK(1);
             stack_pointer[-1] = res;
+            CHECK_EVAL_BREAKER();
             break;
         }
 
@@ -2131,6 +2134,7 @@
             STACK_SHRINK(oparg);
             STACK_SHRINK(1);
             stack_pointer[-1] = res;
+            CHECK_EVAL_BREAKER();
             break;
         }
 
@@ -2243,6 +2247,7 @@
             STACK_SHRINK(oparg);
             STACK_SHRINK(1);
             stack_pointer[-1] = res;
+            CHECK_EVAL_BREAKER();
             break;
         }
 
@@ -2281,6 +2286,7 @@
             STACK_SHRINK(oparg);
             STACK_SHRINK(1);
             stack_pointer[-1] = res;
+            CHECK_EVAL_BREAKER();
             break;
         }
 
@@ -2318,6 +2324,7 @@
             STACK_SHRINK(oparg);
             STACK_SHRINK(1);
             stack_pointer[-1] = res;
+            CHECK_EVAL_BREAKER();
             break;
         }
 
@@ -2496,6 +2503,7 @@
 
         case JUMP_TO_TOP: {
             pc = 0;
+            CHECK_EVAL_BREAKER();
             break;
         }
 
diff --git a/Python/generated_cases.c.h b/Python/generated_cases.c.h
index b2b0aa6..c35b81a 100644
--- a/Python/generated_cases.c.h
+++ b/Python/generated_cases.c.h
@@ -2758,7 +2758,14 @@
             JUMPBY(1-oparg);
             #if ENABLE_SPECIALIZATION
             here[1].cache += (1 << OPTIMIZER_BITS_IN_COUNTER);
-            if (here[1].cache > tstate->interp->optimizer_backedge_threshold) {
+            if (here[1].cache > tstate->interp->optimizer_backedge_threshold &&
+                // Double-check that the opcode isn't instrumented or something:
+                here->op.code == JUMP_BACKWARD &&
+                // _PyOptimizer_BackEdge is going to change frame->prev_instr,
+                // which breaks line event calculations:
+                next_instr->op.code != INSTRUMENTED_LINE
+                )
+            {
                 OBJECT_STAT_INC(optimization_attempts);
                 frame = _PyOptimizer_BackEdge(frame, here, next_instr, stack_pointer);
                 if (frame == NULL) {
diff --git a/Python/optimizer.c b/Python/optimizer.c
index 3d385a1..09120c3 100644
--- a/Python/optimizer.c
+++ b/Python/optimizer.c
@@ -155,6 +155,7 @@ PyUnstable_SetOptimizer(_PyOptimizerObject *optimizer)
 _PyInterpreterFrame *
 _PyOptimizer_BackEdge(_PyInterpreterFrame *frame, _Py_CODEUNIT *src, _Py_CODEUNIT *dest, PyObject **stack_pointer)
 {
+    assert(src->op.code == JUMP_BACKWARD);
     PyCodeObject *code = (PyCodeObject *)frame->f_executable;
     assert(PyCode_Check(code));
     PyInterpreterState *interp = _PyInterpreterState_GET();
diff --git a/Python/pylifecycle.c b/Python/pylifecycle.c
index cf8b437..f91fac9 100644
--- a/Python/pylifecycle.c
+++ b/Python/pylifecycle.c
@@ -1192,7 +1192,11 @@ init_interp_main(PyThreadState *tstate)
         }
         if (enabled) {
             PyObject *opt = PyUnstable_Optimizer_NewUOpOptimizer();
+            if (opt == NULL) {
+                return _PyStatus_ERR("can't initialize optimizer");
+            }
             PyUnstable_SetOptimizer((_PyOptimizerObject *)opt);
+            Py_DECREF(opt);
         }
     }
 
diff --git a/Tools/cases_generator/generate_cases.py b/Tools/cases_generator/generate_cases.py
index 33eff54..a18229a 100644
--- a/Tools/cases_generator/generate_cases.py
+++ b/Tools/cases_generator/generate_cases.py
@@ -1506,6 +1506,8 @@ class Analyzer:
                             self.out.emit("")
                             with self.out.block(f"case {thing.name}:"):
                                 instr.write(self.out, tier=TIER_TWO)
+                                if instr.check_eval_breaker:
+                                    self.out.emit("CHECK_EVAL_BREAKER();")
                                 self.out.emit("break;")
                         # elif instr.kind != "op":
                         #     print(f"NOTE: {thing.name} is not a viable uop")