gh-112383: teach dis how to interpret ENTER_EXECUTOR (#117171)

Author: iritkatriel

Diff for: Lib/dis.py

@@ -17,6 +17,8 @@
     _specialized_opmap,
 )
 
+from _opcode import get_executor
+
 __all__ = ["code_info", "dis", "disassemble", "distb", "disco",
            "findlinestarts", "findlabels", "show_code",
            "get_instructions", "Instruction", "Bytecode"] + _opcodes_all
@@ -205,7 +207,27 @@ def _deoptop(op):
     return _all_opmap[deoptmap[name]] if name in deoptmap else op
 
 def _get_code_array(co, adaptive):
-    return co._co_code_adaptive if adaptive else co.co_code
+    if adaptive:
+        code = co._co_code_adaptive
+        res = []
+        found = False
+        for i in range(0, len(code), 2):
+            op, arg = code[i], code[i+1]
+            if op == ENTER_EXECUTOR:
+                try:
+                    ex = get_executor(co, i)
+                except ValueError:
+                    ex = None
+
+                if ex:
+                    op, arg = ex.get_opcode(), ex.get_oparg()
+                    found = True
+
+            res.append(op.to_bytes())
+            res.append(arg.to_bytes())
+        return code if not found else b''.join(res)
+    else:
+        return co.co_code
 
 def code_info(x):
     """Formatted details of methods, functions, or code."""
@@ -514,8 +536,6 @@ def offset_from_jump_arg(self, op, arg, offset):
             argval = offset + 2 + signed_arg*2
             caches = _get_cache_size(_all_opname[deop])
             argval += 2 * caches
-            if deop == ENTER_EXECUTOR:
-                argval += 2
             return argval
         return None
 
@@ -680,8 +700,7 @@ def _parse_exception_table(code):
 
 def _is_backward_jump(op):
     return opname[op] in ('JUMP_BACKWARD',
-                          'JUMP_BACKWARD_NO_INTERRUPT',
-                          'ENTER_EXECUTOR')
+                          'JUMP_BACKWARD_NO_INTERRUPT')
 
 def _get_instructions_bytes(code, linestarts=None, line_offset=0, co_positions=None,
                             original_code=None, arg_resolver=None):

Diff for: Lib/test/test_capi/test_opt.py

@@ -1,11 +1,11 @@
 import contextlib
-import opcode
 import sys
 import textwrap
 import unittest
 import gc
 import os
 
+import _opcode
 import _testinternalcapi
 
 from test.support import script_helper, requires_specialization
@@ -115,13 +115,11 @@ def testfunc(x):
 def get_first_executor(func):
     code = func.__code__
     co_code = code.co_code
-    JUMP_BACKWARD = opcode.opmap["JUMP_BACKWARD"]
     for i in range(0, len(co_code), 2):
-        if co_code[i] == JUMP_BACKWARD:
-            try:
-                return _testinternalcapi.get_executor(code, i)
-            except ValueError:
-                pass
+        try:
+            return _opcode.get_executor(code, i)
+        except ValueError:
+            pass
     return None
 
 
