Skip to content

Commit a8142bd

Browse files
authored
Merge branch 'main' into gh-158121-monitoring-disable
2 parents 47298eb + 00307b0 commit a8142bd

3 files changed

Lines changed: 90 additions & 57 deletions

File tree

‎Include/internal/pycore_unicodeobject.h‎

Lines changed: 36 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -130,22 +130,39 @@ _PyUnicodeWriter_CanWrite(_PyUnicodeWriter *writer)
130130
#endif
131131

132132
static inline void
133-
_PyUnicodeWriter_Update(_PyUnicodeWriter *writer)
133+
_PyUnicodeWriter_SetBuffer(_PyUnicodeWriter *writer, PyObject *buffer)
134134
{
135-
PyObject *buffer = writer->buffer;
136-
writer->maxchar = PyUnicode_MAX_CHAR_VALUE(buffer);
135+
assert(writer->pos <= PyUnicode_GET_LENGTH(buffer));
136+
137+
// Py_DECREF() the previous buffer (if any)
138+
Py_XSETREF(writer->buffer, buffer);
137139
writer->data = PyUnicode_DATA(buffer);
138140
writer->kind = PyUnicode_KIND(buffer);
141+
writer->maxchar = PyUnicode_MAX_CHAR_VALUE(buffer);
142+
writer->size = PyUnicode_GET_LENGTH(buffer);
143+
writer->readonly = 0;
144+
}
139145

140-
if (!writer->readonly) {
141-
writer->size = PyUnicode_GET_LENGTH(buffer);
142-
}
143-
else {
144-
/* Copy-on-write mode: set buffer size to 0 so
145-
* _PyUnicodeWriter_Prepare() will copy (and enlarge) the buffer on
146-
* next write. */
147-
writer->size = 0;
148-
}
146+
static inline void
147+
_PyUnicodeWriter_SetReadOnly(_PyUnicodeWriter *writer, PyObject *obj,
148+
Py_ssize_t length)
149+
{
150+
assert(writer->buffer == NULL);
151+
assert(writer->pos == 0);
152+
// Micro-optimization: pass length as a parameter, as it's usually known
153+
// by the caller
154+
assert(length == PyUnicode_GET_LENGTH(obj));
155+
156+
writer->buffer = obj;
157+
writer->data = NULL;
158+
/* Set kind and size to 0 to make sure that the next
159+
* _PyUnicodeWriter_Prepare() call allocates a new buffer and copies
160+
* characters. */
161+
writer->kind = 0;
162+
writer->maxchar = PyUnicode_MAX_CHAR_VALUE(obj);
163+
writer->size = 0;
164+
writer->pos = length;
165+
writer->readonly = 1;
149166
}
150167

151168
static inline int
@@ -156,11 +173,9 @@ _PyUnicodeWriter_WriteCharInline(_PyUnicodeWriter *writer, Py_UCS4 ch)
156173
// If the first write is a Latin1 character, use the singleton
157174
// as a read-only object
158175
PyObject *obj = _Py_LATIN1_CHR(ch);
159-
writer->readonly = 1;
160-
writer->buffer = obj; // Py_NewRef() is not need on immortal object
161-
_PyUnicodeWriter_Update(writer);
162-
assert(writer->pos == 0);
163-
writer->pos = 1;
176+
// Py_NewRef() is not need on immortal object
177+
_PyUnicodeWriter_SetReadOnly(writer, obj, 1);
178+
164179
// The next write will create a new buffer and copy the string
165180
return 0;
166181
}
@@ -176,6 +191,10 @@ _PyUnicodeWriter_WriteCharInline(_PyUnicodeWriter *writer, Py_UCS4 ch)
176191
return 0;
177192
}
178193

194+
extern PyObject* _PyUnicodeWriter_FinishWithSize(
195+
_PyUnicodeWriter *writer,
196+
Py_ssize_t size);
197+
179198
/* --- Unicode API -------------------------------------------------------- */
180199

181200
// Export for '_json' shared extension

‎Objects/unicode_writer.c‎

Lines changed: 50 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -178,8 +178,7 @@ _PyUnicodeWriter_InitWithBuffer(_PyUnicodeWriter *writer, PyObject *buffer)
178178
assert(PyUnstable_Object_IsUniquelyReferenced(buffer));
179179

