Skip to content

Commit

Permalink
[mypyc] Fix order of steal/unborrow in tuple unpacking (#18732)
Browse files Browse the repository at this point in the history
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).
  • Loading branch information
ilevkivskyi authored Feb 25, 2025
1 parent 66dde14 commit 34f6f6a
Show file tree
Hide file tree
Showing 4 changed files with 48 additions and 22 deletions.
8 changes: 4 additions & 4 deletions mypyc/ir/ops.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
"""
Expand Down
9 changes: 4 additions & 5 deletions mypyc/irbuild/statement.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
18 changes: 10 additions & 8 deletions mypyc/test-data/irbuild-statements.test
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
35 changes: 30 additions & 5 deletions mypyc/test-data/refcount.test
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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

Expand Down

0 comments on commit 34f6f6a

Please sign in to comment.