gh-131798: Propagate iterator types from GET_ITER in the JIT

Author: lgeiger

Lib/test/test_capi/test_opt.py

@@ -458,10 +458,34 @@ def testfunc(n):
         self.assertIsNotNone(ex)
         uops = get_opnames(ex)
         self.assertIn("_GET_ITER_TRAD", uops)
+        self.assertIn("_ITER_NEXT_INLINE", uops)
+        self.assertNotIn("_GUARD_TYPE_ITER", uops)
         self.assertNotIn("_GET_ITER", uops)
         self.assertNotIn("_GET_ITER_VIRTUAL", uops)
         self.assertNotIn("_GET_ITER_SELF", uops)
 
+    def test_get_iter_trad_set(self):
+        s = set(range(10))
+        def testfunc(n):
+            total = 0
+            while n:
+                n -= 1
+                for value in s:
+                    total += value
+                    break
+            return total
+
+        total = testfunc(TIER2_THRESHOLD)
+        self.assertEqual(total, next(iter(s)) * TIER2_THRESHOLD)
+        ex = get_first_executor(testfunc)
+        self.assertIsNotNone(ex)
+        uops = get_opnames(ex)
+        self.assertIn("_GET_ITER_TRAD", uops)
+        self.assertIn("_ITER_NEXT_INLINE", uops)
+        self.assertNotIn("_GUARD_TYPE_ITER", uops)
+        self.assertNotIn("_GET_ITER", uops)
+        self.assertNotIn("_GET_ITER_VIRTUAL", uops)
+        self.assertNotIn("_GET_ITER_SELF", uops)
 
     def test_for_iter_range(self):
         def testfunc(n):

Python/optimizer_bytecodes.c

@@ -1448,7 +1448,15 @@ dummy_func(void) {
             index_or_null = sym_new_null(ctx);
         }
         else if (is_trad) {
-            iter = sym_new_not_null(ctx);
+            if (tp == &PyDict_Type || tp == &PyFrozenDict_Type) {
+                iter = sym_new_type(ctx, &PyDictIterKey_Type);
+            }
+            else if (tp == &PySet_Type || tp == &PyFrozenSet_Type) {
+                iter = sym_new_type(ctx, &PySetIter_Type);
+            }
+            else {
+                iter = sym_new_not_null(ctx);
+            }
             index_or_null = sym_new_null(ctx);
         }
         else {