Fix optimizer API.

python · markshannon · Jun 19, 2023 · Jun 6, 2023 · Jun 7, 2023 · Jun 7, 2023
commit 2335db629b3172c44e913c0e81e0c53729739e8b
@@ -1917,17 +1917,32 @@ def func():
 
 class TestOptimizerAPI(unittest.TestCase):
 
-    def test_counter_optimizer(self):
+    def test_get_set_optimizer(self):
+        self.assertEqual(_testinternalcapi.get_optimizer(), None)
         opt = _testinternalcapi.get_counter_optimizer()
-        self.assertEqual(opt.get_count(), 0)
-        try:
-            _testinternalcapi.set_optimizer(opt)
-            self.assertEqual(opt.get_count(), 0)
+        _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
-            self.assertEqual(opt.get_count(), 1000)
-        finally:
-            _testinternalcapi.set_optimizer(None)
+
+        for repeat in range(5):
+            opt = _testinternalcapi.get_counter_optimizer()
+            self.assertEqual(opt.get_count(), 0)
+            try:
+                _testinternalcapi.set_optimizer(opt)
+                self.assertEqual(opt.get_count(), 0)
+                loop()
+                self.assertEqual(opt.get_count(), 1000)
+            finally:
+                _testinternalcapi.set_optimizer(None)
+                #Clear executors
+                loop.__code__ = loop.__code__.replace()
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/Modules/_testinternalcapi.c b/Modules/_testinternalcapi.c
@@ -838,6 +838,16 @@ 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 PyMethodDef module_functions[] = {
     {"get_configs", get_configs, METH_NOARGS},
     {"get_recursion_depth", get_recursion_depth, METH_NOARGS},
@@ -866,6 +876,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},
     {NULL, NULL} /* sentinel */

@@ -1465,6 +1465,16 @@ 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)
 {
@@ -1478,6 +1488,20 @@ deopt_code(PyCodeObject *code, _Py_CODEUNIT *instructions)
         }
         i += caches;
     }
+    if (code->co_executors == NULL) {
+        return;
+    }
+    for (int i = 0; i < len; i++) {
+        int opcode = instructions[i].op.code;
+        if (opcode == ENTER_EXECUTOR) {
+            _PyExecutorObject *exec = code->co_executors->executors[instructions[i].op.arg];
+            opcode = instructions[i].op.code = exec->vm_data.opcode;
+            instructions[i].op.arg = exec->vm_data.oparg;
+            assert(instructions[i+1].cache < (1 << OPTIMIZER_BITS_IN_COUNTER));
+        }
+        int caches = _PyOpcode_Caches[opcode];
+        i += caches;
+    }
 }
 
 PyObject *
@@ -1679,10 +1703,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 +2299,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);

@@ -2135,6 +2135,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;
             }
@@ -2143,7 +2144,7 @@ dummy_func(
         }
 
         inst(ENTER_EXECUTOR, (--)) {
-            _PyExecutorObject *executor = (_PyExecutorObject *)frame->f_code->co_executors->executors[oparg];
+            _PyExecutorObject *executor = (_PyExecutorObject *)frame->f_code->co_executors->executors[oparg&255];
             Py_INCREF(executor);
             frame = executor->execute(executor, frame, stack_pointer);
             if (frame == NULL) {