GH-104584: Assorted fixes for the optimizer API.

Lib/test/test_capi/test_misc.py

@@ -2347,17 +2347,67 @@ def func():
 
 class TestOptimizerAPI(unittest.TestCase):
 
-    def test_counter_optimizer(self):
-        opt = _testinternalcapi.get_counter_optimizer()
-        self.assertEqual(opt.get_count(), 0)
+    @contextlib.contextmanager
+    def temporary_optimizer(self, opt):
+        # Code to acquire resource, e.g.:
+        _testinternalcapi.set_optimizer(opt)
         try:
-            _testinternalcapi.set_optimizer(opt)
-            self.assertEqual(opt.get_count(), 0)
-            for _ in range(1000):
-                pass
-            self.assertEqual(opt.get_count(), 1000)
+            yield
         finally:
             _testinternalcapi.set_optimizer(None)
 
+    @contextlib.contextmanager
+    def clear_executors(self, func):
+        try:
+            yield
+        finally:
+            #Clear executors
+            func.__code__ = func.__code__.replace()
+
+    def test_get_set_optimizer(self):
+        self.assertEqual(_testinternalcapi.get_optimizer(), None)
+        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)
+
+    def test_counter_optimizer(self):
+
+        def loop():
+            for _ in range(1000):
+                pass
+
+        for repeat in range(5):
+            opt = _testinternalcapi.get_counter_optimizer()
+            with self.temporary_optimizer(opt):
+                self.assertEqual(opt.get_count(), 0)
+                with self.clear_executors(loop):
+                    loop()
+                self.assertEqual(opt.get_count(), 1000)
+
+    def test_long_loop(self):
+        "Check that we aren't confused by EXTENDED_ARG"
+
+        def nop():
+            pass
+
+        def long_loop():
+            for _ in range(10):
+                nop(); nop(); nop(); nop(); nop(); nop(); nop(); nop();
+                nop(); nop(); nop(); nop(); nop(); nop(); nop(); nop();
+                nop(); nop(); nop(); nop(); nop(); nop(); nop(); nop();
+                nop(); nop(); nop(); nop(); nop(); nop(); nop(); nop();
+                nop(); nop(); nop(); nop(); nop(); nop(); nop(); nop();
+                nop(); nop(); nop(); nop(); nop(); nop(); nop(); nop();
+                nop(); nop(); nop(); nop(); nop(); nop(); nop(); nop();
+
+        opt = _testinternalcapi.get_counter_optimizer()
+        with self.temporary_optimizer(opt):
+            self.assertEqual(opt.get_count(), 0)
+            long_loop()
+            self.assertEqual(opt.get_count(), 10)
+
+
 if __name__ == "__main__":
     unittest.main()

Modules/_testinternalcapi.c

@@ -840,6 +840,15 @@ set_optimizer(PyObject *self, PyObject *opt)
     Py_RETURN_NONE;
 }
 
+static PyObject *
+get_optimizer(PyObject *self, PyObject *Py_UNUSED(ignored))
+{
+    PyObject *opt = (PyObject *)PyUnstable_GetOptimizer();
+    if (opt == NULL) {
+        Py_RETURN_NONE;
+    }
+    return opt;
+}
 
 static int _pending_callback(void *arg)
 {
@@ -982,6 +991,7 @@ static PyMethodDef module_functions[] = {
     {"iframe_getcode", iframe_getcode, METH_O, NULL},
     {"iframe_getline", iframe_getline, METH_O, NULL},
     {"iframe_getlasti", iframe_getlasti, METH_O, NULL},
+    {"get_optimizer", get_optimizer,  METH_NOARGS, NULL},
     {"set_optimizer", set_optimizer,  METH_O, NULL},
     {"get_counter_optimizer", get_counter_optimizer, METH_NOARGS, NULL},
     {"pending_threadfunc", _PyCFunction_CAST(pending_threadfunc),

Objects/codeobject.c

@@ -1465,12 +1465,28 @@ PyCode_GetFreevars(PyCodeObject *code)
     return _PyCode_GetFreevars(code);
 }
 
+static void
+clear_executors(PyCodeObject *co)
+{
+    for (int i = 0; i < co->co_executors->size; i++) {
+        Py_CLEAR(co->co_executors->executors[i]);
+    }
+    PyMem_Free(co->co_executors);
+    co->co_executors = NULL;
+}
+
 static void
 deopt_code(PyCodeObject *code, _Py_CODEUNIT *instructions)
 {
     Py_ssize_t len = Py_SIZE(code);
     for (int i = 0; i < len; i++) {
         int opcode = _Py_GetBaseOpcode(code, i);
+        if (opcode == ENTER_EXECUTOR) {
+            _PyExecutorObject *exec = code->co_executors->executors[instructions[i].op.arg];
+            opcode = exec->vm_data.opcode;
+            instructions[i].op.arg = exec->vm_data.oparg;
+        }
+        assert(opcode != ENTER_EXECUTOR);
         int caches = _PyOpcode_Caches[opcode];
         instructions[i].op.code = opcode;
         for (int j = 1; j <= caches; j++) {
@@ -1679,10 +1695,7 @@ code_dealloc(PyCodeObject *co)
         PyMem_Free(co_extra);
     }
     if (co->co_executors != NULL) {
-        for (int i = 0; i < co->co_executors->size; i++) {
-            Py_CLEAR(co->co_executors->executors[i]);
-        }
-        PyMem_Free(co->co_executors);
+        clear_executors(co);
     }
 
     Py_XDECREF(co->co_consts);
@@ -2278,6 +2291,9 @@ void
 _PyStaticCode_Fini(PyCodeObject *co)
 {
     deopt_code(co, _PyCode_CODE(co));
+    if (co->co_executors != NULL) {
+        clear_executors(co);
+    }
     PyMem_Free(co->co_extra);
     if (co->_co_cached != NULL) {
         Py_CLEAR(co->_co_cached->_co_code);

Python/bytecodes.c

@@ -2157,6 +2157,7 @@ dummy_func(
                     frame = cframe.current_frame;
                     goto error;
                 }
+                assert(frame == cframe.current_frame);
                 here[1].cache &= ((1 << OPTIMIZER_BITS_IN_COUNTER) -1);
                 goto resume_frame;
             }
@@ -2176,7 +2177,7 @@ dummy_func(
 
         inst(ENTER_EXECUTOR, (--)) {
             PyCodeObject *code = _PyFrame_GetCode(frame);
-            _PyExecutorObject *executor = (_PyExecutorObject *)code->co_executors->executors[oparg];
+            _PyExecutorObject *executor = (_PyExecutorObject *)code->co_executors->executors[oparg&255];
             Py_INCREF(executor);
             frame = executor->execute(executor, frame, stack_pointer);
             if (frame == NULL) {