Addressing PR comments

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

Author: matajoh

Include/internal/pycore_dict.h

@@ -32,13 +32,16 @@ typedef struct {
     PyObject *_me_value; /* This field is only meaningful for combined tables */
 } PyDictUnicodeEntry;
 
-#define _PyDictEntry_IsImmutable(entry) (((uintptr_t)((entry)->_me_value)) & 0x1)
-#define _PyDictEntry_SetImmutable(entry) ((entry)->_me_value = (PyObject*)((uintptr_t)(entry)->_me_value | 0x1))
+#define _PyDictEntry_IMMUTABLE_FLAG 0x1
+#define _PyDictEntry_VALUE_MASK ~_PyDictEntry_IMMUTABLE_FLAG
+
+#define _PyDictEntry_IsImmutable(entry) (((uintptr_t)((entry)->_me_value)) & _PyDictEntry_IMMUTABLE_FLAG)
+#define _PyDictEntry_SetImmutable(entry) ((entry)->_me_value = (PyObject*)((uintptr_t)(entry)->_me_value | _PyDictEntry_IMMUTABLE_FLAG))
 #define _PyDictEntry_Hash(entry) ((entry)->me_hash)
 #define _PyDictEntry_Key(entry) ((entry)->me_key)
-#define _PyDictEntry_Value(entry) ((PyObject*)((((uintptr_t)((entry)->_me_value)) >> 1) << 1))
+#define _PyDictEntry_Value(entry) ((PyObject*)(((uintptr_t)((entry)->_me_value)) & _PyDictEntry_VALUE_MASK))
 #define _PyDictEntry_SetValue(entry, value) ((entry)->_me_value = value)
-#define _PyDictEntry_IsEmpty(entry) ((entry)->_me_value == NULL)
+#define _PyDictEntry_IsEmpty(entry) ((((uintptr_t)(entry)->_me_value) & _PyDictEntry_VALUE_MASK) == (uintptr_t)NULL)
 
 extern PyObject *_PyDict_IsKeyImmutable(PyObject* op, PyObject* key);
 

Include/internal/pycore_object.h

@@ -100,6 +100,8 @@ static inline void _Py_SetImmutable(PyObject *op)
 }
 #define _Py_SetImmutable(op) _Py_SetImmutable(_PyObject_CAST(op))
 
+// Check whether an object is writeable.
+// Note that during runtime finalization, all objects must be mutable
 #define Py_CHECKWRITE(op) ((op) && (!_Py_IsImmutable(op) || _Py_IsFinalizing()))
 #define Py_REQUIREWRITE(op, msg) {if (Py_CHECKWRITE(op)) { _PyObject_ASSERT_FAILED_MSG(op, msg); }}
 

Include/object.h

@@ -87,16 +87,16 @@ system also uses the second-to-top bit for managing immutable graphs.
 
 #if SIZEOF_VOID_P > 4
 #define _Py_REFCNT_MASK 0xFFFFFFFF
-#define _Py_IMMUTABLE_MASK 0xC000000000
 #define _Py_IMMUTABLE_FLAG 0x4000000000
 #define _Py_IMMUTABLE_SCC_FLAG 0x8000000000
 #else
-#define _Py_REFCNT_MASK 0x3FFFFFFF
-#define _Py_IMMUTABLE_MASK 0xC0000000
-#define _Py_IMMUTABLE_FLAG 0x40000000
-#define _Py_IMMUTABLE_SCC_FLAG 0x80000000
+#define _Py_REFCNT_MASK 0x1FFFFFFF
+#define _Py_IMMUTABLE_FLAG 0x20000000
+#define _Py_IMMUTABLE_SCC_FLAG 0x40000000
 #endif
 
