[mypyc] Fix semantics for walrus expression in tuple (#21249)

Fixes mypyc/mypyc#1179

Author: hauntsaninja

mypyc/irbuild/expression.py

@@ -1113,7 +1113,12 @@ def transform_tuple_expr(builder: IRBuilder, expr: TupleExpr) -> Value:
     items = []
     for item_expr, item_type in zip(expr.items, types):
         reg = builder.accept(item_expr)
-        items.append(builder.coerce(reg, item_type, item_expr.line))
+        item = builder.coerce(reg, item_type, item_expr.line)
+        if isinstance(item, Register):
+            temp = Register(item.type)
+            builder.assign(temp, item, item_expr.line)
+            item = temp
+        items.append(item)
     return builder.add(TupleSet(items, expr.line))
 
 

mypyc/test-data/irbuild-basic.test

@@ -1665,11 +1665,14 @@ def h() -> Tuple[int, int, int]:
     return f(1, *(2, 3))
 [out]
 def f(a, b, c):
-    a, b, c :: int
-    r0 :: tuple[int, int, int]
+    a, b, c, r0, r1, r2 :: int
+    r3 :: tuple[int, int, int]
 L0:
-    r0 = (a, b, c)
-    return r0
+    r0 = a
+    r1 = b
+    r2 = c
+    r3 = (r0, r1, r2)
+    return r3
 def g():
     r0 :: tuple[int, int, int]
     r1 :: dict
@@ -1724,11 +1727,14 @@ def h() -> Tuple[int, int, int]:
     return f(1, **{'b': 2, 'c': 3})
 [out]
 def f(a, b, c):
-    a, b, c :: int
-    r0 :: tuple[int, int, int]
+    a, b, c, r0, r1, r2 :: int
+    r3 :: tuple[int, int, int]
 L0:
-    r0 = (a, b, c)
-    return r0
+    r0 = a
+    r1 = b
+    r2 = c
+    r3 = (r0, r1, r2)
+    return r3
 def g():
     r0, r1, r2 :: str
     r3, r4, r5 :: object

mypyc/test-data/irbuild-float.test

@@ -483,15 +483,16 @@ def f(x: int) -> None:
     t: tuple[float, float, float] = (x, 2.5, -7)
 [out]
 def f(x):
-    x :: int
-    r0 :: tuple[int, float, int]
-    r1 :: int
-    r2 :: float
-    r3, t :: tuple[float, float, float]
+    x, r0 :: int
+    r1 :: tuple[int, float, int]
+    r2 :: int
+    r3 :: float
+    r4, t :: tuple[float, float, float]
 L0:
-    r0 = (x, 2.5, -14)
-    r1 = r0[0]
-    r2 = CPyFloat_FromTagged(r1)
-    r3 = (r2, 2.5, -7.0)
-    t = r3
+    r0 = x
+    r1 = (r0, 2.5, -14)
+    r2 = r1[0]
+    r3 = CPyFloat_FromTagged(r2)
+    r4 = (r3, 2.5, -7.0)
+    t = r4
     return 1

mypyc/test-data/irbuild-i64.test

@@ -658,11 +658,13 @@ def h() -> i64:
     return x + y + t[0]
 [out]
 def f(x, y):
-    x, y :: i64
-    r0 :: tuple[i64, i64]
+    x, y, r0, r1 :: i64
+    r2 :: tuple[i64, i64]
 L0:
-    r0 = (x, y)
-    return r0
+    r0 = x
+    r1 = y
+    r2 = (r0, r1)
+    return r2
 def g():
     r0 :: tuple[int, int]
     r1 :: tuple[i64, i64]
@@ -2114,33 +2116,34 @@ L3:
     tt = r10
     return 1
 def g(n):
-    n :: int
-    r0 :: tuple[int, int]
-    r1 :: int
-    r2 :: native_int
-    r3 :: bit
-    r4, r5 :: i64
-    r6 :: ptr
-    r7 :: c_ptr
-    r8 :: i64
-    r9, t :: tuple[i64, i64]
+    n, r0 :: int
+    r1 :: tuple[int, int]
+    r2 :: int
+    r3 :: native_int
+    r4 :: bit
+    r5, r6 :: i64
+    r7 :: ptr
+    r8 :: c_ptr
+    r9 :: i64
+    r10, t :: tuple[i64, i64]
 L0:
-    r0 = (2, n)
-    r1 = r0[1]
-    r2 = r1 & 1
-    r3 = r2 == 0
-    if r3 goto L1 else goto L2 :: bool
+    r0 = n
+    r1 = (2, r0)
+    r2 = r1[1]
+    r3 = r2 & 1
+    r4 = r3 == 0
+    if r4 goto L1 else goto L2 :: bool
 L1:
-    r4 = r1 >> 1
-    r5 = r4
+    r5 = r2 >> 1
+    r6 = r5
     goto L3
 L2:
-    r6 = r1 ^ 1
-    r7 = r6
-    r8 = CPyLong_AsInt64(r7)
-    r5 = r8
-    keep_alive r1
+    r7 = r2 ^ 1
+    r8 = r7
+    r9 = CPyLong_AsInt64(r8)
+    r6 = r9
+    keep_alive r2
 L3:
-    r9 = (1, r5)
-    t = r9
+    r10 = (1, r6)
+    t = r10
     return 1

mypyc/test-data/irbuild-statements.test

@@ -381,15 +381,17 @@ def f3(x: int, y: int) -> Tuple[int, int]:
     return (x, y)
 [out]
 def f(x, y):
-    x, y, r0, r1 :: int
-    r2 :: tuple[int, int]
+    x, y, r0, r1, r2, r3 :: int
+    r4 :: tuple[int, int]
 L0:
     r0 = y
     r1 = x
     x = r0
     y = r1
-    r2 = (x, y)
-    return r2
+    r2 = x
+    r3 = y
+    r4 = (r2, r3)
+    return r4
 def f2(x, y, z):
     x :: int
     y :: str
@@ -399,27 +401,34 @@ def f2(x, y, z):
     r2 :: float
     a :: int
     b :: str
-    c :: float
-    r3 :: tuple[float, str, int]
+    c, r3 :: float
+    r4 :: str
+    r5 :: int
+    r6 :: tuple[float, str, int]
 L0:
     r0 = x
     r1 = y
     r2 = z
     a = r0
     b = r1
     c = r2
-    r3 = (c, b, a)
-    return r3
+    r3 = c
+    r4 = b
+    r5 = a
+    r6 = (r3, r4, r5)
+    return r6
 def f3(x, y):
-    x, y, r0, r1 :: int
-    r2 :: tuple[int, int]
+    x, y, r0, r1, r2, r3 :: int
+    r4 :: tuple[int, int]
 L0:
     r0 = y
     r1 = x
     x = r0
     y = r1
-    r2 = (x, y)
-    return r2
+    r2 = x
+    r3 = y
+    r4 = (r2, r3)
+    return r4
 
 [case testMultipleAssignmentBasicUnpacking]
 from typing import Tuple, Any

mypyc/test-data/refcount.test

@@ -617,15 +617,17 @@ def f() -> Tuple[C, C]:
     return a, b
 [out]
 def f():
-    r0, a, r1, b :: __main__.C
-    r2 :: tuple[__main__.C, __main__.C]
+    r0, a, r1, b, r2, r3 :: __main__.C
+    r4 :: tuple[__main__.C, __main__.C]
 L0:
     r0 = C()
     a = r0
     r1 = C()
     b = r1
-    r2 = (a, b)
-    return r2
+    r2 = a
+    r3 = b
+    r4 = (r2, r3)
+    return r4
 
 [case testDecomposeTuple]
 from typing import Tuple

mypyc/test-data/run-tuples.test

@@ -398,6 +398,15 @@ def test_multiply() -> None:
     assert 3 * (1,) == res
     assert multiply((1,), 3) == res
 
+# Ref: https://github.com/mypyc/mypyc/issues/1179
+def test_tuple_with_walrus() -> None:
+    i = 0
+    assert (i, (i := i + 1), i, (i := i + 1), i, (i := i + 1), i) == (0, 1, 1, 2, 2, 3, 3)
+
+def test_nested_tuple_with_walrus() -> None:
+    i = 0
+    assert ((i, i), (i := i + 1)) == ((0, 0), 1)
+
 [case testIsInstance]
 from copysubclass import subc
 def test_built_in() -> None: