@@ -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 ;
@@ -2234,10 +2243,16 @@ set_contains_lock_held(PySetObject *so, PyObject *key)
22342243int
22352244_PySet_Contains (PySetObject * so , PyObject * key )
22362245{
2246+ assert (so );
2247+
22372248 int rv ;
2238- Py_BEGIN_CRITICAL_SECTION (so );
2239- rv = set_contains_lock_held (so , key );
2240- Py_END_CRITICAL_SECTION ();
2249+ if (PyFrozenSet_CheckExact (so )) {
2250+ rv = set_contains_lock_held (so , key );
2251+ } else {
2252+ Py_BEGIN_CRITICAL_SECTION (so );
2253+ rv = set_contains_lock_held (so , key );
2254+ Py_END_CRITICAL_SECTION ();
2255+ }
22412256 return rv ;
22422257}
22432258
0 commit comments