From f5ee4bb0067d5b8a2d64fa4fd9dc09a5d8cfb669 Mon Sep 17 00:00:00 2001
From: Jared Hance <jhance@dropbox.com>
Date: Wed, 22 Sep 2021 10:28:46 -0700
Subject: [PATCH] [mypyc] Coerce Register assignments

Fixes mypyc/mypyc#759 and probably some other issues. Including only
the change to coerce the Register assignments fixes the missing
type coercion for the specialized next-generator case, but causes some
unusual box-then-unbox from short_ints.

These odd coercions are caused by the dict keys in dict literals being
considered short_int, but dict_next_rtuple_single/pair return types bein
annotated as int_primitive. Changing this to short_int_rprimitive removes
the odd coercisions. The remainder of the diff is fixing the irbuild
testcase output to have short_int.
---
 mypyc/ir/rtypes.py                      |  4 +--
 mypyc/irbuild/builder.py                |  2 +-
 mypyc/test-data/irbuild-dict.test       | 12 +++----
 mypyc/test-data/irbuild-lists.test      | 46 +++++++++++++++++++++++++
 mypyc/test-data/irbuild-set.test        |  4 +--
 mypyc/test-data/irbuild-statements.test |  8 ++---
 mypyc/test-data/refcount.test           |  4 +--
 mypyc/test-data/run-lists.test          |  4 +++
 8 files changed, 67 insertions(+), 17 deletions(-)

diff --git a/mypyc/ir/rtypes.py b/mypyc/ir/rtypes.py
index 9c27e425a78a5..4bf71883b15df 100644
--- a/mypyc/ir/rtypes.py
+++ b/mypyc/ir/rtypes.py
@@ -522,11 +522,11 @@ def deserialize(cls, data: JsonDict, ctx: 'DeserMaps') -> 'RTuple':
 # Dictionary iterator tuple: (should continue, internal offset, key, value)
 # See mypyc.irbuild.for_helpers.ForDictionaryCommon for more details.
 dict_next_rtuple_pair = RTuple(
-    [bool_rprimitive, int_rprimitive, object_rprimitive, object_rprimitive]
+    [bool_rprimitive, short_int_rprimitive, object_rprimitive, object_rprimitive]
 )
 # Same as above but just for key or value.
 dict_next_rtuple_single = RTuple(
-    [bool_rprimitive, int_rprimitive, object_rprimitive]
+    [bool_rprimitive, short_int_rprimitive, object_rprimitive]
 )
 
 
diff --git a/mypyc/irbuild/builder.py b/mypyc/irbuild/builder.py
index 5f09f6b3e3b76..a9d72190b272b 100644
--- a/mypyc/irbuild/builder.py
+++ b/mypyc/irbuild/builder.py
@@ -542,7 +542,7 @@ def assign(self,
                rvalue_reg: Value,
                line: int) -> None:
         if isinstance(target, Register):
-            self.add(Assign(target, rvalue_reg))
+            self.add(Assign(target, self.coerce(rvalue_reg, target.type, line)))
         elif isinstance(target, AssignmentTargetRegister):
             rvalue_reg = self.coerce(rvalue_reg, target.type, line)
             self.add(Assign(target.register, rvalue_reg))
diff --git a/mypyc/test-data/irbuild-dict.test b/mypyc/test-data/irbuild-dict.test
index 2cf99ec5b2ac7..69dd42f01dd36 100644
--- a/mypyc/test-data/irbuild-dict.test
+++ b/mypyc/test-data/irbuild-dict.test
@@ -154,8 +154,8 @@ def increment(d):
     r1 :: native_int
     r2 :: short_int
     r3 :: object
-    r4 :: tuple[bool, int, object]
-    r5 :: int
+    r4 :: tuple[bool, short_int, object]
+    r5 :: short_int
     r6 :: bool
     r7 :: object
     r8, k :: str
@@ -232,8 +232,8 @@ def print_dict_methods(d1, d2):
     r1 :: native_int
     r2 :: short_int
     r3 :: object
-    r4 :: tuple[bool, int, object]
-    r5 :: int
+    r4 :: tuple[bool, short_int, object]
+    r5 :: short_int
     r6 :: bool
     r7 :: object
     r8, v :: int
@@ -246,8 +246,8 @@ def print_dict_methods(d1, d2):
     r16 :: native_int
     r17 :: short_int
     r18 :: object
-    r19 :: tuple[bool, int, object, object]
-    r20 :: int
+    r19 :: tuple[bool, short_int, object, object]
+    r20 :: short_int
     r21 :: bool
     r22, r23 :: object
     r24, r25, k :: int
