gh-125196: Use PyUnicodeWriter in symtable.c (#125199)

diff --git a/Python/symtable.c b/Python/symtable.c
index 52be910c0b6a9b1544867b15cab74ca8d9788cc9..709918b27afcc89c3f5376cf2825f520f4ed9ba5 100644 (file)
--- a/Python/symtable.c
+++ b/Python/symtable.c
@@ -3120,33 +3120,30 @@ _Py_Mangle(PyObject *privateobj, PyObject *ident)
     if (ipriv == plen) {
         return Py_NewRef(ident); /* Don't mangle if class is just underscores */
     }
-    plen -= ipriv;
 
-    if (plen + nlen >= PY_SSIZE_T_MAX - 1) {
+    if (nlen + (plen - ipriv) >= PY_SSIZE_T_MAX - 1) {
         PyErr_SetString(PyExc_OverflowError,
                         "private identifier too large to be mangled");
         return NULL;
     }
 
-    Py_UCS4 maxchar = PyUnicode_MAX_CHAR_VALUE(ident);
-    if (PyUnicode_MAX_CHAR_VALUE(privateobj) > maxchar) {
-        maxchar = PyUnicode_MAX_CHAR_VALUE(privateobj);
-    }
-
-    PyObject *result = PyUnicode_New(1 + nlen + plen, maxchar);
-    if (!result) {
+    PyUnicodeWriter *writer = PyUnicodeWriter_Create(1 + nlen + (plen - ipriv));
+    if (!writer) {
         return NULL;
     }
-    /* ident = "_" + priv[ipriv:] + ident # i.e. 1+plen+nlen bytes */
-    PyUnicode_WRITE(PyUnicode_KIND(result), PyUnicode_DATA(result), 0, '_');
-    if (PyUnicode_CopyCharacters(result, 1, privateobj, ipriv, plen) < 0) {
-        Py_DECREF(result);
-        return NULL;
+    // ident = "_" + priv[ipriv:] + ident
+    if (PyUnicodeWriter_WriteChar(writer, '_') < 0) {
+        goto error;
     }
-    if (PyUnicode_CopyCharacters(result, plen+1, ident, 0, nlen) < 0) {
-        Py_DECREF(result);
-        return NULL;
+    if (PyUnicodeWriter_WriteSubstring(writer, privateobj, ipriv, plen) < 0) {
+        goto error;
     }
-    assert(_PyUnicode_CheckConsistency(result, 1));
-    return result;
+    if (PyUnicodeWriter_WriteStr(writer, ident) < 0) {
+        goto error;
+    }
+    return PyUnicodeWriter_Finish(writer);
+
+error:
+    PyUnicodeWriter_Discard(writer);
+    return NULL;
 }