From 8cf122d27a83711a02240ddff45976fe8e549509 Mon Sep 17 00:00:00 2001
From: Lukas Geiger <lukas.geiger94@gmail.com>
Date: Sat, 5 Sep 2026 23:24:23 +0100
Subject: [PATCH] gh-131798: Propagate iterator types from GET_ITER in the JIT

---
 Lib/test/test_capi/test_opt.py | 24 ++++++++++++++++++++++++
 Python/optimizer_bytecodes.c   | 10 +++++++++-
 Python/optimizer_cases.c.h     | 10 +++++++++-
 3 files changed, 42 insertions(+), 2 deletions(-)

diff --git a/Lib/test/test_capi/test_opt.py b/Lib/test/test_capi/test_opt.py
index 36efab51878141..17caa115d64111 100644
--- a/Lib/test/test_capi/test_opt.py
+++ b/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):
diff --git a/Python/optimizer_bytecodes.c b/Python/optimizer_bytecodes.c
index 5246e50633461b..71d40cce5c87b2 100644
--- a/Python/optimizer_bytecodes.c
+++ b/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 {
diff --git a/Python/optimizer_cases.c.h b/Python/optimizer_cases.c.h
index 21f275f27cafe0..58a9b8c87e6a71 100644
--- a/Python/optimizer_cases.c.h
+++ b/Python/optimizer_cases.c.h
@@ -3605,7 +3605,15 @@
                 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 {