TRegions: Support Region References.

Include/cpython/weakrefobject.h

@@ -2,6 +2,18 @@
 #  error "this header file must not be included directly"
 #endif
 
+/* A region reference is a weak reference that survives its target's region
+ * being closed. Instead of keeping the region open it checks on every
+ * dereference whether this interpreter may reach the target, and opens the
+ * region tree if it may. The metadata carrying that information lives in
+ * `pycore_regionref.h`; it is opaque here.
+ */
+PyAPI_DATA(PyTypeObject) _PyRegionref_RefType;
+
+#define _PyRegionRef_CheckExact(op) Py_IS_TYPE((op), &_PyRegionref_RefType)
+
+struct _PyRegionRefMetadata;
+
 /* PyWeakReference is the base struct for the Python ReferenceType, ProxyType,
  * and CallableProxyType.
  */
@@ -43,12 +55,23 @@ struct _PyWeakReference {
      */
     PyMutex *weakrefs_lock;
 #endif
+
+    /* The ownership domain of `wr_object`, or NULL if this object doesn't have an ownership
+     * domain. This can happen if this is a normal weakref or if the object is immutable.
+     */
+    struct _PyRegionRefMetadata *region_ref;
 };
 
 PyAPI_FUNC(void) _PyWeakref_ClearRef(PyWeakReference *self);
 
+/* Region references reuse this struct but are deliberately not a subtype of
+ * `_PyWeakref_RefType`, so that `PyWeakref_Check()` stays false for them and
+ * the region close trace does not follow them. */
+#define _PyWeakrefOrRegionRef_Check(op) \
+    (PyWeakref_Check(op) || _PyRegionRef_CheckExact(op))
+
 #define _PyWeakref_CAST(op) \
-    (assert(PyWeakref_Check(op)), _Py_CAST(PyWeakReference*, (op)))
+    (assert(_PyWeakrefOrRegionRef_Check(op)), _Py_CAST(PyWeakReference*, (op)))
 
 // Test if a weak reference is dead.
 PyAPI_FUNC(int) PyWeakref_IsDead(PyObject *ref);

Include/internal/pycore_cown.h

@@ -22,6 +22,16 @@ typedef uint64_t _PyCown_thread_id_t;
 PyAPI_FUNC(_PyCown_ipid_t) _PyCown_ThisInterpreterId(void);
 PyAPI_FUNC(_PyCown_thread_id_t) _PyCown_ThisThreadId(void);
 
+/* The interpreter currently owning the cown, or `_PyCown_ReleasedIpid()` when
+ * no interpreter does. Safe to call from any interpreter. */
+PyAPI_FUNC(_PyCown_ipid_t) _PyCown_Owner(PyObject *cown);
+PyAPI_FUNC(_PyCown_ipid_t) _PyCown_ReleasedIpid(void);
+
+/* The thread that acquired the cown, or `_PyCown_UnsetThreadId()` when it was
+ * acquired without the GIL. Not enforced, only reported. */
+PyAPI_FUNC(_PyCown_thread_id_t) _PyCown_LockingThread(PyObject *cown);
+PyAPI_FUNC(_PyCown_thread_id_t) _PyCown_UnsetThreadId(void);
+
 
 #ifdef __cplusplus
 }

Include/internal/pycore_immutability.h

