Modifying the semantics of function freezing

Signed-off-by: Matthew A Johnson <matjoh@microsoft.com>

Author: matajoh

Lib/test/test_freeze.py

@@ -425,5 +425,6 @@ def test_global_dict_mutation(self):
         self.assertTrue(isimmutable(f1))
         self.assertRaises(NotWriteableError, f1)
 
+
 if __name__ == '__main__':
     unittest.main()

Python/immutability.c

@@ -80,18 +80,6 @@ static bool is_c_wrapper(PyObject* obj){
     } \
 } while(0)
 
-static PyObject* make_global_immutable(PyObject* globals, PyObject* name)
-{
-    PyObject* value = PyDict_GetItem(globals, name); // value.rc = x
-
-    _PyDict_SetKeyImmutable((PyDictObject*)globals, name);
-
-    if(!_Py_IsImmutable(value)){
-        return value;
-    }else{
-        Py_RETURN_NONE;
-    }
-}
 
 /**
  * Special function for walking the reachable graph of a function object.
@@ -106,27 +94,26 @@ static PyObject* make_global_immutable(PyObject* globals, PyObject* name)
  */
 static PyObject* walk_function(PyObject* op, stack* frontier)
 {
-    PyObject* builtins;
-    PyObject* globals;
-    PyObject* module;
-    PyObject* module_dict;
-    PyFunctionObject* f;
-    PyObject* f_ptr;
-    PyCodeObject* f_code;
+    PyObject* builtins = NULL;
+    PyObject* frozen_builtins = NULL;
+    PyObject* globals = NULL;
+    PyObject* frozen_globals = NULL;
+    PyObject* module = NULL;
+    PyObject* module_dict = NULL;
+    PyFunctionObject* f = NULL;
+    PyObject* f_ptr = NULL;
+    PyCodeObject* f_code = NULL;
     Py_ssize_t size;
-    stack* f_stack;
+    stack* f_stack = NULL;
     bool check_globals = false;
+
     _PyObject_ASSERT(op, PyFunction_Check(op));
 
     _Py_SetImmutable(op);
 
 
     f = (PyFunctionObject*)op;
 
-    // TODO find a way to use traverse to avoid having to manually walk
-    // the function's members
-    // f->func_code needs special treatment (see below)
-    // func_globals, func_builtins, and func_module can stay mutable, but depending on code we may need to make some keys immutable
     globals = f->func_globals;
     builtins = f->func_builtins;
 
@@ -163,8 +150,17 @@ static PyObject* walk_function(PyObject* op, stack* frontier)
 
     f_ptr = f->func_code;
     if(stack_push(f_stack, f_ptr)){
-        stack_free(f_stack);
-        return PyErr_NoMemory();
+        goto nomemory;
+    }
+
+    frozen_builtins = PyDict_New();
+    if(frozen_builtins == NULL){
+        goto nomemory;
+    }
+
+    frozen_globals = PyDict_New();
+    if(frozen_globals == NULL){
+        goto nomemory;
     }
 
     while(!stack_empty(f_stack)){
@@ -188,21 +184,20 @@ static PyObject* walk_function(PyObject* op, stack* frontier)
             }
 
             if(PyDict_Contains(globals, name)){
-                PyObject* value = make_global_immutable(globals, name);
-                if(!Py_IsNone(value)){
-                    if(stack_push(frontier, value)){
-                        stack_free(f_stack);
-                        // frontier freed by the caller
-                        return PyErr_NoMemory();
-                    }
+                PyObject* value = PyDict_GetItem(globals, name);
+                if(PyDict_SetItem(frozen_globals, name, value)){
+                    Py_DECREF(frozen_builtins);
+                    Py_DECREF(frozen_globals);
+                    stack_free(f_stack);
+                    return NULL;
                 }
             }else if(PyDict_Contains(builtins, name)){
-
-                _PyDict_SetKeyImmutable((PyDictObject*)builtins, name);
-
-                PyObject* value = PyDict_GetItem(builtins, name); // value.rc = x
-                if(!_Py_IsImmutable(value)){
-                    _Py_SetImmutable(value);
+                PyObject* value = PyDict_GetItem(builtins, name);
+                if(PyDict_SetItem(frozen_builtins, name, value)){
+                    Py_DECREF(frozen_builtins);
+                    Py_DECREF(frozen_globals);
+                    stack_free(f_stack);
+                    return NULL;
                 }
             }else if(PyDict_Contains(module_dict, name)){
                 PyObject* value = PyDict_GetItem(module_dict, name); // value.rc = x
@@ -211,9 +206,7 @@ static PyObject* walk_function(PyObject* op, stack* frontier)
 
                 if(!_Py_IsImmutable(value)){
                     if(stack_push(frontier, value)){
-                        stack_free(f_stack);
-                        // frontier freed by the caller
-                        return PyErr_NoMemory();
+                        goto nomemory;
                     }
                 }
             }
@@ -227,29 +220,24 @@ static PyObject* walk_function(PyObject* op, stack* frontier)
                     _Py_SetImmutable(value);
 
                     if(stack_push(f_stack, value)){
-                        stack_free(f_stack);
-                        // frontier freed by the caller
-                        return PyErr_NoMemory();
+                        goto nomemory;
                     }
                 }else{
                     if(stack_push(frontier, value)){
-                        stack_free(f_stack);
-                        // frontier freed by the caller
-                        return PyErr_NoMemory();
+                        goto nomemory;
                     }
                 }
             }
 
             if(check_globals && PyUnicode_Check(value)){
                 PyObject* name = value;
                 if(PyDict_Contains(globals, name)){
-                    value = make_global_immutable(globals, name);
-                    if(!Py_IsNone(value)){
-                        if(stack_push(frontier, value)){
-                            stack_free(f_stack);
-                            // frontier freed by the caller
-                            return PyErr_NoMemory();
-                        }
+                    value = PyDict_GetItem(globals, name);
+                    if(PyDict_SetItem(frozen_globals, name, value)){
+                        Py_DECREF(frozen_builtins);
+                        Py_DECREF(frozen_globals);
+                        stack_free(f_stack);
+                        return NULL;
                     }
                 }
             }
@@ -270,29 +258,39 @@ static PyObject* walk_function(PyObject* op, stack* frontier)
             if(PyUnicode_Check(value)){
                 PyObject* name = value;
                 if(PyDict_Contains(globals, name)){
-                    value = make_global_immutable(globals, name);
-                    if(!Py_IsNone(value)){
-                        if(stack_push(frontier, value)){
-                            stack_free(f_stack);
-                            // frontier freed by the caller
-                            return PyErr_NoMemory();
-                        }
+                    value = PyDict_GetItem(globals, name);
+                    if(PyDict_SetItem(frozen_globals, name, value)){
+                        Py_DECREF(frozen_builtins);
+                        Py_DECREF(frozen_globals);
+                        return NULL;
                     }
                 }
             }
         }
     }
 
+    if(stack_push(frontier, frozen_globals)){
+        goto nomemory;
+    }
+
+    f->func_globals = frozen_globals;
+    Py_DECREF(globals);
+
+    if(stack_push(frontier, frozen_builtins)){
+        goto nomemory;
+    }
+
+    f->func_builtins = frozen_builtins;
+    Py_DECREF(builtins);
+
     Py_RETURN_NONE;
-}
 
