optimize set-like iterables

HarryLHW · Jul 14, 2024 · 0822d3c · 0822d3c
1 parent c2c8a93
commit 0822d3c
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 0 deletions.
diff --git a/Include/dictobject.h b/Include/dictobject.h
@@ -95,6 +95,18 @@ PyAPI_DATA(PyTypeObject) PyDictRevIterKey_Type;
 PyAPI_DATA(PyTypeObject) PyDictRevIterItem_Type;
 PyAPI_DATA(PyTypeObject) PyDictRevIterValue_Type;
 
+#define PyDictIterKey_Check(op) PyObject_TypeCheck((op), &PyDictIterKey_Type)
+#define PyDictIterValue_Check(op) PyObject_TypeCheck((op), &PyDictIterValue_Type)
+#define PyDictIterItem_Check(op) PyObject_TypeCheck((op), &PyDictIterItem_Type)
+
+#define PyDictRevIterKey_Check(op) PyObject_TypeCheck((op), &PyDictRevIterKey_Type)
+#define PyDictRevIterItem_Check(op) PyObject_TypeCheck((op), &PyDictRevIterItem_Type)
+#define PyDictRevIterValue_Check(op) PyObject_TypeCheck((op), &PyDictRevIterValue_Type)
+
+#define PyDictViewSetIter_Check(op) \
+    (PyDictItems_Check(op) || PyDictIterKey_Check(op) || PyDictIterItem_Check(op) || \
+     PyDictRevIterKey_Check(op) || PyDictRevIterItem_Check(op))
+
 
 #ifndef Py_LIMITED_API
 #  define Py_CPYTHON_DICTOBJECT_H

diff --git a/Include/setobject.h b/Include/setobject.h
@@ -37,6 +37,8 @@ PyAPI_FUNC(Py_ssize_t) PySet_Size(PyObject *anyset);
     (Py_IS_TYPE((ob), &PySet_Type) || \
     PyType_IsSubtype(Py_TYPE(ob), &PySet_Type))
 
+#define PySetIter_Check(ob) Py_IS_TYPE((ob), &PySetIter_Type)
+
 #ifndef Py_LIMITED_API
 #  define Py_CPYTHON_SETOBJECT_H
 #  include "cpython/setobject.h"

diff --git a/Objects/setobject.c b/Objects/setobject.c
@@ -967,6 +967,26 @@ set_update_iterable_lock_held(PySetObject *so, PyObject *other)
     }
 
     PyObject *key;
+
+    /* If our table is empty, and other is a set-like iterable, we can use set_insert_clean() */
+    if (so->fill == 0) {
+        if (PyRange_Check(so) | PyDictViewSetIter_Check(so) || PySetIter_Check(so)) {
+            setentry *newtable = so->table;
+            size_t newmask = (size_t)so->mask;
+            while ((key = PyIter_Next(it)) != NULL) {
+                Py_hash_t hash = _PyObject_HashFast(key);
+                if (hash == -1) {
+                    return -1;
+                }
+                set_insert_clean(newtable, newmask, key, hash);
+            }
+            Py_DECREF(it);
+            if (PyErr_Occurred())
+                return -1;
+            return 0;
+        }
+    }
+
     while ((key = PyIter_Next(it)) != NULL) {
         if (set_add_key(so, key)) {
             Py_DECREF(it);