From 5ae23c1c39ed3bc69cf0672802f9b4828c62bb17 Mon Sep 17 00:00:00 2001
From: Xuanda Yang <th3charlie@gmail.com>
Date: Sun, 11 Oct 2020 07:18:06 +0800
Subject: [PATCH] multiple assignment

---
 mypyc/irbuild/statement.py         | 16 +++++++++++-
 mypyc/test-data/irbuild-basic.test | 42 ++++++++++++++++++++++++++++++
 2 files changed, 57 insertions(+), 1 deletion(-)

diff --git a/mypyc/irbuild/statement.py b/mypyc/irbuild/statement.py
index 5e1bf6e914012..3f3c282f6155b 100644
--- a/mypyc/irbuild/statement.py
+++ b/mypyc/irbuild/statement.py
@@ -12,7 +12,7 @@
 from mypy.nodes import (
     Block, ExpressionStmt, ReturnStmt, AssignmentStmt, OperatorAssignmentStmt, IfStmt, WhileStmt,
     ForStmt, BreakStmt, ContinueStmt, RaiseStmt, TryStmt, WithStmt, AssertStmt, DelStmt,
-    Expression, StrExpr, TempNode, Lvalue, Import, ImportFrom, ImportAll
+    Expression, StrExpr, TempNode, Lvalue, Import, ImportFrom, ImportAll, TupleExpr
 )
 
 from mypyc.ir.ops import (
@@ -79,6 +79,20 @@ def transform_assignment_stmt(builder: IRBuilder, stmt: AssignmentStmt) -> None:
         builder.get_assignment_target(lvalue, stmt.line)
         return
 
+    # multiple assignment
+    if (isinstance(lvalue, TupleExpr) and isinstance(stmt.rvalue, TupleExpr)
+            and len(lvalue.items) == len(stmt.rvalue.items)):
+        temps = []
+        for right in stmt.rvalue.items:
+            rvalue_reg = builder.accept(right)
+            temp = builder.alloc_temp(rvalue_reg.type)
+            builder.assign(temp, rvalue_reg, stmt.line)
+            temps.append(temp)
+        for (left, temp) in zip(lvalue.items, temps):
+            assignment_target = builder.get_assignment_target(left)
+            builder.assign(assignment_target, temp, stmt.line)
+        return
+
     line = stmt.rvalue.line
     rvalue_reg = builder.accept(stmt.rvalue)
     if builder.non_function_scope() and stmt.is_final_def:
diff --git a/mypyc/test-data/irbuild-basic.test b/mypyc/test-data/irbuild-basic.test
index 62a93e62c829c..8c7f8aeb989ad 100644
--- a/mypyc/test-data/irbuild-basic.test
+++ b/mypyc/test-data/irbuild-basic.test
@@ -3624,3 +3624,45 @@ L0:
     r0 = PyObject_IsTrue(x)
     r1 = truncate r0: int32 to builtins.bool
     return r1
+
+[case testMultipleAssignment]
+from typing import Tuple
+
+def f(x: int, y: int) -> Tuple[int, int]:
+    x, y = y, x
+    return (x, y)
+
+def f2(x: int, y: str, z: float) -> Tuple[float, str, int]:
+    a, b, c = x, y, z
+    return (c, b, a)
+[out]
+def f(x, y):
+    x, y, r0, r1 :: int
+    r2 :: tuple[int, int]
+L0:
+    r0 = y
+    r1 = x
+    x = r0
+    y = r1
+    r2 = (x, y)
+    return r2
+def f2(x, y, z):
+    x :: int
+    y :: str
+    z :: float
+    r0 :: int
+    r1 :: str
+    r2 :: float
+    a :: int
+    b :: str
+    c :: float
+    r3 :: tuple[float, str, int]
+L0:
+    r0 = x
+    r1 = y
+    r2 = z
+    a = r0
+    b = r1
+    c = r2
+    r3 = (c, b, a)
+    return r3