-#define _Py_FREEZE_CALL(f, item, frontier) do { \
-    PyObject* err = f((item), (frontier));             \
-    if(!Py_IsNone(err)){                               \
-        stack_free((frontier));                        \
-        return err;                                    \
-    }                                                  \
-} while(0)
+nomemory:
+    Py_XDECREF(frozen_builtins);
+    Py_XDECREF(frozen_globals);
+    stack_free(f_stack);
+    return PyErr_NoMemory();
+}
 
 static int _freeze_visit(PyObject* obj, void* frontier)
 {
@@ -309,7 +307,6 @@ static int _freeze_visit(PyObject* obj, void* frontier)
     }
 
     if(!_Py_IsImmutable(obj)){
-        //printf("Pushing object %p tp=%p rc=0x%lx\n", (void *)obj, obj->ob_type, obj->ob_refcnt);
         if(stack_push((stack*)frontier, obj)){
             PyErr_NoMemory();
             return -1;
@@ -321,9 +318,6 @@ static int _freeze_visit(PyObject* obj, void* frontier)
 
 PyObject* _Py_Freeze(PyObject* obj)
 {
-    //printf("\nFreezing object: ");
-    //PyObject_Print(obj, stdout, 0);
-    //printf(" rc=0x%lx\n", obj->ob_refcnt);
     if(_Py_IsImmutable(obj) && _Py_IsImmutable(Py_TYPE(obj))){
         Py_RETURN_NONE;
     }
@@ -340,32 +334,29 @@ PyObject* _Py_Freeze(PyObject* obj)
 
     while(!stack_empty(frontier)){
         PyObject* item = stack_pop(frontier); // item.rc = x + 1
-        //printf("Current item %p tp=%p rc=0x%lx: ", (void *)item, item->ob_type, item->ob_refcnt);
-        //fflush(stdout);
-        //PyObject_Print(item, stdout, 0);
-        //printf("\n");
 
         PyTypeObject* type = Py_TYPE(item);
         traverseproc traverse;
         PyObject* type_op = NULL;
 
         if(_Py_IsImmutable(item)){
-            // Direct access like this is not recommended, but will be removed in the future as
-            // this is just for debugging purposes.
             goto handle_type;
         }
 
         _Py_SetImmutable(item);
 
-        //printf("Frozen object %p tp=%p rc=0x%lx\n", (void *)item, item->ob_type, item->ob_refcnt);
 
         if(is_c_wrapper(item)) {
             // C functions are not mutable, so we can skip them.
             continue;
         }
 
         if(PyFunction_Check(item)){
-            _Py_FREEZE_CALL(walk_function, item, frontier);
+            PyObject* err = walk_function(item, frontier);
+            if(!Py_IsNone(err)){
+                stack_free(frontier);
+                return err;
+            }
             goto handle_type;
         }
 
@@ -380,7 +371,6 @@ PyObject* _Py_Freeze(PyObject* obj)
 handle_type:
         type_op = _PyObject_CAST(item->ob_type);
         if (!_Py_IsImmutable(type_op)){
-            // Previously this included a check for is_leaf_type, but
             if (stack_push(frontier, type_op))
             {
                 stack_free(frontier);