SET_ADD should not lock

Author: DinoV

Include/internal/pycore_setobject.h

@@ -33,6 +33,13 @@ PyAPI_FUNC(int) _PySet_Contains(PySetObject *so, PyObject *key);
 // Clears the set without acquiring locks. Used by _PyCode_Fini.
 extern void _PySet_ClearInternal(PySetObject *so);
 
+PyAPI_FUNC(int) _PySet_AddTakeRef(PySetObject *so, PyObject *key);
+
+PyAPI_FUNC(PyObject *)
+_PySet_FromStackRefSteal(const _PyStackRef *src, Py_ssize_t n);
+
+
+
 #ifdef __cplusplus
 }
 #endif

Objects/setobject.c

@@ -122,7 +122,7 @@ set_lookkey(PySetObject *so, PyObject *key, Py_hash_t hash)
 static int set_table_resize(PySetObject *, Py_ssize_t);
 
 static int
-set_add_entry(PySetObject *so, PyObject *key, Py_hash_t hash)
+set_add_entry_takeref(PySetObject *so, PyObject *key, Py_hash_t hash)
 {
     setentry *table;
     setentry *freeslot;
@@ -133,12 +133,6 @@ set_add_entry(PySetObject *so, PyObject *key, Py_hash_t hash)
     int probes;
     int cmp;
 
-    _Py_CRITICAL_SECTION_ASSERT_OBJECT_LOCKED(so);
-
-    /* Pre-increment is necessary to prevent arbitrary code in the rich
-       comparison from deallocating the key just before the insertion. */
-    Py_INCREF(key);
-
   restart:
 
     mask = so->mask;
@@ -209,6 +203,57 @@ set_add_entry(PySetObject *so, PyObject *key, Py_hash_t hash)
     return -1;
 }
 
+static int
+set_add_entry(PySetObject *so, PyObject *key, Py_hash_t hash)
+{
+    _Py_CRITICAL_SECTION_ASSERT_OBJECT_LOCKED(so);
+
+    /* Pre-increment is necessary to prevent arbitrary code in the rich
+       comparison from deallocating the key just before the insertion. */
+    Py_INCREF(key);
+
+    return set_add_entry_takeref(so, key, hash);
+}
+
+int
+_PySet_AddTakeRef(PySetObject *so, PyObject *key)
+{
+    Py_hash_t hash = _PyObject_HashFast(key);
+    if (hash == -1) {
+        return -1;
+    }
+    // We don't pre-increment here, the caller holds a strong
+    // reference to the object which we are stealing.
+    return set_add_entry_takeref(so, key, hash);
+}
+
+PyObject *
+_PySet_FromStackRefSteal(const _PyStackRef *src, Py_ssize_t n)
+{
+    PySetObject *set = (PySetObject *)PySet_New(NULL);
+    if (n == 0) {
+        return (PyObject *)set;
+    }
+
+    if (set_table_resize(set, n*2) != 0) {
+        return NULL;
+    }
+
+    int err = 0;
+    for (Py_ssize_t i = 0; i < n; i++) {
+        if (err == 0) {
+            err = _PySet_AddTakeRef(set, PyStackRef_AsPyObjectSteal(src[i]));
+        } else {
+            PyStackRef_CLOSE(src[i]);
+        }
+    }
+    if (err) {
+        Py_CLEAR(set);
+    }
+
+    return (PyObject *)set;
+}
+
 /*
 Internal routine used by set_table_resize() to insert an item which is
 known to be absent from the set.  Besides the performance benefit,

Python/bytecodes.c

@@ -975,9 +975,8 @@ dummy_func(
         }
 
         inst(SET_ADD, (set, unused[oparg-1], v -- set, unused[oparg-1])) {
-            int err = PySet_Add(PyStackRef_AsPyObjectBorrow(set),
-                                PyStackRef_AsPyObjectBorrow(v));
-            PyStackRef_CLOSE(v);
+            int err = _PySet_AddTakeRef((PySetObject *)PyStackRef_AsPyObjectBorrow(set),
+                                PyStackRef_AsPyObjectSteal(v));
             ERROR_IF(err, error);
         }
 
@@ -1911,22 +1910,11 @@ dummy_func(
         }
 
         inst(BUILD_SET, (values[oparg] -- set)) {
-            PyObject *set_o = PySet_New(NULL);
+            PyObject *set_o = _PySet_FromStackRefSteal(values, oparg);
             if (set_o == NULL) {
-                DECREF_INPUTS();
-                ERROR_IF(true, error);
-            }
-            int err = 0;
-            for (int i = 0; i < oparg; i++) {
-                if (err == 0) {
-                    err = PySet_Add(set_o, PyStackRef_AsPyObjectBorrow(values[i]));
-                }
-            }
-            DECREF_INPUTS();
-            if (err != 0) {
-                Py_DECREF(set_o);
                 ERROR_IF(true, error);
             }
+            INPUTS_DEAD();
             set = PyStackRef_FromPyObjectStealMortal(set_o);
         }