Immutable: Add invariant for debugging (#59)

xFrednet · web-flow · commit 88f13d231328 · 2025-07-17T13:52:19.000+01:00
* Ownership: Add debugging invariant with `--with-ownership-invariant`

* Review comments &lt;3

* sudo CI=green
diff --git a/Include/internal/pycore_interp_structs.h b/Include/internal/pycore_interp_structs.h
@@ -11,6 +11,7 @@ extern "C" {
 #include "pycore_immutability.h"  // struct _immutability_runtime_state
 #include "pycore_llist.h"         // struct llist_node
 #include "pycore_opcode_utils.h"  // NUM_COMMON_CONSTANTS
+#include "pycore_ownership.h"     // struct _Py_ownership_state
 #include "pycore_pymath.h"        // _PY_SHORT_FLOAT_REPR
 #include "pycore_structs.h"       // PyHamtObject
 #include "pycore_tstate.h"        // _PyThreadStateImpl
@@ -937,6 +938,7 @@ struct _is {
     struct _Py_exc_state exc_state;
     struct _Py_immutability_state immutability;
     struct _Py_mem_interp_free_queue mem_free_queue;
+    _Py_ownership_state ownership;
 
     struct ast_state ast;
     struct types_state types;
diff --git a/Include/internal/pycore_ownership.h b/Include/internal/pycore_ownership.h
@@ -0,0 +1,79 @@
+#ifndef Py_INTERNAL_OWNERSHIP_H
+#define Py_INTERNAL_OWNERSHIP_H
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+#ifndef Py_BUILD_CORE
+#  error "Py_BUILD_CORE must be defined to include this header"
+#endif
+
+#include "exports.h"
+
+typedef struct _Py_ownership_state {
+    /* Temporary value until the state always has a field to indicate this. */
+    int is_initialized;
+#ifdef Py_OWNERSHIP_INVARIANT
+    /* Tracks the state of the ownership invariant. Some ownership-related
+     * operations may temporarily violate the invariant. To handle this safely,
+     * the invariant must be suspended during such operations and only resumed
+     * once all of them complete. This is necessary to support re-entrancy.
+     *
+     * For example, during freezing, the object graph is traversed and objects
+     * are marked as immutable — even while they may still reference mutable
+     * objects. If the invariant were enforced mid-way, it would raise a
+     * (premature) error, despite the state being corrected as the operation
+     * completes. To avoid this, the invariant must be paused during the freeze.
+     *
+     * States:
+     *  -1 => The invariant is disabled.
+     *   0 => The invariant is active and enforced.
+     *   N => The invariant is temporarily paused. The value indicates the
+     *        number of suspensions yet to be resumed (this supports nesting).
+     */
+    int invariant_state;
+#endif
+} _Py_ownership_state;
+
+/* This function returns true for C wrappers around functions, types and
+ * all kinds of wrappers around C with immutable state. For ownership these
+ * can be seen as immutable, meaning they can be referenced from immutable
+ * objects and from inside regions.
+ */
+PyAPI_FUNC(int) _PyOwnership_is_c_wrapper(PyObject *obj);
+
+/* This function calls the `visit` function for the fields of the `obj`
+ * which should be effected by ownership. The `data` pointer will be
+ * passed along as the second argument to `visit`.
+ */
+PyAPI_FUNC(int) _PyOwnership_traverse_obj(PyObject *obj, visitproc visit, void *data);
+
+#ifdef Py_OWNERSHIP_INVARIANT
+
+#include "object.h" // PyObject, visitproc
+#include "pytypedefs.h" // PyThreadState
+
+#define Py_OWNERSHIP_INVARIANT_DISABLED -1
+#define Py_OWNERSHIP_INVARIANT_ENABLED 0
+
+/* This function validates that the current heap follows the ownership
+ * rules. This is a slow operation and should only be done for debugging.
+ *
+ * 0 indicates a valid heap, -1 will be returned if an error was thrown.
+ */
+PyAPI_FUNC(int) _PyOwnership_check_invariant(PyThreadState *tstate);
+
+PyAPI_FUNC(int) _PyOwnership_invariant_enable(void);
+PyAPI_FUNC(int) _PyOwnership_invariant_pause(void);
+PyAPI_FUNC(int) _PyOwnership_invariant_resume(void);
+
+#else
+#   define _PyOwnership_invariant_enable() 0 /* success */
+#   define _PyOwnership_invariant_pause() 0 /* success */
+#   define _PyOwnership_invariant_resume() 0 /* success */
+#endif
+
+#ifdef __cplusplus
+}
+#endif
+#endif /* !Py_INTERNAL_OWNERSHIP_H */
diff --git a/Lib/test/test_freeze/test_core.py b/Lib/test/test_freeze/test_core.py
@@ -465,15 +465,6 @@ def test_weakref(self):
         # self.assertTrue(c.val() is obj)
         self.assertIsNone(c.val())
 
-class TestStackCapture(unittest.TestCase):
-     def test_stack_capture(self):
-         import sys
-         x = {}
-         x["frame"] = sys._getframe()
-         freeze(x)
-         self.assertTrue(isfrozen(x))
-         self.assertTrue(isfrozen(x["frame"]))
-
 global_test_dict = 0
 class TestGlobalDictMutation(unittest.TestCase):
     def g():
diff --git a/Lib/test/test_freeze/test_gc.py b/Lib/test/test_freeze/test_gc.py
@@ -1,6 +1,6 @@
 from gc import collect
 import unittest
-from immutable import freeze, NotFreezable, isfrozen
+from immutable import freeze
 
 class GCInteropTest(unittest.TestCase):
   def test_collect(self):
@@ -11,4 +11,4 @@ def test_collect(self):
     # Freeze it
     freeze(a)
     # f
-    collect()
+    collect()
diff --git a/Makefile.pre.in b/Makefile.pre.in
@@ -474,6 +474,7 @@ PYTHON_OBJS=	\
 		Python/optimizer.o \
 		Python/optimizer_analysis.o \
 		Python/optimizer_symbols.o \
+		Python/ownership.o \
 		Python/parking_lot.o \
 		Python/pathconfig.o \
 		Python/preconfig.o \
@@ -1356,6 +1357,7 @@ PYTHON_HEADERS= \
 		$(srcdir)/Include/internal/pycore_opcode_metadata.h \
 		$(srcdir)/Include/internal/pycore_opcode_utils.h \
 		$(srcdir)/Include/internal/pycore_optimizer.h \
+		$(srcdir)/Include/internal/pycore_ownership.h \
 		$(srcdir)/Include/internal/pycore_parking_lot.h \
 		$(srcdir)/Include/internal/pycore_parser.h \
 		$(srcdir)/Include/internal/pycore_pathconfig.h \
diff --git a/PCbuild/_freeze_module.vcxproj b/PCbuild/_freeze_module.vcxproj
@@ -243,6 +243,7 @@
     <ClCompile Include="..\Python\optimizer.c" />
     <ClCompile Include="..\Python\optimizer_analysis.c" />
     <ClCompile Include="..\Python\optimizer_symbols.c" />
+    <ClCompile Include="..\Python\ownership.c" />
     <ClCompile Include="..\Python\parking_lot.c" />
     <ClCompile Include="..\Python\pathconfig.c" />
     <ClCompile Include="..\Python\perf_trampoline.c" />
diff --git a/PCbuild/_freeze_module.vcxproj.filters b/PCbuild/_freeze_module.vcxproj.filters
@@ -328,6 +328,9 @@
     <ClCompile Include="..\Python\optimizer_symbols.c">
       <Filter>Python</Filter>
     </ClCompile>
+    <ClCompile Include="..\Python\ownership.c">
+      <Filter>Source Files</Filter>
+    </ClCompile>
     <ClCompile Include="..\Parser\parser.c">
       <Filter>Source Files</Filter>
     </ClCompile>
diff --git a/PCbuild/pythoncore.vcxproj b/PCbuild/pythoncore.vcxproj
@@ -285,6 +285,7 @@
     <ClInclude Include="..\Include\internal\pycore_obmalloc.h" />
     <ClInclude Include="..\Include\internal\pycore_obmalloc_init.h" />
     <ClInclude Include="..\Include\internal\pycore_optimizer.h" />
+    <ClInclude Include="..\Include\internal\pycore_ownership.h" />
     <ClInclude Include="..\Include\internal\pycore_parking_lot.h" />
     <ClInclude Include="..\Include\internal\pycore_pathconfig.h" />
     <ClInclude Include="..\Include\internal\pycore_pyarena.h" />
@@ -640,6 +641,7 @@
     <ClCompile Include="..\Python\optimizer.c" />
     <ClCompile Include="..\Python\optimizer_analysis.c" />
     <ClCompile Include="..\Python\optimizer_symbols.c" />
+    <ClCompile Include="..\Python\ownership.c" />
     <ClCompile Include="..\Python\parking_lot.c" />
     <ClCompile Include="..\Python\pathconfig.c" />
     <ClCompile Include="..\Python\perf_trampoline.c" />
diff --git a/PCbuild/pythoncore.vcxproj.filters b/PCbuild/pythoncore.vcxproj.filters
@@ -694,6 +694,8 @@
       <Filter>Include\internal</Filter>
     </ClInclude>
     <ClInclude Include="..\Include\internal\pycore_index_pool.h">
+      <Filter>Include\internal</Filter>
+    </ClInclude>
     <ClInclude Include="..\Include\internal\pycore_immutability.h">
       <Filter>Include\internal</Filter>
     </ClInclude>
@@ -1427,6 +1429,8 @@
       <Filter>Python</Filter>
     </ClCompile>
     <ClCompile Include="..\Python\index_pool.c">
+      <Filter>Python</Filter>
+    </ClCompile>
     <ClCompile Include="..\Python\immutability.c">
       <Filter>Python</Filter>
     </ClCompile>
@@ -1478,6 +1482,9 @@
     <ClCompile Include="..\Python\optimizer_symbols.c">
       <Filter>Python</Filter>
     </ClCompile>
+    <ClCompile Include="..\Python\ownership.c">
+      <Filter>Python</Filter>
+    </ClCompile>
     <ClCompile Include="..\Python\parking_lot.c">
       <Filter>Python</Filter>
     </ClCompile>
diff --git a/Python/ceval.c b/Python/ceval.c
@@ -29,6 +29,7 @@
 #include "pycore_opcode_metadata.h" // EXTRA_CASES
 #include "pycore_opcode_utils.h"  // MAKE_FUNCTION_*
 #include "pycore_optimizer.h"     // _PyUOpExecutor_Type
+#include "pycore_ownership.h"     // _PyOwnership_check_invariant
 #include "pycore_pyatomic_ft_wrappers.h" // FT_ATOMIC_*
 #include "pycore_pyerrors.h"      // _PyErr_GetRaisedException()
 #include "pycore_pystate.h"       // _PyInterpreterState_GET()
diff --git a/Python/ceval_gil.c b/Python/ceval_gil.c
@@ -7,7 +7,7 @@
 #include "pycore_pylifecycle.h"   // _PyErr_Print()
 #include "pycore_pystats.h"       // _Py_PrintSpecializationStats()
 #include "pycore_runtime.h"       // _PyRuntime
-
+#include "pycore_ownership.h"     // _PyOwnership_check_invariant
 
 /*
    Notes about the implementation:
@@ -1367,6 +1367,13 @@ _Py_HandlePending(PyThreadState *tstate)
         _PyThreadState_Attach(tstate);
     }
 
+#ifdef Py_OWNERSHIP_INVARIANT
+    /* Check the region invariant if required. */
+    if (_PyOwnership_check_invariant(tstate) != 0) {
+        return -1;
+    }
+#endif
+
     /* Pending signals */
     if ((breaker & _PY_SIGNALS_PENDING_BIT) != 0) {
         if (handle_signals(tstate) != 0) {
diff --git a/Python/ceval_macros.h b/Python/ceval_macros.h
@@ -150,13 +150,26 @@ do { \
 
 
 /* Do interpreter dispatch accounting for tracing and instrumentation */
-#define DISPATCH() \
+#ifdef Py_OWNERSHIP_INVARIANT
+#  define DISPATCH() \
+    { \
+        if (_PyOwnership_check_invariant(tstate) != 0) { \
+            JUMP_TO_LABEL(error); \
+        } \
+        assert(frame->stackpointer == NULL); \
+        NEXTOPARG(); \
+        PRE_DISPATCH_GOTO(); \
+        DISPATCH_GOTO(); \
+    }
+#else
+#  define DISPATCH() \
     { \
         assert(frame->stackpointer == NULL); \
         NEXTOPARG(); \
         PRE_DISPATCH_GOTO(); \
         DISPATCH_GOTO(); \
     }
+#endif
 
 #define DISPATCH_SAME_OPARG() \
     { \
diff --git a/Python/immutability.c b/Python/immutability.c
@@ -5,8 +5,9 @@
 #include <stdio.h>
 #include "pycore_descrobject.h"
 #include "pycore_gc.h"
-#include "pycore_object.h"
 #include "pycore_immutability.h"
+#include "pycore_object.h"
+#include "pycore_ownership.h"
 #include "pycore_list.h"
 
 
@@ -152,10 +153,6 @@ static PyObject* pop(PyObject* s){
     return item;
 }
 
-static bool is_c_wrapper(PyObject* obj){
-    return PyCFunction_Check(obj) || Py_IS_TYPE(obj, &_PyMethodWrapper_Type) || Py_IS_TYPE(obj, &PyWrapperDescr_Type);
-}
-
 // Lifted from Python/gc.c
 //******************************** */
 #ifndef Py_GIL_DISABLED
@@ -775,7 +772,7 @@ int _Py_DecRef_Immutable(PyObject *op)
 
 int traverse_freeze(PyObject* obj, PyObject* dfs)
 {
-    if(is_c_wrapper(obj)) {
+    if(_PyOwnership_is_c_wrapper(obj)) {
         // C functions are not mutable
         // Types are manually traversed
         return 0;
@@ -788,33 +785,9 @@ int traverse_freeze(PyObject* obj, PyObject* dfs)
         SUCCEEDS(shadow_function_globals(obj));
     }
 
-    if(PyType_Check(obj)){
-        // TODO(Immutable): mjp: Special case for types not sure if required. We should review.
-        PyTypeObject* type = (PyTypeObject*)obj;
-
-        SUCCEEDS(freeze_visit(type->tp_dict, dfs));
-        SUCCEEDS(freeze_visit(type->tp_mro, dfs));
-        // We need to freeze the tuple object, even though the types
-        // within will have been frozen already.
-        SUCCEEDS(freeze_visit(type->tp_bases, dfs));
-    }
-    else
-    {
-        traverseproc traverse = Py_TYPE(obj)->tp_traverse;
-        if(traverse != NULL){
-            SUCCEEDS(traverse(obj, (visitproc)freeze_visit, dfs));
-        }
-    }
-
-    // The default tp_traverse will not visit the type object if it is
-    // not heap allocated, so we need to do that manually here to freeze
-    // the statically allocated types that are reachable.
-    if (!(Py_TYPE(obj)->tp_flags & Py_TPFLAGS_HEAPTYPE)) {
-        SUCCEEDS(freeze_visit(_PyObject_CAST(Py_TYPE(obj)), dfs));
-    }
+    SUCCEEDS(_PyOwnership_traverse_obj(obj, (visitproc)freeze_visit, (void*)dfs));
 
     return 0;
-
 error:
     return -1;
 }
@@ -827,6 +800,21 @@ int _PyImmutability_Freeze(PyObject* obj)
     }
     int result = 0;
 
+#ifdef Py_DEBUG
+    // This has to be declared early to support the `Py_XDECREF` if any of the
+    // `SUCCEEDS` fails
+    PyObject* freeze_location = NULL;
+#endif
+
+    // Enable the invariant. It has to be enabled at the beginning to allow
+    // reentry and failure in internal calls.
+    SUCCEEDS(_PyOwnership_invariant_enable());
+    // This function incrementally marks new objects as frozen. During this
+    // process it is possible that frozen objects point to mutable ones. This
+    // therefore needs to pause the invariant. Otherwise we might get an
+    // exception when freezing calls into Python and triggers the invariant.
+    SUCCEEDS(_PyOwnership_invariant_pause());
+
     struct FreezeState freeze_state;
     // Initialize the freeze state
     SUCCEEDS(init_freeze_state(&freeze_state));
@@ -836,9 +824,7 @@ int _PyImmutability_Freeze(PyObject* obj)
         goto error;
     }
 
-
 #ifdef Py_DEBUG
-    PyObject* freeze_location = NULL;
     // In debug mode, we can set a freeze location for debugging purposes.
     // Get a traceback object to use as the freeze location.
     if (state->traceback_func == NULL) {
@@ -900,5 +886,11 @@ int _PyImmutability_Freeze(PyObject* obj)
 #ifdef Py_DEBUG
     Py_XDECREF(freeze_location);
 #endif
+    // Indicate that this funciton no longer requires the invariant to be paused.
+    // This can't use the `SUCCEEDS` macro, since that one would jump to the
+    // `error` label above.
+    if (_PyOwnership_invariant_resume() != 0) {
+        result = -1;
+    }
     return result;
 }
diff --git a/Python/ownership.c b/Python/ownership.c
diff --git a/Python/pystate.c b/Python/pystate.c
diff --git a/configure b/configure
diff --git a/configure.ac b/configure.ac
diff --git a/pyconfig.h.in b/pyconfig.h.in
