[mypyc] Access tuple attributes under critical section

mypyc/codegen/emit.py

@@ -471,7 +471,14 @@ def _emit_attr_bitmap_update(
         index = cl.bitmap_attrs.index(attr)
         mask = 1 << (index & (BITMAP_BITS - 1))
         bitmap = self.attr_bitmap_expr(obj, cl, index)
-        if clear:
+        if self.target_python.free_threaded and not cl.attrs_are_thread_confined():
+            # Scalar and tuple attributes share bitmap words. All concurrent
+            # loads and RMWs must be atomic.
+            if clear:
+                self.emit_line(f"_Py_atomic_and_uint32(&{bitmap}, ~{mask}U);")
+            else:
+                self.emit_line(f"_Py_atomic_or_uint32(&{bitmap}, {mask});")
+        elif clear:
             self.emit_line(f"{bitmap} &= ~{mask};")
         else:
             self.emit_line(f"{bitmap} |= {mask};")
@@ -497,9 +504,82 @@ def emit_undefined_attr_check(
             bit = 1 << (index & (BITMAP_BITS - 1))
             attr = self.bitmap_field(index)
             obj_expr = f"({cl.struct_name(self.names)} *){obj}"
-            check = f"{check} && !(({obj_expr})->{attr} & {bit})"
+            bitmap = f"({obj_expr})->{attr}"
+            if self.target_python.free_threaded and not cl.attrs_are_thread_confined():
+                bitmap = f"_Py_atomic_load_uint32(&{bitmap})"
+            check = f"{check} && !({bitmap} & {bit})"
         self.emit_line(f"if ({check}) {{")
 
+    def emit_tuple_attr_load_locked(
+        self,
+        dest: str,
+        attr_expr: str,
+        obj: str,
+        rtype: RTuple,
+        cl: ClassIR,
+        attr: str,
+        *,
+        borrow: bool = False,
+    ) -> str:
+        """Snapshot a tuple attribute and its defined state while holding the owner lock.
+
+        Secure references before unlocking: tuple fields can otherwise tear or be
+        reclaimed by a concurrent setter. The bitmap must be read under the same lock.
+        """
+        undefined = self.temp_name()
+        self.emit_line(f"char {undefined} = 0;")
+        self.emit_line(f"Py_BEGIN_CRITICAL_SECTION({obj});")
+        self.emit_line(f"{dest} = {attr_expr};")
+        if not cl.is_always_defined(attr) or rtype.is_refcounted:
+            self.emit_undefined_attr_check(rtype, dest, "==", obj, attr, cl)
+            self.emit_line(f"{undefined} = 1;")
+            self.emit_line("}")
+        if rtype.is_refcounted and not borrow:
+            self.emit_line(f"if (!{undefined}) {{")
+            self.emit_inc_ref(dest, rtype)
+            self.emit_line("}")
+        self.emit_line("Py_END_CRITICAL_SECTION();")
+        return undefined
+
+    def emit_tuple_attr_store_locked(
+        self,
+        src: str,
+        attr_expr: str,
+        obj: str,
+        rtype: RTuple,
+        cl: ClassIR,
+        attr: str,
+        *,
+        delete: bool = False,
+        is_init: bool = False,
+    ) -> None:
+        """Assign src to a tuple attribute attr under the owner lock.
+
+        References to the old value of the attribute are released outside the lock.
+        """
+        release_old = rtype.is_refcounted and not is_init
+        old = self.temp_name()
+        if release_old:
+            self.emit_line(f"{self.ctype(rtype)} {old};")
+        self.emit_line(f"Py_BEGIN_CRITICAL_SECTION({obj});")
+        if release_old:
+            self.emit_line(f"{old} = {attr_expr};")
+        if delete:
+            self.set_undefined_value(attr_expr, rtype)
+        else:
+            self.emit_line(f"{attr_expr} = {src};")
+        if rtype.error_overlap and not cl.is_always_defined(attr):
+            if delete:
+                self.emit_attr_bitmap_clear(obj, rtype, cl, attr)
+            else:
+                self.emit_attr_bitmap_set(src, obj, rtype, cl, attr)
+        self.emit_line("Py_END_CRITICAL_SECTION();")
+        if release_old:
+            # Refcounted tuple types have a disjoint undefined value (no bitmap).
+            self.emit_line(f"if ({self.error_value_check(rtype, old, '!=')}) {{")
+            self.emit_dec_ref(old, rtype)
+            self.emit_line("}")
+
     def error_value_check(self, rtype: RType, value: str, compare: str) -> str:
         if isinstance(rtype, RTuple):
             return self.tuple_undefined_check_cond(

mypyc/codegen/emitclass.py

@@ -1424,14 +1424,30 @@ def generate_getter(cl: ClassIR, attr: str, rtype: RType, emitter: Emitter) -> N
     #       values is benign.
     always_defined = cl.is_always_defined(attr) and not rtype.is_refcounted
 
+    tuple_snapshot = (
+        emitter.target_python.free_threaded
+        and isinstance(rtype, RTuple)
+        and attr not in cl.final_attributes
+    )
+    if tuple_snapshot:
+        assert isinstance(rtype, RTuple)
+        emitter.emit_line(f"{emitter.ctype(rtype)} tmp;")
+        undefined = emitter.emit_tuple_attr_load_locked("tmp", attr_expr, "self", rtype, cl, attr)
+        attr_expr = "tmp"
     if not always_defined:
-        emitter.emit_undefined_attr_check(rtype, attr_expr, "==", "self", attr, cl, unlikely=True)
+        if tuple_snapshot:
+            emitter.emit_line(f"if (unlikely({undefined})) {{")
+        else:
+            emitter.emit_undefined_attr_check(
+                rtype, attr_expr, "==", "self", attr, cl, unlikely=True
+            )
         emitter.emit_line("PyErr_SetString(PyExc_AttributeError,")
         emitter.emit_line(f'    "attribute {repr(attr)} of {repr(cl.name)} undefined");')
         emitter.emit_line("return NULL;")
         emitter.emit_line("}")
-    emitter.emit_inc_ref(f"self->{attr_field}", rtype)
-    emitter.emit_box(f"self->{attr_field}", "retval", rtype, declare_dest=True)
+    if not tuple_snapshot:
+        emitter.emit_inc_ref(attr_expr, rtype)
+    emitter.emit_box(attr_expr, "retval", rtype, declare_dest=True)
     emitter.emit_line("return retval;")
     emitter.emit_line("}")
 
@@ -1542,17 +1558,25 @@ def emit_decref_old_value() -> None:
     # Take ownership of the incoming value before releasing the old one. In
     # particular, a failed cast must leave the attribute and its reference intact.
     emitter.emit_inc_ref("tmp", rtype)
-    emit_decref_old_value()
-    emitter.emit_line(f"self->{attr_field} = tmp;")
-    if rtype.error_overlap and not always_defined:
-        emitter.emit_attr_bitmap_set("tmp", "self", rtype, cl, attr)
+    if emitter.target_python.free_threaded and isinstance(rtype, RTuple):
+        emitter.emit_tuple_attr_store_locked("tmp", f"self->{attr_field}", "self", rtype, cl, attr)
+    else:
+        emit_decref_old_value()
+        emitter.emit_line(f"self->{attr_field} = tmp;")
+        if rtype.error_overlap and not always_defined:
+            emitter.emit_attr_bitmap_set("tmp", "self", rtype, cl, attr)
 
     if deletable:
         emitter.emit_line("} else {")
-        emit_decref_old_value()
-        emitter.set_undefined_value(f"self->{attr_field}", rtype)
-        if rtype.error_overlap:
-            emitter.emit_attr_bitmap_clear("self", rtype, cl, attr)
+        if emitter.target_python.free_threaded and isinstance(rtype, RTuple):
+            emitter.emit_tuple_attr_store_locked(
+                "", f"self->{attr_field}", "self", rtype, cl, attr, delete=True
+            )
+        else:
+            emit_decref_old_value()
+            emitter.set_undefined_value(f"self->{attr_field}", rtype)
+            if rtype.error_overlap:
+                emitter.emit_attr_bitmap_clear("self", rtype, cl, attr)
         emitter.emit_line("}")
     emitter.emit_line("return 0;")
     emitter.emit_line("}")

mypyc/codegen/emitfunc.py

@@ -526,15 +526,30 @@ def visit_get_attr(self, op: GetAttr) -> None:
         else:
             # Otherwise, use direct or offset struct access.
             attr_expr = self.get_attr_expr(obj, op, decl_cl)
-            use_get_attr_ref = self.emit_load_attr_take_ref(
-                dest, obj, op, cl, attr_rtype, attr_expr
-            )
+            tuple_undefined = None
+            if (
+                self.emitter.target_python.free_threaded
+                and isinstance(attr_rtype, RTuple)
+                and not cl.is_final_attr(op.attr)
+                and not cl.attrs_are_thread_confined()
+            ):
+                tuple_undefined = self.emitter.emit_tuple_attr_load_locked(
+                    dest, attr_expr, obj, attr_rtype, cl, op.attr, borrow=op.is_borrowed
+                )
+                use_get_attr_ref = True
+            else:
+                use_get_attr_ref = self.emit_load_attr_take_ref(
+                    dest, obj, op, cl, attr_rtype, attr_expr
+                )
             always_defined = cl.is_always_defined(op.attr)
             merged_branch = None
             if not always_defined:
-                self.emitter.emit_undefined_attr_check(
-                    attr_rtype, dest, "==", obj, op.attr, cl, unlikely=True
-                )
+                if tuple_undefined is not None:
+                    self.emitter.emit_line(f"if (unlikely({tuple_undefined})) {{")
+                else:
+                    self.emitter.emit_undefined_attr_check(
+                        attr_rtype, dest, "==", obj, op.attr, cl, unlikely=True
+                    )
                 branch = self.next_branch()
                 if branch is not None:
                     if (
@@ -580,6 +595,16 @@ def get_attr_with_allow_error_value(self, op: GetAttr) -> None:
         attr_rtype, decl_cl = cl.attr_details(op.attr)
 
         attr_expr = self.get_attr_expr(obj, op, decl_cl)
+        if (
+            self.emitter.target_python.free_threaded
+            and isinstance(attr_rtype, RTuple)
+            and not cl.is_final_attr(op.attr)
+            and not cl.attrs_are_thread_confined()
+        ):
+            self.emitter.emit_tuple_attr_load_locked(
+                dest, attr_expr, obj, attr_rtype, cl, op.attr, borrow=op.is_borrowed
+            )
+            return
         # On free-threaded builds this takes a new reference atomically (see
         # emit_load_attr_take_ref). CPy_GetAttrRef returns NULL and CPyTagged_GetAttr
         # returns CPY_INT_TAG when the field is undefined, which is precisely the error
@@ -671,6 +696,17 @@ def visit_set_attr(self, op: SetAttr) -> None:
                 )
             if op.error_kind == ERR_FALSE:
                 self.emitter.emit_line(f"{dest} = 1;")
+        elif (
+            self.emitter.target_python.free_threaded
+            and isinstance(attr_rtype, RTuple)
+            and not cl.attrs_are_thread_confined()
+        ):
+            attr_expr = self.get_attr_expr(obj, op, decl_cl)
+            self.emitter.emit_tuple_attr_store_locked(
+                src, attr_expr, obj, attr_rtype, cl, op.attr, is_init=op.is_init
+            )
+            if op.error_kind == ERR_FALSE:
+                self.emitter.emit_line(f"{dest} = 1;")
         else:
             # ...and struct access for normal attributes.
             attr_expr = self.get_attr_expr(obj, op, decl_cl)

mypyc/test-data/run-classes.test

@@ -9904,3 +9904,53 @@ def test_trait() -> None:
     set_k(w, BIG + 1)
     assert get_k(w) == BIG + 1
     assert w.k == BIG + 1
+
+
+[case testFinalTupleAttributeOwnedBorrowedInherited]
+from typing import Final
+
+class Base:
+    def __init__(self, item: object, number: int) -> None:
+        self.refs: Final[tuple[object, tuple[object, int]]] = (item, (item, number))
+        self.numbers: Final[tuple[float, float]] = (1.0, 2.0)
+
+class Child(Base):
+    pass
+
+def owned_base(box: Base) -> tuple[object, tuple[object, int]]:
+    return box.refs
+
+def owned_child(box: Child) -> tuple[object, tuple[object, int]]:
+    return box.refs
+
+def borrowed_base(box: Base) -> object:
+    return box.refs[1][0]
+
+def borrowed_child(box: Child) -> object:
+    return box.refs[1][0]
+
+def read_numbers(box: Child, count: int) -> float:
+    result = 0.0
+    for _ in range(count):
+        result += box.numbers[0] + box.numbers[1]
+    return result
+
+[file driver.py]
+import sys
+from native import Base, Child, owned_base, owned_child, borrowed_base, borrowed_child, read_numbers
+
+item = object()
+number = 1 << 100
+for box in [Base(item, number), Child(item, number)]:
+    before = sys.getrefcount(item)
+    for _ in range(100):
+        assert owned_base(box) == (item, (item, number))
+        assert borrowed_base(box) is item
+        # Interpreted access uses the defining-class Final getset descriptor.
+        assert box.refs == (item, (item, number))
+    assert sys.getrefcount(item) == before
+child = Child(item, number)
+for _ in range(100):
+    assert owned_child(child) == (item, (item, number))
+    assert borrowed_child(child) is item
+assert read_numbers(child, 100) == 300.0