LRC correction for set iteration

Author: xFrednet

Lib/test/test_regions/test_set.py

@@ -0,0 +1,19 @@
+import unittest
+from regions import Region
+
+class TestSet(unittest.TestCase):
+    def test_for(self):
+        class A: pass
+
+        # Setup
+        r = Region()
+        r.obj1=A()
+        r.obj2=A()
+        s = {r.obj1, r.obj2}
+
+        base_lrc = r._lrc
+        for i in s:
+            pass
+        i = None
+
+        self.assertEqual(base_lrc, r._lrc)

Objects/object.c

@@ -1595,7 +1595,7 @@ _PyObject_GetDictPtr(PyObject *obj)
 PyObject *
 PyObject_SelfIter(PyObject *obj)
 {
-    return Py_NewRef(obj);
+    return PyRegion_NewRef(obj);
 }
 
 /* Helper used when the __next__ method is removed from a type:

Objects/setobject.c

@@ -110,8 +110,12 @@ set_lookkey(PySetObject *so, PyObject *key, Py_hash_t hash)
                         return NULL;
                 } else {
                     // incref startkey because it can be removed from the set by the compare
+                    if (PyRegion_AddLocalRef(startkey)) {
+                        return NULL;
+                    }
                     Py_INCREF(startkey);
                     cmp = PyObject_RichCompareBool(startkey, key, Py_EQ);
+                    PyRegion_RemoveLocalRef(startkey);
                     Py_DECREF(startkey);
                     if (cmp < 0)
                         return NULL;
@@ -855,7 +859,7 @@ setiter_dealloc(PyObject *self)
     setiterobject *si = (setiterobject*)self;
     /* bpo-31095: UnTrack is needed before calling any callbacks */
     _PyObject_GC_UNTRACK(si);
-    Py_XDECREF(si->si_set);
+    PyRegion_CLEAR(si, si->si_set);
     PyObject_GC_Del(si);
 }
 
@@ -886,11 +890,14 @@ setiter_reduce(PyObject *op, PyObject *Py_UNUSED(ignored))
 
     /* copy the iterator state */
     setiterobject tmp = *si;
+    if (PyRegion_AddLocalRef(si)) {
+        return NULL;
+    }
     Py_XINCREF(tmp.si_set);
 
     /* iterate the temporary into a list */
     PyObject *list = PySequence_List((PyObject*)&tmp);
-    Py_XDECREF(tmp.si_set);
+    PyRegion_RemoveLocalRef(tmp.si_set);
     if (list == NULL) {
         return NULL;
     }
@@ -935,13 +942,16 @@ static PyObject *setiter_iternext(PyObject *self)
         i++;
     }
     if (i <= mask) {
-        key = Py_NewRef(entry[i].key);
+        key = entry[i].key;
+        if (PyRegion_AddLocalRef(key)) {
+            return NULL;
+        }
+        Py_INCREF(key);
     }
     Py_END_CRITICAL_SECTION();
     si->si_pos = i+1;
     if (key == NULL) {
-        si->si_set = NULL;
-        Py_DECREF(so);
+        PyRegion_CLEAR(si, si->si_set);
         return NULL;
     }
     si->len--;