From 5f237e6498d408988be9d444738bf9363143b0d9 Mon Sep 17 00:00:00 2001
From: Jukka Lehtosalo <jukka.lehtosalo@iki.fi>
Date: Sat, 24 Oct 2020 16:35:10 +0100
Subject: [PATCH 1/2] [mypyc] Fix type of for loop index register in for over
 range

It sometimes was short int even though it should have been int.
---
 mypyc/irbuild/for_helpers.py            |  8 +++--
 mypyc/test-data/irbuild-statements.test | 41 +++++++++++++++++++++++++
 2 files changed, 47 insertions(+), 2 deletions(-)

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..d86a24cd5d089 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 :: int64
+    r3 :: bit
+    r4 :: int64
+    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):

From 4b9525552f57d77d94bf9521943713fc166cfd24 Mon Sep 17 00:00:00 2001
From: Xuanda Yang <th3charlie@gmail.com>
Date: Sun, 25 Oct 2020 16:57:27 +0800
Subject: [PATCH 2/2] Update irbuild-statements.test

---
 mypyc/test-data/irbuild-statements.test | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/mypyc/test-data/irbuild-statements.test b/mypyc/test-data/irbuild-statements.test
index d86a24cd5d089..d824bedb206f4 100644
--- a/mypyc/test-data/irbuild-statements.test
+++ b/mypyc/test-data/irbuild-statements.test
@@ -37,9 +37,9 @@ def f(a: int) -> None:
 def f(a):
     a, r0, i :: int
     r1 :: bool
-    r2 :: int64
+    r2 :: native_int
     r3 :: bit
-    r4 :: int64
+    r4 :: native_int
     r5, r6, r7, r8 :: bit
     r9 :: int
 L0: