From 0df87cf6300cd478597383b987fc5fe6dd085034 Mon Sep 17 00:00:00 2001
From: Ivan Levkivskyi <levkivskyi@gmail.com>
Date: Tue, 25 Feb 2025 10:55:27 +0000
Subject: [PATCH] Fix order of steal/unborrow in tuple unpacking

---
 mypyc/ir/ops.py                         |  8 +++---
 mypyc/irbuild/statement.py              |  9 +++----
 mypyc/test-data/irbuild-statements.test | 18 +++++++------
 mypyc/test-data/refcount.test           | 35 +++++++++++++++++++++----
 4 files changed, 48 insertions(+), 22 deletions(-)

diff --git a/mypyc/ir/ops.py b/mypyc/ir/ops.py
index 6a2e70aee6d7..0323d31d0605 100644
--- a/mypyc/ir/ops.py
+++ b/mypyc/ir/ops.py
@@ -1532,12 +1532,12 @@ class Unborrow(RegisterOp):
       # t is a 2-tuple
       r0 = borrow t[0]
       r1 = borrow t[1]
+      keep_alive steal t
       r2 = unborrow r0
       r3 = unborrow r1
-      # now (r2, r3) represent the tuple as separate items, and the
-      # original tuple can be considered dead and available to be
-      # stolen
-      keep_alive steal t
+      # now (r2, r3) represent the tuple as separate items, that are
+      # managed again. (Note we need to steal before unborrow, to avoid
+      # refcount briefly touching zero if r2 or r3 are unused.)
 
     Be careful with this -- this can easily cause double freeing.
     """
diff --git a/mypyc/irbuild/statement.py b/mypyc/irbuild/statement.py
index cdc1d54589eb..f5b65bedbbca 100644
--- a/mypyc/irbuild/statement.py
+++ b/mypyc/irbuild/statement.py
@@ -211,12 +211,11 @@ def transform_assignment_stmt(builder: IRBuilder, stmt: AssignmentStmt) -> None:
         and any(t.is_refcounted for t in rvalue_reg.type.types)
     ):
         n = len(first_lvalue.items)
-        for i in range(n):
-            target = builder.get_assignment_target(first_lvalue.items[i])
-            rvalue_item = builder.add(TupleGet(rvalue_reg, i, borrow=True))
-            rvalue_item = builder.add(Unborrow(rvalue_item))
-            builder.assign(target, rvalue_item, line)
+        borrows = [builder.add(TupleGet(rvalue_reg, i, borrow=True)) for i in range(n)]
         builder.builder.keep_alive([rvalue_reg], steal=True)
+        for lvalue_item, rvalue_item in zip(first_lvalue.items, borrows):
+            rvalue_item = builder.add(Unborrow(rvalue_item))
+            builder.assign(builder.get_assignment_target(lvalue_item), rvalue_item, line)
         builder.flush_keep_alives()
         return
 
diff --git a/mypyc/test-data/irbuild-statements.test b/mypyc/test-data/irbuild-statements.test
index d5df984cfe4b..1f9336d32140 100644
--- a/mypyc/test-data/irbuild-statements.test
+++ b/mypyc/test-data/irbuild-statements.test
@@ -492,19 +492,21 @@ def from_any(a: Any) -> None:
 [out]
 def from_tuple(t):
     t :: tuple[int, object]
-    r0, r1 :: int
-    r2, x, r3, r4 :: object
+    r0 :: int
+    r1 :: object
+    r2 :: int
+    r3, x, r4 :: object
     r5, y :: int
 L0:
     r0 = borrow t[0]
-    r1 = unborrow r0
-    r2 = box(int, r1)
-    x = r2
-    r3 = borrow t[1]
-    r4 = unborrow r3
+    r1 = borrow t[1]
+    keep_alive steal t
+    r2 = unborrow r0
+    r3 = box(int, r2)
+    x = r3
+    r4 = unborrow r1
     r5 = unbox(int, r4)
     y = r5
-    keep_alive steal t
     return 1
 def from_any(a):
     a, r0, r1 :: object
diff --git a/mypyc/test-data/refcount.test b/mypyc/test-data/refcount.test
index c311f042ad5e..22153cff5a91 100644
--- a/mypyc/test-data/refcount.test
+++ b/mypyc/test-data/refcount.test
@@ -642,15 +642,15 @@ def g() -> Tuple[C, C]:
 [out]
 def f():
     r0 :: tuple[__main__.C, __main__.C]
-    r1, r2, x, r3, r4, y :: __main__.C
+    r1, r2, r3, x, r4, y :: __main__.C
     r5, r6, r7 :: int
 L0:
     r0 = g()
     r1 = borrow r0[0]
-    r2 = unborrow r1
-    x = r2
-    r3 = borrow r0[1]
-    r4 = unborrow r3
+    r2 = borrow r0[1]
+    r3 = unborrow r1
+    x = r3
+    r4 = unborrow r2
     y = r4
     r5 = borrow x.a
     r6 = borrow y.a
@@ -800,6 +800,31 @@ L2:
 L3:
     return 1
 
+[case testTupleUnpackUnused]
+from typing import Tuple
+
+def f(x: Tuple[str, int]) -> int:
+    a, xi = x
+    return 0
+[out]
+def f(x):
+    x :: tuple[str, int]
+    r0 :: str
+    r1 :: int
+    r2, a :: str
+    r3, xi :: int
+L0:
+    r0 = borrow x[0]
+    r1 = borrow x[1]
+    inc_ref x
+    r2 = unborrow r0
+    a = r2
+    dec_ref a
+    r3 = unborrow r1
+    xi = r3
+    dec_ref xi :: int
+    return 0
+
 [case testGetElementPtrLifeTime]
 from typing import List