#1569291: speed up array.repeat() by making only O(log n) memcpy() calls; the code follows unicode_repeat.

birkenfeld · birkenfeld · commit c29cc6a8f27b · 2010-12-04T11:02:04.000Z
diff --git a/Lib/test/test_array.py b/Lib/test/test_array.py
@@ -504,6 +504,12 @@ def test_mul(self):
             array.array(self.typecode)
         )
 
+        a = 5 * array.array(self.typecode, self.example[:1])
+        self.assertEqual(
+            a,
+            array.array(self.typecode, [a[0]] * 5)
+        )
+
         self.assertRaises(TypeError, a.__mul__, "bad")
 
     def test_imul(self):
diff --git a/Misc/NEWS b/Misc/NEWS
@@ -45,6 +45,8 @@ Core and Builtins
 Library
 -------
 
+- Issue #1569291: Speed up array.repeat().
+
 - Provide an interface to set the optimization level of compilation in
   py_compile, compileall and zipfile.PyZipFile.
 
diff --git a/Modules/arraymodule.c b/Modules/arraymodule.c
@@ -674,11 +674,9 @@ array_concat(arrayobject *a, PyObject *bb)
 static PyObject *
 array_repeat(arrayobject *a, Py_ssize_t n)
 {
-    Py_ssize_t i;
     Py_ssize_t size;
     arrayobject *np;
-    char *p;
-    Py_ssize_t nbytes;
+    Py_ssize_t oldbytes, newbytes;
     if (n < 0)
         n = 0;
     if ((Py_SIZE(a) != 0) && (n > PY_SSIZE_T_MAX / Py_SIZE(a))) {
@@ -688,13 +686,23 @@ array_repeat(arrayobject *a, Py_ssize_t n)
     np = (arrayobject *) newarrayobject(&Arraytype, size, a->ob_descr);
     if (np == NULL)
         return NULL;
-    p = np->ob_item;
-    nbytes = Py_SIZE(a) * a->ob_descr->itemsize;
-    for (i = 0; i < n; i++) {
-        memcpy(p, a->ob_item, nbytes);
-        p += nbytes;
+    if (n == 0)
+        return (PyObject *)np;
+    oldbytes = Py_SIZE(a) * a->ob_descr->itemsize;
+    newbytes = oldbytes * n;
+    /* this follows the code in unicode_repeat */
+    if (oldbytes == 1) {
+        memset(np->ob_item, a->ob_item[0], newbytes);
+    } else {
+        Py_ssize_t done = oldbytes;
+        Py_MEMCPY(np->ob_item, a->ob_item, oldbytes);
+        while (done < newbytes) {
+            Py_ssize_t ncopy = (done <= newbytes-done) ? done : newbytes-done;
+            Py_MEMCPY(np->ob_item+done, np->ob_item, ncopy);
+            done += ncopy;
+        }
     }
-    return (PyObject *) np;
+    return (PyObject *)np;
 }
 
 static int
