diff --git a/Doc/c-api/list.rst b/Doc/c-api/list.rst index 21769c4c5c869c..4b85e5833cde93 100644 --- a/Doc/c-api/list.rst +++ b/Doc/c-api/list.rst @@ -227,3 +227,96 @@ List Objects Return a new tuple object containing the contents of *list*; equivalent to ``tuple(list)``. + + +.. _pylistwriter: + +PyListWriter +------------ + +The :c:type:`PyListWriter` API can be used to create a Python :class:`list` +object. + +On Free Threading, it's more efficient than the ``PyList`` C API to create a +:class:`list`, because it can avoid the critical section on the created list. + +.. versionadded:: next + + +.. c:type:: PyListWriter + + A list writer object. + + The API is **not thread safe**. To share a writer with multiple threads, a + critical section or a lock is needed. + + The instance must be destroyed by :c:func:`PyListWriter_Finish` on + success, or :c:func:`PyListWriter_Discard` on error. + + +.. c:function:: PyListWriter* PyListWriter_Create(Py_ssize_t size) + + Create a :c:type:`PyListWriter` and preallocate *size* items. + + On error, set an exception and return ``NULL``. + + *size* must be positive or zero. + +.. c:function:: PyObject* PyListWriter_Finish(PyListWriter *writer) + + Finish a :c:type:`PyListWriter` created by + :c:func:`PyListWriter_Create`. + + On success, return a Python :class:`list` object. + On error, set an exception and return ``NULL``. + + The writer instance is invalid after the call in any case. + No API can be called on the writer after :c:func:`PyListWriter_Finish`. + +.. c:function:: void PyListWriter_Discard(PyListWriter *writer) + + Discard a :c:type:`PyListWriter` created by :c:func:`PyListWriter_Create`. + + Do nothing if *writer* is ``NULL``. + + The writer instance is invalid after the call. + No API can be called on the writer after :c:func:`PyListWriter_Discard`. + +.. c:function:: int PyListWriter_Append(PyListWriter *writer, PyObject *item) + + Append an item to the list writer. + Similar to :c:func:`PyList_Append`. + + On success, return ``0``. + On error, set an exception and return ``-1``. + +.. c:function:: int PyListWriter_AppendSteal(PyListWriter *writer, PyObject *item) + + Similar to :c:func:`PyListWriter_Append`, but :term:`steals ` the + reference to *item*, instead of creating a new :term:`strong reference`. + +.. c:function:: int PyListWriter_Extend(PyListWriter *writer, PyObject *iterable) + + Extend the list writer with the contents of *iterable*. + Similar to :c:func:`PyList_Extend`. + + On success, return ``0``. + On error, set an exception and return ``-1``. + +.. c:function:: int PyListWriter_ExtendArray(PyListWriter *writer, PyObject *const *items, Py_ssize_t nitem) + + Extend the list writer with the contents of *items*: copy references from + *items* to the list writer. + + On success, return ``0``. + On error, set an exception and return ``-1``. + +.. c:function:: int PyListWriter_ExtendArraySteal(PyListWriter *writer, PyObject *const *items, Py_ssize_t nitem) + + Similar to :c:func:`PyListWriter_ExtendArray`, but :term:`steal` references + instead of copying references. + + On success, return ``0``. + On error, set an exception and return ``-1``. + + On error, all *items* references are also cleared. diff --git a/Doc/whatsnew/3.16.rst b/Doc/whatsnew/3.16.rst index 09477b22d95413..0b91987aace7ba 100644 --- a/Doc/whatsnew/3.16.rst +++ b/Doc/whatsnew/3.16.rst @@ -1310,6 +1310,22 @@ New features * Add :c:func:`Py_HashBuffer` to the limited C API. (Contributed by Victor Stinner in :gh:`148233`.) +* Add a new :ref:`PyListWriter ` C API to create a :class:`list` + object. On Free Threading, it's more efficient than the ``PyList`` C API to + create a :class:`list`, because it can avoid the critical section on the + created list. Add functions: + + * :c:func:`PyListWriter_Create` + * :c:func:`PyListWriter_Finish` + * :c:func:`PyListWriter_Discard` + * :c:func:`PyListWriter_Append` + * :c:func:`PyListWriter_AppendSteal` + * :c:func:`PyListWriter_Extend` + * :c:func:`PyListWriter_ExtendArray` + * :c:func:`PyListWriter_ExtendArraySteal` + + (Contributed by Victor Stinner in :gh:`158660`.) + * Add :c:macro:`Py_SETREF` and :c:macro:`Py_XSETREF` macros to the limited C API. (Contributed by Nathan Goldbaum and Victor Stinner in :gh:`158842`.) diff --git a/Include/cpython/listobject.h b/Include/cpython/listobject.h index 49f5e8d6d1a0d6..a0843bf185e514 100644 --- a/Include/cpython/listobject.h +++ b/Include/cpython/listobject.h @@ -51,3 +51,29 @@ PyList_SET_ITEM(PyObject *op, Py_ssize_t index, PyObject *value) { PyAPI_FUNC(int) PyList_Extend(PyObject *self, PyObject *iterable); PyAPI_FUNC(int) PyList_Clear(PyObject *self); + + +/* --- PyListWriter API --------------------------------------------------- */ + +typedef struct PyListWriter PyListWriter; + +PyAPI_FUNC(PyListWriter*) PyListWriter_Create(Py_ssize_t size); +PyAPI_FUNC(PyObject*) PyListWriter_Finish(PyListWriter *writer); +PyAPI_FUNC(void) PyListWriter_Discard(PyListWriter *writer); +PyAPI_FUNC(int) PyListWriter_Append( + PyListWriter *writer, + PyObject *item); +PyAPI_FUNC(int) PyListWriter_AppendSteal( + PyListWriter *writer, + PyObject *item); +PyAPI_FUNC(int) PyListWriter_Extend( + PyListWriter *writer, + PyObject *iterable); +PyAPI_FUNC(int) PyListWriter_ExtendArray( + PyListWriter *writer, + PyObject *const *items, + Py_ssize_t nitem); +PyAPI_FUNC(int) PyListWriter_ExtendArraySteal( + PyListWriter *writer, + PyObject *const *items, + Py_ssize_t nitem); diff --git a/Include/internal/pycore_list.h b/Include/internal/pycore_list.h index df0d00f752573b..6db69da20c9eaf 100644 --- a/Include/internal/pycore_list.h +++ b/Include/internal/pycore_list.h @@ -12,6 +12,8 @@ extern "C" { #include "pycore_stackref.h" #endif +#define _Py_LIST_SMALL_ALLOCATED 32 + PyAPI_FUNC(PyObject*) _PyList_Extend(PyListObject *, PyObject *); PyAPI_FUNC(PyObject) *_PyList_SliceSubscript(PyObject*, PyObject*); PyAPI_FUNC(PyObject *) _PyList_BinarySlice(PyObject *, PyObject *, PyObject *); diff --git a/Lib/test/test_capi/test_list.py b/Lib/test/test_capi/test_list.py index b95b8ba960bd8b..6cd1aa10b18c09 100644 --- a/Lib/test/test_capi/test_list.py +++ b/Lib/test/test_capi/test_list.py @@ -1,12 +1,18 @@ +import threading import unittest -from test.support import import_helper from collections import UserList +from test import support +from test.support import import_helper +from test.support import threading_helper + _testcapi = import_helper.import_module('_testcapi') _testlimitedcapi = import_helper.import_module('_testlimitedcapi') + NULL = None PY_SSIZE_T_MIN = _testcapi.PY_SSIZE_T_MIN PY_SSIZE_T_MAX = _testcapi.PY_SSIZE_T_MAX +LIST_SMALL_ALLOCATED = _testcapi.LIST_SMALL_ALLOCATED class ListSubclass(list): pass @@ -355,5 +361,170 @@ def test_uninitialized_list_repr(self): self.assertEqual(repr(lst), '[, , ]') +class ListWriterTests(unittest.TestCase): + def create_writer(self, size): + # Test PyListWriter_Create() + return _testcapi.PyListWriter(size) + + def list_equal(self, result, expected): + if not isinstance(expected, list): + expected = list(expected) + self.assertEqual(result, expected) + self.assertEqual(type(result), list) + + def create_items(self, nitem): + # Call object() to create unique objects. Don't use small integers + # which are immortal (and so would not check for reference leaks). + return tuple(object() for i in range(nitem)) + + def check_append(self, method): + items = self.create_items(3) + for size in (0, len(items), len(items) * 10): + with self.subTest(size=size): + writer = self.create_writer(size) + append = getattr(writer, method) + for item in items: + append(item) + self.list_equal(writer.finish(), items) + + # test MemoryError + writer = self.create_writer(0) + append = getattr(writer, method) + with self.assertRaises(MemoryError): + with support.inject_memory_error_cm(): + append("alloc failed") + for item in items: + append(item) + self.list_equal(writer.finish(), items) + + def test_append(self): + # Test PyListWriter_Append() + self.check_append('append') + + def test_append_steal(self): + # Test PyListWriter_AppendSteal() + self.check_append('append_steal') + + def check_extend(self, method): + items = self.create_items(5) + for size in (0, len(items), len(items) * 10): + with self.subTest(size=size): + writer = self.create_writer(size) + extend = getattr(writer, method) + extend(items) + self.list_equal(writer.finish(), items) + + # test MemoryError + writer = self.create_writer(0) + extend = getattr(writer, method) + items_error = ("alloc failed",) + with self.assertRaises(MemoryError): + with support.inject_memory_error_cm(): + extend(items_error) + extend(items) + self.list_equal(writer.finish(), items) + + def test_extend(self): + # Test PyListWriter_Extend() + self.check_extend('extend') + + def test_extend_array(self): + # Test PyListWriter_ExtendArray() + self.check_extend('extend_array') + + def test_extend_array_steal(self): + # Test PyListWriter_ExtendArraySteal() + self.check_extend('extend_array_steal') + + def test_empty(self): + list_get_allocated = _testcapi._list_get_allocated + empty_tuple = () + + for size in (0, LIST_SMALL_ALLOCATED, 1024): + with self.subTest(size=size): + writer = self.create_writer(size) + writer.extend(empty_tuple) + writer.extend_array(empty_tuple) + writer.extend_array_steal(empty_tuple) + result = writer.finish() + self.list_equal(result, []) + + if size == 0 or size > LIST_SMALL_ALLOCATED: + expected = 0 + else: + expected = size + self.assertEqual(list_get_allocated(result), expected) + + def test_finish_shrink(self): + # PyListWriter_Finish() shrinks the list if its allocated size + # is way too big + list_get_allocated = _testcapi._list_get_allocated + + def compute_allocated(size): + return (size + (size >> 3) + 6) & ~3 + + def check_allocated(items, alloc, expected): + writer = self.create_writer(alloc) + # Don't use PyListWriter_Extend() since it can shrink the list + for item in items: + writer.append(item) + result = writer.finish() + self.list_equal(result, list(items)) + self.assertEqual(list_get_allocated(result), expected) + + # small list + items = [1, 2, 3] + check_allocated(items, + 1024, + compute_allocated(len(items))) + + alloc = LIST_SMALL_ALLOCATED + check_allocated(items, alloc, alloc) + + # larger list + items = tuple(range(100)) + check_allocated(items, + (len(items) + 1) * 2, + compute_allocated(len(items))) + + alloc = len(items) * 2 + check_allocated(items, alloc, alloc) + + # do not shrink on MemoryError + alloc = 1024 + writer = self.create_writer(alloc) + writer.append(1) + with support.inject_memory_error_cm(): + # list_resize() never fails when shrinking + result = writer.finish() + self.list_equal(result, [1]) + self.assertEqual(list_get_allocated(result), alloc) + + def test_discard(self): + # Test PyListWriter_Discard(): must not leak memory + writer = self.create_writer(0) + writer.extend([2, 5, 7]) + writer.discard() + + @threading_helper.requires_working_threading() + def test_thread(self): + # PyListWriter can be used by multiple threads: it's up to the caller + # to implement a lock to prevent concurrent accesses. + writer = self.create_writer(0) + + def thread_func(writer): + writer.append('a') + writer.append_steal('b') + writer.extend(('c',)) + writer.extend_array(('d',)) + writer.extend_array_steal(('e',)) + + thread = threading.Thread(target=thread_func, args=(writer,)) + thread.start() + threading_helper.join_thread(thread) + + self.list_equal(writer.finish(), list('abcde')) + + if __name__ == "__main__": unittest.main() diff --git a/Misc/NEWS.d/next/C_API/2026-10-21-16-35-42.gh-issue-158660.KvixNp.rst b/Misc/NEWS.d/next/C_API/2026-10-21-16-35-42.gh-issue-158660.KvixNp.rst new file mode 100644 index 00000000000000..bbedca24001651 --- /dev/null +++ b/Misc/NEWS.d/next/C_API/2026-10-21-16-35-42.gh-issue-158660.KvixNp.rst @@ -0,0 +1,15 @@ +Add a new :ref:`PyListWriter ` C API to create a :class:`list` +object. On Free Threading, it's more efficient than the ``PyList`` C API to +create a :class:`list`, because it can avoid the critical section on the +created list. Add functions: + +* :c:func:`PyListWriter_Create` +* :c:func:`PyListWriter_Finish` +* :c:func:`PyListWriter_Discard` +* :c:func:`PyListWriter_Append` +* :c:func:`PyListWriter_AppendSteal` +* :c:func:`PyListWriter_Extend` +* :c:func:`PyListWriter_ExtendArray` +* :c:func:`PyListWriter_ExtendArraySteal` + +Patch by Victor Stinner. diff --git a/Modules/_testcapi/list.c b/Modules/_testcapi/list.c index 530b47780ac94e..5f82f5c304476b 100644 --- a/Modules/_testcapi/list.c +++ b/Modules/_testcapi/list.c @@ -1,6 +1,11 @@ +// Use pycore_list.h +#define PYTESTCAPI_NEED_INTERNAL_API + #include "parts.h" #include "util.h" +#include "pycore_list.h" // _Py_LIST_SMALL_ALLOCATED + static PyObject * list_get_size(PyObject *Py_UNUSED(module), PyObject *obj) @@ -103,6 +108,243 @@ test_list_api(PyObject *self, PyObject *Py_UNUSED(ignored)) } +static PyObject * +_list_get_allocated(PyObject *Py_UNUSED(module), PyObject *obj) +{ + NULLABLE(obj); + PyListObject *list = (PyListObject*)obj; + return PyLong_FromSsize_t(list->allocated); +} + + +typedef struct { + PyObject_HEAD + PyListWriter *writer; +} WriterObject; + + +static inline int +writer_check(WriterObject *self) +{ + if (self->writer == NULL) { + PyErr_SetString(PyExc_ValueError, + "operation on finished/discarded writer"); + return -1; + } + return 0; +} + + +static PyObject * +writer_new(PyTypeObject *type, PyObject *args, PyObject *kwargs) +{ + WriterObject *self = (WriterObject *)type->tp_alloc(type, 0); + if (!self) { + return NULL; + } + self->writer = NULL; + return (PyObject*)self; +} + + +// Test PyListWriter_Create() +static int +writer_init(PyObject *self_raw, PyObject *args, PyObject *kwargs) +{ + if (kwargs && PyDict_GET_SIZE(kwargs)) { + PyErr_Format(PyExc_TypeError, + "PyListWriter() takes no keyword arguments"); + return -1; + } + + Py_ssize_t size = 0; + if (!PyArg_ParseTuple(args, "|n", &size)) { + return -1; + } + + WriterObject *self = (WriterObject *)self_raw; + PyListWriter_Discard(self->writer); + + self->writer = PyListWriter_Create(size); + if (self->writer == NULL) { + return -1; + } + + return 0; +} + + +// Test PyListWriter_Discard() +static void +writer_dealloc(PyObject *self_raw) +{ + WriterObject *self = (WriterObject *)self_raw; + PyTypeObject *tp = Py_TYPE(self); + if (self->writer) { + PyListWriter_Discard(self->writer); + } + tp->tp_free(self); + Py_DECREF(tp); +} + + +// Test PyListWriter_Append() +static PyObject* +writer_append(PyObject *self_raw, PyObject *item) +{ + WriterObject *self = (WriterObject *)self_raw; + if (writer_check(self) < 0) { + return NULL; + } + + if (PyListWriter_Append(self->writer, item) < 0) { + return NULL; + } + Py_RETURN_NONE; +} + + +// Test PyListWriter_AppendSteal() +static PyObject* +writer_append_steal(PyObject *self_raw, PyObject *item) +{ + WriterObject *self = (WriterObject *)self_raw; + if (writer_check(self) < 0) { + return NULL; + } + + if (PyListWriter_AppendSteal(self->writer, Py_NewRef(item)) < 0) { + return NULL; + } + Py_RETURN_NONE; +} + + +// Test PyListWriter_Extend() +static PyObject* +writer_extend(PyObject *self_raw, PyObject *iterable) +{ + WriterObject *self = (WriterObject *)self_raw; + if (writer_check(self) < 0) { + return NULL; + } + + if (PyListWriter_Extend(self->writer, iterable) < 0) { + return NULL; + } + Py_RETURN_NONE; +} + + +// Test PyListWriter_ExtendArray() +static PyObject* +writer_extend_array(PyObject *self_raw, PyObject *tuple) +{ + WriterObject *self = (WriterObject *)self_raw; + if (writer_check(self) < 0) { + return NULL; + } + + if (!PyTuple_Check(tuple)) { + PyErr_Format(PyExc_TypeError, + "expect tuple, got %T", tuple); + return NULL; + } + PyObject **items = &PyTuple_GET_ITEM(tuple, 0); + Py_ssize_t len = PyTuple_GET_SIZE(tuple); + + if (PyListWriter_ExtendArray(self->writer, items, len) < 0) { + return NULL; + } + Py_RETURN_NONE; +} + + +// Test PyListWriter_ExtendArraySteal() +static PyObject* +writer_extend_array_steal(PyObject *self_raw, PyObject *tuple) +{ + WriterObject *self = (WriterObject *)self_raw; + if (writer_check(self) < 0) { + return NULL; + } + + if (!PyTuple_Check(tuple)) { + PyErr_Format(PyExc_TypeError, + "expect tuple, got %T", tuple); + return NULL; + } + PyObject **items = &PyTuple_GET_ITEM(tuple, 0); + Py_ssize_t len = PyTuple_GET_SIZE(tuple); + + for (Py_ssize_t i = 0; i < len; i++) { + Py_INCREF(items[i]); + } + + if (PyListWriter_ExtendArraySteal(self->writer, items, len) < 0) { + return NULL; + } + Py_RETURN_NONE; +} + + +// Test PyListWriter_Finish() +static PyObject* +writer_finish(PyObject *self_raw, PyObject *Py_UNUSED(args)) +{ + WriterObject *self = (WriterObject *)self_raw; + if (writer_check(self) < 0) { + return NULL; + } + + PyObject *list = PyListWriter_Finish(self->writer); + self->writer = NULL; + return list; +} + + +// Test PyListWriter_Discard() +static PyObject* +writer_discard(PyObject *self_raw, PyObject *Py_UNUSED(args)) +{ + WriterObject *self = (WriterObject *)self_raw; + if (writer_check(self) < 0) { + return NULL; + } + + PyListWriter_Discard(self->writer); + self->writer = NULL; + Py_RETURN_NONE; +} + + +static PyMethodDef writer_methods[] = { + {"append", _PyCFunction_CAST(writer_append), METH_O}, + {"append_steal", _PyCFunction_CAST(writer_append_steal), METH_O}, + {"extend", _PyCFunction_CAST(writer_extend), METH_O}, + {"extend_array", _PyCFunction_CAST(writer_extend_array), METH_O}, + {"extend_array_steal", _PyCFunction_CAST(writer_extend_array_steal), METH_O}, + {"finish", _PyCFunction_CAST(writer_finish), METH_NOARGS}, + {"discard", _PyCFunction_CAST(writer_discard), METH_NOARGS}, + {NULL, NULL} /* sentinel */ +}; + +static PyType_Slot Writer_Type_slots[] = { + {Py_tp_new, writer_new}, + {Py_tp_init, writer_init}, + {Py_tp_dealloc, writer_dealloc}, + {Py_tp_methods, writer_methods}, + {0, 0}, /* sentinel */ +}; + +static PyType_Spec Writer_spec = { + .name = "_testcapi.PyListWriter", + .basicsize = sizeof(WriterObject), + .flags = Py_TPFLAGS_DEFAULT, + .slots = Writer_Type_slots, +}; + + static PyMethodDef test_methods[] = { {"list_get_size", list_get_size, METH_O}, {"list_get_item", list_get_item, METH_VARARGS}, @@ -110,11 +352,30 @@ static PyMethodDef test_methods[] = { {"list_clear", list_clear, METH_O}, {"list_extend", list_extend, METH_VARARGS}, {"test_list_api", test_list_api, METH_NOARGS}, + {"_list_get_allocated", _list_get_allocated, METH_O}, {NULL}, }; int _PyTestCapi_Init_List(PyObject *m) { - return PyModule_AddFunctions(m, test_methods); + if (PyModule_AddFunctions(m, test_methods) < 0) { + return -1; + } + + PyTypeObject *writer_type = (PyTypeObject *)PyType_FromSpec(&Writer_spec); + if (writer_type == NULL) { + return -1; + } + if (PyModule_AddType(m, writer_type) < 0) { + Py_DECREF(writer_type); + return -1; + } + Py_DECREF(writer_type); + + if (PyModule_AddIntConstant(m, "LIST_SMALL_ALLOCATED", _Py_LIST_SMALL_ALLOCATED) < 0) { + return -1; + } + + return 0; } diff --git a/Objects/listobject.c b/Objects/listobject.c index 5ca8cdf09d3160..23165c5bc70fa3 100644 --- a/Objects/listobject.c +++ b/Objects/listobject.c @@ -91,7 +91,7 @@ ensure_shared_on_resize(PyListObject *self) #endif } -#define LIST_SMALL_ALLOCATED 32 +#define LIST_SMALL_ALLOCATED _Py_LIST_SMALL_ALLOCATED static int py_list_resize(PyListObject *self, Py_ssize_t newsize); @@ -299,30 +299,48 @@ PyList_New(Py_ssize_t size) return (PyObject *) op; } +// Similar to PyList_New() but don't initialize memory and do not track +// the newly created list in the garbage collector. +// +// The list is initialized with a size of 0. +// +// size must be >= 0. static PyObject * list_new_prealloc(Py_ssize_t size) { - assert(size > 0); - PyListObject *op = (PyListObject *) PyList_New(0); + // The caller is responsible to check the size + assert(size >= 0); + + PyListObject *op = _Py_FREELIST_POP(PyListObject, lists); if (op == NULL) { - return NULL; + op = PyObject_GC_New(PyListObject, &PyList_Type); + if (op == NULL) { + return NULL; + } } - assert(op->ob_item == NULL); -#ifdef Py_GIL_DISABLED - _PyListArray *array = list_allocate_array(size); - if (array == NULL) { - Py_DECREF(op); - return PyErr_NoMemory(); + + if (size <= 0) { + op->ob_item = NULL; } - op->ob_item = array->ob_item; + else { +#ifdef Py_GIL_DISABLED + _PyListArray *array = list_allocate_array(size); + if (array == NULL) { + Py_DECREF(op); + return PyErr_NoMemory(); + } + op->ob_item = array->ob_item; #else - op->ob_item = PyMem_New(PyObject *, size); - if (op->ob_item == NULL) { - Py_DECREF(op); - return PyErr_NoMemory(); - } + op->ob_item = (PyObject **) PyMem_Malloc(size * sizeof(PyObject *)); #endif + if (op->ob_item == NULL) { + Py_DECREF(op); + return PyErr_NoMemory(); + } + } + Py_SET_SIZE(op, 0); op->allocated = size; + return (PyObject *) op; } @@ -766,6 +784,7 @@ list_slice_lock_held(PyListObject *a, Py_ssize_t ilow, Py_ssize_t ihigh) dest[i] = Py_NewRef(v); } Py_SET_SIZE(np, len); + _PyObject_GC_TRACK(np); return (PyObject *)np; } @@ -848,6 +867,7 @@ list_concat_lock_held(PyListObject *a, PyListObject *b) dest[i] = Py_NewRef(v); } Py_SET_SIZE(np, size); + _PyObject_GC_TRACK(np); return (PyObject *)np; } @@ -918,6 +938,7 @@ list_repeat_impl(PyListObject *a, Py_ssize_t n) } Py_SET_SIZE(np, output_size); + _PyObject_GC_TRACK(np); return (PyObject *) np; } @@ -1507,6 +1528,7 @@ list_extend_dictitems(PyListObject *self, PyDictObject *dict) return 0; } +// If you update this function, update also PyListWriter_Extend(). static int _list_extend(PyListObject *self, PyObject *iterable) { @@ -3722,6 +3744,7 @@ list_slice_step_lock_held(PyListObject *a, Py_ssize_t start, Py_ssize_t step, Py dest[i] = Py_NewRef(v); } Py_SET_SIZE(np, len); + _PyObject_GC_TRACK(np); return (PyObject *)np; } @@ -4387,3 +4410,252 @@ listiter_reduce_general(void *_it, int forward) } return Py_BuildValue("N(N)", iter, list); } + + +/* --- PyListWriter API --------------------------------------------------- */ + +#ifndef NDEBUG +static int +listwriter_check_consistency(PyListWriter *writer) +{ + assert(writer != NULL); + assert(PyList_CheckExact(_PyObject_CAST(writer))); + PyListObject *list = (PyListObject*)writer; + assert(Py_SIZE(list) >= 0); + assert(Py_SIZE(list) <= list->allocated); + // Don't test _PyObject_IsUniquelyReferenced() since a list writer can + // be shared between multiple threads with external locking. + assert(Py_REFCNT(list) == 1); + assert(!_PyObject_GC_IS_TRACKED(list)); + return 1; +} +#endif + + +PyListWriter* +PyListWriter_Create(Py_ssize_t size) +{ + if (size < 0) { + PyErr_SetString(PyExc_ValueError, "size must not be negative"); + return NULL; + } + + PyObject *list = list_new_prealloc(size); + if (list == NULL) { + return NULL; + } + assert(Py_SIZE(list) == 0); + // In release mode, list items are uninitialized pointers which is safe, + // since the list object is not tracked by the GC, and the writer caller + // must initialize all items. + +#ifdef Py_DEBUG + // In debug mode, set all items to NULL to be able to check that they are + // non-NULL in PyListWriter_Finish(). + if (1 <= size) { + memset(_PyList_CAST(list)->ob_item, 0, size * sizeof(PyObject*)); + } +#endif + + PyListWriter *writer = (PyListWriter*)list; + assert(listwriter_check_consistency(writer)); + return writer; +} + + +void +PyListWriter_Discard(PyListWriter *writer) +{ + if (writer == NULL) { + return; + } + + assert(listwriter_check_consistency(writer)); + PyObject *list = _PyObject_CAST(writer); + Py_DECREF(list); +} + + +PyObject* +PyListWriter_Finish(PyListWriter *writer) +{ + assert(listwriter_check_consistency(writer)); + PyListObject *list = _PyList_CAST(writer); + + // Shrink the list if the allocated size is greater than the double size. + // For example, if PyListWriter_Create() argument overestimated the final + // size. + if (Py_SIZE(list) < (list->allocated >> 1) + && list->allocated > LIST_SMALL_ALLOCATED) + { + if (list_resize(list, Py_SIZE(list)) < 0) { + // This case cannot happen: list_resize() never fails when + // shrinking + PyListWriter_Discard(writer); + return NULL; + } + // Note: cannot test that list->allocated changed in debug mode, since + // it's left unchanged on memory allocation failure. + } + +#ifdef Py_DEBUG + // Make sure that all items have been initialized + Py_ssize_t size = Py_SIZE(list); + PyObject **list_items = list->ob_item; + for (Py_ssize_t i = 0; i < size; i++) { + assert(list_items[i] != NULL); + } +#endif + + // Now that the list is fully initialized, it's safe to track it in the GC + _PyObject_GC_TRACK(list); + return _PyObject_CAST(list); +} + + +static int +listwriter_resize(PyListWriter *writer, Py_ssize_t new_size) +{ + PyListObject *list = _PyList_CAST(writer); + + // On Free Threading, it can call _PyObject_GC_SET_SHARED() on the list + return list_resize(list, new_size); +} + + +int +PyListWriter_AppendSteal(PyListWriter *writer, PyObject *item) +{ + assert(listwriter_check_consistency(writer)); + PyListObject *list = _PyList_CAST(writer); + + Py_ssize_t len = Py_SIZE(list); + Py_ssize_t allocated = list->allocated; + if (len < allocated) { + list->ob_item[len] = item; + Py_SET_SIZE(list, len + 1); + return 0; + } + + if (listwriter_resize(writer, len + 1) < 0) { + Py_DECREF(item); + return -1; + } + list->ob_item[len] = item; + return 0; +} + + +int +PyListWriter_Append(PyListWriter *writer, PyObject *item) +{ + return PyListWriter_AppendSteal(writer, Py_NewRef(item)); +} + + +// Similar to _list_extend() but don't lock the writer list. +// If you update this function, update also _list_extend(). +int +PyListWriter_Extend(PyListWriter *writer, PyObject *iterable) +{ + assert(listwriter_check_consistency(writer)); + PyListObject *self = _PyList_CAST(writer); + assert(_PyObject_CAST(self) != iterable); + + int res; + if (PyList_CheckExact(iterable)) { + Py_BEGIN_CRITICAL_SECTION(iterable); + res = list_extend_lock_held(self, iterable); + Py_END_CRITICAL_SECTION(); + } + else if (PyTuple_CheckExact(iterable)) { + res = list_extend_lock_held(self, iterable); + } + else if (PyAnySet_CheckExact(iterable)) { + Py_BEGIN_CRITICAL_SECTION(iterable); + res = list_extend_set(self, (PySetObject *)iterable); + Py_END_CRITICAL_SECTION(); + } + else if (PyDict_CheckExact(iterable)) { + Py_BEGIN_CRITICAL_SECTION(iterable); + res = list_extend_dict(self, (PyDictObject *)iterable, 0 /*keys*/); + Py_END_CRITICAL_SECTION(); + } + else if (Py_IS_TYPE(iterable, &PyDictKeys_Type)) { + PyDictObject *dict = ((_PyDictViewObject *)iterable)->dv_dict; + Py_BEGIN_CRITICAL_SECTION(dict); + res = list_extend_dict(self, dict, 0 /*keys*/); + Py_END_CRITICAL_SECTION(); + } + else if (Py_IS_TYPE(iterable, &PyDictValues_Type)) { + PyDictObject *dict = ((_PyDictViewObject *)iterable)->dv_dict; + Py_BEGIN_CRITICAL_SECTION(dict); + res = list_extend_dict(self, dict, 1 /*values*/); + Py_END_CRITICAL_SECTION(); + } + else if (Py_IS_TYPE(iterable, &PyDictItems_Type)) { + PyDictObject *dict = ((_PyDictViewObject *)iterable)->dv_dict; + Py_BEGIN_CRITICAL_SECTION(dict); + res = list_extend_dictitems(self, dict); + Py_END_CRITICAL_SECTION(); + } + else { + res = list_extend_iter_lock_held(self, iterable); + } + return res; +} + + +static int +listwriter_extend_array(PyListWriter *writer, + PyObject *const *items, Py_ssize_t nitem, + int steal) +{ + assert(listwriter_check_consistency(writer)); + assert(nitem >= 0); + + PyListObject *list = _PyList_CAST(writer); + Py_ssize_t size = Py_SIZE(list); + Py_ssize_t allocated = list->allocated; + Py_ssize_t new_size = size + nitem; + if (new_size > allocated) { + if (listwriter_resize(writer, new_size) < 0) { + if (steal) { + for (Py_ssize_t i = 0; i < nitem; i++) { + Py_DECREF(items[i]); + } + } + return -1; + } + } + else { + Py_SET_SIZE(list, new_size); + } + + PyObject **list_items = &list->ob_item[size]; + if (steal) { + memcpy(list_items, items, nitem * sizeof(PyObject*)); + } + else { + for (Py_ssize_t i = 0; i < nitem; i++) { + list_items[i] = Py_NewRef(items[i]); + } + } + return 0; +} + + +int +PyListWriter_ExtendArray(PyListWriter *writer, + PyObject *const *items, Py_ssize_t nitem) +{ + return listwriter_extend_array(writer, items, nitem, 0); +} + + +int +PyListWriter_ExtendArraySteal(PyListWriter *writer, + PyObject *const *items, Py_ssize_t nitem) +{ + return listwriter_extend_array(writer, items, nitem, 1); +}