Skip to content

Commit 34f6f6a

Browse files
authored
[mypyc] Fix order of steal/unborrow in tuple unpacking (#18732)
Currently, although globally the refcount is correct, it may briefly touch 0 if a target of unpacking in unused, e.g. `_, _, last = some_tuple`. This can be prevented by placing steal before unborrow (which IMO should be the recommended way, if I understand the logic of these terms correctly).
1 parent 66dde14 commit 34f6f6a

File tree

4 files changed

+48
-22
lines changed

4 files changed

+48
-22
lines changed

mypyc/ir/ops.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1532,12 +1532,12 @@ class Unborrow(RegisterOp):
15321532
# t is a 2-tuple
15331533
r0 = borrow t[0]
15341534
r1 = borrow t[1]
1535+
keep_alive steal t
15351536
r2 = unborrow r0
15361537
r3 = unborrow r1
1537-
# now (r2, r3) represent the tuple as separate items, and the
1538-
# original tuple can be considered dead and available to be
1539-
# stolen
1540-
keep_alive steal t
1538+
# now (r2, r3) represent the tuple as separate items, that are
1539+
# managed again. (Note we need to steal before unborrow, to avoid
1540+
# refcount briefly touching zero if r2 or r3 are unused.)
15411541
15421542
Be careful with this -- this can easily cause double freeing.
15431543
"""

mypyc/irbuild/statement.py

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -211,12 +211,11 @@ def transform_assignment_stmt(builder: IRBuilder, stmt: AssignmentStmt) -> None:
211211
and any(t.is_refcounted for t in rvalue_reg.type.types)
212212
):
213213
n = len(first_lvalue.items)
214-
for i in range(n):
215-
target = builder.get_assignment_target(first_lvalue.items[i])
216-
rvalue_item = builder.add(TupleGet(rvalue_reg, i, borrow=True))
217-
rvalue_item = builder.add(Unborrow(rvalue_item))
218-
builder.assign(target, rvalue_item, line)
214+
borrows = [builder.add(TupleGet(rvalue_reg, i, borrow=True)) for i in range(n)]
219215
builder.builder.keep_alive([rvalue_reg], steal=True)
216+
for lvalue_item, rvalue_item in zip(first_lvalue.items, borrows):
217+
rvalue_item = builder.add(Unborrow(rvalue_item))
218+
builder.assign(builder.get_assignment_target(lvalue_item), rvalue_item, line)
220219
builder.flush_keep_alives()
221220
return
222221

mypyc/test-data/irbuild-statements.test

Lines changed: 10 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -492,19 +492,21 @@ def from_any(a: Any) -> None:
492492
[out]
493493
def from_tuple(t):
494494
t :: tuple[int, object]
495-
r0, r1 :: int
496-
r2, x, r3, r4 :: object
495+
r0 :: int
496+
r1 :: object
497+
r2 :: int
498+
r3, x, r4 :: object
497499
r5, y :: int
498500
L0:
499501
r0 = borrow t[0]
500-
r1 = unborrow r0
501-
r2 = box(int, r1)
502-
x = r2
503-
r3 = borrow t[1]
504-
r4 = unborrow r3
502+
r1 = borrow t[1]
503+
keep_alive steal t
504+
r2 = unborrow r0
505+
r3 = box(int, r2)
506+
x = r3
507+
r4 = unborrow r1
505508
r5 = unbox(int, r4)
506509
y = r5
507-
keep_alive steal t
508510
return 1
509511
def from_any(a):
510512
a, r0, r1 :: object

mypyc/test-data/refcount.test

Lines changed: 30 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -642,15 +642,15 @@ def g() -> Tuple[C, C]:
642642
[out]
643643
def f():
644644
r0 :: tuple[__main__.C, __main__.C]
645-
r1, r2, x, r3, r4, y :: __main__.C
645+
r1, r2, r3, x, r4, y :: __main__.C
646646
r5, r6, r7 :: int
647647
L0:
648648
r0 = g()
649649
r1 = borrow r0[0]
650-
r2 = unborrow r1
651-
x = r2
652-
r3 = borrow r0[1]
653-
r4 = unborrow r3
650+
r2 = borrow r0[1]
651+
r3 = unborrow r1
652+
x = r3
653+
r4 = unborrow r2
654654
y = r4
655655
r5 = borrow x.a
656656
r6 = borrow y.a
@@ -800,6 +800,31 @@ L2:
800800
L3:
801801
return 1
802802

803+
[case testTupleUnpackUnused]
804+
from typing import Tuple
805+
806+
def f(x: Tuple[str, int]) -> int:
807+
a, xi = x
808+
return 0
809+
[out]
810+
def f(x):
811+
x :: tuple[str, int]
812+
r0 :: str
813+
r1 :: int
814+
r2, a :: str
815+
r3, xi :: int
816+
L0:
817+
r0 = borrow x[0]
818+
r1 = borrow x[1]
819+
inc_ref x
820+
r2 = unborrow r0
821+
a = r2
822+
dec_ref a
823+
r3 = unborrow r1
824+
xi = r3
825+
dec_ref xi :: int
826+
return 0
827+
803828
[case testGetElementPtrLifeTime]
804829
from typing import List
805830

0 commit comments

Comments
 (0)