[pypy-commit] pypy default: (arigo, fijal) Implement hacks for array(X, [0]) * number to be a good

[fijal]

Author: Maciej Fijalkowski <fijall at gmail.com>
Branch: 
Changeset: r55895:1f69aec5550f
Date: 2012-07-03 17:38 +0200
http://bitbucket.org/pypy/pypy/changeset/1f69aec5550f/

Log:	(arigo, fijal) Implement hacks for array(X, [0]) * number to be a
	good initializer for arrays of known length and unknown content

diff --git a/pypy/module/array/interp_array.py b/pypy/module/array/interp_array.py
--- a/pypy/module/array/interp_array.py
+++ b/pypy/module/array/interp_array.py
@@ -227,7 +227,7 @@
             # length
             self.setlen(0)
 
-        def setlen(self, size):
+        def setlen(self, size, zero=False):
             if size > 0:
                 if size > self.allocated or size < self.allocated / 2:
                     if size < 9:
@@ -236,11 +236,17 @@
                         some = 6
                     some += size >> 3
                     self.allocated = size + some
-                    new_buffer = lltype.malloc(mytype.arraytype,
-                                               self.allocated, flavor='raw',
-                                               add_memory_pressure=True)
-                    for i in range(min(size, self.len)):
-                        new_buffer[i] = self.buffer[i]
+                    if zero:
+                        new_buffer = lltype.malloc(mytype.arraytype,
+                                                   self.allocated, flavor='raw',
+                                                   add_memory_pressure=True,
+                                                   zero=True)
+                    else:
+                        new_buffer = lltype.malloc(mytype.arraytype,
+                                                   self.allocated, flavor='raw',
+                                                   add_memory_pressure=True)
+                        for i in range(min(size, self.len)):
+                            new_buffer[i] = self.buffer[i]
                 else:
                     self.len = size
                     return
@@ -487,46 +493,50 @@
         return self
 
     def mul__Array_ANY(space, self, w_repeat):
+        return _mul_helper(space, self, w_repeat, False)
+
+    def mul__ANY_Array(space, w_repeat, self):
+        return _mul_helper(space, self, w_repeat, False)
+
+    def inplace_mul__Array_ANY(space, self, w_repeat):
+        return _mul_helper(space, self, w_repeat, True)
+
+    def _mul_helper(space, self, w_repeat, is_inplace):
         try:
             repeat = space.getindex_w(w_repeat, space.w_OverflowError)
         except OperationError, e:
             if e.match(space, space.w_TypeError):
                 raise FailedToImplement
             raise
-        a = mytype.w_class(space)
         repeat = max(repeat, 0)
         try:
             newlen = ovfcheck(self.len * repeat)
         except OverflowError:
             raise MemoryError
+        oldlen = self.len
+        if is_inplace:
+            a = self
+            start = 1
+        else:
+            a = mytype.w_class(space)
+            start = 0
+        # <a performance hack>
+        if oldlen == 1:
+            if self.buffer[0] == rffi.cast(mytype.itemtype, 0):
+                a.setlen(newlen, zero=True)
+                return a
+            a.setlen(newlen)
+            item = self.buffer[0]
+            for r in range(start, repeat):
+                a.buffer[r] = item
+            return a
+        # </a performance hack>
         a.setlen(newlen)
-        for r in range(repeat):
-            for i in range(self.len):
-                a.buffer[r * self.len + i] = self.buffer[i]
+        for r in range(start, repeat):
+            for i in range(oldlen):
+                a.buffer[r * oldlen + i] = self.buffer[i]
         return a
 
-    def mul__ANY_Array(space, w_repeat, self):
-        return mul__Array_ANY(space, self, w_repeat)
-
-    def inplace_mul__Array_ANY(space, self, w_repeat):
-        try:
-            repeat = space.getindex_w(w_repeat, space.w_OverflowError)
-        except OperationError, e:
-            if e.match(space, space.w_TypeError):
-                raise FailedToImplement
-            raise
-        oldlen = self.len
-        repeat = max(repeat, 0)
-        try:
-            newlen = ovfcheck(self.len * repeat)
-        except OverflowError:
-            raise MemoryError
-        self.setlen(newlen)
-        for r in range(1, repeat):
-            for i in range(oldlen):
-                self.buffer[r * oldlen + i] = self.buffer[i]
-        return self
-
     # Convertions
 
     def array_tolist__Array(space, self):
diff --git a/pypy/module/array/test/test_array.py b/pypy/module/array/test/test_array.py
--- a/pypy/module/array/test/test_array.py
+++ b/pypy/module/array/test/test_array.py
@@ -890,6 +890,46 @@
         a[::-1] = a
         assert a == self.array('b', [3, 2, 1, 0])
 
+    def test_array_multiply(self):
+        a = self.array('b', [0])
+        b = a * 13
+        assert b[12] == 0
+        b = 13 * a
+        assert b[12] == 0
+        a *= 13
+        assert a[12] == 0
+        a = self.array('b', [1])
+        b = a * 13
+        assert b[12] == 1
+        b = 13 * a
+        assert b[12] == 1
+        a *= 13
+        assert a[12] == 1
+        a = self.array('i', [0])
+        b = a * 13
+        assert b[12] == 0
+        b = 13 * a
+        assert b[12] == 0
+        a *= 13
+        assert a[12] == 0
+        a = self.array('i', [1])
+        b = a * 13
+        assert b[12] == 1
+        b = 13 * a
+        assert b[12] == 1
+        a *= 13
+        assert a[12] == 1
+        a = self.array('i', [0, 0])
+        b = a * 13
+        assert len(b) == 26
+        assert b[22] == 0
+        b = 13 * a
+        assert len(b) == 26
+        assert b[22] == 0
+        a *= 13
+        assert a[22] == 0
+        assert len(a) == 26        
+
 
 class AppTestArrayBuiltinShortcut(AppTestArray):
     OPTIONS = {'objspace.std.builtinshortcut': True}