+#define _Py_IMMUTABLE_MASK (_Py_IMMUTABLE_SCC_FLAG | _Py_IMMUTABLE_FLAG)
+
 /*
 Immortalization:
 
@@ -275,7 +275,7 @@ static inline int Py_IS_TYPE(PyObject *ob, PyTypeObject *type) {
 
 static inline Py_ALWAYS_INLINE int _Py_IsImmutable(PyObject *op)
 {
-    return (op->ob_refcnt & _Py_IMMUTABLE_FLAG) > 0;
+    return (op->ob_refcnt & _Py_IMMUTABLE_MASK) > 0;
 }
 #define _Py_IsImmutable(op) _Py_IsImmutable(_PyObject_CAST(op))
 

Lib/test/test_freeze.py

@@ -293,6 +293,8 @@ def d():
     def test_globals_copy(self):
         def f():
             global global0
+            ref_1 = global0
+            ref_2 = global0
             return global0
 
         expected = f()

Python/immutability.c

@@ -19,11 +19,13 @@ static PyObject* pop(PyObject* s){
         return NULL;
     }
 
+    Py_INCREF(item);
+
     if(PyList_SetSlice(s, size - 1, size, NULL)){
+        Py_DECREF(item);
         return NULL;
     }
 
-    Py_INCREF(item);
     return item;
 }
 
@@ -44,12 +46,14 @@ static bool is_c_wrapper(PyObject* obj){
  * Special function for walking the reachable graph of a function object.
  *
  * This is necessary because the function object has a pointer to the global
- * object, and this is problematic because freezing any function will make the
- * global object immutable, which is not always the desired behaviour.
+ * dictionary, and this is problematic because freezing any function directly
+ * (as we do with other objects) would make all globals immutable.
  *
- * This function attempts to find the globals that a function will use, and freeze
- * just those, and prevent those keys from being updated in the global dictionary
- * from this point onwards.
+ * 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 frozen
+ * copies containing just those globals and builtins we were able to determine
+ * the function uses.
  */
 static PyObject* walk_function(PyObject* op, PyObject* frontier)
 {
@@ -129,15 +133,15 @@ static PyObject* walk_function(PyObject* op, PyObject* frontier)
     }
 
     while(PyList_Size(f_stack) != 0){
-        f_ptr = pop(f_stack); // fp.rc = x + 1
+        f_ptr = pop(f_stack);
         _PyObject_ASSERT(f_ptr, PyCode_Check(f_ptr));
         f_code = (PyCodeObject*)f_ptr;
 
         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); // name.rc = x