diff --git a/mypyc/test-data/irbuild-lists.test b/mypyc/test-data/irbuild-lists.test
index 2d9f24fe8e487..0d29f256cc597 100644
--- a/mypyc/test-data/irbuild-lists.test
+++ b/mypyc/test-data/irbuild-lists.test
@@ -381,3 +381,49 @@ L7:
 L8:
     b = r16
     return 1
+[case testGeneratorNext]
+from typing import List, Optional
+
+def test(x: List[int]) -> None:
+    res = next((i for i in x), None)
+[out]
+def test(x):
+    x :: list
+    r0 :: short_int
+    r1 :: ptr
+    r2 :: native_int
+    r3 :: short_int
+    r4 :: bit
+    r5 :: object
+    r6, i :: int
+    r7 :: object
+    r8 :: union[int, None]
+    r9 :: short_int
+    r10 :: object
+    res :: union[int, None]
+L0:
+    r0 = 0
+L1:
+    r1 = get_element_ptr x ob_size :: PyVarObject
+    r2 = load_mem r1 :: native_int*
+    keep_alive x
+    r3 = r2 << 1
+    r4 = r0 < r3 :: signed
+    if r4 goto L2 else goto L4 :: bool
+L2:
+    r5 = CPyList_GetItemUnsafe(x, r0)
+    r6 = unbox(int, r5)
+    i = r6
+    r7 = box(int, i)
+    r8 = r7
+    goto L5
+L3:
+    r9 = r0 + 2
+    r0 = r9
+    goto L1
+L4:
+    r10 = box(None, 1)
+    r8 = r10
+L5:
+    res = r8
+    return 1
diff --git a/mypyc/test-data/irbuild-set.test b/mypyc/test-data/irbuild-set.test
index aab70ecd0e762..f620039b76557 100644
--- a/mypyc/test-data/irbuild-set.test
+++ b/mypyc/test-data/irbuild-set.test
@@ -173,8 +173,8 @@ def test3():
     r9 :: native_int
     r10 :: short_int
     r11 :: object
-    r12 :: tuple[bool, int, object]
-    r13 :: int
+    r12 :: tuple[bool, short_int, object]
+    r13 :: short_int
     r14 :: bool
     r15 :: object
     r16, x, r17 :: int
diff --git a/mypyc/test-data/irbuild-statements.test b/mypyc/test-data/irbuild-statements.test
index 0975ac301db94..8b85e494b07db 100644
--- a/mypyc/test-data/irbuild-statements.test
+++ b/mypyc/test-data/irbuild-statements.test
@@ -342,8 +342,8 @@ def f(d):
     r1 :: native_int
     r2 :: short_int
     r3 :: object
-    r4 :: tuple[bool, int, object]
-    r5 :: int
+    r4 :: tuple[bool, short_int, object]
+    r5 :: short_int
     r6 :: bool
     r7 :: object
     r8, key :: int
@@ -394,8 +394,8 @@ def sum_over_even_values(d):
     r1 :: native_int
     r2 :: short_int
     r3 :: object
-    r4 :: tuple[bool, int, object]
-    r5 :: int
+    r4 :: tuple[bool, short_int, object]
+    r5 :: short_int
     r6 :: bool
     r7 :: object
     r8, key :: int
diff --git a/mypyc/test-data/refcount.test b/mypyc/test-data/refcount.test
index 64b38f33a26bc..c48ae1acce098 100644
--- a/mypyc/test-data/refcount.test
+++ b/mypyc/test-data/refcount.test
@@ -724,8 +724,8 @@ def f(d):
     r1 :: native_int
     r2 :: short_int
     r3 :: object
-    r4 :: tuple[bool, int, object]
-    r5 :: int
+    r4 :: tuple[bool, short_int, object]
+    r5 :: short_int
     r6 :: bool
     r7 :: object
     r8, key :: int
diff --git a/mypyc/test-data/run-lists.test b/mypyc/test-data/run-lists.test
index 1366b1fd857e7..c98ab91711238 100644
--- a/mypyc/test-data/run-lists.test
+++ b/mypyc/test-data/run-lists.test
@@ -379,4 +379,8 @@ def test() -> None:
     source_str = "abcd"
     f = list("str:" + x for x in source_str)
     assert f == ["str:a", "str:b", "str:c", "str:d"]
+[case testNextBug]
+from typing import List, Optional
 
+def test(x: List[int]) -> None:
+    res = next((i for i in x), None)