python · markshannon · Jul 4, 2023 · Jul 3, 2023 · Jul 3, 2023 · Jul 4, 2023
diff --git a/Include/cpython/optimizer.h b/Include/cpython/optimizer.h
@@ -38,6 +38,8 @@ PyAPI_FUNC(void) PyUnstable_SetOptimizer(_PyOptimizerObject* optimizer);
 
 PyAPI_FUNC(_PyOptimizerObject *) PyUnstable_GetOptimizer(void);
 
+PyAPI_FUNC(_PyExecutorObject *)PyUnstable_GetExecutor(PyCodeObject *code, int offset);
+
 struct _PyInterpreterFrame *
 _PyOptimizer_BackEdge(struct _PyInterpreterFrame *frame, _Py_CODEUNIT *src, _Py_CODEUNIT *dest, PyObject **stack_pointer);
 

diff --git a/Include/internal/pycore_opcode.h b/Include/internal/pycore_opcode.h
@@ -2415,5 +2415,28 @@ def long_loop():
             self.assertEqual(opt.get_count(), 10)
 
 
+class TestUops(unittest.TestCase):
+
+    def test_basic_loop(self):
+
+        def testfunc(x):
+            i = 0
+            while i < x:
+                i += 1
+
+        testfunc(1000)
+
+        ex = None
+        for offset in range(0, 100, 2):
+            try:
+                ex = _testinternalcapi.get_executor(testfunc.__code__, offset)
+                break
+            except ValueError:
+                pass
+        if ex is None:
+            return
+        self.assertIn("SAVE_IP", str(ex))
+
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/Modules/_testinternalcapi.c b/Modules/_testinternalcapi.c
@@ -858,6 +858,26 @@ 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 int _pending_callback(void *arg)
 {
     /* we assume the argument is callable object to which we own a reference */
@@ -1326,6 +1346,7 @@ 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},
     {"get_counter_optimizer", get_counter_optimizer, METH_NOARGS, NULL},
     {"get_uop_optimizer", get_uop_optimizer, METH_NOARGS, NULL},
     {"pending_threadfunc", _PyCFunction_CAST(pending_threadfunc),

diff --git a/Python/opcode_metadata.h b/Python/opcode_metadata.h
@@ -188,6 +188,23 @@ _PyOptimizer_BackEdge(_PyInterpreterFrame *frame, _Py_CODEUNIT *src, _Py_CODEUNI
     return frame;
 }
 
+_PyExecutorObject *
+PyUnstable_GetExecutor(PyCodeObject *code, int offset)
+{
+    int code_len = (int)Py_SIZE(code);
+    for (int i = 0 ; i < code_len;) {
+        if (_PyCode_CODE(code)[i].op.code == ENTER_EXECUTOR && i*2 == offset) {
+            int oparg = _PyCode_CODE(code)[i].op.arg;
+            _PyExecutorObject *res = code->co_executors->executors[oparg];
+            Py_INCREF(res);
+            return res;
+        }
+        i += _PyInstruction_GetLength(code, i);
+    }
+    PyErr_SetString(PyExc_ValueError, "no executor at given offset");
+    return NULL;
+}
+
 /** Test support **/
 
 
@@ -287,13 +304,66 @@ uop_dealloc(_PyUOpExecutorObject *self) {
     PyObject_Free(self);
 }
 
+static const char *
+uop_name(int index) {
+    if (index < EXIT_TRACE) {
+        return _PyOpcode_OpName[index];
+    }
+    return _PyOpcode_uop_name[index];
+}
+
+static Py_ssize_t
+uop_len(_PyUOpExecutorObject *self)
+{
+    int count = 1;
+    for (; count < _Py_UOP_MAX_TRACE_LENGTH; count++) {
+        if (self->trace[count-1].opcode == EXIT_TRACE) {
+            break;
+        }
+    }
+    return count;
+}
+
+static PyObject *
+uop_item(_PyUOpExecutorObject *self, Py_ssize_t index)
+{
+    for (int i = 0; i < _Py_UOP_MAX_TRACE_LENGTH; i++) {
+        if (self->trace[i].opcode == EXIT_TRACE) {
+            break;
+        }
+        if (i != index) {
+            continue;
+        }
+        const char *name = uop_name(self->trace[i].opcode);
+        PyObject *oname = _PyUnicode_FromASCII(name, strlen(name));
+        if (oname == NULL) {
+            return NULL;
+        }
+        PyObject *operand = PyLong_FromUnsignedLongLong(self->trace[i].operand);
+        if (operand == NULL) {
+            Py_DECREF(oname);
+            return NULL;
+        }
+        PyObject *args[2] = { oname, operand };
+        return _PyTuple_FromArraySteal(args, 2);
+    }
+    PyErr_SetNone(PyExc_IndexError);
+    return NULL;
+}
+
+PySequenceMethods uop_as_sequence = {
+    .sq_length = (lenfunc)uop_len,
+    .sq_item = (ssizeargfunc)uop_item,
+};
+
 static PyTypeObject UOpExecutor_Type = {
     PyVarObject_HEAD_INIT(&PyType_Type, 0)
     .tp_name = "uop_executor",
     .tp_basicsize = sizeof(_PyUOpExecutorObject),
     .tp_itemsize = 0,
     .tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_DISALLOW_INSTANTIATION,
     .tp_dealloc = (destructor)uop_dealloc,
+    .tp_as_sequence = &uop_as_sequence,
 };
 
 static int

diff --git a/Tools/build/generate_opcode_h.py b/Tools/build/generate_opcode_h.py
@@ -184,14 +184,15 @@ def main(opcode_py,
         fobj.write(f"#define ENABLE_SPECIALIZATION {int(ENABLE_SPECIALIZATION)}")
 
         iobj.write("\n")
-        iobj.write("#ifdef Py_DEBUG\n")
-        iobj.write(f"static const char *const _PyOpcode_OpName[{NUM_OPCODES}] = {{\n")
+        iobj.write(f"\nextern const char *const _PyOpcode_OpName[{NUM_OPCODES}];\n")
+        iobj.write("\n#ifdef NEED_OPCODE_TABLES\n")
+        iobj.write(f"const char *const _PyOpcode_OpName[{NUM_OPCODES}] = {{\n")
         for op, name in enumerate(opname_including_specialized):
             if name[0] != "<":
                 op = name
             iobj.write(f'''    [{op}] = "{name}",\n''')
         iobj.write("};\n")
-        iobj.write("#endif\n")
+        iobj.write("#endif   // NEED_OPCODE_TABLES\n")
 
         iobj.write("\n")
         iobj.write("#define EXTRA_CASES \\\n")

@@ -1224,9 +1224,7 @@ def write_metadata(self) -> None:
             self.out.emit("#ifndef NEED_OPCODE_METADATA")
             self.out.emit("extern const struct opcode_metadata _PyOpcode_opcode_metadata[512];")
             self.out.emit("extern const struct opcode_macro_expansion _PyOpcode_macro_expansion[256];")
-            self.out.emit("#ifdef Py_DEBUG")
             self.out.emit("extern const char * const _PyOpcode_uop_name[512];")
-            self.out.emit("#endif")
             self.out.emit("#else")
 
             self.out.emit("const struct opcode_metadata _PyOpcode_opcode_metadata[512] = {")
@@ -1273,10 +1271,10 @@ def write_metadata(self) -> None:
                         case _:
                             typing.assert_never(thing)
 
-            self.out.emit("#ifdef Py_DEBUG")
+            self.out.emit("#ifdef NEED_OPCODE_METADATA")
             with self.out.block("const char * const _PyOpcode_uop_name[512] =", ";"):
                 self.write_uop_items(lambda name, counter: f"[{counter}] = \"{name}\",")
-            self.out.emit("#endif")
+            self.out.emit("#endif // NEED_OPCODE_METADATA")
 
             self.out.emit("#endif")