@@ -760,17 +758,16 @@ def test_promote_globals_to_constants(self):
         result = script_helper.run_python_until_end('-c', textwrap.dedent("""
         import _testinternalcapi
         import opcode
+        import _opcode
 
         def get_first_executor(func):
             code = func.__code__
             co_code = code.co_code
-            JUMP_BACKWARD = opcode.opmap["JUMP_BACKWARD"]
             for i in range(0, len(co_code), 2):
-                if co_code[i] == JUMP_BACKWARD:
-                    try:
-                        return _testinternalcapi.get_executor(code, i)
-                    except ValueError:
-                        pass
+                try:
+                    return _opcode.get_executor(code, i)
+                except ValueError:
+                    pass
             return None
 
         def get_opnames(ex):

Diff for: Lib/test/test_dis.py

@@ -1201,19 +1201,10 @@ def test_call_specialize(self):
     @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)
+        self.code_quicken(loop_test, 4)
         got = self.get_disassembly(loop_test, adaptive=True)
         expected = dis_loop_test_quickened_code
-        if _testinternalcapi.get_optimizer():
-            # We *may* see ENTER_EXECUTOR in the disassembly. This is a
-            # temporary hack to keep the test working until dis is able to
-            # handle the instruction correctly (GH-112383):
-            got = got.replace(
-                "ENTER_EXECUTOR          16",
-                "JUMP_BACKWARD           16 (to L1)",
-            )
         self.do_disassembly_compare(got, expected)
 
     @cpython_only

Diff for: Misc/NEWS.d/next/Library/2024-03-23-13-40-13.gh-issue-112383.XuHf3G.rst

@@ -0,0 +1 @@
+Fix :mod:`dis` module's handling of ``ENTER_EXECUTOR`` instructions.

Diff for: Modules/_opcode.c

@@ -347,6 +347,28 @@ _opcode_get_intrinsic2_descs_impl(PyObject *module)
     return list;
 }
 
+/*[clinic input]
+
+_opcode.get_executor
+
+  code: object
+  offset: int
+
+Return the executor object at offset in code if exists, None otherwise.
+[clinic start generated code]*/
+
+static PyObject *
+_opcode_get_executor_impl(PyObject *module, PyObject *code, int offset)
+/*[clinic end generated code: output=c035c7a47b16648f input=85eff93ea7aac282]*/
+{
+    if (!PyCode_Check(code)) {
+        PyErr_Format(PyExc_TypeError,
+                     "expected a code object, not '%.100s'",
+                     Py_TYPE(code)->tp_name);
+        return NULL;
+    }
+    return (PyObject *)PyUnstable_GetExecutor((PyCodeObject *)code, offset);
+}
 
 static PyMethodDef
 opcode_functions[] =  {
@@ -363,6 +385,7 @@ opcode_functions[] =  {
     _OPCODE_GET_NB_OPS_METHODDEF
     _OPCODE_GET_INTRINSIC1_DESCS_METHODDEF
     _OPCODE_GET_INTRINSIC2_DESCS_METHODDEF
+    _OPCODE_GET_EXECUTOR_METHODDEF
     {NULL, NULL, 0, NULL}
 };
 

Diff for: Modules/_testinternalcapi.c

@@ -991,26 +991,6 @@ get_optimizer(PyObject *self, PyObject *Py_UNUSED(ignored))
     return opt;
 }
 
-static PyObject *
-get_executor(PyObject *self, PyObject *const *args, Py_ssize_t nargs)
-{
-
-    if (!_PyArg_CheckPositional("get_executor", nargs, 2, 2)) {
-        return NULL;
-    }
-    PyObject *code = args[0];
-    PyObject *offset = args[1];
-    long ioffset = PyLong_AsLong(offset);
-    if (ioffset == -1 && PyErr_Occurred()) {
-        return NULL;
-    }
-    if (!PyCode_Check(code)) {
-         PyErr_SetString(PyExc_TypeError, "first argument must be a code object");
-        return NULL;
-    }
-    return (PyObject *)PyUnstable_GetExecutor((PyCodeObject *)code, ioffset);
-}
-
 static PyObject *
 add_executor_dependency(PyObject *self, PyObject *args)
 {
@@ -1836,7 +1816,6 @@ static PyMethodDef module_functions[] = {
     {"iframe_getlasti", iframe_getlasti, METH_O, NULL},
     {"get_optimizer", get_optimizer,  METH_NOARGS, NULL},
     {"set_optimizer", set_optimizer,  METH_O, NULL},
-    {"get_executor", _PyCFunction_CAST(get_executor),  METH_FASTCALL, NULL},
     {"new_counter_optimizer", new_counter_optimizer, METH_NOARGS, NULL},
     {"new_uop_optimizer", new_uop_optimizer, METH_NOARGS, NULL},
     {"add_executor_dependency", add_executor_dependency, METH_VARARGS, NULL},