bpo-34151: Improve performance of some list operations (GH-8332)

1 files changed, 36 insertions, 15 deletions

diff --git a/Objects/listobject.c b/Objects/listobject.c
index 308ae83..3d4a187 100644
--- a/Objects/listobject.c
+++ b/Objects/listobject.c
@@ -180,6 +180,23 @@ PyList_New(Py_ssize_t size)
     return (PyObject *) op;
 }
 
+static PyObject *
+list_new_prealloc(Py_ssize_t size)
+{
+    PyListObject *op = (PyListObject *) PyList_New(0);
+    if (size == 0 || op == NULL) {
+        return (PyObject *) op;
+    }
+    assert(op->ob_item == NULL);
+    op->ob_item = PyMem_New(PyObject *, size);
+    if (op->ob_item == NULL) {
+        Py_DECREF(op);
+        return PyErr_NoMemory();
+    }
+    op->allocated = size;
+    return (PyObject *) op;
+}
+
 Py_ssize_t
 PyList_Size(PyObject *op)
 {
@@ -438,7 +455,7 @@ list_slice(PyListObject *a, Py_ssize_t ilow, Py_ssize_t ihigh)
     else if (ihigh > Py_SIZE(a))
         ihigh = Py_SIZE(a);
     len = ihigh - ilow;
-    np = (PyListObject *) PyList_New(len);
+    np = (PyListObject *) list_new_prealloc(len);
     if (np == NULL)
         return NULL;
 
@@ -449,6 +466,7 @@ list_slice(PyListObject *a, Py_ssize_t ilow, Py_ssize_t ihigh)
         Py_INCREF(v);
         dest[i] = v;
     }
+    Py_SIZE(np) = len;
     return (PyObject *)np;
 }
 
@@ -479,7 +497,7 @@ list_concat(PyListObject *a, PyObject *bb)
     if (Py_SIZE(a) > PY_SSIZE_T_MAX - Py_SIZE(b))
         return PyErr_NoMemory();
     size = Py_SIZE(a) + Py_SIZE(b);
-    np = (PyListObject *) PyList_New(size);
+    np = (PyListObject *) list_new_prealloc(size);
     if (np == NULL) {
         return NULL;
     }
@@ -497,6 +515,7 @@ list_concat(PyListObject *a, PyObject *bb)
         Py_INCREF(v);
         dest[i] = v;
     }
+    Py_SIZE(np) = size;
     return (PyObject *)np;
 #undef b
 }
@@ -516,28 +535,30 @@ list_repeat(PyListObject *a, Py_ssize_t n)
     size = Py_SIZE(a) * n;
     if (size == 0)
         return PyList_New(0);
-    np = (PyListObject *) PyList_New(size);
+    np = (PyListObject *) list_new_prealloc(size);
     if (np == NULL)
         return NULL;
 
-    items = np->ob_item;
     if (Py_SIZE(a) == 1) {
+        items = np->ob_item;
         elem = a->ob_item[0];
         for (i = 0; i < n; i++) {
             items[i] = elem;
             Py_INCREF(elem);
         }
-        return (PyObject *) np;
-    }
-    p = np->ob_item;
-    items = a->ob_item;
-    for (i = 0; i < n; i++) {
-        for (j = 0; j < Py_SIZE(a); j++) {
-            *p = items[j];
-            Py_INCREF(*p);
-            p++;
+    }
+    else {
+        p = np->ob_item;
+        items = a->ob_item;
+        for (i = 0; i < n; i++) {
+            for (j = 0; j < Py_SIZE(a); j++) {
+                *p = items[j];
+                Py_INCREF(*p);
+                p++;
+            }
         }
     }
+    Py_SIZE(np) = size;
     return (PyObject *) np;
 }
 
@@ -2738,7 +2759,7 @@ list_subscript(PyListObject* self, PyObject* item)
             return list_slice(self, start, stop);
         }
         else {
-            result = PyList_New(slicelength);
+            result = list_new_prealloc(slicelength);
             if (!result) return NULL;
 
             src = self->ob_item;
@@ -2749,7 +2770,7 @@ list_subscript(PyListObject* self, PyObject* item)
                 Py_INCREF(it);
                 dest[i] = it;
             }
-
+            Py_SIZE(result) = slicelength;
             return result;
         }
     }