diff --git a/mypyc/irbuild/for_helpers.py b/mypyc/irbuild/for_helpers.py index 62da773f68ada..94c11c4d1356b 100644 --- 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 index 225b93c1c50dd..d824bedb206f4 100644 --- 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):