From 0b19e8664075b3b9f72be827eeb14113097517da Mon Sep 17 00:00:00 2001 From: Victor Stinner Date: Mon, 5 Oct 2026 03:54:59 +0200 Subject: [PATCH] gh-158585: Simplify PyBytesWriter_Finish() PyBytesWriter_Finish() doesn't need to check the size. It's known to be valid. PyBytesWriter_Create() now only calls byteswriter_resize() if the size is larger than the small buffer. Remove byteswriter_data() function: call static inline _PyBytesWriter_GetData() function directly. --- Modules/_struct.c | 3 +- Objects/bytesobject.c | 101 +++++++++++++++++++++------------------- Objects/unicodeobject.c | 7 +-- Python/assemble.c | 2 +- 4 files changed, 57 insertions(+), 56 deletions(-) diff --git a/Modules/_struct.c b/Modules/_struct.c index fb2f40ab9766c5b..1e5d52b758f7e29 100644 --- a/Modules/_struct.c +++ b/Modules/_struct.c @@ -2498,7 +2498,8 @@ Struct_pack_impl(PyStructObject *self, PyObject * const *values, return NULL; } - return PyBytesWriter_FinishWithSize(writer, self->s_size); + assert(PyBytesWriter_GetSize(writer) == self->s_size); + return PyBytesWriter_Finish(writer); } /*[clinic input] diff --git a/Objects/bytesobject.c b/Objects/bytesobject.c index af5dce763e1f25d..9bd9afc9b51ba4c 100644 --- a/Objects/bytesobject.c +++ b/Objects/bytesobject.c @@ -3731,12 +3731,6 @@ _PyBytes_RepeatBuffer(char* dest, Py_ssize_t len_dest, // one extra NUL byte which is a common error. #define PyBytesWriter_CANARY_BYTE PYMEM_DEADBYTE -static inline char* -byteswriter_data(PyBytesWriter *writer) -{ - return _PyBytesWriter_GetData(writer); -} - static inline Py_ssize_t byteswriter_allocated(PyBytesWriter *writer) @@ -3758,7 +3752,7 @@ byteswriter_allocated(PyBytesWriter *writer) static void byteswriter_write_canary_byte(PyBytesWriter *writer) { - unsigned char *data = (unsigned char*)byteswriter_data(writer); + unsigned char *data = (unsigned char*)_PyBytesWriter_GetData(writer); data[writer->size] = PyBytesWriter_CANARY_BYTE; } @@ -3770,7 +3764,7 @@ byteswriter_reset_trailing_byte(PyBytesWriter *writer) // bytes/bytearray expects the last byte to be a null byte. // Reset the last byte to null for bytes/bytearray. Py_ssize_t allocated = byteswriter_allocated(writer); - char *data = byteswriter_data(writer); + char *data = _PyBytesWriter_GetData(writer); data[allocated] = '\0'; } #endif @@ -3801,7 +3795,7 @@ byteswriter_check_consistency(PyBytesWriter *writer) } #ifdef Py_DEBUG - const unsigned char *data = (const unsigned char*)byteswriter_data(writer); + const unsigned char *data = (const unsigned char*)_PyBytesWriter_GetData(writer); unsigned char canary = data[writer->size]; if (canary != PyBytesWriter_CANARY_BYTE) { _Py_FatalErrorFormat(__func__, @@ -3895,7 +3889,7 @@ byteswriter_resize(PyBytesWriter *writer, Py_ssize_t new_size, int resize) if (resize) { Py_ssize_t old_size = writer->size; assert(allocated > old_size); - memset(byteswriter_data(writer) + old_size, PyBytesWrite_NEW_BYTE, + memset(_PyBytesWriter_GetData(writer) + old_size, PyBytesWrite_NEW_BYTE, allocated - old_size); } #endif @@ -3921,24 +3915,25 @@ byteswriter_create(Py_ssize_t size, int use_bytearray) } } writer->obj = NULL; - writer->size = 0; + writer->size = size; writer->use_bytearray = use_bytearray; writer->overallocate = !use_bytearray; - if (size >= 1) { + Py_ssize_t allocated = sizeof(writer->small_buffer) - 1; + if (size > allocated) { if (byteswriter_resize(writer, size, 0) < 0) { #ifdef Py_DEBUG // Write the canary byte so byteswriter_check_consistency() // doesn't fail in PyBytesWriter_Discard() + writer->size = 0; byteswriter_write_canary_byte(writer); #endif PyBytesWriter_Discard(writer); return NULL; } - writer->size = size; } #ifdef Py_DEBUG - memset(byteswriter_data(writer), PyBytesWrite_NEW_BYTE, + memset(_PyBytesWriter_GetData(writer), PyBytesWrite_NEW_BYTE, byteswriter_allocated(writer)); byteswriter_write_canary_byte(writer); #endif @@ -3967,34 +3962,23 @@ PyBytesWriter_Discard(PyBytesWriter *writer) } assert(byteswriter_check_consistency(writer)); -#ifdef Py_DEBUG if (writer->obj != NULL) { +#ifdef Py_DEBUG byteswriter_reset_trailing_byte(writer); - } #endif - - Py_XDECREF(writer->obj); + Py_DECREF(writer->obj); + } _Py_FREELIST_FREE(bytes_writers, writer, PyMem_Free); } -PyObject* -PyBytesWriter_FinishWithSize(PyBytesWriter *writer, Py_ssize_t size) +static inline PyObject* +byteswriter_finish_with_size(PyBytesWriter *writer, Py_ssize_t final_size) { assert(byteswriter_check_consistency(writer)); - if (size < 0) { - PyErr_Format(PyExc_ValueError, "size must be positive"); - goto error; - } - - if (size > writer->size) { - PyErr_SetString(PyExc_ValueError, "size larger than allocated size"); - goto error; - } - PyObject *result; - if (size == 0 && !writer->use_bytearray) { + if (final_size == 0 && !writer->use_bytearray) { result = bytes_get_empty(); if (writer->obj != NULL) { #ifdef Py_DEBUG @@ -4011,19 +3995,19 @@ PyBytesWriter_FinishWithSize(PyBytesWriter *writer, Py_ssize_t size) #endif if (writer->use_bytearray) { - if (size != PyByteArray_GET_SIZE(writer->obj)) { - if (PyByteArray_Resize(writer->obj, size)) { + if (final_size != PyByteArray_GET_SIZE(writer->obj)) { + if (PyByteArray_Resize(writer->obj, final_size)) { goto error; } } } else { - if (size == 1) { + if (final_size == 1) { unsigned char ch = PyBytes_AS_STRING(writer->obj)[0]; Py_SETREF(writer->obj, bytes_get_char(ch)); } - else if (size != PyBytes_GET_SIZE(writer->obj)) { - if (bytes_resize_inplace(&writer->obj, size)) { + else if (final_size != PyBytes_GET_SIZE(writer->obj)) { + if (bytes_resize_inplace(&writer->obj, final_size)) { goto error; } } @@ -4036,18 +4020,18 @@ PyBytesWriter_FinishWithSize(PyBytesWriter *writer, Py_ssize_t size) // Create an object from the small buffer const char *buffer = (const char *)writer->small_buffer; if (writer->use_bytearray) { - result = PyByteArray_FromStringAndSize(buffer, size); + result = PyByteArray_FromStringAndSize(buffer, final_size); } else { - if (size == 1) { + if (final_size == 1) { result = bytes_get_char((uint8_t)buffer[0]); } else { - result = bytes_alloc(size); + result = bytes_alloc(final_size); if (result == NULL) { goto error; } - memcpy(PyBytes_AS_STRING(result), buffer, size); + memcpy(PyBytes_AS_STRING(result), buffer, final_size); } } } @@ -4061,17 +4045,36 @@ PyBytesWriter_FinishWithSize(PyBytesWriter *writer, Py_ssize_t size) return NULL; } +PyObject* +PyBytesWriter_FinishWithSize(PyBytesWriter *writer, Py_ssize_t size) +{ + if (size < 0) { + PyErr_Format(PyExc_ValueError, "size must be positive"); + goto error; + } + if (size > writer->size) { + PyErr_SetString(PyExc_ValueError, "size larger than allocated size"); + goto error; + } + return byteswriter_finish_with_size(writer, size); + +error: + PyBytesWriter_Discard(writer); + return NULL; +} + PyObject* PyBytesWriter_Finish(PyBytesWriter *writer) { - return PyBytesWriter_FinishWithSize(writer, writer->size); + return byteswriter_finish_with_size(writer, writer->size); } PyObject* PyBytesWriter_FinishWithPointer(PyBytesWriter *writer, void *buf) { - Py_ssize_t size = (char*)buf - byteswriter_data(writer); + Py_ssize_t size = (char*)buf - _PyBytesWriter_GetData(writer); + // Call PyBytesWriter_FinishWithSize() to check size return PyBytesWriter_FinishWithSize(writer, size); } @@ -4081,7 +4084,7 @@ PyBytesWriter_GetData(PyBytesWriter *writer) { assert(byteswriter_check_consistency(writer)); - return byteswriter_data(writer); + return _PyBytesWriter_GetData(writer); } @@ -4125,11 +4128,11 @@ static void* _PyBytesWriter_ResizeAndUpdatePointer(PyBytesWriter *writer, Py_ssize_t size, void *data) { - Py_ssize_t pos = (char*)data - byteswriter_data(writer); + Py_ssize_t pos = (char*)data - _PyBytesWriter_GetData(writer); if (PyBytesWriter_Resize(writer, size) < 0) { return NULL; } - return byteswriter_data(writer) + pos; + return _PyBytesWriter_GetData(writer) + pos; } @@ -4176,11 +4179,11 @@ void* PyBytesWriter_GrowAndUpdatePointer(PyBytesWriter *writer, Py_ssize_t size, void *buf) { - Py_ssize_t pos = (char*)buf - byteswriter_data(writer); + Py_ssize_t pos = (char*)buf - _PyBytesWriter_GetData(writer); if (PyBytesWriter_Grow(writer, size) < 0) { return NULL; } - return byteswriter_data(writer) + pos; + return _PyBytesWriter_GetData(writer) + pos; } @@ -4201,7 +4204,7 @@ PyBytesWriter_WriteBytes(PyBytesWriter *writer, if (PyBytesWriter_Grow(writer, size) < 0) { return -1; } - char *buf = byteswriter_data(writer); + char *buf = _PyBytesWriter_GetData(writer); memcpy(buf + pos, bytes, size); assert(byteswriter_check_consistency(writer)); @@ -4232,7 +4235,7 @@ PyBytesWriter_Format(PyBytesWriter *writer, const char *format, ...) return -1; } - Py_ssize_t size = buf - byteswriter_data(writer); + Py_ssize_t size = buf - _PyBytesWriter_GetData(writer); return PyBytesWriter_Resize(writer, size); } diff --git a/Objects/unicodeobject.c b/Objects/unicodeobject.c index 157e005c3c0b600..ded13cc37f69f17 100644 --- a/Objects/unicodeobject.c +++ b/Objects/unicodeobject.c @@ -4880,7 +4880,7 @@ _PyUnicode_EncodeUTF7(PyObject *str, else { /* not in a shift sequence */ if (ch == '+') { *out++ = '+'; - *out++ = '-'; + *out++ = '-'; } else if (ENCODE_DIRECT(ch)) { *out++ = (char) ch; @@ -8644,10 +8644,7 @@ _PyUnicode_EncodeIconv(const char *encoding, PyObject *unicode, iconv(cd, NULL, NULL, NULL, NULL); } - if (PyBytesWriter_Resize(writer, out - (char *)PyBytesWriter_GetData(writer)) < 0) { - goto done; - } - result = PyBytesWriter_Finish(writer); + result = PyBytesWriter_FinishWithPointer(writer, out); writer = NULL; done: diff --git a/Python/assemble.c b/Python/assemble.c index 2c8abb116b2e336..8f68e233ab7b93c 100644 --- a/Python/assemble.c +++ b/Python/assemble.c @@ -102,7 +102,7 @@ assemble_free(struct assembler *a) static inline void write_except_byte(struct assembler *a, int byte) { - unsigned char *p = (unsigned char *) PyBytesWriter_GetData(a->a_except_table_writer); + unsigned char *p = PyBytesWriter_GetData(a->a_except_table_writer); p[a->a_except_table_off++] = byte; }