simple implementation of _PyType_SharedKeysMaxSize (no inheritance lo…

…okup)
iritkatriel · Mar 23, 2024 · 1bb4ad6 · 1bb4ad6
1 parent add1804
commit 1bb4ad6
Show file tree

Hide file tree

Showing 4 changed files with 20 additions and 8 deletions.
diff --git a/Include/internal/pycore_dict.h b/Include/internal/pycore_dict.h
@@ -256,7 +256,7 @@ PyAPI_FUNC(PyObject *)_PyDict_FromItems(
 static inline void
 _PyDictValues_AddToInsertionOrder(PyDictValues *values, Py_ssize_t ix)
 {
-    assert(ix < SHARED_KEYS_MAX_SIZE);
+    // assert(ix < SHARED_KEYS_MAX_SIZE);
     uint8_t *size_ptr = ((uint8_t *)values)-2;
     int size = *size_ptr;
     assert(size+2 < DICT_VALUES_SIZE(values));

diff --git a/Include/internal/pycore_object.h b/Include/internal/pycore_object.h
@@ -237,6 +237,7 @@ Py_ssize_t _Py_ExplicitMergeRefcount(PyObject *op, Py_ssize_t extra);
 extern int _PyType_CheckConsistency(PyTypeObject *type);
 extern int _PyDict_CheckConsistency(PyObject *mp, int check_content);
 
+
 /* Update the Python traceback of an object. This function must be called
    when a memory block is reused from a free list.
 
@@ -255,6 +256,8 @@ extern PyStatus _PyObject_InitState(PyInterpreterState *interp);
 extern void _PyObject_FiniState(PyInterpreterState *interp);
 extern bool _PyRefchain_IsTraced(PyInterpreterState *interp, PyObject *obj);
 
+extern Py_ssize_t _PyType_SharedKeysMaxSize(PyTypeObject *type);
+
 /* Inline functions trading binary compatibility for speed:
    _PyObject_Init() is the fast version of PyObject_Init(), and
    _PyObject_InitVar() is the fast version of PyObject_InitVar().

diff --git a/Objects/dictobject.c b/Objects/dictobject.c
@@ -623,7 +623,7 @@ static PyDictKeysObject empty_keys_struct = {
 static inline int
 get_index_from_order(PyDictObject *mp, Py_ssize_t i)
 {
-    assert(mp->ma_used <= SHARED_KEYS_MAX_SIZE);
+    // assert(mp->ma_used <= SHARED_KEYS_MAX_SIZE);
     assert(i < (((char *)mp->ma_values)[-2]));
     return ((char *)mp->ma_values)[-3-i];
 }
@@ -671,7 +671,7 @@ _PyDict_CheckConsistency(PyObject *op, int check_content)
     }
     else {
         CHECK(keys->dk_kind == DICT_KEYS_SPLIT);
-        CHECK(mp->ma_used <= SHARED_KEYS_MAX_SIZE);
+        // CHECK(mp->ma_used <= SHARED_KEYS_MAX_SIZE);
     }
 
     if (check_content) {
@@ -1615,7 +1615,7 @@ insert_into_splitdictkeys(PyDictKeysObject *keys, PyObject *name, Py_hash_t hash
         ep->me_key = Py_NewRef(name);
         split_keys_entry_added(keys);
     }
-    assert (ix < SHARED_KEYS_MAX_SIZE);
+    // assert (ix < SHARED_KEYS_MAX_SIZE);
     return ix;
 }
 

diff --git a/Objects/typeobject.c b/Objects/typeobject.c
@@ -7760,9 +7760,18 @@ type_ready_set_new(PyTypeObject *type, int rerunbuiltin)
     return 0;
 }
 
-static Py_ssize_t
-expected_number_of_shared_keys(PyTypeObject *type)
-{
+Py_ssize_t
+_PyType_SharedKeysMaxSize(PyTypeObject *type)
+{
+    if (PyObject_HasAttr((PyObject*)type, &_Py_ID(__expected_attributes__))) {
+        PyObject *attrs = PyObject_GetAttr((PyObject*)type, &_Py_ID(__expected_attributes__));
+        if (attrs != NULL && PyTuple_Check(attrs)) {
+            Py_ssize_t num_keys = PyTuple_GET_SIZE(attrs);
+            if (num_keys > SHARED_KEYS_MAX_SIZE) {
+                return 1 + num_keys;
+            }
+        }
+    }
     return SHARED_KEYS_MAX_SIZE;
 }
 
@@ -7781,7 +7790,7 @@ type_ready_managed_dict(PyTypeObject *type)
     }
     PyHeapTypeObject* et = (PyHeapTypeObject*)type;
     if (et->ht_cached_keys == NULL) {
-        et->ht_cached_keys = _PyDict_NewKeysForClass(expected_number_of_shared_keys(type));
+        et->ht_cached_keys = _PyDict_NewKeysForClass(_PyType_SharedKeysMaxSize(type));
         if (et->ht_cached_keys == NULL) {
             PyErr_NoMemory();
             return -1;