@@ -86,6 +86,8 @@ set_lookkey(PySetObject *so, PyObject *key, Py_hash_t hash)
8686 int probes ;
8787 int cmp ;
8888
89+ int frozenset = PyFrozenSet_CheckExact (so );
90+
8991 while (1 ) {
9092 entry = & so -> table [i ];
9193 probes = (i + LINEAR_PROBES <= mask ) ? LINEAR_PROBES : 0 ;
@@ -102,13 +104,20 @@ set_lookkey(PySetObject *so, PyObject *key, Py_hash_t hash)
102104 && unicode_eq (startkey , key ))
103105 return entry ;
104106 table = so -> table ;
105- Py_INCREF (startkey );
106- cmp = PyObject_RichCompareBool (startkey , key , Py_EQ );
107- Py_DECREF (startkey );
108- if (cmp < 0 )
109- return NULL ;
110- if (table != so -> table || entry -> key != startkey )
111- return set_lookkey (so , key , hash );
107+ if (frozenset ) {
108+ cmp = PyObject_RichCompareBool (startkey , key , Py_EQ );
109+ if (cmp < 0 )
110+ return NULL ;
111+ } else {
112+ // incref startkey because it can be removed from the set by the compare
113+ Py_INCREF (startkey );
114+ cmp = PyObject_RichCompareBool (startkey , key , Py_EQ );
115+ Py_DECREF (startkey );
116+ if (cmp < 0 )
117+ return NULL ;
118+ if (table != so -> table || entry -> key != startkey )
119+ return set_lookkey (so , key , hash );
120+ }
112121 if (cmp > 0 )
113122 return entry ;
114123 mask = so -> mask ;
@@ -2235,10 +2244,16 @@ set_contains_lock_held(PySetObject *so, PyObject *key)
22352244int
22362245_PySet_Contains (PySetObject * so , PyObject * key )
22372246{
2247+ assert (so );
2248+
22382249 int rv ;
2239- Py_BEGIN_CRITICAL_SECTION (so );
2240- rv = set_contains_lock_held (so , key );
2241- Py_END_CRITICAL_SECTION ();
2250+ if (PyFrozenSet_CheckExact (so )) {
2251+ rv = set_contains_lock_held (so , key );
2252+ } else {
2253+ Py_BEGIN_CRITICAL_SECTION (so );
2254+ rv = set_contains_lock_held (so , key );
2255+ Py_END_CRITICAL_SECTION ();
2256+ }
22422257 return rv ;
22432258}
22442259
0 commit comments