180180
memset(writer, 0, sizeof(*writer));
181-
writer->buffer = buffer;
182-
_PyUnicodeWriter_Update(writer);
181+
_PyUnicodeWriter_SetBuffer(writer, buffer);
183182
writer->min_length = writer->size;
184183
assert(_PyUnicodeWriter_CanWrite(writer));
185184
}
@@ -204,16 +203,19 @@ _PyUnicodeWriter_PrepareInternal(_PyUnicodeWriter *writer,
204203

205204
maxchar = Py_MAX(maxchar, writer->min_char);
206205

207-
PyObject *newbuffer;
206+
PyObject *new_buffer;
208207
if (writer->buffer == NULL) {
209208
assert(!writer->readonly);
209+
210210
// Do not overallocate at the first allocation, but use min_length
211-
if (alloc < writer->min_length)
211+
if (alloc < writer->min_length) {
212212
alloc = writer->min_length;
213+
}
213214

214-
writer->buffer = PyUnicode_New(alloc, maxchar);
215-
if (writer->buffer == NULL)
215+
new_buffer = PyUnicode_New(alloc, maxchar);
216+
if (new_buffer == NULL) {
216217
return -1;
218+
}
217219
}
218220
else if (alloc > writer->size) {
219221
// Do not overallocate at the first allocation, but use min_length
@@ -223,38 +225,42 @@ _PyUnicodeWriter_PrepareInternal(_PyUnicodeWriter *writer,
223225
/* overallocate to limit the number of realloc() */
224226
alloc += alloc / OVERALLOCATE_FACTOR;
225227
}
226-
if (alloc < writer->min_length)
228+
if (alloc < writer->min_length) {
227229
alloc = writer->min_length;
230+
}
228231

229232
if (maxchar > writer->maxchar || writer->readonly) {
230233
/* resize + widen */
231234
maxchar = Py_MAX(maxchar, writer->maxchar);
232-
newbuffer = PyUnicode_New(alloc, maxchar);
233-
if (newbuffer == NULL)
235+
new_buffer = PyUnicode_New(alloc, maxchar);
236+
if (new_buffer == NULL) {
234237
return -1;
235-
_PyUnicode_FastCopyCharacters(newbuffer, 0,
238+
}
239+
_PyUnicode_FastCopyCharacters(new_buffer, 0,
236240
writer->buffer, 0, writer->pos);
237-
writer->readonly = 0;
238-
Py_DECREF(writer->buffer);
239-
writer->buffer = newbuffer;
240241
}
241242
else {
242-
newbuffer = _PyUnicode_ResizeCompact(writer->buffer, alloc);
243-
if (newbuffer == NULL)
243+
new_buffer = _PyUnicode_ResizeCompact(writer->buffer, alloc);
244+
if (new_buffer == NULL) {
244245
return -1;
245-
writer->buffer = newbuffer;
246+
}
247+
// Do not DECREF the old buffer
248+
writer->buffer = NULL;
246249
}
247250
}
248-
else if (maxchar > writer->maxchar) {
251+
else {
252+
assert(maxchar > writer->maxchar);
249253
assert(!writer->readonly);
250-
newbuffer = PyUnicode_New(writer->size, maxchar);
251-
if (newbuffer == NULL)
254+
255+
new_buffer = PyUnicode_New(writer->size, maxchar);
256+
if (new_buffer == NULL) {
252257
return -1;
253-
_PyUnicode_FastCopyCharacters(newbuffer, 0,
258+
}
259+
_PyUnicode_FastCopyCharacters(new_buffer, 0,
254260
writer->buffer, 0, writer->pos);
255-
Py_SETREF(writer->buffer, newbuffer);
256261
}
257-
_PyUnicodeWriter_Update(writer);
262+
263+
_PyUnicodeWriter_SetBuffer(writer, new_buffer);
258264
return 0;
259265

260266
#undef OVERALLOCATE_FACTOR
@@ -316,11 +322,7 @@ _PyUnicodeWriter_WriteStr(_PyUnicodeWriter *writer, PyObject *str)
316322
if (maxchar > writer->maxchar || len > writer->size - writer->pos) {
317323
if (writer->buffer == NULL && PyUnicode_CheckExact(str)) {
318324
assert(_PyUnicode_CheckConsistency(str, 1));
319-
writer->readonly = 1;
320-
writer->buffer = Py_NewRef(str);
321-
_PyUnicodeWriter_Update(writer);
322-
writer->pos += len;
323-
// The next write will create a new buffer and copy the string
325+
_PyUnicodeWriter_SetReadOnly(writer, Py_NewRef(str), len);
324326
return 0;
325327
}
326328
if (_PyUnicodeWriter_PrepareInternal(writer, len, maxchar) == -1)
@@ -457,10 +459,7 @@ _PyUnicodeWriter_WriteASCIIString(_PyUnicodeWriter *writer,
457459
if (str == NULL)
458460
return -1;
459461

460-
writer->readonly = 1;
461-
writer->buffer = str;
462-
_PyUnicodeWriter_Update(writer);
463-
writer->pos += len;
462+
_PyUnicodeWriter_SetReadOnly(writer, str, len);
464463
return 0;
465464
}
466465

@@ -639,6 +638,26 @@ _PyUnicodeWriter_Finish(_PyUnicodeWriter *writer)
639638
}
640639

641640

641+
PyObject *
642+
_PyUnicodeWriter_FinishWithSize(_PyUnicodeWriter *writer, Py_ssize_t size)
643+
{
644+
assert(0 <= size);
645+
if (writer->buffer != NULL) {
646+
assert(size <= writer->pos);
647+
assert(size <= PyUnicode_GET_LENGTH(writer->buffer));
648+
if (size < writer->pos) {
649+
// Truncate the string: we may need to adjust the string kind
650+
writer->recheck_maxchar = 1;
651+
}
652+
}
653+
else {
654+
assert(size == 0);
655+
}
656+
writer->pos = size;
657+
return _PyUnicodeWriter_Finish(writer);
658+
}
659+
660+
642661
PyObject*
643662
PyUnicodeWriter_Finish(PyUnicodeWriter *writer)
644663
{

‎Objects/unicodeobject.c‎

Lines changed: 4 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -4761,15 +4761,10 @@ PyUnicode_DecodeUTF7Stateful(const char *s,
47614761
if (consumed) {
47624762
if (inShift) {
47634763
*consumed = startinpos;
4764-
if (writer.pos != shiftOutStart && writer.maxchar > 127) {
4765-
PyObject *result = PyUnicode_FromKindAndData(
4766-
writer.kind, writer.data, shiftOutStart);
4767-
Py_XDECREF(errorHandler);
4768-
Py_XDECREF(exc);
4769-
_PyUnicodeWriter_Dealloc(&writer);
4770-
return result;
4771-
}
4772-
writer.pos = shiftOutStart; /* back off output */
4764+
4765+
Py_XDECREF(errorHandler);
4766+
Py_XDECREF(exc);
4767+
return _PyUnicodeWriter_FinishWithSize(&writer, shiftOutStart);
47734768
}
47744769
else {
47754770
*consumed = s-starts;

0 commit comments

Comments
 (0)