Add pre-freeze hook

Include/cpython/object.h

@@ -242,6 +242,9 @@ struct _typeobject {
 
     /* call function for all referenced objects (includes non-cyclic refs) */
     traverseproc tp_reachable;
+
+    /* A callback called before a type is frozen. */
+    prefreezeproc tp_prefreeze;
 };
 
 #define _Py_ATTR_CACHE_UNUSED (30000)  // (see tp_versions_used)

Include/internal/pycore_gc.h

@@ -205,6 +205,16 @@ static inline void _PyGC_CLEAR_FINALIZED(PyObject *op) {
 #endif
 }
 
+static inline void _PyGC_CLEAR_COLLECTING(PyObject *op) {
+#ifdef Py_GIL_DISABLED
+    // TODO(immutable): Does NoGil have a collecting flag? If so, how do we
+    // clear it?
+#else
+    PyGC_Head *gc = _Py_AS_GC(op);
+    gc->_gc_prev &= ~_PyGC_PREV_MASK_COLLECTING;
+#endif
+}
+
 
 /* Tell the GC to track this object.
  *

Include/internal/pycore_immutability.h

@@ -23,6 +23,9 @@ struct _Py_immutability_state {
     _Py_hashtable_t *shallow_immutable_types;
     PyObject *destroy_cb;
     _Py_hashtable_t *warned_types;
+    // With the pre-freeze hook it can happen that freeze calls are
+    // nested. This is stack of the enclosing freeze states.
+    struct FreezeState *freeze_stack;
 #ifdef Py_DEBUG
     PyObject *traceback_func;  // For debugging purposes, can be NULL
 #endif

Include/object.h

@@ -358,6 +358,7 @@ typedef int(*objobjargproc)(PyObject *, PyObject *, PyObject *);
 typedef int (*objobjproc)(PyObject *, PyObject *);
 typedef int (*visitproc)(PyObject *, void *);
 typedef int (*traverseproc)(PyObject *, visitproc, void *);
+typedef int (*prefreezeproc)(PyObject *);
 
 
 typedef void (*freefunc)(void *);
@@ -641,6 +642,7 @@ given type object has a specified feature.
 #if defined(Py_GIL_DISABLED) && defined(Py_DEBUG)
 #define _Py_TYPE_REVEALED_FLAG (1 << 3)
 #endif
+#define _Py_PREFREEZE_RAN_FLAG (1 << 8)
 
 #define Py_CONSTANT_NONE 0
 #define Py_CONSTANT_FALSE 1

Lib/test/test_freeze/test_prefreeze.py

@@ -0,0 +1,156 @@
+import unittest
+
+from immutable import freeze, isfrozen, set_freezable, FREEZABLE_NO
+
+
+class TestPreFreezeHook(unittest.TestCase):
+    def test_prefreeze_hook_is_called(self):
+        class C:
+            def __init__(self):
+                self.hook_calls = 0
+
+            def __pre_freeze__(self):
+                self.hook_calls += 1
+
+        obj = C()
+        freeze(obj)
+
+        self.assertEqual(obj.hook_calls, 1)
+        self.assertTrue(isfrozen(obj))
+
+    def test_prefreeze_hook_runs_before_object_is_frozen(self):
+        class C:
+            def __init__(self):
+                self.was_frozen_inside_hook = None
+
+            def __pre_freeze__(self):
+                self.was_frozen_inside_hook = isfrozen(obj)
+
+        obj = C()
+        freeze(obj)
+
+        self.assertIs(obj.was_frozen_inside_hook, False)
+        self.assertTrue(isfrozen(obj))
+
+    def test_prefreeze_hook_remains_called_after_failure(self):
+        class C:
+            def __init__(self):
+                self.hook_calls = 0
+                self.child = {}
+                set_freezable(self.child, FREEZABLE_NO)
+
+            def __pre_freeze__(self):
+                self.hook_calls += 1
+
+        obj = C()
+
+        with self.assertRaises(TypeError):
+            freeze(obj)
+        with self.assertRaises(TypeError):
+            freeze(obj)
+
+        self.assertEqual(obj.hook_calls, 1)
+        self.assertFalse(isfrozen(obj))
+
+    def test_nested_freeze(self):
+        class A:
+            def __init__(self, field):
+                self.field = field
+            def __pre_freeze__(self):
+                freeze(self.field)
+
+        a = A(A(None))
+
+        # Freezing A should succeed even with nested `freeze()` calls
+        freeze(a)
+
+        # Objects frozen by nested freeze calls should remain frozen
+        self.assertTrue(isfrozen(a))
+        self.assertTrue(isfrozen(a.field))
+        self.assertTrue(isfrozen(a.field.field))
+
+    def test_nested_cycle(self):
+        class A:
+            def __init__(self, next):
+                self.next = next
+            def __pre_freeze__(self):
+                freeze(self.next)
+
+        # Create a cycle of pre-freezes
+        a = A(None)
+        b = A(a)
+        c = A(b)
+        d = A(c)
+        e = A(d)
+        a.next = e
+
+        # Freezing should succeed even with the cycle of pre-freezes
+        freeze(a)
+
+        # Check the objects are frozen
+        self.assertTrue(isfrozen(a))
+        self.assertTrue(isfrozen(b))
+        self.assertTrue(isfrozen(c))
+        self.assertTrue(isfrozen(d))
+        self.assertTrue(isfrozen(e))
+
+    def test_nested_freeze_stays_frozen_on_fail(self):
+        class A:
+            def __init__(self):
+                self.freezable = {}
+                self.unfreezable = {}
+                set_freezable(self.unfreezable, FREEZABLE_NO)
+
+            def __pre_freeze__(self):
+                freeze(self.freezable)
+
+        a = A()
+
+        # Freezing A should succeed even with nested `freeze()` calls
+        with self.assertRaises(TypeError):
+            freeze(a)
+
+        # Objects frozen by nested freeze calls should remain frozen
+        self.assertFalse(isfrozen(a))
+        self.assertTrue(isfrozen(a.freezable))
+
+    def test_pre_freeze_can_stop_freezing(self):
+        class A:
+            def __init__(self, fail):
+                self.fail = fail
+            def __pre_freeze__(self):
+                if self.fail:
+                    raise ValueError(2)
+
+        # This should fail, since the pre-freeze throws a ValueError
+        a = A(True)
+        with self.assertRaises(ValueError):
+            freeze(a)
+        self.assertFalse(isfrozen(a))
+
+        # This should succeed, since the pre-freeze succeeds
+        a = A(False)
+        freeze(a)
+        self.assertTrue(isfrozen(a))
+
+    def test_pre_freeze_self(self):
+        class A:
+            def __pre_freeze__(self):
+                freeze(self)
+
+        # This assumes that list items are visited in order
+        a = A()
+        b = A()
+        set_freezable(b, FREEZABLE_NO)
+        lst = [a, b]
+
+        # Freezing lst will fail due to b
+        with self.assertRaises(TypeError):
+            freeze(lst)
+
+        # a should remain frozen due to its pre-freeze
+        self.assertTrue(isfrozen(a))
+        self.assertFalse(isfrozen(b))
+
+if __name__ == "__main__":
+    unittest.main()

Lib/test/test_sys.py

@@ -1782,7 +1782,7 @@ def delx(self): del self.__x
         check((1,2,3), vsize('') + self.P + 3*self.P)
         # type
         # static type: PyTypeObject
-        fmt = 'P2nPI13Pl4Pn9Pn12PIPcP'
+        fmt = 'P2nPI13Pl4Pn9Pn12PIPcPP'
         s = vsize(fmt)
         check(int, s)
         typeid = 'n' if support.Py_GIL_DISABLED else ''

Objects/funcobject.c

@@ -1205,6 +1205,115 @@ func_descr_get(PyObject *func, PyObject *obj, PyObject *type)
     return PyMethod_New(func, obj);
 }
 
+/**
+ * Special function for replacing globals and builtins with a copy of just what they use.
+ *
+ * This is necessary because the function object has a pointer to the global
+ * dictionary, and this is problematic because freezing any function directly
+ * (as we do with other objects) would make all globals immutable.
+ *
+ * Instead, we walk the function and find any places where it references
+ * global variables or builtins, and then freeze just those objects. The globals
+ * and builtins dictionaries for the function are then replaced with
+ * copies containing just those globals and builtins we were able to determine
+ * the function uses.
+ */
+static int func_prefreeze_shadow_captures(PyObject* op)
+{
+    _PyObject_ASSERT(op, PyFunction_Check(op));
+
+    PyObject* shadow_builtins = NULL;
+    PyObject* shadow_globals = NULL;
+    PyObject* shadow_closure = NULL;
+    Py_ssize_t size;
+
+    PyFunctionObject* f = _PyFunction_CAST(op);
+    PyObject* globals = f->func_globals;
+    PyObject* builtins = f->func_builtins;
+
+    shadow_builtins = PyDict_New();
+    if(shadow_builtins == NULL){
+        goto nomemory;
+    }
+
+    shadow_globals = PyDict_New();
+    if(shadow_globals == NULL){
+        goto nomemory;
+    }
+
+    if (PyDict_SetItemString(shadow_globals, "__builtins__", Py_NewRef(shadow_builtins))) {
+        goto error;
+    }
+
+    _PyObject_ASSERT(f->func_code, PyCode_Check(f->func_code));
+    PyCodeObject* f_code = (PyCodeObject*)f->func_code;
+
+    size = 0;
+    if (f_code->co_names != NULL)
+        size = PySequence_Fast_GET_SIZE(f_code->co_names);
+    for (Py_ssize_t i = 0; i < size; i++) {
+        PyObject* name = PySequence_Fast_GET_ITEM(f_code->co_names, i);
+
+        if( PyDict_Contains(globals, name)){
+            PyObject* value = PyDict_GetItem(globals, name);
+            if (PyDict_SetItem(shadow_globals, Py_NewRef(name), Py_NewRef(value))) {
+                goto error;
+            }
+        } else if (PyDict_Contains(builtins, name)) {
+            PyObject* value = PyDict_GetItem(builtins, name);
+            if (PyDict_SetItem(shadow_builtins, Py_NewRef(name), Py_NewRef(value))) {
+                goto error;
+            }
+        }
+    }
+
+    // Shadow cells with a new frozen cell to warn on reassignments in the
+    // capturing function.
+    size = 0;
+    if(f->func_closure != NULL) {
+        size = PyTuple_Size(f->func_closure);
+        if (size == -1) {
+            goto error;
+        }
+        shadow_closure = PyTuple_New(size);
+        if (shadow_closure == NULL) {
+            goto error;
+        }
+    }
+    for(Py_ssize_t i=0; i < size; ++i){
+        PyObject* cellvar = PyTuple_GET_ITEM(f->func_closure, i);
+        PyObject* value = PyCell_GET(cellvar);
+
+        PyObject* shadow_cellvar = PyCell_New(value);
+        if(PyTuple_SetItem(shadow_closure, i, shadow_cellvar) == -1){
+            goto error;
+        }
+    }
+
+    if (f->func_annotations == NULL) {
+        PyObject* new_annotations = PyDict_New();
+        if (new_annotations == NULL) {
+            goto nomemory;
+        }
+        f->func_annotations = new_annotations;
+    }
+
+    // Only assign them at the end when everything succeeded
+    Py_XSETREF(f->func_closure, shadow_closure);
+    Py_SETREF(f->func_globals, shadow_globals);
+    Py_SETREF(f->func_builtins, shadow_builtins);
+
+    return 0;
+
+nomemory:
+    PyErr_NoMemory();
+error:
+    Py_XDECREF(shadow_builtins);
+    Py_XDECREF(shadow_globals);
+    Py_XDECREF(shadow_closure);
+    return -1;
+}
+
 PyTypeObject PyFunction_Type = {
     PyVarObject_HEAD_INIT(&PyType_Type, 0)
     "function",
@@ -1247,6 +1356,7 @@ PyTypeObject PyFunction_Type = {
     0,                                          /* tp_alloc */
     func_new,                                   /* tp_new */
     .tp_reachable = _PyObject_ReachableVisitTypeAndTraverse,
+    .tp_prefreeze = func_prefreeze_shadow_captures,
 };