fix bug && add more tests

Author: cocolato

Lib/test/test_capi/test_opt.py

@@ -1290,21 +1290,7 @@ def testfunc(n):
         self.assertIsNotNone(ex)
         self.assertIn("_RETURN_GENERATOR", get_opnames(ex))
 
-    def test_make_heap_safe_optimized(self):
-        def returns_owned(x):
-            return x + 1
-        def testfunc(n):
-            a = 0
-            for _ in range(n):
-                a = returns_owned(a)
-            return a
-        res, ex = self._run_with_optimizer(testfunc, TIER2_THRESHOLD)
-        self.assertEqual(res, TIER2_THRESHOLD)
-        self.assertIsNotNone(ex)
-        uops = get_opnames(ex)
-        self.assertNotIn("_MAKE_HEAP_SAFE", uops)
-        self.assertIn("_RETURN_VALUE", uops)
-
+    def test_make_heap_safe_optimized_immortal(self):
         def returns_immortal():
             return None
         def testfunc(n):
@@ -1319,12 +1305,12 @@ def testfunc(n):
         self.assertNotIn("_MAKE_HEAP_SAFE", uops)
         self.assertIn("_RETURN_VALUE", uops)
 
-    def test_make_heap_safe_yield_correctness(self):
+    def test_make_heap_safe_optimized_yield(self):
         def gen(n):
-            for i in range(n):
-                yield i + i
+            for _ in range(n):
+                yield 1
         def testfunc(n):
-            for x in gen(n):
+            for _ in gen(n):
                 pass
         testfunc(TIER2_THRESHOLD * 2)
         gen_ex = get_first_executor(gen)
@@ -1333,6 +1319,21 @@ def testfunc(n):
         self.assertNotIn("_MAKE_HEAP_SAFE", uops)
         self.assertIn("_YIELD_VALUE", uops)
 
+    def test_make_heap_safe_not_optimized_for_owned(self):
+        def returns_owned(x):
+            return x + 1
+        def testfunc(n):
+            a = 0
+            for _ in range(n):
+                a = returns_owned(a)
+            return a
+        res, ex = self._run_with_optimizer(testfunc, TIER2_THRESHOLD)
+        self.assertEqual(res, TIER2_THRESHOLD)
+        self.assertIsNotNone(ex)
+        uops = get_opnames(ex)
+        self.assertIn("_MAKE_HEAP_SAFE", uops)
+        self.assertIn("_RETURN_VALUE", uops)
+
     def test_for_iter(self):
         def testfunc(n):
             t = 0

Python/optimizer_bytecodes.c

@@ -87,9 +87,8 @@ dummy_func(void) {
 // BEGIN BYTECODES //
 
     op(_MAKE_HEAP_SAFE, (value -- value)) {
-        // If the value is not borrowed, or is immortal, it's heap-safe.
-        if (!PyJitRef_IsBorrowed(value) ||
-            sym_is_immortal(PyJitRef_Unwrap(value))) {
+        // eliminate _MAKE_HEAP_SAFE when we *know* the value is immortal
+        if (sym_is_immortal(PyJitRef_Unwrap(value))) {
             ADD_OP(_NOP, 0, 0);
         }
         value = PyJitRef_StripReferenceInfo(value);