+            PyObject* name = PySequence_Fast_GET_ITEM(f_code->co_names, i);
 
             if(PyUnicode_CompareWithASCIIString(name, "globals") == 0){
                 // if the code calls the globals() builtin, then any
@@ -165,7 +169,7 @@ static PyObject* walk_function(PyObject* op, PyObject* frontier)
                     return NULL;
                 }
             }else if(PyDict_Contains(module_dict, name)){
-                PyObject* value = PyDict_GetItem(module_dict, name); // value.rc = x
+                PyObject* value = PyDict_GetItem(module_dict, name);
 
                 _PyDict_SetKeyImmutable((PyDictObject*)module_dict, name);
 
@@ -179,7 +183,7 @@ static PyObject* walk_function(PyObject* op, PyObject* frontier)
 
         size = PySequence_Fast_GET_SIZE(f_code->co_consts);
         for(Py_ssize_t i = 0; i < size; i++){
-            PyObject* value = PySequence_Fast_GET_ITEM(f_code->co_consts, i); // value.rc = x
+            PyObject* value = PySequence_Fast_GET_ITEM(f_code->co_consts, i);
             if(!_Py_IsImmutable(value)){
                 if(PyCode_Check(value)){
                     _Py_SetImmutable(value);
@@ -195,6 +199,11 @@ static PyObject* walk_function(PyObject* op, PyObject* frontier)
             }
 
             if(check_globals && PyUnicode_Check(value)){
+                // if the code calls the globals() builtin, then any
+                // cellvar or const in the function could, potentially, refer to
+                // a global variable. As such, we need to check if the globals
+                // dictionary contains that key and then make it immutable
+                // from this point forwards.
                 PyObject* name = value;
                 if(PyDict_Contains(globals, name)){
                     value = PyDict_GetItem(globals, name);
@@ -212,13 +221,20 @@ static PyObject* walk_function(PyObject* op, PyObject* frontier)
     Py_DECREF(f_stack);
 
     if(check_globals){
+        // if the code calls the globals() builtin, then any
+        // cellvar or const in the function could, potentially, refer to
+        // a global variable. As such, we need to check if the globals
+        // dictionary contains that key and then make it immutable
+        // from this point forwards.
+        // we need to check the closure for any cellvars that are not
+        // referenced in the code object, but are still used in the function
         size = 0;
         if(f->func_closure != NULL)
             size = PySequence_Fast_GET_SIZE(f->func_closure);
 
         for(Py_ssize_t i=0; i < size; ++i){
-            PyObject* cellvar = PySequence_Fast_GET_ITEM(f->func_closure, i); // cellvar.rc = x
-            PyObject* value = PyCell_GET(cellvar); // value.rc = x
+            PyObject* cellvar = PySequence_Fast_GET_ITEM(f->func_closure, i);
+            PyObject* value = PyCell_GET(cellvar);
 
             if(PyUnicode_Check(value)){
                 PyObject* name = value;
@@ -271,7 +287,7 @@ static int freeze_visit(PyObject* obj, void* frontier)
 
 PyObject* _Py_Freeze(PyObject* obj)
 {
-    if(_Py_IsImmutable(obj) && _Py_IsImmutable(Py_TYPE(obj))){
+    if(_Py_IsImmutable(obj)){
         Py_RETURN_NONE;
     }
 
@@ -309,7 +325,7 @@ PyObject* _Py_Freeze(PyObject* obj)
     Py_DECREF(frozen_importlib);
 
     while(PyList_Size(frontier) != 0){
-        PyObject* item = pop(frontier); // item.rc = x + 1
+        PyObject* item = pop(frontier);
 
         if(item == blocking_on ||
            item == module_locks){
@@ -323,12 +339,11 @@ PyObject* _Py_Freeze(PyObject* obj)
         PyObject* type_op = NULL;
 
         if(_Py_IsImmutable(item)){
-            goto handle_type;
+            continue;
         }
 
         _Py_SetImmutable(item);
 
-
         if(is_c_wrapper(item)) {
             // C functions are not mutable, so we can skip them.
             continue;
@@ -342,20 +357,20 @@ PyObject* _Py_Freeze(PyObject* obj)
                 Py_DECREF(frontier);
                 return err;
             }
-            goto handle_type;
         }
-
-        traverse = type->tp_traverse;
-        if(traverse != NULL){
-            if(traverse(item, (visitproc)freeze_visit, frontier)){
-                Py_DECREF(blocking_on);
-                Py_DECREF(module_locks);
-                Py_DECREF(frontier);
-                return NULL;
+        else
+        {
+            traverse = type->tp_traverse;
+            if(traverse != NULL){
+                if(traverse(item, (visitproc)freeze_visit, frontier)){
+                    Py_DECREF(blocking_on);
+                    Py_DECREF(module_locks);
+                    Py_DECREF(frontier);
+                    return NULL;
+                }
             }
         }
 
-handle_type:
         type_op = _PyObject_CAST(item->ob_type);
         if (!_Py_IsImmutable(type_op)){
             if (PyList_Append(frontier, type_op))

Tools/gdb/libpython.py

@@ -783,6 +783,7 @@ def iteritems(self):
         for i in safe_range(nentries):
             ep = entries[i]
             value_ptr = ep['_me_value'].reinterpret_cast(gdb.lookup_type('uintptr_t'))
+            # Clear the immutability flag
             pyop_value = PyObjectPtr.from_pyobject_ptr((value_ptr >> 1) << 1)
             if not pyop_value.is_null():
                 pyop_key = PyObjectPtr.from_pyobject_ptr(ep['me_key'])