[mypyc] Fix type of for loop index register in for over range (python#9634)

JukkaL · web-flow · commit af99ebc03cb2 · 2020-10-25T20:52:34.000+08:00
It sometimes was short int even though it should have been int.
diff --git a/mypyc/irbuild/for_helpers.py b/mypyc/irbuild/for_helpers.py
@@ -17,7 +17,7 @@
 )
 from mypyc.ir.rtypes import (
     RType, is_short_int_rprimitive, is_list_rprimitive, is_sequence_rprimitive,
-    RTuple, is_dict_rprimitive, short_int_rprimitive
+    RTuple, is_dict_rprimitive, short_int_rprimitive, int_rprimitive
 )
 from mypyc.primitives.registry import CFunctionDescription
 from mypyc.primitives.dict_ops import (
@@ -605,7 +605,11 @@ def init(self, start_reg: Value, end_reg: Value, step: int) -> None:
         self.end_reg = end_reg
         self.step = step
         self.end_target = builder.maybe_spill(end_reg)
-        index_reg = builder.alloc_temp(start_reg.type)
+        if is_short_int_rprimitive(start_reg.type) and is_short_int_rprimitive(end_reg.type):
+            index_type = short_int_rprimitive
+        else:
+            index_type = int_rprimitive
+        index_reg = builder.alloc_temp(index_type)
         builder.assign(index_reg, start_reg, -1)
         self.index_reg = builder.maybe_spill_assignable(index_reg)
         # Initialize loop index to 0. Assert that the index target is assignable.
diff --git a/mypyc/test-data/irbuild-statements.test b/mypyc/test-data/irbuild-statements.test
@@ -29,6 +29,47 @@ L3:
 L4:
     return 1
 
+[case testForInRangeVariableEndIndxe]
+def f(a: int) -> None:
+    for i in range(a):
+        pass
+[out]
+def f(a):
+    a, r0, i :: int
+    r1 :: bool
+    r2 :: native_int
+    r3 :: bit
+    r4 :: native_int
+    r5, r6, r7, r8 :: bit
+    r9 :: int
+L0:
+    r0 = 0
+    i = r0
+L1:
+    r2 = r0 & 1
+    r3 = r2 == 0
+    r4 = a & 1
+    r5 = r4 == 0
+    r6 = r3 & r5
+    if r6 goto L2 else goto L3 :: bool
+L2:
+    r7 = r0 < a :: signed
+    r1 = r7
+    goto L4
+L3:
+    r8 = CPyTagged_IsLt_(r0, a)
+    r1 = r8
+L4:
+    if r1 goto L5 else goto L7 :: bool
+L5:
+L6:
+    r9 = CPyTagged_Add(r0, 2)
+    r0 = r9
+    i = r9
+    goto L1
+L7:
+    return 1
+
 [case testForInNegativeRange]
 def f() -> None:
     for i in range(10, 0, -1):
