[WIP] Add basic 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/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,58 @@
+import unittest
+
+from immutable import NotFreezable, freeze, isfrozen
+
+
+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 = NotFreezable()
+
+            def __pre_freeze__(self):
+                self.hook_calls += 1
+
+        obj = C()
+
+        with self.assertRaises(TypeError):
+            freeze(obj)
+        with self.assertRaises(TypeError):
+            freeze(obj)
+        with self.assertRaises(TypeError):
+            freeze(obj)
+
+        self.assertEqual(obj.hook_calls, 1)
+        self.assertFalse(isfrozen(obj))
+
+
+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

@@ -1212,6 +1212,115 @@ func_reachable(PyObject *self, visitproc visit, void *arg)
     return func_traverse(self, visit, arg);
 }
 
+/**
+ * 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",
@@ -1254,6 +1363,7 @@ PyTypeObject PyFunction_Type = {
     0,                                          /* tp_alloc */
     func_new,                                   /* tp_new */
     .tp_reachable = func_reachable,
+    .tp_prefreeze = func_prefreeze_shadow_captures,
 };
 
 

Objects/moduleobject.c

@@ -864,54 +864,6 @@ _PyModule_ClearDict(PyObject *d)
 
 }
 
-int _Py_module_freeze_hook(PyObject *self) {
-    // Use cast, since we want this exact object
-    PyModuleObject *m = _PyModule_CAST(self);
-
-    if (m->md_frozen) {
-        return 0;
-    }
-
-    // Get the interpreter state early to make error handling easy
-    PyInterpreterState* ip = PyInterpreterState_Get();
-    if (ip == NULL) {
-        PyErr_Format(PyExc_RuntimeError, "Well, this is a problem", Py_None);
-        return -1;
-    }
-
-    // Create a new module module
-    PyModuleObject *mut_state = new_module_notrack(&PyModule_Type);
-    if (mut_state == NULL) {
-        PyErr_NoMemory();
-        return -1;
-    }
-    track_module(mut_state);
-
-    // Insert our mutable module into `sys.mut_modules`
-    if (PyDict_SetItem(ip->mutable_modules, m->md_name, _PyObject_CAST(mut_state))) {
-        // Make sure failure keeps self intact
-        Py_DECREF(mut_state);
-        return -1;
-    }
-
-    // Copy mutable state
-    mut_state->md_name = Py_NewRef(m->md_name);
-    mut_state->md_dict = m->md_dict;
-    mut_state->md_def = m->md_def;
-    mut_state->md_state = m->md_state;
-    mut_state->md_weaklist = m->md_weaklist;
-
-    // Clear the state to freeze the module
-    m->md_dict = NULL;
-    m->md_def = NULL;
-    m->md_state = NULL;
-    m->md_weaklist = NULL;
-    m->md_frozen = true;
-    m->ob_base.ob_type = &_PyImmModule_Type;
-
-    return 0;
-}
-
 /*[clinic input]
 class module "PyModuleObject *" "&PyModule_Type"
 [clinic start generated code]*/
@@ -1553,6 +1505,61 @@ module_reachable(PyObject *self, visitproc visit, void *arg)
     return module_traverse(self, visit, arg);
 }
 
+int module_make_immutable_proxy(PyObject *self) {
+    // Use cast, since we want this exact object
+    PyModuleObject *m = _PyModule_CAST(self);
+
+    if (m->md_frozen) {
+        return 0;
+    }
+
+    // Get the interpreter state early to make error handling easy
+    PyInterpreterState* ip = PyInterpreterState_Get();
+    if (ip == NULL) {
+        PyErr_Format(PyExc_RuntimeError, "Well, this is a problem", Py_None);
+        return -1;
+    }
+
+    // Create a new module module
+    PyModuleObject *mut_state = new_module_notrack(&PyModule_Type);
+    if (mut_state == NULL) {
+        PyErr_NoMemory();
+        return -1;
+    }
+    track_module(mut_state);
+
+    // Insert our mutable module into `sys.mut_modules`
+    if (PyDict_SetItem(ip->mutable_modules, m->md_name, _PyObject_CAST(mut_state))) {
+        // Make sure failure keeps self intact
+        Py_DECREF(mut_state);
+        return -1;
+    }
+
+    // Copy mutable state
+    mut_state->md_name = Py_NewRef(m->md_name);
+    mut_state->md_dict = m->md_dict;
+    mut_state->md_def = m->md_def;
+    mut_state->md_state = m->md_state;
+    mut_state->md_weaklist = m->md_weaklist;
+
+    // Clear the state to freeze the module
+    m->md_dict = NULL;
+    m->md_def = NULL;
+    m->md_state = NULL;
+    m->md_weaklist = NULL;
+    m->md_frozen = true;
+    m->ob_base.ob_type = &_PyImmModule_Type;
+
+    return 0;
+}
+
+int module_prefreeze(PyObject *self) {
+    // TODO(immutable): Check if the module defines a custom pre-freeze hook:
+
+    // TODO(immutable): Check if the module wants to be a proxy first:
+    return module_make_immutable_proxy(self);
+}
+
 PyTypeObject PyModule_Type = {
     PyVarObject_HEAD_INIT(&PyType_Type, 0)
     "module",                                   /* tp_name */
@@ -1595,6 +1602,7 @@ PyTypeObject PyModule_Type = {
     new_module,                                 /* tp_new */
     PyObject_GC_Del,                            /* tp_free */
     .tp_reachable = module_reachable,
+    .tp_prefreeze = module_prefreeze,
 };
 
 PyTypeObject _PyImmModule_Type = {

Objects/weakrefobject.c

@@ -158,6 +158,8 @@ static int
 weakref_reachable(PyObject *op, visitproc visit, void *arg)
 {
     Py_VISIT(_PyObject_CAST(Py_TYPE(op)));
+    // FIXME(immutable): Currently we manually visit the weak references object.
+    // we might want to do this here to remove (some) special handling.
     return gc_traverse(op, visit, arg);
 }