@@ -8,9 +8,32 @@ extern "C" {
 #  error "Py_BUILD_CORE must be defined to include this header"
 #endif
 
+struct _PyRegionRefMetadata;
+
 PyAPI_DATA(PyTypeObject) _PyTracingRegion_Type;
 PyAPI_FUNC(int) _PyTracingRegion_Close(PyObject* region);
 PyAPI_FUNC(int) _PyTracingRegion_IsClosed(PyObject* region);
+PyAPI_FUNC(void) _PyTracingRegion_Open(PyObject* region);
+
+/* Returns the region's metadata node, allocating it if this is the first
+ * region reference the current close has found. Borrowed, and only valid while
+ * the region stays closed. The caller must hold `_PyWeakref_Lock`. */
+PyAPI_FUNC(struct _PyRegionRefMetadata*)
+    _PyTracingRegion_MetaLockHeld(PyObject* region);
+
+/* Hands a closed region's node to `cown`, used when a cown takes ownership of
+ * the region. Does nothing for an open region. */
+PyAPI_FUNC(void) _PyTracingRegion_SetMetaCown(PyObject* region, PyObject* cown);
+
+/* Records who owns a closed region, used when it leaves the cown that owned it.
+ * The owner is NOT necessarily the calling interpreter: a cown is immutable and
+ * may be deallocated by anyone holding a reference, including an interpreter
+ * that never owned it. Pass `_PyCown_ReleasedIpid()` when nobody owns it.
+ * Does nothing for an open region. */
+// FIXME: The deallocation will be fixed in a follow-up, then we can remove the
+// owner argument and assert that it's always local.
+PyAPI_FUNC(void) _PyTracingRegion_SetMetaOwner(
+    PyObject* region, uint64_t owner);
 
 struct _Py_immutability_state {
     int late_init_done;

Include/internal/pycore_regionref.h

@@ -0,0 +1,86 @@
+#ifndef Py_INTERNAL_REGIONREF_H
+#define Py_INTERNAL_REGIONREF_H
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+#ifndef Py_BUILD_CORE
+#  error "Py_BUILD_CORE must be defined to include this header"
+#endif
+
+#include "pycore_cown.h"          // _PyCown_ipid_t
+#include "pycore_hashtable.h"     // _Py_hashtable_t
+
+/* Every `RegionRef` points at a `_PyRegionRefMetadata` node, and those nodes
+ * form a tree that mirrors the region hierarchy: a nested region's node
+ * delegates to the node of its parent, and the outermost node names the owner,
+ * either an interpreter or the cown holding the region. A cown node looks the
+ * owner up on the cown, which is what makes acquiring and releasing free;
+ * moving a closed region between owners restamps a single node.
+ *
+ * A reference to an object that is in no region at all owns its own node.
+ */
+
+typedef enum {
+    /* A close is in progress. The referenced data is unavailable but should be
+     * soon; dereferencing fails and the caller may retry. Only ever reachable
+     * through a child's `parent`, never directly from a reference. */
+    _Py_REGION_REF_WIP,
+    /* Delegates to `value.parent`, the node of the enclosing region. */
+    _Py_REGION_REF_META,
+    /* Terminal. The region is held by `value.cown`, on which the owner is
+     * looked up dynamically. */
+    _Py_REGION_REF_COWN,
+    /* Terminal, owned by one interpreter. */
+    _Py_REGION_REF_IPID,
+} _PyRegionRefKind;
+
+typedef struct _PyRegionRefMetadata {
+    uint32_t rc;
+    /* A `_PyRegionRefKind`. */
+    uint8_t kind;
+    /* Borrowed. Set only on region nodes, and only while that region is
+     * closed. Names the region a dereference has to open on its way down.
+     *
+     * Borrowing is safe because the pointer is only followed after the terminal
+     * check established that this interpreter owns the region, and a region can
+     * only be deallocated by its owner. */
+    PyObject *region;
+    union {
+        struct _PyRegionRefMetadata *parent;  /* META */
+        PyObject *cown;                       /* COWN, borrowed */
+        _PyCown_ipid_t ipid;                  /* IPID */
+    } value;
+} _PyRegionRefMetadata;
+
+/* Creates the node of a region that is being closed. Returns a new reference. */
+extern _PyRegionRefMetadata *_PyRegionRef_NewRegionMetaLockHeld(PyObject *region);
+
+extern void _PyRegionRef_MetaDecref(_PyRegionRefMetadata *meta);
+
+// Ownership transitions.
+extern void _PyRegionRef_MetaSetParentLockHeld(_PyRegionRefMetadata *meta,
+                                               _PyRegionRefMetadata *parent);
+/* Hands the node to `cown`, which is borrowed. The owner is from then on
+ * whoever holds the cown. */
+extern void _PyRegionRef_MetaSetCown(_PyRegionRefMetadata *meta, PyObject *cown);
+/* Stamps an explicit owner, which need not be the current interpreter and may
+ * be `_PyCown_ReleasedIpid()` to mean nobody owns the region. */
+extern void _PyRegionRef_MetaSetIpid(_PyRegionRefMetadata *meta,
+                                     _PyCown_ipid_t ipid);
+extern void _PyRegionRef_MetaRegionOpened(_PyRegionRefMetadata *meta);
+extern void _PyRegionRef_MetaResolveWip(_PyRegionRefMetadata *meta);
+
+/* Re-homes every `RegionRef` pointing at `obj` onto `region`'s node, allocating
+ * that node if this is the first reference the close has found. Every other
+ * weak reference to `obj` is cleared unless it is listed in `keep`.
+ *
+ * This is the region close hook; it replaces `_PyWeakref_ClearWeakRefsExcept()`
+ * for objects that are being closed into a region. */
+extern void _PyRegionRef_CloseWeakRefs(PyObject *obj, _Py_hashtable_t *keep,
+                                       PyObject *region);
+
+#ifdef __cplusplus
+}
+#endif
+#endif /* !Py_INTERNAL_REGIONREF_H */

Include/internal/pycore_weakref.h

@@ -13,6 +13,15 @@ extern "C" {
 #include "pycore_object.h"           // _Py_REF_IS_MERGED()
 #include "pycore_pyatomic_ft_wrappers.h"
 
+/* Guards weakrefs to immutable objects, and all region reference metadata.
+ * Declared for both builds because `_PyRegionRefMetadata` uses it either way,
+ * while the weakref lists themselves are striped in free-threaded builds. */
+extern PyMutex _PyWeakref_Lock;
+
+#define LOCK_REGION_REF_META() \
+    PyMutex_LockFlags(&_PyWeakref_Lock, _Py_LOCK_DONT_DETACH)
+#define UNLOCK_REGION_REF_META() PyMutex_Unlock(&_PyWeakref_Lock)
+
 #ifdef Py_GIL_DISABLED
 
 #define WEAKREF_LIST_LOCK(obj) \
@@ -37,9 +46,6 @@ extern "C" {
 
 #else
 
-// Lock used for weakrefs to immutable objects
-extern PyMutex _PyWeakref_Lock;
-
 #define LOCK_WEAKREFS(obj) PyMutex_LockFlags(&_PyWeakref_Lock, _Py_LOCK_DONT_DETACH)
 #define UNLOCK_WEAKREFS(obj) PyMutex_Unlock(&_PyWeakref_Lock)
 
@@ -108,7 +114,7 @@ static inline PyObject* get_ref_lock_held(PyWeakReference *ref, PyObject *obj)
 
 static inline PyObject* _PyWeakref_GET_REF(PyObject *ref_obj)
 {
-    assert(PyWeakref_Check(ref_obj));
+    assert(_PyWeakrefOrRegionRef_Check(ref_obj));
     PyWeakReference *ref = _Py_CAST(PyWeakReference*, ref_obj);
 
     PyObject *obj = _Py_atomic_load_ptr(&ref->wr_object);
@@ -125,7 +131,7 @@ static inline PyObject* _PyWeakref_GET_REF(PyObject *ref_obj)
 
 static inline int _PyWeakref_IS_DEAD(PyObject *ref_obj)
 {
-    assert(PyWeakref_Check(ref_obj));
+    assert(_PyWeakrefOrRegionRef_Check(ref_obj));
     int ret = 0;
     PyWeakReference *ref = _Py_CAST(PyWeakReference*, ref_obj);
     PyObject *obj = FT_ATOMIC_LOAD_PTR(ref->wr_object);

Lib/immutable.py

@@ -23,6 +23,7 @@
 SharedField = _c.SharedField
 TracingRegion = _c.TracingRegion
 Cown = _c.Cown
+RegionRef = _c.RegionRef
 
 # FIXME(immutable): For the longest time we used the name `isfrozen`
 # without the underscore. This keeps the function name for now, but
@@ -141,6 +142,9 @@ def __enter__(self):
     "FREEZABLE_PROXY",
     "InterpreterLocal",
     "SharedField",
+    "TracingRegion",
+    "Cown",
+    "RegionRef",
     "freezable",
     "unfreezable",
     "explicitlyFreezable",