pythongh-104584: Clean up and fix uops tests and fix crash (python#106492)

gvanrossum · gvanrossum · commit 62ef3d883d37 · 2023-07-06T16:25:42.000-07:00
The uops test wasn't testing anything by default,
and was failing when run with -Xuops.

Made the two executor-related context managers global,
so TestUops can use them (notably `with temporary_optimizer(opt)`).

Made clear_executor() a little more thorough.

Fixed a crash upon finalizing a uop optimizer,
by adding a `tp_dealloc` handler.
diff --git a/Lib/test/test_capi/test_misc.py b/Lib/test/test_capi/test_misc.py
@@ -2343,23 +2343,26 @@ def func():
         names = ["func", "outer", "outer", "inner", "inner", "outer", "inner"]
         self.do_test(func, names)
 
-class TestOptimizerAPI(unittest.TestCase):
 
-    @contextlib.contextmanager
-    def temporary_optimizer(self, opt):
-        _testinternalcapi.set_optimizer(opt)
-        try:
-            yield
-        finally:
-            _testinternalcapi.set_optimizer(None)
+@contextlib.contextmanager
+def temporary_optimizer(opt):
+    _testinternalcapi.set_optimizer(opt)
+    try:
+        yield
+    finally:
+        _testinternalcapi.set_optimizer(None)
+
+@contextlib.contextmanager
+def clear_executors(func):
+    # Clear executors in func before and after running a block
+    func.__code__ = func.__code__.replace()
+    try:
+        yield
+    finally:
+        func.__code__ = func.__code__.replace()
 
-    @contextlib.contextmanager
-    def clear_executors(self, func):
-        try:
-            yield
-        finally:
-            #Clear executors
-            func.__code__ = func.__code__.replace()
+
+class TestOptimizerAPI(unittest.TestCase):
 
     def test_get_set_optimizer(self):
         self.assertEqual(_testinternalcapi.get_optimizer(), None)
@@ -2381,9 +2384,9 @@ def loop():
 
         for repeat in range(5):
             opt = _testinternalcapi.get_counter_optimizer()
-            with self.temporary_optimizer(opt):
+            with temporary_optimizer(opt):
                 self.assertEqual(opt.get_count(), 0)
-                with self.clear_executors(loop):
+                with clear_executors(loop):
                     loop()
                 self.assertEqual(opt.get_count(), 1000)
 
@@ -2409,11 +2412,37 @@ def long_loop():
         long_loop = ns['long_loop']
 
         opt = _testinternalcapi.get_counter_optimizer()
-        with self.temporary_optimizer(opt):
+        with temporary_optimizer(opt):
             self.assertEqual(opt.get_count(), 0)
             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
+
+        opt = _testinternalcapi.get_uop_optimizer()
+
+        with temporary_optimizer(opt):
+            testfunc(1000)
+
+        ex = None
+        for offset in range(0, len(testfunc.__code__.co_code), 2):
+            try:
+                ex = _testinternalcapi.get_executor(testfunc.__code__, offset)
+                break
+            except ValueError:
+                pass
+        self.assertIsNotNone(ex)
+        uops = {opname for opname, _ in ex}
+        self.assertIn("SAVE_IP", uops)
+        self.assertIn("LOAD_FAST", uops)
+
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/Python/optimizer.c b/Python/optimizer.c
@@ -458,7 +458,7 @@ uop_optimize(
         return trace_length;
     }
     OBJECT_STAT_INC(optimization_traces_created);
-    _PyUOpExecutorObject *executor = (_PyUOpExecutorObject *)_PyObject_New(&UOpExecutor_Type);
+    _PyUOpExecutorObject *executor = PyObject_New(_PyUOpExecutorObject, &UOpExecutor_Type);
     if (executor == NULL) {
         return -1;
     }
@@ -468,18 +468,24 @@ uop_optimize(
     return 1;
 }
 
+static void
+uop_opt_dealloc(PyObject *self) {
+    PyObject_Free(self);
+}
+
 static PyTypeObject UOpOptimizer_Type = {
     PyVarObject_HEAD_INIT(&PyType_Type, 0)
     .tp_name = "uop_optimizer",
     .tp_basicsize = sizeof(_PyOptimizerObject),
     .tp_itemsize = 0,
     .tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_DISALLOW_INSTANTIATION,
+    .tp_dealloc = uop_opt_dealloc,
 };
 
 PyObject *
 PyUnstable_Optimizer_NewUOpOptimizer(void)
 {
-    _PyOptimizerObject *opt = (_PyOptimizerObject *)_PyObject_New(&UOpOptimizer_Type);
+    _PyOptimizerObject *opt = PyObject_New(_PyOptimizerObject, &UOpOptimizer_Type);
     if (opt == NULL) {
         return NULL;
     }
