Optimize _Py_strhex_impl() (GH-19535)

vstinner · web-flow · commit 455df9779873 · 2020-04-15T14:05:24.000+02:00
Avoid a temporary buffer to create a bytes string: use
PyBytes_FromStringAndSize() to directly allocate a bytes object.

Cleanup also the code: PEP 7 formatting, move variable definitions
closer to where they are used. Fix assertion checking "j" index.
diff --git a/Python/pystrhex.c b/Python/pystrhex.c
@@ -8,12 +8,9 @@ static PyObject *_Py_strhex_impl(const char* argbuf, const Py_ssize_t arglen,
                                  const PyObject* sep, int bytes_per_sep_group,
                                  const int return_bytes)
 {
-    PyObject *retval;
-    Py_UCS1* retbuf;
-    Py_ssize_t i, j, resultlen = 0;
-    Py_UCS1 sep_char = 0;
-    unsigned int abs_bytes_per_sep;
+    assert(arglen >= 0);
 
+    Py_UCS1 sep_char = 0;
     if (sep) {
         Py_ssize_t seplen = PyObject_Length((PyObject*)sep);
         if (seplen < 0) {
@@ -31,22 +28,25 @@ static PyObject *_Py_strhex_impl(const char* argbuf, const Py_ssize_t arglen,
                 return NULL;
             }
             sep_char = PyUnicode_READ_CHAR(sep, 0);
-        } else if (PyBytes_Check(sep)) {
+        }
+        else if (PyBytes_Check(sep)) {
             sep_char = PyBytes_AS_STRING(sep)[0];
-        } else {
+        }
+        else {
             PyErr_SetString(PyExc_TypeError, "sep must be str or bytes.");
             return NULL;
         }
         if (sep_char > 127 && !return_bytes) {
             PyErr_SetString(PyExc_ValueError, "sep must be ASCII.");
             return NULL;
         }
-    } else {
+    }
+    else {
         bytes_per_sep_group = 0;
     }
 
-    assert(arglen >= 0);
-    abs_bytes_per_sep = abs(bytes_per_sep_group);
+    unsigned int abs_bytes_per_sep = abs(bytes_per_sep_group);
+    Py_ssize_t resultlen = 0;
     if (bytes_per_sep_group && arglen > 0) {
         /* How many sep characters we'll be inserting. */
         resultlen = (arglen - 1) / abs_bytes_per_sep;
@@ -62,26 +62,32 @@ static PyObject *_Py_strhex_impl(const char* argbuf, const Py_ssize_t arglen,
         abs_bytes_per_sep = 0;
     }
 
+    PyObject *retval;
+    Py_UCS1 *retbuf;
     if (return_bytes) {
         /* If _PyBytes_FromSize() were public we could avoid malloc+copy. */
-        retbuf = (Py_UCS1*) PyMem_Malloc(resultlen);
-        if (!retbuf)
-            return PyErr_NoMemory();
-        retval = NULL;  /* silence a compiler warning, assigned later. */
-    } else {
+        retval = PyBytes_FromStringAndSize(NULL, resultlen);
+        if (!retval) {
+            return NULL;
+        }
+        retbuf = (Py_UCS1 *)PyBytes_AS_STRING(retval);
+    }
+    else {
         retval = PyUnicode_New(resultlen, 127);
-        if (!retval)
+        if (!retval) {
             return NULL;
+        }
         retbuf = PyUnicode_1BYTE_DATA(retval);
     }
 
     /* Hexlify */
+    Py_ssize_t i, j;
     for (i=j=0; i < arglen; ++i) {
-        assert(j < resultlen);
+        assert((j + 1) < resultlen);
         unsigned char c;
-        c = (argbuf[i] >> 4) & 0xf;
+        c = (argbuf[i] >> 4) & 0x0f;
         retbuf[j++] = Py_hexdigits[c];
-        c = argbuf[i] & 0xf;
+        c = argbuf[i] & 0x0f;
         retbuf[j++] = Py_hexdigits[c];
         if (bytes_per_sep_group && i < arglen - 1) {
             Py_ssize_t anchor;
@@ -93,12 +99,8 @@ static PyObject *_Py_strhex_impl(const char* argbuf, const Py_ssize_t arglen,
     }
     assert(j == resultlen);
 
-    if (return_bytes) {
-        retval = PyBytes_FromStringAndSize((const char *)retbuf, resultlen);
-        PyMem_Free(retbuf);
-    }
 #ifdef Py_DEBUG
-    else {
+    if (!return_bytes) {
         assert(_PyUnicode_CheckConsistency(retval, 1));
     }
 #endif
