gh-141510: Add frozendict fast-path to the set type (#144912)

diff --git a/Lib/test/test_set.py b/Lib/test/test_set.py
index 554716aed1e98b7cd8a96277e35036f439acce66..9bfd4bc7d636699c5a17de07c276ed79b5f446f2 100644 (file)
--- a/Lib/test/test_set.py
+++ b/Lib/test/test_set.py
@@ -188,7 +188,10 @@ class TestJointOps:
         self.assertEqual(type(i), self.basetype)
         self.assertRaises(PassThru, self.s.symmetric_difference, check_pass_thru())
         self.assertRaises(TypeError, self.s.symmetric_difference, [[]])
-        for C in set, frozenset, dict.fromkeys, str, list, tuple:
+        constructors = (set, frozenset,
+                        dict.fromkeys, frozendict.fromkeys,
+                        str, list, tuple)
+        for C in constructors:
             self.assertEqual(self.thetype('abcba').symmetric_difference(C('cdc')), set('abd'))
             self.assertEqual(self.thetype('abcba').symmetric_difference(C('efgfe')), set('abcefg'))
             self.assertEqual(self.thetype('abcba').symmetric_difference(C('ccb')), set('a'))
@@ -1591,6 +1594,14 @@ class TestOnlySetsDict(TestOnlySetsInBinaryOps, unittest.TestCase):
 
 #------------------------------------------------------------------------------
 
+class TestOnlySetsFrozenDict(TestOnlySetsInBinaryOps, unittest.TestCase):
+    def setUp(self):
+        self.set   = set((1, 2, 3))
+        self.other = frozendict({1:2, 3:4})
+        self.otherIsIterable = True
+
+#------------------------------------------------------------------------------
+
 class TestOnlySetsOperator(TestOnlySetsInBinaryOps, unittest.TestCase):
     def setUp(self):
         self.set   = set((1, 2, 3))

diff --git a/Objects/setobject.c b/Objects/setobject.c
index f8713bf3d1a4321a6a10e6e119fd61df7b8ee203..ae6c1d1248d2fc9e8076ef3b2961d21d200bbc6d 100644 (file)
--- a/Objects/setobject.c
+++ b/Objects/setobject.c
@@ -1186,10 +1186,14 @@ set_iter(PyObject *so)
 static int
 set_update_dict_lock_held(PySetObject *so, PyObject *other)
 {
-    assert(PyDict_CheckExact(other));
+    assert(PyAnyDict_CheckExact(other));
 
     _Py_CRITICAL_SECTION_ASSERT_OBJECT_LOCKED(so);
-    _Py_CRITICAL_SECTION_ASSERT_OBJECT_LOCKED(other);
+#ifdef Py_DEBUG
+    if (!PyFrozenDict_CheckExact(other)) {
+        _Py_CRITICAL_SECTION_ASSERT_OBJECT_LOCKED(other);
+    }
+#endif
 
     /* Do one big resize at the start, rather than
     * incrementally resizing as we insert new keys.  Expect
@@ -1245,7 +1249,7 @@ set_update_lock_held(PySetObject *so, PyObject *other)
     if (PyAnySet_Check(other)) {
         return set_merge_lock_held(so, other);
     }
-    else if (PyDict_CheckExact(other)) {
+    else if (PyAnyDict_CheckExact(other)) {
         return set_update_dict_lock_held(so, other);
     }
     return set_update_iterable_lock_held(so, other);
@@ -1270,6 +1274,9 @@ set_update_local(PySetObject *so, PyObject *other)
         Py_END_CRITICAL_SECTION();
         return rv;
     }
+    else if (PyFrozenDict_CheckExact(other)) {
+        return set_update_dict_lock_held(so, other);
+    }
     return set_update_iterable_lock_held(so, other);
 }
 
@@ -1293,6 +1300,13 @@ set_update_internal(PySetObject *so, PyObject *other)
         Py_END_CRITICAL_SECTION2();
         return rv;
     }
+    else if (PyFrozenDict_CheckExact(other)) {
+        int rv;
+        Py_BEGIN_CRITICAL_SECTION(so);
+        rv = set_update_dict_lock_held(so, other);
+        Py_END_CRITICAL_SECTION();
+        return rv;
+    }
     else {
         int rv;
         Py_BEGIN_CRITICAL_SECTION(so);
@@ -2033,7 +2047,7 @@ set_difference(PySetObject *so, PyObject *other)
     if (PyAnySet_Check(other)) {
         other_size = PySet_GET_SIZE(other);
     }
-    else if (PyDict_CheckExact(other)) {
+    else if (PyAnyDict_CheckExact(other)) {
         other_size = PyDict_GET_SIZE(other);
     }
     else {
@@ -2050,7 +2064,7 @@ set_difference(PySetObject *so, PyObject *other)
     if (result == NULL)
         return NULL;
 
-    if (PyDict_CheckExact(other)) {
+    if (PyAnyDict_CheckExact(other)) {
         while (set_next(so, &pos, &entry)) {
             key = entry->key;
             hash = entry->hash;
@@ -2172,7 +2186,11 @@ static int
 set_symmetric_difference_update_dict(PySetObject *so, PyObject *other)
 {
     _Py_CRITICAL_SECTION_ASSERT_OBJECT_LOCKED(so);
-    _Py_CRITICAL_SECTION_ASSERT_OBJECT_LOCKED(other);
+#ifdef Py_DEBUG
+    if (!PyFrozenDict_CheckExact(other)) {
+        _Py_CRITICAL_SECTION_ASSERT_OBJECT_LOCKED(other);
+    }
+#endif
 
     Py_ssize_t pos = 0;
     PyObject *key, *value;
@@ -2246,6 +2264,11 @@ set_symmetric_difference_update_impl(PySetObject *so, PyObject *other)
         rv = set_symmetric_difference_update_dict(so, other);
         Py_END_CRITICAL_SECTION2();
     }
+    else if (PyFrozenDict_CheckExact(other)) {
+        Py_BEGIN_CRITICAL_SECTION(so);
+        rv = set_symmetric_difference_update_dict(so, other);
+        Py_END_CRITICAL_SECTION();
+    }
     else if (PyAnySet_Check(other)) {
         Py_BEGIN_CRITICAL_SECTION2(so, other);
         rv = set_symmetric_difference_update_set(so, (PySetObject *)other);