gh-116621: Specialize list.extend for dict keys/values (gh-116816)

This commit is contained in:
Donghee Na 2024-03-15 23:48:34 +09:00 committed by GitHub
parent 59e30f41ed
commit 8da83f3386
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 37 additions and 0 deletions

View File

@ -3,6 +3,7 @@
#include "Python.h" #include "Python.h"
#include "pycore_abstract.h" // _PyIndex_Check() #include "pycore_abstract.h" // _PyIndex_Check()
#include "pycore_ceval.h" // _PyEval_GetBuiltin() #include "pycore_ceval.h" // _PyEval_GetBuiltin()
#include "pycore_dict.h" // _PyDictViewObject
#include "pycore_pyatomic_ft_wrappers.h" #include "pycore_pyatomic_ft_wrappers.h"
#include "pycore_interp.h" // PyInterpreterState.list #include "pycore_interp.h" // PyInterpreterState.list
#include "pycore_list.h" // struct _Py_list_freelist, _PyListIterObject #include "pycore_list.h" // struct _Py_list_freelist, _PyListIterObject
@ -1295,6 +1296,30 @@ list_extend_set(PyListObject *self, PySetObject *other)
return 0; return 0;
} }
static int
list_extend_dict(PyListObject *self, PyDictObject *dict, int which_item)
{
// which_item: 0 for keys and 1 for values
Py_ssize_t m = Py_SIZE(self);
Py_ssize_t n = PyDict_GET_SIZE(dict);
if (list_resize(self, m + n) < 0) {
return -1;
}
PyObject **dest = self->ob_item + m;
Py_ssize_t pos = 0;
PyObject *keyvalue[2];
while (_PyDict_Next((PyObject *)dict, &pos, &keyvalue[0], &keyvalue[1], NULL)) {
PyObject *obj = keyvalue[which_item];
Py_INCREF(obj);
*dest = obj;
dest++;
}
Py_SET_SIZE(self, m + n);
return 0;
}
static int static int
_list_extend(PyListObject *self, PyObject *iterable) _list_extend(PyListObject *self, PyObject *iterable)
{ {
@ -1322,6 +1347,18 @@ _list_extend(PyListObject *self, PyObject *iterable)
res = list_extend_set(self, (PySetObject *)iterable); res = list_extend_set(self, (PySetObject *)iterable);
Py_END_CRITICAL_SECTION2(); Py_END_CRITICAL_SECTION2();
} }
else if (Py_IS_TYPE(iterable, &PyDictKeys_Type)) {
PyDictObject *dict = ((_PyDictViewObject *)iterable)->dv_dict;
Py_BEGIN_CRITICAL_SECTION2(self, dict);
res = list_extend_dict(self, dict, 0 /*keys*/);
Py_END_CRITICAL_SECTION2();
}
else if (Py_IS_TYPE(iterable, &PyDictValues_Type)) {
PyDictObject *dict = ((_PyDictViewObject *)iterable)->dv_dict;
Py_BEGIN_CRITICAL_SECTION2(self, dict);
res = list_extend_dict(self, dict, 1 /*values*/);
Py_END_CRITICAL_SECTION2();
}
else { else {
Py_BEGIN_CRITICAL_SECTION(self); Py_BEGIN_CRITICAL_SECTION(self);
res = list_extend_iter_lock_held(self, iterable); res = list_extend_iter_lock_held(self, iterable);