Immutability: minor fixes

Lib/test/test_freeze/test_implicit.py

@@ -1,3 +1,4 @@
+import sys
 import unittest
 from immutable import freeze, is_frozen
 
@@ -139,6 +140,26 @@ def test_deeply_nested_no_stack_overflow(self):
             obj = (obj,)
         self.assertTrue(is_frozen(obj))
 
+    def test_abandoned_walk_keeps_references(self):
+        """An aborted walk must not drop references it never took.
+
+        The walk pushes objects onto a worklist without increfing them, so
+        anything still on the worklist when a mutable object aborts the walk
+        used to be decrefed when the worklist was released. That freed the
+        object while its real owners were still pointing at it, which showed
+        up much later as a negative refcount.
+        """
+        # Built at runtime so it is neither interned nor immortal, which makes
+        # its reference count fully accounted for by this test.
+        item = "".join(["abandoned", "-", "worklist", "-", "entry"])
+        # Tuples are traversed back to front, so `item` reaches the worklist
+        # before the dict aborts the walk.
+        obj = ({"mutable": 1}, item)
+
+        before = sys.getrefcount(item)
+        self.assertFalse(is_frozen(obj))
+        self.assertEqual(sys.getrefcount(item), before)
+
 
 if __name__ == '__main__':
     unittest.main()

Modules/_interpretersmodule.c

@@ -484,13 +484,18 @@ _interp_call_pack(PyThreadState *tstate, struct interp_call *call,
                       "expected a callable, got %R", func);
         return -1;
     }
-    if (_PyFunction_GetXIData(tstate, func, &call->_preallocated.func) < 0) {
-        PyObject *exc = _PyErr_GetRaisedException(tstate);
-        if (_PyPickle_GetXIData(tstate, func, &call->_preallocated.func) < 0) {
-            _PyErr_SetRaisedException(tstate, exc);
-            return -1;
+    // If func is immutable (e.g. frozen), share it directly instead of
+    // marshaling its code.
+    if (_PyObject_GetXIDataNoFallback(tstate, func, &call->_preallocated.func) < 0) {
+        _PyErr_Clear(tstate);
+        if (_PyFunction_GetXIData(tstate, func, &call->_preallocated.func) < 0) {
+            PyObject *exc = _PyErr_GetRaisedException(tstate);
+            if (_PyPickle_GetXIData(tstate, func, &call->_preallocated.func) < 0) {
+                _PyErr_SetRaisedException(tstate, exc);
+                return -1;
+            }
+            Py_DECREF(exc);
         }
-        Py_DECREF(exc);
     }
     call->func = &call->_preallocated.func;
     // Handle the args.

Python/immutability.c

@@ -675,7 +675,7 @@
     return true;
 }
 
 static PyObject* get_next(PyObject* obj, struct FreezeState *freeze_state)
 {
     (void)freeze_state;
     PyObject* next = scc_next(obj);
@@ -1828,6 +1828,14 @@
     }
 
     _Py_hashtable_destroy(state.visited);
+
+    // We can't call the destructor directly as we didn't newref the objects
+    // on push. Breaking out of the loop above leaves the remaining objects
+    // on the worklist, so drain it here. This is a slow path if there are
+    // still objects in the stack, so there is no need to optimize it.
+    while (PyList_Size(state.worklist) > 0) {
+        pop(state.worklist);
+    }
     Py_DECREF(state.worklist);
 
     if (result < 0) {
@@ -2130,7 +2138,7 @@
         return 0;
     }
 #else
 #error "Immutability currently only works on 64bit platforms"
 #endif
 
     // Mark pre-freeze hook as completed. This has to be set before calling